MTP support in per-layer fused export

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

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())

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.