fix(export): stop the resume dir growing per layer, and close four review gaps

The auto-derived resume dir was never opted into, and kept one full activation
set per layer -- on a 40-layer model that dwarfs the checkpoint it sits beside.
Prune every boundary but the committed one, after the manifest commits it so a
crash mid-write still resumes. Scoped to the per-layer export path, leaving the
explicit checkpoint_dir semantics untouched.

Fuse the sibling experts the probe never routed to, mirroring the replay
requantize_resmooth_fused_llm_layers does for the same reason: sync_moe_gate_up_amax
covers weight_quantizer.amax but not a static quantizer's global_amax, so unrouted
gate/up pairs would keep unmerged scales.

Drop the format re-check latch -- it inspected only one layer, so a recipe applying
a pre-quant-scale format to a subset still shipped unfused pre_quant_scale.

Refuse MTP checkpoints before calibration on the normal load path: only the FSDP2
loader flags the prefixes early, so the run used to write a complete-looking
checkpoint and fail afterwards.

Also: no MoE gate/up warning at finalize (the sync happens inside
transient_module_state, so it fired on every MoE run and reported a miss that did
not happen); delete out-of-range shards from a longer previous run; null-safe
manifest read plus a num_layers check; drop __all__, which advertised names the
package never exported.

The resume test now interrupts for real rather than rewinding a finished run's
manifest -- a state a crash cannot produce.

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

examples/hf_ptq/hf_ptq.py

@@ -1363,6 +1363,11 @@ def _layerwise_get(cfg, key, default=None):
         # Complementary to recipe `*mtp*` wildcards (name-match); this catches MTP layers
         # identified by index.
         mtp_layer_prefixes = getattr(full_model, "_mtp_layer_prefixes", None)
+        if args.layerwise_export and not mtp_layer_prefixes:
+            # Only the FSDP2 loader flags these before quantization. Per-layer export has
+            # to refuse *before* calibration, or the run writes a complete-looking
+            # checkpoint and only then discovers it is missing the MTP weights.
+            mtp_layer_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path)
         if mtp_layer_prefixes:
             quant_cfg = copy.deepcopy(quant_cfg)
             for prefix in mtp_layer_prefixes:

modelopt/torch/export/layerwise_export.py

@@ -18,6 +18,7 @@
 import contextlib
 import hashlib
 import json
+import re
 import warnings
 from collections.abc import Callable
 from pathlib import Path
@@ -29,13 +30,6 @@
 from .model_config import FUSION_FREE_FORMATS, QUANTIZATION_NVFP4
 from .quant_utils import get_quant_config, get_quantization_format
 
-__all__ = [
-    "LayerwiseExporter",
-    "assert_layerwise_export_supported",
-    "layer_shard_name",
-    "transient_module_state",
-]
-
 # Fusing formats this path can handle itself. The groups _fuse_shared_input_modules works
 # on -- q/k/v behind input_layernorm, gate/up behind post_attention_layernorm -- live
 # inside one decoder layer, so export_layer rediscovers them per layer instead of needing
@@ -72,6 +66,40 @@ def _is_quantized_module(module: nn.Module) -> bool:
     )
 
 
+def _fuse_unrouted_experts(layer_module: nn.Module, fused_linears: dict[str, list[str]]) -> None:
+    """Fuse the sibling experts the probe never routed a token to.
+
+    The probe runs one real batch, so on a 256-expert layer it activates a handful and the
+    rest would keep unmerged gate/up scales -- ``sync_moe_gate_up_amax`` covers
+    ``weight_quantizer.amax`` but not the static quantizer's ``global_amax``. Mirrors the
+    replay ``requantize_resmooth_fused_llm_layers`` does for the same reason, scoped to
+    this layer because experts never span one.
+    """
+    from .quant_utils import preprocess_linear_fusion
+
+    names = {name for name, _ in layer_module.named_modules()}
+    for group, members in fused_linears.items():
+        if not re.search(r"experts?\.\d+", group):
+            continue
+        expert_id = 0
+        while True:
+            sibling = re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", group, count=1)
+            if sibling in fused_linears:  # the probe already fused this one
+                expert_id += 1
+                continue
+            if sibling not in names:
+                break
+            preprocess_linear_fusion(
+                [
+                    layer_module.get_submodule(
+                        re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", member)
+                    )
+                    for member in members
+                ]
+            )
+            expert_id += 1
+
+
 def _module_formats(model: nn.Module) -> set:
     """Every distinct format present. ``get_quantization_format`` stops at the first."""
     return {
@@ -279,7 +307,6 @@ def __init__(
             else raw_tied_keys
         )
 
-        self._calibrated_format_checked = False
         self._bind_identity(quant_config)
 
     def export_layer(
@@ -335,12 +362,10 @@ def _assert_calibrated_format_supported(self, layer_module: nn.Module) -> None:
 
         AWQ and SVDQuant are only distinguishable after calibration -- their discriminators
         (``_pre_quant_scale``, ``svdquant_lora_a``) are registered by the calibrator -- so
-        the constructor's gate sees plain NVFP4 and passes. Checking again on the first
-        exported layer fails loudly instead of shipping unfused pre_quant_scale.
+        the constructor's gate sees plain NVFP4 and passes. Every layer, not just the
+        first: a recipe may apply such a format to a subset, and a resumed run starts
+        part-way through. One walk per layer is one whole-model walk in total.
         """
-        if self._calibrated_format_checked:
-            return
-        self._calibrated_format_checked = True
         unsupported = sorted(str(f) for f in _module_formats(layer_module) - SUPPORTED_FORMATS)
         if unsupported:
             raise NotImplementedError(
@@ -377,9 +402,10 @@ def _fuse_shared_inputs(
         input_to_linear, _ = collect_shared_input_modules(
             layer_module, lambda: probe_forward(layer_module)
         )
-        _fuse_shared_input_modules(
+        fused = _fuse_shared_input_modules(
             self._ctx.model, input_to_linear, quantization_format=layer_format
         )
+        _fuse_unrouted_experts(layer_module, fused)
 
     def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> dict:
         """Export the tail, write the config artifacts, and index all shards.
@@ -395,7 +421,6 @@ def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> d
         from .unified_export_hf import (
             _add_mtp_exclusions,
             _dispatch_export_handler,
-            _warn_on_unsynced_moe_gate_up,
             _write_hf_export_config,
             save_non_weight_artifacts,
         )
@@ -409,7 +434,9 @@ def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> d
         model = self._ctx.model
         quant_config = get_quant_config(model, is_modelopt_qlora=self._ctx.is_modelopt_qlora)
         _add_mtp_exclusions(model, quant_config)
-        _warn_on_unsynced_moe_gate_up(model)
+        # No _warn_on_unsynced_moe_gate_up: export_layer syncs each layer inside
+        # transient_module_state, so the shards are synced but the live model is not --
+        # the check would fire on every MoE run and report a miss that did not happen.
         if getattr(model, "hf_quantizer", None) is not None:
             model.hf_quantizer = None
         # Config module names must match the tensors', or a loader treats an excluded
@@ -585,6 +612,12 @@ def _write_index(self) -> None:
         """
         from safetensors import safe_open
 
+        # A longer previous run's shards are already out of the index; delete them too,
+        # so the directory *is* the checkpoint rather than the checkpoint plus leftovers.
+        for stale in self._export_dir.glob("model-layer-*.safetensors"):
+            if int(stale.stem.rsplit("-", 1)[1]) >= len(self._layers):
+                stale.unlink()
+
         shards = [self._export_dir / layer_shard_name(i) for i in range(len(self._layers))]
         shards.append(self._export_dir / _TAIL_SHARD)
 

modelopt/torch/quantization/model_calib.py

@@ -2120,8 +2120,15 @@ def layerwise_calibrate(
         # detect_resume_point returns None once the manifest is complete, which would put
         # start_layer back at 0 and recalibrate everything. The shards are already on disk,
         # so a completed manifest means finalize-only.
-        manifest = _read_manifest(checkpoint_dir)
-        if manifest is not None and manifest.get("last_completed_layer", -1) + 1 >= num_layers:
+        manifest = _read_manifest(checkpoint_dir) or {}
+        last, total = manifest.get("last_completed_layer"), manifest.get("num_layers")
+        if total is not None and total != num_layers:
+            raise ValueError(
+                f"Layerwise checkpoint at {checkpoint_dir} was written for {total} layers "
+                f"but this model has {num_layers}. Use a fresh checkpoint_dir."
+            )
+        # None, not -1: a truncated or hand-edited manifest must not arithmetic-error here.
+        if last is not None and last + 1 >= num_layers:
             exporter.assert_shards_present(num_layers)
             exporter.finalize()
             print_rank_0(f"Layerwise export: finalized existing shards in {export_dir}")

modelopt/torch/quantization/utils/layerwise_calib.py

@@ -737,6 +737,15 @@ def full_restore(self, layers: nn.ModuleList, model: nn.Module) -> None:
 
         print_rank_0(f"Checkpoint: restored {self.start_layer} previously calibrated layers")
 
+    def _prune_stale_next_inputs(self, keep: int) -> None:
+        """Drop every layer's cached activations but the committed boundary's."""
+        for idx in range(self.num_layers):
+            if idx == keep:
+                continue
+            stale = os.path.join(_layer_dir(self.checkpoint_dir, idx), "next_inputs.pt")
+            if os.path.exists(stale):
+                os.remove(stale)
+
     def save(
         self,
         layer_idx: int,
@@ -803,6 +812,12 @@ def save(
             calib_mutates_weights=self.calib_mutates_weights,
             save_layer_state=self.save_layer_state,
         )
+        # Per-layer export only, whose resume dir is auto-derived and so was never opted
+        # into: without this it keeps one activation set per layer, which dwarfs the
+        # checkpoint it sits beside. After the manifest commits the boundary, so a crash
+        # mid-write still leaves the previous one resumable.
+        if not self.save_layer_state:
+            self._prune_stale_next_inputs(keep=layer_idx)
         window_start = self._last_saved_layer + 1
         self._last_saved_layer = layer_idx
         window_size = layer_idx - window_start + 1

tests/gpu/torch/export/test_layerwise_export.py

@@ -17,14 +17,15 @@
 
 import copy
 import json
-import shutil
+from unittest.mock import patch
 
 import pytest
 import torch
-from _test_utils.torch.transformers_models import get_tiny_llama
+from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe
 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.unified_export_hf import export_hf_checkpoint
 
 NUM_LAYERS = 4
@@ -215,28 +216,41 @@ def test_layerwise_export_replaces_resume_artifacts(tmp_path):
 
     assert not list(checkpoint_dir.rglob("weights.pt"))
     assert not list(checkpoint_dir.rglob("quantizer_buffers.pt"))
-    # next_inputs and output_meta are not reconstructible from exported weights, so they stay.
+    # output_meta is not reconstructible from exported weights, so it stays.
     assert list(checkpoint_dir.rglob("output_meta.pt"))
+    # The cached activations are the bulk of the resume dir, and only the committed
+    # boundary's are resumable -- keeping one per layer would dwarf the checkpoint.
+    assert not list(checkpoint_dir.rglob("next_inputs.pt")), (
+        "a completed run has nothing to resume from, so no activation cache should remain"
+    )
 
 
 def test_resume_skips_exported_layers(tmp_path, baseline_checkpoint):
     """A run resuming mid-model must still produce the full, correct checkpoint."""
     export_dir = tmp_path / "fused"
     checkpoint_dir = tmp_path / "ckpt"
-    mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib)
+    # Die partway, the way a lost GPU session would: only the committed boundary is
+    # resumable, so rewinding a *finished* run's manifest would not reproduce this state.
+    real_export_layer = LayerwiseExporter.export_layer
+
+    def die_at_layer_2(self, layer_idx, *args, **kwargs):
+        if layer_idx == 2:
+            raise RuntimeError("simulated interruption")
+        return real_export_layer(self, layer_idx, *args, **kwargs)
+
+    with (
+        patch.object(LayerwiseExporter, "export_layer", die_at_layer_2),
+        pytest.raises(RuntimeError, match="simulated interruption"),
+    ):
+        mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib)
 
-    # Rewind the manifest so the next run believes only layers 0..1 finished; their shards
-    # are on disk and must be reused rather than recalculated.
-    manifest_path = checkpoint_dir / "manifest.json"
-    manifest = json.loads(manifest_path.read_text())
-    manifest["last_completed_layer"] = 1
-    manifest_path.write_text(json.dumps(manifest))
+    assert (export_dir / layer_shard_name(1)).is_file(), "layer 1 was never committed"
+    assert not (export_dir / layer_shard_name(2)).exists(), "layer 2 should not have landed"
 
-    resumed_dir = tmp_path / "resumed"
-    shutil.copytree(export_dir, resumed_dir)
-    mtq.quantize(_build_model(), _layerwise_cfg(resumed_dir, checkpoint_dir), _calib)
+    # Shards 0..1 are on disk and must be reused rather than recalculated.
+    mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib)
 
-    _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(resumed_dir))
+    _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(export_dir))
 
 
 def test_resume_without_matching_shards_fails_fast(tmp_path):
@@ -347,6 +361,33 @@ def test_identity_without_shards_does_not_block_a_rerun(tmp_path):
     assert _load_checkpoint(export_dir), "rerun produced no checkpoint"
 
 
+def _build_moe_model():
+    torch.manual_seed(0)
+    model = get_tiny_qwen3_moe(num_experts=16, num_experts_per_tok=1).cuda().eval()
+    model.config.architectures = ["Qwen3MoeForCausalLM"]
+    return model
+
+
+def test_moe_export_matches(tmp_path):
+    """MoE layers take a different path: fused expert inputs and gate/up amax sync.
+
+    This fixture uses the fused expert representation (one ``mlp.experts`` module), so it
+    does not cover the per-expert sibling replay in ``_fuse_unrouted_experts`` -- that
+    needs a checkpoint whose experts are separate ``experts.N.gate_proj`` modules.
+    """
+    baseline_dir = tmp_path / "baseline"
+    base = _nvfp4_cfg()
+    base["algorithm"] = {"method": "max", "layerwise": {"enable": True}}
+    export_hf_checkpoint(mtq.quantize(_build_moe_model(), base, _calib), export_dir=baseline_dir)
+
+    export_dir = tmp_path / "fused"
+    mtq.quantize(
+        _build_moe_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()), _calib
+    )
+
+    _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir))
+
+
 def test_export_does_not_mutate_the_model(tmp_path):
     """Exporting a layer must leave the model exactly as calibration left it.