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

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. Supports FP8 and NVFP4 on single-process models, resident or accelerate-offloaded; fusing formats work because each layer rediscovers its own q/k/v and gate/up scale-fusion groups. AWQ and SVDQuant (whose pre-quant-scale steps are still whole-model), 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 field implies ``--skip_generate``.
 
 *Misc*
 

examples/hf_ptq/example_utils.py

@@ -1166,6 +1166,23 @@ 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)
+    algorithm = quant_cfg.get("algorithm")
+    # A recipe may carry one algorithm or a list of them; detection accepts both, so the
+    # retarget has to as well or a list-shaped recipe dies on a str index.
+    for entry in algorithm if isinstance(algorithm, list) else [algorithm]:
+        if isinstance(entry, dict) and isinstance(entry.get("layerwise"), dict):
+            entry["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,61 @@ 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."
+        )
+
+    # Mirrors the branch conditions in export_quantized: anything that routes to a second
+    # exporter would write over the --export_path layerwise calibration already populated.
+    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()"),
+        ("--qformat int8_sq", "int8_sq" in args.qformat, "export_tensorrt_llm_checkpoint()"),
+        (
+            "an encoder-decoder model_type (t5/bart/whisper)",
+            getattr(full_model.config, "model_type", None) in ("t5", "bart", "whisper"),
+            "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 +971,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(
@@ -1176,17 +1243,46 @@ def quantize_main(
         aq_config = None
         fixed_quantize_config = None
 
-    def _is_layerwise(obj):
+    def _layerwise_cfg(obj):
+        """The recipe's ``layerwise`` block, or None.
+
+        An algorithm parsed from YAML arrives as a plain dict, while the deprecated
+        ``--auto_quantize_*`` path builds config objects, so both shapes reach here.
+        """
         if isinstance(obj, ModelOptPTQRecipe):
-            return _is_layerwise(obj.quantize.algorithm)
+            return _layerwise_cfg(obj.quantize.algorithm)
         if isinstance(obj, ModelOptAutoQuantizeRecipe):
-            return obj.quantize is not None and _is_layerwise(obj.quantize.algorithm)
+            return _layerwise_cfg(obj.quantize.algorithm) if obj.quantize is not None else None
         if isinstance(obj, list):
-            return any(_is_layerwise(a) for a in obj)
-        layerwise = getattr(obj, "layerwise", None)
-        return bool(getattr(layerwise, "enable", False))
-
-    is_layerwise = _is_layerwise(recipe)
+            return next((cfg for cfg in map(_layerwise_cfg, obj) if cfg is not None), None)
+        if isinstance(obj, dict):
+            return obj.get("layerwise")
+        return getattr(obj, "layerwise", None)
+
+    def _layerwise_get(cfg, key, default=None):
+        if cfg is None:
+            return default
+        return cfg.get(key, default) if isinstance(cfg, dict) else getattr(cfg, key, default)
+
+    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:
+        if isinstance(recipe, ModelOptAutoQuantizeRecipe):
+            # Only the mono-quantize path retargets export_dir and runs the compatibility
+            # refusals. Reaching auto_quantize with the flag set would export to the
+            # recipe's placeholder directory and skip export_hf_checkpoint(), leaving
+            # --export_path with no weights and no error.
+            raise NotImplementedError(
+                "layerwise.export_dir is not supported with an AutoQuantize recipe; "
+                "use a PTQ recipe, or drop export_dir and export afterwards."
+            )
+        # 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
@@ -1316,6 +1412,15 @@ def _is_layerwise(obj):
                 quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False})
                 print(f"Excluding MTP layer from quantization: {pattern}")
 
+        # Retarget export_dir before resolving checkpoint_dir: the resolved name hashes the
+        # config, so hashing it with the recipe's placeholder would make two runs to
+        # different --export_path values share one checkpoint dir and resume against the
+        # wrong shards.
+        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 needs_checkpoint_path_update(quant_cfg):
             quant_cfg, resolved_dir = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path)
             print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}")