diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 66bc31f2cfa..c3563d8e4a3 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -77,6 +77,7 @@ has_spec_opt, save_expert_token_count_table, ) +from modelopt.torch.export.layerwise_export import MTP_EXTRA_STATE_ATTR from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights @@ -770,7 +771,7 @@ def mono_quantize( warnings.warn("Skipping quantization: model is already quantized.") -def assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) -> None: +def assert_layerwise_export_compatible(args, full_model) -> None: """Refuse layerwise export before calibration starts, not after it writes a checkpoint. Layerwise export writes the finished checkpoint during calibration, so anything that @@ -785,13 +786,6 @@ def assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) -> "overwrite config.json with the unquantized source config." ) - if mtp_layer_prefixes: - raise NotImplementedError( - f"layerwise.export_dir does not support models with MTP layers {mtp_layer_prefixes}: " - "their exclusions and any orphaned MTP weights are applied after calibration, by " - "which point every shard and the quant config are already written." - ) - if has_spec_opt(full_model): raise NotImplementedError( "layerwise.export_dir does not support speculative-decoding models: " @@ -932,12 +926,13 @@ def export_quantized( full_model._mtp_layer_prefixes = mtp_layer_prefixes if args.layerwise_export: - if mtp_state_dict: + staged = getattr(full_model, MTP_EXTRA_STATE_ATTR, None) or {} + if mtp_state_dict.keys() - staged.keys(): raise NotImplementedError( - "layerwise.export_dir does not support models with MTP weights: " - "they are loaded after calibration has already written every " - "shard, so they would be missing from the checkpoint. Export " - "without layerwise.export_dir." + "layerwise.export_dir found MTP weights it did not stage before " + f"calibration: {sorted(mtp_state_dict.keys() - staged.keys())[:4]}. " + "Every shard is already written, so they cannot be added now. " + "Export without layerwise.export_dir." ) # Calibration already wrote every shard, the index and the configs. print(f"Layerwise export already wrote the checkpoint to {export_path}") @@ -1378,10 +1373,20 @@ def _layerwise_get(cfg, key, default=None): quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) print(f"Excluding MTP layer from quantization: {pattern}") + if args.layerwise_export and mtp_layer_prefixes: + # Per-layer export writes the checkpoint *during* calibration, so the MTP + # weights have to be in place first. load_mtp_weights only fills existing slots + # and hands back the rest, so running it early is safe; the orphans are stashed + # for finalize(), which owns the tail shard. + _, mtp_state_dict = load_mtp_weights(full_model, args.pyt_ckpt_path) + if mtp_state_dict: + setattr(full_model, MTP_EXTRA_STATE_ATTR, mtp_state_dict) + print(f"Layerwise export: staged {len(mtp_state_dict)} orphaned MTP tensors") + # Before resolve_checkpoint_dir, which hashes the config: with the placeholder # still in it, two --export_path values would share one checkpoint dir. if args.layerwise_export: - assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) + assert_layerwise_export_compatible(args, full_model) quant_cfg = set_layerwise_export_dir(quant_cfg, args.export_path) print(f"Layerwise export enabled: writing quantized shards to {args.export_path}") # The shards are only a resume artifact if the manifest that names the resume diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 62ba5a64776..a5672feb5d2 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -42,6 +42,12 @@ _INDEX_FILE = "model.safetensors.index.json" _IDENTITY_FILE = ".layerwise_export.json" +#: Set by the caller on the model, holding MTP tensors that have no slot in +#: ``state_dict()``. finalize() runs inside calibration, so the caller cannot pass them as +#: an argument; this follows the ``_mtp_layer_prefixes`` convention already used to hand +#: MTP information across the same boundary. +MTP_EXTRA_STATE_ATTR = "_mtp_extra_state_dict" + def layer_shard_name(layer_idx: int) -> str: """Shard filename for one decoder layer. @@ -451,7 +457,9 @@ def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> d self._collect(tail, name, tensor) # Tensors the model never held (e.g. orphaned MTP weights), already in export - # form, so only the hub-name reversal applies. + # form, so only the hub-name reversal applies. The stash is the layerwise route: + # calibration owns the finalize() call, so the caller cannot pass them directly. + extra_state_dict = extra_state_dict or getattr(model, MTP_EXTRA_STATE_ATTR, None) for name, tensor in (extra_state_dict or {}).items(): mapped = self._name_mapper(name) if self._name_mapper is not None else name tail.setdefault(mapped, tensor.detach().contiguous().cpu()) diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index 34f70f5502f..b46c35c9bca 100644 --- a/tests/gpu/torch/export/test_layerwise_export.py +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -26,7 +26,11 @@ from safetensors.torch import load_file import modelopt.torch.quantization as mtq -from modelopt.torch.export.layerwise_export import LayerwiseExporter, layer_shard_name +from modelopt.torch.export.layerwise_export import ( + MTP_EXTRA_STATE_ATTR, + LayerwiseExporter, + layer_shard_name, +) from modelopt.torch.export.unified_export_hf import export_hf_checkpoint NUM_LAYERS = 4 @@ -415,6 +419,32 @@ def test_moe_export_matches(tmp_path): _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) +def test_orphaned_mtp_tensors_reach_the_tail_shard(tmp_path): + """MTP weights with no slot in state_dict() must still land in the checkpoint. + + finalize() runs inside calibration, so hf_ptq cannot pass them as an argument; it + stashes them on the model and the exporter picks them up. Without that they are + silently absent from a checkpoint that otherwise looks complete. + """ + export_dir = tmp_path / "fused" + model = _build_model() + orphans = { + "mtp.layers.0.weight": torch.ones(4, 4, dtype=torch.bfloat16), + "mtp.norm.weight": torch.ones(4, dtype=torch.bfloat16), + } + setattr(model, MTP_EXTRA_STATE_ATTR, orphans) + + mtq.quantize(model, _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib) + + exported = _load_checkpoint(export_dir) + for key, value in orphans.items(): + assert key in exported, f"{key} missing from the exported checkpoint" + assert torch.equal(exported[key].cpu(), value) + + weight_map = json.loads((export_dir / "model.safetensors.index.json").read_text())["weight_map"] + assert set(orphans) <= set(weight_map), "orphans written but left out of the index" + + def test_export_does_not_mutate_the_model(tmp_path): """Exporting a layer must leave the model exactly as calibration left it.