diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 47ad628cf89..f38f643978a 100755 --- a/CHANGELOG.rst +++ b/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* diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 740a09b2267..c4e4b42132f 100755 --- a/examples/hf_ptq/example_utils.py +++ b/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( diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 7a2328d10f7..105f34477a4 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/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}") diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py new file mode 100644 index 00000000000..5f13e478439 --- /dev/null +++ b/modelopt/torch/export/layerwise_export.py @@ -0,0 +1,565 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Write each decoder layer's quantized checkpoint shard as soon as it is calibrated.""" + +import contextlib +import json +import warnings +from collections.abc import Callable +from pathlib import Path + +import torch +import torch.nn as nn +from safetensors.torch import save_file + +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 +# the whole-model forward. AWQ and SVDQuant are excluded: they additionally need +# requantize_resmooth_fused_llm_layers' pre-quant-scale steps, which are still model-wide. +_PER_LAYER_FUSABLE_FORMATS = frozenset({QUANTIZATION_NVFP4}) + +SUPPORTED_FORMATS = FUSION_FREE_FORMATS | _PER_LAYER_FUSABLE_FORMATS + +_TAIL_SHARD = "model-tail.safetensors" +_INDEX_FILE = "model.safetensors.index.json" + + +def layer_shard_name(layer_idx: int) -> str: + """Shard filename for one decoder layer. + + Derived from the index rather than a running counter so that re-exporting a layer + overwrites its shard instead of leaving a stale copy behind for the index to pick up. + """ + return f"model-layer-{layer_idx:05d}.safetensors" + + +def _is_quantized_module(module: nn.Module) -> bool: + """Whether this module carries quantizers of its own. + + Detected by type, not attribute name: fused-expert modules name theirs after the weight + (``gate_up_proj_weight_quantizer``), so a name-based test misses whole MoE blocks. + """ + from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer + + return any( + isinstance(child, (TensorQuantizer, SequentialQuantizer)) for child in module.children() + ) + + +def _module_formats(model: nn.Module) -> set: + """Every distinct quantization format present, not just the first one found. + + ``get_quantization_format(model)`` stops at the first quantized child, so gating on it + would let an NVFP4 layer slip through a check meant to exclude it. + """ + return { + get_quantization_format(module) + for _, module in model.named_modules() + if _is_quantized_module(module) + } + + +def _tied_quantized_modules(model: nn.Module) -> list[str]: + """Names of quantized modules that share a weight tensor with another quantized module. + + The whole-model export merges their amaxes via ``sync_tied_input_amax``; a per-layer + pass cannot, since a tie partner may be uncalibrated or already written. + """ + by_ptr: dict[int, list[str]] = {} + for name, module in model.named_modules(): + weight = getattr(module, "weight", None) + if weight is None or not _is_quantized_module(module) or weight.is_meta: + continue + ptr = weight.data_ptr() + # data_ptr() is 0 for meta tensors and DTensors, grouping unrelated modules. + if ptr: + by_ptr.setdefault(ptr, []).append(name) + return sorted(n for names in by_ptr.values() if len(names) > 1 for n in names) + + +def assert_layerwise_export_supported(model: nn.Module) -> None: + """Raise ``NotImplementedError`` unless per-layer export is valid for this model. + + Each case would otherwise produce a checkpoint differing from a whole-model export + without failing, so all are rejected before the first shard is written. + + The central one is the format gate. Both other export paths begin with + ``requantize_resmooth_fused_llm_layers``, which this one never calls: its core step + discovers modules sharing an input via a dummy forward over the *whole* model, and + there is no such forward here. Restricting to ``FUSION_FREE_FORMATS`` is what makes + that omission invisible -- for those formats all three of its steps are no-ops, since + pre-quant-scale fusion requires ``nvfp4_awq`` and MoE expert resmoothing requires AWQ + or SVDQuant. Any other format would silently lose them. + + .. todo:: + Support the fusing formats (NVFP4 above all) by making + ``collect_shared_input_modules`` operate on a single decoder layer rather than the + whole model. The groups it finds -- q/k/v, gate/up -- are intra-layer, so a dummy + forward over one layer can discover them; the whole-model scope is an artifact of + how the batch exporter happens to call it, not a requirement. AWQ and SVDQuant + additionally need the pre-quant-scale steps made per-layer. + """ + from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload + from modelopt.torch.utils import distributed as dist + + unsupported = sorted(str(f) for f in _module_formats(model) - SUPPORTED_FORMATS) + if unsupported: + raise NotImplementedError( + f"layerwise export does not support quantization format(s) {unsupported}: they " + "need requantize_resmooth_fused_llm_layers' pre-quant-scale steps, which are " + "still whole-model. Supported today: " + f"{sorted(str(f) for f in SUPPORTED_FORMATS if f)}." + ) + + tied = _tied_quantized_modules(model) + if has_accelerate_offload(model) and not tied: + # Offloaded weights sit on meta between windows, so the grouping above inspected + # nothing; an empty result is not a clean bill of health. Resolving ties by name + # instead would survive weight moves -- the same fix ExportContext.__post_init__ + # already carries a TODO for. + warnings.warn( + "Weight-tie detection is skipped for offloaded models: data_ptr() cannot group " + "weights that are not resident. A model with tied quantized modules would " + "export per-module input_scale instead of the merged value." + ) + + if tied: + raise NotImplementedError( + f"layerwise export does not support weight-tied quantized modules {tied[:6]}: " + "the whole-model path merges their input_quantizer amaxes via " + "sync_tied_input_amax so both sides share one input_scale, which a per-layer " + "pass cannot do because a tie partner may be uncalibrated or already written." + ) + + if dist.is_initialized() and dist.size() > 1: + raise NotImplementedError( + "layerwise export does not support multi-process jobs (e.g. FSDP2): every rank " + "would write the same shard files. Use single-process calibration." + ) + + +@contextlib.contextmanager +def transient_module_state(module: nn.Module): + """Undo everything export does to ``module``, so calibration can continue through it. + + Export is destructive -- packed weights, new scale buffers, grafted per-expert + submodules. An offloaded model discards that when its materialization window closes; a + resident one has no window, and calibration still has every later layer to run. + + Restoring the dicts suffices, and costs references rather than a deep copy, because + export rebinds them instead of mutating tensors in place. + """ + snapshot = [ + (m, dict(m._parameters), dict(m._buffers), dict(m._modules)) for m in module.modules() + ] + try: + yield + finally: + for m, params, buffers, children in snapshot: + m._parameters.clear() + m._parameters.update(params) + m._buffers.clear() + m._buffers.update(buffers) + m._modules.clear() + m._modules.update(children) + + +class LayerwiseExporter: + """Writes one decoder layer's quantized shard per call, then the tail and index. + + Constructed before calibration begins, driven once per layer from inside the window + calibration already opens, and finalized after the last one:: + + exporter = LayerwiseExporter(model, export_dir) + ... + with persistent_materialization(layer, writeback=False): + calib_func(layer, ...) + exporter.export_layer(layer_idx, layer) + ... + quant_config = exporter.finalize(extra_state_dict=mtp_state_dict) + + ``finalize()`` rebuilds the index from the shards present on disk, so layers exported + by an earlier run that this one skipped are picked up without being re-exported. + """ + + def __init__( + self, + model: nn.Module, + export_dir: Path | str, + dtype: torch.dtype | None = None, + is_modelopt_qlora: bool = False, + ) -> None: + """Validate support and capture model-level state, before calibration runs. + + Only quantizer *configuration* is read here, which ``mtq.quantize`` fixes when it + swaps modules; anything amax-dependent belongs in :meth:`finalize`. + """ + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + from .layer_utils import is_moe + from .quant_aware_conversion import build_reverse_name_mapper + from .registry import ExportContext, PrepareMoEInputsRegistry + from .unified_export_hf import _resolve_export_dtype + from .unified_export_hf_streaming import _assert_no_split_rules + + assert_layerwise_export_supported(model) + # Splits regroup tensors across the whole state dict; no per-layer pass can reverse it. + _assert_no_split_rules(model) + + for _, sub_module in model.named_modules(): + if ( + is_moe(sub_module) + and hasattr(sub_module, "experts") + and PrepareMoEInputsRegistry.match(sub_module.experts) is None + ): + raise NotImplementedError( + f"MoE model with experts type '{type(sub_module.experts).__name__}' is " + "not supported in export." + ) + + layers = LayerActivationCollector.get_decoder_layers(model) + if layers is None: + raise RuntimeError( + "Layerwise export requires discoverable decoder layers. The model " + "architecture is not supported by LayerActivationCollector." + ) + # The same call calibration uses, so layer_idx means the same thing on both sides. + self._layers = layers + layer_ids = {id(m): i for i, m in enumerate(layers)} + self._layer_names: dict[int, str] = {} + for name, module in model.named_modules(): + idx = layer_ids.get(id(module)) + if idx is not None: + self._layer_names[idx] = name + # Descendants too: the tail pass must skip anything a layer shard already covered. + self._decoder_owned_ids = {id(m) for layer in layers for m in layer.modules()} + # Materialization dispatch rebuilds this map per call when not supplied; only the + # handful of tail modules reach it, but the map is model-sized either way. + self._name_to_module = dict(model.named_modules()) + + self._model = model + self._export_dir = Path(export_dir) + self._export_dir.mkdir(parents=True, exist_ok=True) + self._is_modelopt_qlora = is_modelopt_qlora + self._dtype = _resolve_export_dtype(model, dtype) + # Not get_kv_cache_dtype(model): it does not recurse, so given the root it always + # answers None, which then trips the KV assert in the per-tensor pass. + self._kv_cache_format = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora)[ + "quantization" + ]["kv_cache_quant_algo"] + self._finalized = False + + self._name_mapper = None + try: + self._name_mapper = build_reverse_name_mapper(model) + except Exception as exc: + warnings.warn( + f"Reverse name mapper unavailable ({exc}); exported tensor names may not " + "match the original HF hub checkpoint." + ) + # By name, not data_ptr: layers and tail are separate passes, so there is never a + # whole-dict view to compare pointers across. + raw_tied_keys: set[str] = ( + set(getattr(model, "_tied_weights_keys", None) or []) + if getattr(model.config, "tie_word_embeddings", False) + else set() + ) + self._tied_alias_keys: set[str] = ( + {self._name_mapper(k) for k in raw_tied_keys} + if self._name_mapper is not None + else raw_tied_keys + ) + + # Dedup off for the reason the offload path turns it off (registry.py + # __post_init__): data_ptr() cannot identify a tensor across an export that keeps + # rolling packed weights back. Ties are refused anyway, and with both caches None + # the context is immutable, so one instance serves every pass. + self._ctx = ExportContext( + model=model, + dtype=self._dtype, + is_modelopt_qlora=is_modelopt_qlora, + tied_cache=None, + moe_tied_cache=None, + ) + + def export_layer( + self, + layer_idx: int, + layer_module: nn.Module, + probe_forward: Callable[[nn.Module], None] | None = None, + ) -> None: + """Pack one calibrated layer into its shard, leaving the layer itself untouched. + + ``probe_forward`` runs the layer once on real activations; a fusing format needs it + to rediscover which modules share an input. Omitting it is only valid when no + format present fuses. + """ + from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + + from .layer_utils import sync_moe_gate_up_amax + from .unified_export_hf import _dispatch_export_handler, _prepare_moe_inputs + + assert not self._finalized, "export_layer() called after finalize()" + assert layer_module is self._layers[layer_idx], ( + f"layer_idx {layer_idx} does not match the module passed; calibration and export " + "disagree on decoder layer order." + ) + + layer_name = self._layer_names[layer_idx] + tensors: dict[str, torch.Tensor] = {} + with transient_module_state(layer_module): + # Per-block, so neither fits a whole-model prep pass: earlier there is no amax + # yet, later the layer is already written. + _prepare_moe_inputs(layer_module, self._dtype, self._is_modelopt_qlora) + self._fuse_shared_inputs(layer_module, probe_forward) + sync_moe_gate_up_amax(layer_module) + + for sub_name, sub_mod in layer_module.named_modules(): + full_name = f"{layer_name}.{sub_name}" if sub_name else layer_name + _dispatch_export_handler(full_name, sub_mod, self._ctx) + _reconstruct_fused_moe_linear(layer_module) + + prefix = f"{layer_name}." if layer_name else "" + for key, tensor in layer_module.state_dict().items(): + self._collect(tensors, prefix + key, tensor) + + save_file(_copy_storage_aliases(tensors), str(self._export_dir / layer_shard_name(layer_idx))) + + def _fuse_shared_inputs( + self, layer_module: nn.Module, probe_forward: Callable[[nn.Module], None] | None + ) -> None: + """Unify scales across the modules of this layer that share an input. + + The whole-model exporters get these groups from one forward over the entire model. + Rediscovering them per layer is equivalent because the groups never cross a layer + boundary, and it uses the layer's real activations rather than a synthetic probe. + """ + from .quant_utils import get_quantization_format + from .unified_export_hf import _fuse_shared_input_modules, collect_shared_input_modules + + # Per-module scan, not get_quantization_format(layer_module): that returns the first + # format found, so a layer with FP8 attention and NVFP4 experts reports fp8 and + # would skip fusing its NVFP4 groups. _fuse_shared_input_modules re-evaluates the + # format per group, so the value passed below is only a fallback. + if not (_module_formats(layer_module) - FUSION_FREE_FORMATS): + return + layer_format = get_quantization_format(layer_module) + if probe_forward is None: + raise RuntimeError( + f"layer format {layer_format!r} needs input-sharing groups to fuse its " + "scales, but no probe_forward was supplied to rediscover them." + ) + + input_to_linear, _ = collect_shared_input_modules( + layer_module, lambda: probe_forward(layer_module) + ) + _fuse_shared_input_modules(self._model, input_to_linear, quantization_format=layer_format) + + def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> dict: + """Export the tail, write every config artifact, and index all shards. + + Leaves ``export_dir`` a complete, loadable checkpoint, so no separate + ``export_hf_checkpoint()`` call is needed. Returns the quant config. + """ + from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + requires_weight_materialization, + ) + + from .quant_aware_conversion import revert_quant_config_names + 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, + ) + + assert not self._finalized, "finalize() called twice" + self._finalized = True + + model = self._model + quant_config = get_quant_config(model, is_modelopt_qlora=self._is_modelopt_qlora) + _add_mtp_exclusions(model, quant_config) + _warn_on_unsynced_moe_gate_up(model) + if getattr(model, "hf_quantizer", None) is not None: + model.hf_quantizer = None + # Module references in the config must use the same hub names the tensors were + # written under, or a loader will treat an excluded BF16 layer as quantized. + if self._name_mapper is not None and quant_config: + with contextlib.suppress(Exception): + revert_quant_config_names(quant_config.get("quantization", {}), self._name_mapper) + + tail: dict[str, torch.Tensor] = {} + seen_keys: set[str] = set() + handled_ids: set[int] = set() + # Decoder tensors are already in their own shards. + skip_prefixes = tuple(f"{n}." for n in self._layer_names.values() if n) + + # Non-decoder modules whose weights are not directly readable -- embeddings, norms, + # lm_head on an offloaded model. Each needs its own materialization window, or its + # tensors are still on meta here and _collect drops them silently. Containers are + # skipped: their children get their own window. + for name, module in model.named_modules(): + if id(module) in self._decoder_owned_ids: + continue + if not requires_weight_materialization(module, model, self._name_to_module): + continue + with enable_weight_access_and_writeback( + module, model, self._name_to_module, writeback=False + ): + for sub_name, sub_mod in module.named_modules(): + full_name = f"{name}.{sub_name}" if sub_name else name + _dispatch_export_handler(full_name, sub_mod, self._ctx) + handled_ids.add(id(sub_mod)) + prefix = f"{name}." if name else "" + for key, tensor in module.state_dict().items(): + seen_keys.add(prefix + key) + self._collect(tail, prefix + key, tensor) + + # Everything already resident. On a model with no offload this is the whole tail. + for name, module in model.named_modules(): + if id(module) in self._decoder_owned_ids or id(module) in handled_ids: + continue + if _holds_meta_tensor(module): + # requires_weight_materialization said no window was needed, yet the weights + # are not here. Packing would raise deep inside the export handler; skipping + # would drop the tensor silently. Neither is acceptable. + raise RuntimeError( + f"{name!r} holds meta tensors but was not offered a materialization " + "window, so its weights cannot be exported. Export without export_dir " + "and use export_hf_checkpoint() for this model." + ) + _dispatch_export_handler(name, module, self._ctx) + for name, tensor in model.state_dict().items(): + if name.startswith(skip_prefixes) or name in seen_keys: + continue + self._collect(tail, name, tensor) + + # Tensors the model never held -- e.g. MTP weights, which HF leaves orphaned because + # it only builds num_hidden_layers decoders. Already materialized and already in + # export form, so only the hub-name reversal applies. + 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()) + + save_file(_copy_storage_aliases(tail), str(self._export_dir / _TAIL_SHARD)) + self._write_index() + save_non_weight_artifacts(model, self._export_dir) + _write_hf_export_config(model, quant_config, self._export_dir) + return quant_config + + def assert_shards_present(self, upto: int) -> None: + """Require shards for layers ``[0, upto)``, which a resume intends to skip. + + Calibration resumes from its own checkpoint directory, which knows nothing about + what was exported. If the two were produced by different runs, the skipped layers + have no shards and the gap would only surface at :meth:`finalize`, after the whole + calibration had run. Fail before any of that work instead. + """ + missing = [i for i in range(upto) if not (self._export_dir / layer_shard_name(i)).exists()] + if missing: + raise RuntimeError( + f"Resuming calibration at layer {upto} would skip layers {missing}, but " + f"their shards are missing from {self._export_dir}. The checkpoint and " + "export directories are from different runs; delete one and restart." + ) + + def _collect(self, out: dict[str, torch.Tensor], full_key: str, tensor: torch.Tensor) -> None: + """Apply per-tensor export postprocessing and hub-name reversal, or drop the tensor.""" + from .quant_utils import _postprocess_single_tensor + + if tensor is None or tensor.is_meta: + return + new_key, new_value = _postprocess_single_tensor( + full_key, tensor, 448, self._kv_cache_format, self._is_modelopt_qlora + ) + if new_key is None or new_value is None: + return + if self._name_mapper is not None: + new_key = self._name_mapper(new_key) + if new_key in self._tied_alias_keys: + return + out[new_key] = new_value.detach().contiguous().cpu() + + def _write_index(self) -> None: + """Build ``model.safetensors.index.json`` by reading back the shards on disk. + + Read from disk rather than accumulated in memory, because shards this run resumed + past were never seen by this process. Enumerated from the layer count rather than + globbed, so leftovers from a longer previous run cannot leak into the index. + """ + from safetensors import safe_open + + shards = [self._export_dir / layer_shard_name(i) for i in range(len(self._layers))] + shards.append(self._export_dir / _TAIL_SHARD) + + weight_map: dict[str, str] = {} + total_size = 0 + for shard in shards: + with safe_open(str(shard), framework="pt") as f: + for key in f.keys(): # noqa: SIM118 -- safe_open has no __iter__ + weight_map[key] = shard.name + total_size += _shard_data_bytes(shard) + index = {"metadata": {"total_size": total_size}, "weight_map": weight_map} + (self._export_dir / _INDEX_FILE).write_text(json.dumps(index, indent=2)) + + +def _holds_meta_tensor(module: nn.Module) -> bool: + """Whether this module's own parameters or buffers are still on meta.""" + return any( + t is not None and t.is_meta + for t in (*module._parameters.values(), *module._buffers.values()) + ) + + +def _copy_storage_aliases(tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Copy tensors that share storage with an earlier key. + + ``save_file`` rejects two keys backed by the same storage, and ``_collect``'s ``.cpu()`` + is a no-op rather than a copy when the tensor is already there. Copy rather than drop: + every key has to survive. Mirrors ``_StreamingShardWriter.add``. + """ + seen: set[int] = set() + for key, tensor in tensors.items(): + if tensor.data_ptr() in seen: + tensors[key] = tensor.clone() + else: + seen.add(tensor.data_ptr()) + return tensors + + +def _shard_data_bytes(path: Path) -> int: + """Payload size of a safetensors file, excluding its header. + + Layout is an 8-byte little-endian header length, that much JSON, then tensor data. + Subtracting is exact and avoids a dtype-size table that would have to track every + safetensors dtype name. + """ + with open(path, "rb") as f: + header_len = int.from_bytes(f.read(8), "little") + return path.stat().st_size - 8 - header_len diff --git a/modelopt/torch/export/model_config.py b/modelopt/torch/export/model_config.py index 5f92cc2e5dc..ebafc1d9d7a 100755 --- a/modelopt/torch/export/model_config.py +++ b/modelopt/torch/export/model_config.py @@ -44,6 +44,11 @@ QUANTIZATION_FP8_PB_WO = "fp8_pb_wo" QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" +# Formats whose scales are purely per-module, so export never merges them across the q/k/v +# and gate/up groups that share an input. Every other format unifies input_amax (and, for +# NVFP4, weight_scale_2) across such a group, which only a whole-model forward can discover. +FUSION_FREE_FORMATS = frozenset({QUANTIZATION_FP8, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL}) + KV_CACHE_FP8 = "FP8" KV_CACHE_INT8 = "INT8" KV_CACHE_NVFP4 = "NVFP4" diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 2a605ed6d9e..1700cbd782a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,8 +15,10 @@ """Code that export quantized Hugging Face models for deployment.""" +import contextlib import json import re +import shutil import tempfile import warnings from builtins import ValueError @@ -80,6 +82,7 @@ sync_moe_gate_up_amax, ) from .model_config import ( + FUSION_FREE_FORMATS, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PC_PT, @@ -379,11 +382,7 @@ def _fuse_shared_input_modules( # (must be re-evaluated per group as different modules may have different formats) group_quant_format = get_quantization_format(modules[0]) if modules else quantization_format - if len(modules) > 1 and group_quant_format not in [ - QUANTIZATION_FP8, - QUANTIZATION_NONE, - QUANTIZATION_FP8_PB_REAL, - ]: + if len(modules) > 1 and group_quant_format not in FUSION_FREE_FORMATS: if qkv_only: # Filter to only include QKV projection layers (diffusion models) qkv_modules = [m for m in modules if is_qkv_projection(getattr(m, "name", ""))] @@ -1481,6 +1480,38 @@ def _sanitize_generation_config_for_save(model: torch.nn.Module) -> None: gc.do_sample = True +def save_non_weight_artifacts(model: nn.Module, export_dir: Path) -> None: + """Write config.json, generation_config.json, and trust_remote_code modeling files. + + The ``*.py`` files are what ``trust_remote_code`` checkpoints (e.g. NemotronH) need to + load at all, so they are copied across from the source directory. + + For exporters that stream weights out themselves and so never hand a state dict to + ``save_pretrained``. Calling ``save_pretrained(state_dict={})`` instead is not an + option: MoE models (e.g. DSR1) have expert weights that share underlying storage across + layers, and safetensors' shared-tensor check fires even when the dict is empty -- + crashing the export after every shard is already written correctly. + """ + _sanitize_generation_config_for_save(model) + # transformers' own revert_weight_conversion cannot handle quantized state dicts. + patches = _patch_revert_weight_conversion() + try: + model.config.save_pretrained(str(export_dir)) + finally: + _unpatch_revert_weight_conversion(patches) + + if getattr(model, "generation_config", None) is not None: + with contextlib.suppress(Exception): + model.generation_config.save_pretrained(str(export_dir)) + + src_dir = Path(getattr(model.config, "_name_or_path", "") or "") + if src_dir.is_dir(): + for py_file in src_dir.glob("*.py"): + dst = export_dir / py_file.name + if not dst.exists(): + shutil.copy2(py_file, dst) + + def export_speculative_decoding( model: torch.nn.Module, dtype: torch.dtype | None = None, diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 5e2d682c770..4daefe726d5 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -21,10 +21,8 @@ lazily to keep the dependency acyclic. """ -import contextlib import itertools import json -import shutil import warnings from pathlib import Path from typing import Any @@ -39,13 +37,11 @@ from .unified_export_hf import ( _add_mtp_exclusions, _dispatch_export_handler, - _patch_revert_weight_conversion, _prepare_moe_inputs, _resolve_export_dtype, - _sanitize_generation_config_for_save, - _unpatch_revert_weight_conversion, _warn_on_unsynced_moe_gate_up, requantize_resmooth_fused_llm_layers, + save_non_weight_artifacts, ) __all__ = ["_export_transformers_checkpoint_streaming"] @@ -418,28 +414,6 @@ def _stream_tensor(full_key: str, tensor: torch.Tensor) -> None: writer.finalize() - # Write non-weight artifacts: config.json, generation_config.json, and the custom - # modeling *.py files that trust_remote_code models (e.g. NemotronH) need. - # We avoid model.save_pretrained(state_dict={}) here because MoE models (e.g. DSR1) - # have expert weights that share underlying storage across layers; safetensors' shared- - # tensor check fires even when saving an empty state dict, crashing the export after - # all shards are already written correctly. - _sanitize_generation_config_for_save(model) - _patches = _patch_revert_weight_conversion() - try: - model.config.save_pretrained(str(export_dir)) - finally: - _unpatch_revert_weight_conversion(_patches) - if hasattr(model, "generation_config") and model.generation_config is not None: - with contextlib.suppress(Exception): - model.generation_config.save_pretrained(str(export_dir)) - - # Copy custom modeling *.py files for trust_remote_code checkpoints. - _src_dir = Path(getattr(model.config, "_name_or_path", "") or "") - if _src_dir.is_dir(): - for _py in _src_dir.glob("*.py"): - _dst = export_dir / _py.name - if not _dst.exists(): - shutil.copy2(_py, _dst) + save_non_weight_artifacts(model, export_dir) return None, quant_config diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 88927fc7895..04dad1c2015 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -761,6 +761,23 @@ class LayerwiseConfig(ModeloptBaseConfig): ), ) + export_dir: str | None = ModeloptField( + default=None, + title="Export each layer's quantized checkpoint as soon as it is calibrated.", + description=( + "If set, each decoder layer is written to a quantized HF checkpoint shard in " + "this directory the moment its calibration finishes, leaving a complete, " + "loadable checkpoint when the last layer lands. Removes the separate " + "``export_hf_checkpoint()`` pass and its full-precision intermediate. " + "Combined with ``checkpoint_dir``, an interrupted run resumes without " + "re-exporting finished layers. Supports FP8 and NVFP4 on single-process " + "models, resident or accelerate-offloaded; AWQ, SVDQuant, multi-process jobs, " + "weight-tied quantized modules, multimodal and MTP models raise " + "NotImplementedError. The model left in memory afterwards is not valid for " + "inference if the run resumed." + ), + ) + calib_mutates_weights: bool = ModeloptField( default=True, title="Whether layerwise calibration mutates layer weights.", @@ -870,12 +887,13 @@ def _coerce_layerwise(cls, value): @model_validator(mode="after") def validate_layerwise_checkpoint_dir(self): - """Raise if layerwise.checkpoint_dir is set but layerwise.enable is False.""" - if self.layerwise.checkpoint_dir is not None and not self.layerwise.enable: - raise ValueError( - "layerwise.checkpoint_dir requires layerwise.enable=True. " - "Set layerwise.enable=True or remove layerwise.checkpoint_dir." - ) + """Raise if a layerwise directory is set but layerwise.enable is False.""" + for field in ("checkpoint_dir", "export_dir"): + if getattr(self.layerwise, field) is not None and not self.layerwise.enable: + raise ValueError( + f"layerwise.{field} requires layerwise.enable=True. " + f"Set layerwise.enable=True or remove layerwise.{field}." + ) return self @model_validator(mode="after") diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index c096aaeb00e..db7704e89b5 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -230,6 +230,7 @@ def wrapped_calib_func( layerwise_cfg = kwargs.pop("layerwise", None) or {} layerwise = layerwise_cfg.get("enable", False) checkpoint_dir = layerwise_cfg.get("checkpoint_dir") + export_dir = layerwise_cfg.get("export_dir") qdq_from_prev = layerwise_cfg.get("get_qdq_activations_from_prev_layer", False) save_every = layerwise_cfg.get("save_every", 1) calib_mutates_weights = layerwise_cfg.get("calib_mutates_weights", True) @@ -265,6 +266,7 @@ def wrapped_calib_func( forward_loop=forward_loop, calib_func=func, checkpoint_dir=checkpoint_dir, + export_dir=export_dir, get_qdq_activations_from_prev_layer=qdq_from_prev, save_every=save_every, calib_mutates_weights=calib_mutates_weights, diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index c85e97a104d..909b05aa408 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2064,12 +2064,20 @@ def layerwise_calibrate( are saved after each layer completes. On restart, calibration resumes from the last completed layer. + If ``export_dir`` is passed, each layer is additionally written to a quantized HF + checkpoint shard as soon as it is calibrated, leaving a complete checkpoint when the + last layer lands and removing the need for a separate ``export_hf_checkpoint()`` pass. + Those shards then serve as the resume artifact, so the per-layer weight and quantizer + files are not written and finished layers are skipped rather than restored -- which + also means a resumed run leaves the in-memory model unusable for inference. + ``get_qdq_activations_from_prev_layer`` (via ``calib_kwargs``) controls whether the cached inputs handed to layer N+1 come from a forward through the just-calibrated layer with quantizers active (True; e.g. GPTQ) or temporarily disabled (False; matches non-layerwise max-calib semantics). """ checkpoint_dir = calib_kwargs.pop("checkpoint_dir", None) + export_dir = calib_kwargs.pop("export_dir", None) qdq_from_prev = calib_kwargs.pop("get_qdq_activations_from_prev_layer", False) save_every = calib_kwargs.pop("save_every", 1) calib_mutates_weights = calib_kwargs.pop("calib_mutates_weights", True) @@ -2090,13 +2098,24 @@ def layerwise_calibrate( num_layers = len(transformer_layers) print_rank_0(f"Layerwise calibration: Found {num_layers} transformer layers") + # Before any calibration, so unsupported models fail immediately rather than after + # hours of work with nothing exportable. + exporter = None + if export_dir is not None: + from modelopt.torch.export.layerwise_export import LayerwiseExporter + + exporter = LayerwiseExporter(model, export_dir) + ckpt = _CheckpointState.from_folder( checkpoint_dir, num_layers, save_every=save_every, calib_mutates_weights=calib_mutates_weights, + save_layer_state=exporter is None, ) start_layer = ckpt.start_layer if ckpt else 0 + if exporter is not None and start_layer > 0: + exporter.assert_shards_present(start_layer) layer_pbar = tqdm( total=num_layers, @@ -2171,6 +2190,27 @@ def _layer_forward_loop(m, _inputs=layer_inputs): elif is_last: next_inputs = None + # After the next-layer capture in both orderings, so the shard reflects + # the layer's final state. + if exporter is not None: + # One real batch, so a fusing format can rediscover which modules share + # an input without the whole-model forward this path never runs. + def _fusion_probe(m, _inputs=layer_inputs): + args, kwargs_input = _inputs[0] + # Same reset _layer_forward_loop does: these tuples were already + # replayed once, so the cache holds this layer's keys and the probe + # would see kv_len twice the mask width. + cache = kwargs_input.get("past_key_values") + if cache is not None: + kwargs_input = dict(kwargs_input) + if hasattr(cache, "reset"): + cache.reset() + else: + kwargs_input["past_key_values"] = None + m(*args, **kwargs_input) + + exporter.export_layer(layer_idx, layer, _fusion_probe) + if ckpt: ckpt.save(layer_idx, model, transformer_layers, next_inputs) @@ -2185,6 +2225,16 @@ def _layer_forward_loop(m, _inputs=layer_inputs): if ckpt: ckpt.full_restore(transformer_layers, model) + if exporter is not None: + exporter.finalize() + print_rank_0(f"Layerwise export: wrote quantized checkpoint to {export_dir}") + if start_layer > 0: + warn_rank_0( + f"This run resumed at layer {start_layer}, so layers 0..{start_layer - 1} " + "were never re-calibrated: the exported checkpoint is complete, but the " + "in-memory model is not and must not be used for inference." + ) + print_rank_0("Layerwise calibration completed") diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 070ee521cd5..7518e58eb31 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -494,6 +494,7 @@ def _write_manifest( num_layers: int, save_every: int, calib_mutates_weights: bool, + save_layer_state: bool, ) -> None: """Atomically write manifest.json. Config keys are persisted so resume can detect drift.""" path = os.path.join(checkpoint_dir, "manifest.json") @@ -505,6 +506,7 @@ def _write_manifest( "num_layers": num_layers, "save_every": save_every, "calib_mutates_weights": calib_mutates_weights, + "save_layer_state": save_layer_state, }, f, ) @@ -519,7 +521,7 @@ def _save_layer_files( checkpoint_dir: str, idx: int, weights: dict | None, - qstate: dict, + qstate: dict | None, quantizer_buffers: dict | None, output_meta: tuple, ) -> None: @@ -527,7 +529,9 @@ def _save_layer_files( Exactly one of ``weights`` (full layer state_dict) or ``quantizer_buffers`` (just the TensorQuantizer state_dict slice, used when calibration does not mutate weights) - is written; ``full_restore`` falls back to whichever is present. + is written; ``full_restore`` falls back to whichever is present. Both may be None, + along with ``qstate``, when per-layer export already captured the layer durably and + resume will skip it rather than restore it. ``next_inputs.pt`` and ``manifest.json`` are deferred to window boundaries in :meth:`_CheckpointState.save`. """ @@ -539,7 +543,8 @@ def _save_layer_files( torch.save(weights, os.path.join(d, "weights.pt")) elif quantizer_buffers is not None: torch.save(quantizer_buffers, os.path.join(d, "quantizer_buffers.pt")) - torch.save(qstate, os.path.join(d, "quantizer_state.pt")) + if qstate is not None: + torch.save(qstate, os.path.join(d, "quantizer_state.pt")) torch.save(output_meta, os.path.join(d, "output_meta.pt")) @@ -580,6 +585,7 @@ def __init__( start_layer: int = 0, save_every: int = 1, calib_mutates_weights: bool = True, + save_layer_state: bool = True, ): if dist.is_initialized() and dist.size() > 1: raise RuntimeError( @@ -593,6 +599,9 @@ def __init__( self.start_layer = start_layer self.save_every = save_every self.calib_mutates_weights = calib_mutates_weights + # False when per-layer export runs alongside: its shards already hold each layer's + # result, so resume skips the layer instead of restoring it. + self.save_layer_state = save_layer_state # Tracks the most recent saved layer so save() can window-save the layers # since the last save event. Initialized to start_layer - 1 so the first # save event after resume covers the new work only. @@ -605,6 +614,7 @@ def from_folder( num_layers: int, save_every: int = 1, calib_mutates_weights: bool = True, + save_layer_state: bool = True, ) -> _CheckpointState | None: """Create from folder. Detects resume point. Returns None if no checkpoint_dir.""" if not checkpoint_dir: @@ -617,6 +627,9 @@ def from_folder( ("num_layers", num_layers), ("save_every", save_every), ("calib_mutates_weights", calib_mutates_weights), + # Else resuming an export-mode checkpoint without export_dir recalibrates + # everything, then fails in full_restore on files that were never written. + ("save_layer_state", save_layer_state), ): ckpt_value = manifest.get(key) if ckpt_value is not None and ckpt_value != new_value: @@ -635,6 +648,7 @@ def from_folder( start_layer=start, save_every=save_every, calib_mutates_weights=calib_mutates_weights, + save_layer_state=save_layer_state, ) def setup_resume(self, layers: nn.ModuleList) -> list | None: @@ -675,7 +689,7 @@ def full_restore(self, layers: nn.ModuleList, model: nn.Module) -> None: set_quantizer_state_dict, ) - if self.start_layer == 0: + if self.start_layer == 0 or not self.save_layer_state: return dummy_config = QuantizeConfig() @@ -747,14 +761,14 @@ def save( _cpu = torch.device("cpu") layer = layers[layer_idx] - with enable_weight_access_and_writeback(layer, model, writeback=False): - qstate = _move_to_device(quantizer_state(layer), _cpu) - if self.calib_mutates_weights: - weights = _move_to_device(layer.state_dict(), _cpu) - quantizer_buffers = None - else: - weights = None - quantizer_buffers = _move_to_device(get_quantizer_state_dict(layer), _cpu) + qstate = weights = quantizer_buffers = None + if self.save_layer_state: + with enable_weight_access_and_writeback(layer, model, writeback=False): + qstate = _move_to_device(quantizer_state(layer), _cpu) + if self.calib_mutates_weights: + weights = _move_to_device(layer.state_dict(), _cpu) + else: + quantizer_buffers = _move_to_device(get_quantizer_state_dict(layer), _cpu) output_meta = getattr(layer._layerwise_calib, "output_meta", None) if output_meta is None: @@ -787,6 +801,7 @@ def save( self.num_layers, save_every=self.save_every, calib_mutates_weights=self.calib_mutates_weights, + save_layer_state=self.save_layer_state, ) window_start = self._last_saved_layer + 1 self._last_saved_layer = layer_idx diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export.yaml new file mode 100644 index 00000000000..0d444c0aaf3 --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export.yaml @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: > + NVFP4 static weight and dynamic activation for expert layers only (W4A4), FP8 KV cache, + max layerwise calibration, exporting each decoder layer to a quantized checkpoint shard + as soon as it is calibrated. Setting layerwise.export_dir is what enables that; hf_ptq.py + rewrites it to --export_path, the way it already rewrites layerwise.checkpoint_dir. + Paired with checkpoint_dir so an interrupted run resumes without recalibrating or + re-exporting finished layers -- the point of the combination for a PTQ run that outlasts + its GPU session. Export rediscovers the q/k/v and gate/up scale-fusion groups per layer, + so the checkpoint matches a whole-model export byte for byte. + + Resident (non-offloaded), single-process models only. A resumed run never recalibrates + the layers it skipped, so the exported checkpoint is complete but the in-memory model is + not and must not be used for inference. +quantize: + algorithm: + method: max + layerwise: + enable: true + # max only updates _amax, so the exported shard stays valid for its layer. + calib_mutates_weights: false + checkpoint_dir: /tmp/modelopt_layerwise_ckpt + # Presence enables per-layer export; the value is replaced with --export_path. + export_dir: /tmp/modelopt_layerwise_export + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*block_sparse_moe*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py new file mode 100644 index 00000000000..288d5c96bd6 --- /dev/null +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-layer export must match whole-model export, and refuse what it cannot match.""" + +import copy +import json +import shutil + +import pytest +import torch +from _test_utils.torch.transformers_models import get_tiny_llama +from safetensors.torch import load_file + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.unified_export_hf import export_hf_checkpoint + +NUM_LAYERS = 4 +CALIB_BATCHES = [torch.randint(0, 32, (1, 16)) for _ in range(2)] + + +def _calib(model): + for batch in CALIB_BATCHES: + model(batch.cuda()) + + +def _build_model(): + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=NUM_LAYERS).cuda().eval() + # get_tiny_llama leaves this unset, but export reads it to detect multimodal models. + model.config.architectures = ["LlamaForCausalLM"] + return model + + +def _layerwise_cfg(export_dir, checkpoint_dir, base=None): + cfg = copy.deepcopy(base or mtq.FP8_DEFAULT_CFG) + cfg["algorithm"] = { + "method": "max", + "layerwise": { + "enable": True, + "export_dir": str(export_dir), + "checkpoint_dir": str(checkpoint_dir), + # max is amax-only, so the layer weights the shard captured stay valid. + "calib_mutates_weights": False, + }, + } + return cfg + + +def _load_checkpoint(export_dir): + index = export_dir / "model.safetensors.index.json" + shards = ( + set(json.loads(index.read_text())["weight_map"].values()) + if index.exists() + else ["model.safetensors"] + ) + tensors = {} + for shard in shards: + tensors.update(load_file(str(export_dir / shard))) + return tensors + + +def _assert_same_checkpoint(expected, actual): + assert set(expected) == set(actual), ( + f"key mismatch: missing={sorted(set(expected) - set(actual))}, " + f"extra={sorted(set(actual) - set(expected))}" + ) + for key, want in expected.items(): + got = actual[key] + assert got.dtype == want.dtype and got.shape == want.shape, f"{key}: dtype/shape differs" + assert torch.equal(got.float(), want.float()), f"{key}: values differ" + + +@pytest.fixture(scope="module") +def baseline_checkpoint(tmp_path_factory): + """A normal layerwise calibration followed by a separate whole-model export.""" + export_dir = tmp_path_factory.mktemp("baseline") + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + model = mtq.quantize(_build_model(), cfg, _calib) + export_hf_checkpoint(model, export_dir=export_dir) + return _load_checkpoint(export_dir) + + +def test_layerwise_export_matches_whole_model_export(tmp_path, baseline_checkpoint): + """Exporting per layer during calibration must yield the same checkpoint.""" + export_dir = tmp_path / "fused" + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib) + + _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(export_dir)) + # The directory must be loadable on its own, with no follow-up export call. + for artifact in ("config.json", "hf_quant_config.json", "model.safetensors.index.json"): + assert (export_dir / artifact).is_file(), f"{artifact} missing" + + +def test_layerwise_export_replaces_resume_artifacts(tmp_path): + """The shards are the resume artifact, so per-layer weight copies are not written.""" + checkpoint_dir = tmp_path / "ckpt" + mtq.quantize(_build_model(), _layerwise_cfg(tmp_path / "fused", checkpoint_dir), _calib) + + 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. + assert list(checkpoint_dir.rglob("output_meta.pt")) + + +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) + + # 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)) + + resumed_dir = tmp_path / "resumed" + shutil.copytree(export_dir, resumed_dir) + mtq.quantize(_build_model(), _layerwise_cfg(resumed_dir, checkpoint_dir), _calib) + + _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(resumed_dir)) + + +def test_resume_without_matching_shards_fails_fast(tmp_path): + """Mismatched checkpoint/export dirs must fail before recalibrating, not at the end.""" + checkpoint_dir = tmp_path / "ckpt" + mtq.quantize(_build_model(), _layerwise_cfg(tmp_path / "fused", checkpoint_dir), _calib) + + 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)) + + with pytest.raises(RuntimeError, match="shards are missing"): + mtq.quantize( + _build_model(), _layerwise_cfg(tmp_path / "empty_export", checkpoint_dir), _calib + ) + + +def test_kv_cache_quantized_export_matches(tmp_path): + """KV-cache scales must survive: the format has to be read off the whole quant config. + + Deriving it from the root module alone yields None, which makes the per-tensor pass + assert on the first ``*_bmm_quantizer._amax`` it sees. + """ + kv_cfg = mtq.update_quant_cfg_with_kv_cache_quant( + copy.deepcopy(mtq.FP8_DEFAULT_CFG), copy.deepcopy(mtq.FP8_KV_CFG["quant_cfg"]) + ) + + baseline_dir = tmp_path / "baseline" + base = copy.deepcopy(kv_cfg) + base["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + export_hf_checkpoint(mtq.quantize(_build_model(), base, _calib), export_dir=baseline_dir) + + export_dir = tmp_path / "fused" + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt", base=kv_cfg), _calib) + + exported = _load_checkpoint(export_dir) + assert any(k.endswith(("k_scale", "v_scale")) for k in exported), ( + "no KV cache scales were exported" + ) + _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) + + +def _nvfp4_cfg(): + """NVFP4 with o_proj left unquantized. + + Layerwise calibration leaves ``self_attn.o_proj``'s input amax at 0 on every layer but + the last, so a full-NVFP4 model cannot be exported by *any* path -- a pre-existing bug + unrelated to per-layer export. The shipped NVFP4 layerwise recipes are experts-only and + never quantize o_proj, which is why it has gone unnoticed. Excluding it here keeps this + test on the behaviour it is meant to cover: q/k/v and gate/up scale fusion. + """ + cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG) + cfg["quant_cfg"].append({"quantizer_name": "*o_proj*", "enable": False}) + return cfg + + +def test_nvfp4_export_matches(tmp_path): + """NVFP4 fuses q/k/v and gate/up scales; per-layer rediscovery must match.""" + baseline_dir = tmp_path / "baseline" + base = _nvfp4_cfg() + base["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + export_hf_checkpoint(mtq.quantize(_build_model(), base, _calib), export_dir=baseline_dir) + + export_dir = tmp_path / "fused" + mtq.quantize( + _build_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()), _calib + ) + + exported = _load_checkpoint(export_dir) + assert any(k.endswith("weight_scale_2") for k in exported), "no NVFP4 global scales exported" + _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) + + +def _mixed_fp8_nvfp4_cfg(): + """FP8 attention, NVFP4 MLP -- a layer whose format depends on where you look. + + ``get_quantization_format`` returns the first format found, so gating fusion on it + reports fp8 here and silently skips fusing the NVFP4 groups. o_proj stays unquantized + for the reason in :func:`_nvfp4_cfg`. + """ + nvfp4 = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG) + numerics = next( + e["cfg"] for e in nvfp4["quant_cfg"] if e.get("quantizer_name") == "*weight_quantizer" + ) + fp8 = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + fp8_numerics = next( + e["cfg"] for e in fp8["quant_cfg"] if e.get("quantizer_name") == "*weight_quantizer" + ) + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*self_attn*weight_quantizer", "cfg": copy.deepcopy(fp8_numerics)}, + {"quantizer_name": "*self_attn*input_quantizer", "cfg": copy.deepcopy(fp8_numerics)}, + {"quantizer_name": "*mlp*weight_quantizer", "cfg": copy.deepcopy(numerics)}, + {"quantizer_name": "*mlp*input_quantizer", "cfg": copy.deepcopy(numerics)}, + {"quantizer_name": "*o_proj*", "enable": False}, + ] + } + + +def test_mixed_format_export_matches(tmp_path): + """A layer holding two formats must still fuse the one that needs it.""" + baseline_dir = tmp_path / "baseline" + base = _mixed_fp8_nvfp4_cfg() + base["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + export_hf_checkpoint(mtq.quantize(_build_model(), base, _calib), export_dir=baseline_dir) + + export_dir = tmp_path / "fused" + mtq.quantize( + _build_model(), + _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_mixed_fp8_nvfp4_cfg()), + _calib, + ) + _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) + + +def test_awq_is_refused(tmp_path): + """AWQ needs the pre-quant-scale steps, which are still whole-model.""" + cfg = _layerwise_cfg(tmp_path / "fused", tmp_path / "ckpt", base=mtq.INT4_AWQ_CFG) + with pytest.raises(NotImplementedError, match="awq"): + mtq.quantize(_build_model(), cfg, _calib) diff --git a/tests/unit/torch/quantization/test_config_validation.py b/tests/unit/torch/quantization/test_config_validation.py index 4b969d3259c..0d94af38229 100644 --- a/tests/unit/torch/quantization/test_config_validation.py +++ b/tests/unit/torch/quantization/test_config_validation.py @@ -694,6 +694,7 @@ def test_default_dump_shape(self): "get_qdq_activations_from_prev_layer": False, "checkpoint_dir": None, "save_every": 1, + "export_dir": None, "calib_mutates_weights": True, } assert "layerwise_checkpoint_dir" not in dumped