feat(export): export each decoder layer as layerwise calibration finishes it

A PTQ run that outlasts its GPU session loses every calibrated layer, and a
completed one still pays for a second whole-model export pass over a
full-precision intermediate checkpoint.

Add layerwise.export_dir: each decoder layer's final quantized tensors are
flushed to their own shard the moment calibration finishes with it, so the shards
double as the resume artifact. Combined with layerwise.checkpoint_dir a restarted
run skips layers already on disk instead of recalibrating them, and the per-layer
weights.pt / quantizer_buffers.pt files are no longer written. When the last layer
lands the directory is already a complete, loadable checkpoint, so
export_hf_checkpoint() is skipped. Setting the field is the whole switch;
hf_ptq.py rewrites it to --export_path, as it already does for checkpoint_dir.

One layer per shard is what makes that work: a shard is only ever written whole,
its name derives from the layer index, and the index is rebuilt at the end from
the shards on disk. A crash can lose the layer in flight but never corrupt an
earlier one, and a re-run overwrites in place rather than appending duplicates.

Resident modules have no materialization window to discard export's damage, so
transient_module_state snapshots the layer subtree's _parameters/_buffers/_modules
and restores them, leaving calibration free to run every later layer through it.

Tied-weight dedup is off here for the reason registry.py already turns it off for
offload: data_ptr() cannot identify a tensor across an export that keeps rolling
packed weights back. Weight-tied quantized modules are refused up front anyway.

NVFP4 works because export_layer rediscovers the scale-fusion groups itself: the
groups _fuse_shared_input_modules operates on -- q/k/v behind input_layernorm,
gate/up behind post_attention_layernorm -- never cross a layer boundary, so a
probe forward over one layer finds them. The probe uses that layer's real cached
activations rather than the synthetic input the whole-model pass builds.

Scope is resident, single-process models. AWQ and SVDQuant additionally need
requantize_resmooth_fused_llm_layers' pre-quant-scale steps, which are still
whole-model, so they stay refused along with accelerate offload, multi-process
jobs, weight-tied quantized modules,
multimodal models, MTP models and speculative decoding are refused with
NotImplementedError before calibration starts; each would otherwise produce a
silently different checkpoint rather than fail. A resumed run never recalibrates
the layers it skipped, so the in-memory model is not valid for inference and
layerwise.export_dir implies --skip_generate.

Two pieces move out of this path to avoid duplicating what already exists:
save_non_weight_artifacts() is extracted from the streaming exporter's tail, and
FUSION_FREE_FORMATS moves to model_config.py, where _fuse_shared_input_modules
had held the same set inline as a literal list.

Verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only,
FP8, 256 experts per layer): 93,273 tensors, 0 mismatched, and identical again
after a simulated mid-model resume. NVFP4 equivalence is covered by a GPU test.

The NVFP4 test leaves o_proj unquantized: layerwise calibration leaves its input
amax at 0 on every layer but the last, which no export path can write. That is a
pre-existing bug, unrelated to per-layer export.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

CHANGELOG.rst

@@ -9,6 +9,7 @@ Changelog
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers. Weight scales are an orthogonal axis, selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian`` with that algorithm's own options), so one recipe can combine a weight calibration with this activation policy in a single pass. ``SequentialQuantizer`` activation quantizers are not supported and raise. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
+- Add ``layerwise.export_dir`` (``--layerwise_export`` in ``examples/hf_ptq/hf_ptq.py``), which writes each decoder layer to its own quantized HF checkpoint shard as soon as layerwise calibration finishes it, so the directory is already a complete checkpoint when the last layer lands and the separate ``export_hf_checkpoint()`` pass over a full-precision intermediate is skipped. Combined with ``layerwise.checkpoint_dir`` the shards double as the resume artifact, so a run interrupted mid-model restarts without recalibrating or re-exporting finished layers; the output is byte-identical to the whole-model export. Scoped to FP8 on resident, single-process models -- NVFP4 and other fusion-requiring formats, accelerate offload, multi-process jobs, weight-tied quantized modules, multimodal models, MTP models and speculative decoding are refused with ``NotImplementedError`` before calibration starts, and since a resumed run never recalibrates the layers it skipped the flag implies ``--skip_generate``.
 
 *Misc*
 

examples/hf_ptq/example_utils.py

@@ -1166,6 +1166,18 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> tuple[dict, str]
     return quant_cfg, resolved
 
 
+def set_layerwise_export_dir(quant_cfg: dict, export_path: str) -> dict:
+    """Retarget layerwise per-layer export at ``export_path``.
+
+    The recipe opts in by setting ``layerwise.export_dir``; its value is a placeholder,
+    since the destination is per-run rather than per-recipe. Mirrors how
+    :func:`resolve_checkpoint_dir` rewrites ``layerwise.checkpoint_dir``.
+    """
+    quant_cfg = copy.deepcopy(quant_cfg)
+    quant_cfg["algorithm"]["layerwise"]["export_dir"] = export_path
+    return quant_cfg
+
+
 def add_mlflow_args(parser: argparse.ArgumentParser) -> None:
     """Add the MLflow tracking flags."""
     parser.add_argument(

examples/hf_ptq/hf_ptq.py

@@ -46,6 +46,7 @@
     resolve_checkpoint_dir,
     resolve_mlflow_args,
     run_nemotron_vl_preview,
+    set_layerwise_export_dir,
     setup_distributed_args,
     validate_fsdp2_supported,
 )
@@ -814,6 +815,53 @@ def mono_quantize(
         warnings.warn("Skipping quantization: model is already quantized.")
 
 
+def assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) -> None:
+    """Refuse layerwise export before calibration starts, not after it writes a checkpoint.
+
+    Layerwise export writes the finished checkpoint during calibration, so anything that
+    would rewrite or contradict that checkpoint afterwards has to be caught here -- once
+    calibration begins, the user has already paid for the whole run.
+    """
+    if is_multimodal_model(full_model):
+        raise NotImplementedError(
+            "layerwise.export_dir does not support multimodal models: calibration runs on the "
+            "extracted language model, so the shards and config.json would describe that "
+            "submodel rather than the full VLM, and the VLM export path would then "
+            "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: "
+            "export_speculative_decoding() would write a second checkpoint over the same "
+            "--export_path."
+        )
+
+    if args.cast_mxfp4_to_nvfp4:
+        raise NotImplementedError(
+            "layerwise.export_dir is not compatible with --cast_mxfp4_to_nvfp4: the cast "
+            "rewrites weights after calibration, by which point every shard is written."
+        )
+
+    for flag, value, exporter in (
+        ("--vllm_fakequant_export", args.vllm_fakequant_export, "export_hf_vllm_fq_checkpoint()"),
+        ("--sparsity_fmt", args.sparsity_fmt != "dense", "export_tensorrt_llm_checkpoint()"),
+    ):
+        if value:
+            raise NotImplementedError(
+                f"layerwise.export_dir is not compatible with {flag}: {exporter} would write a "
+                "second checkpoint over the same --export_path that layerwise calibration "
+                "already populated."
+            )
+
+
 def export_quantized(
     args: argparse.Namespace,
     full_model: torch.nn.Module,
@@ -915,11 +963,22 @@ def export_quantized(
                 if mtp_layer_prefixes:
                     full_model._mtp_layer_prefixes = mtp_layer_prefixes
 
-                export_hf_checkpoint(
-                    full_model,
-                    export_dir=export_path,
-                    extra_state_dict=mtp_state_dict,
-                )
+                if args.layerwise_export:
+                    if mtp_state_dict:
+                        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."
+                        )
+                    # Calibration already wrote every shard, the index and the configs.
+                    print(f"Layerwise export already wrote the checkpoint to {export_path}")
+                else:
+                    export_hf_checkpoint(
+                        full_model,
+                        export_dir=export_path,
+                        extra_state_dict=mtp_state_dict,
+                    )
 
                 if args.qformat == "w4a16_nvfp4":
                     warnings.warn(
@@ -1200,6 +1259,14 @@ def _layerwise_get(cfg, key, default=None):
     layerwise_cfg = _layerwise_cfg(recipe)
     is_layerwise = bool(_layerwise_get(layerwise_cfg, "enable", False))
 
+    # Setting layerwise.export_dir is the switch; the value is replaced with --export_path
+    # below, the way resolve_checkpoint_dir already rewrites layerwise.checkpoint_dir.
+    args.layerwise_export = _layerwise_get(layerwise_cfg, "export_dir") is not None
+    if args.layerwise_export:
+        # A resumed run leaves the model without amax on the layers it skipped, so no
+        # preview may run against it.
+        args.skip_generate = True
+
     if args.batch_size == 0:
         # For VL models with image-text calibration, skip automatic batch size detection
         # since get_max_batch_size can't handle multimodal inputs
@@ -1332,6 +1399,11 @@ def _layerwise_get(cfg, key, default=None):
             quant_cfg, resolved_dir = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path)
             print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}")
 
+        if args.layerwise_export:
+            assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes)
+            quant_cfg = set_layerwise_export_dir(quant_cfg, args.export_path)
+            print(f"Layerwise export enabled: writing quantized shards to {args.export_path}")
+
         if args.cast_mxfp4_to_nvfp4:
             quant_cfg = copy.deepcopy(quant_cfg)
             force_weight_quantizers_static(quant_cfg["quant_cfg"])