From 5b32946bd8c110c9b8879d4d919116ce07be94f9 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Sun, 30 Aug 2026 04:47:14 +0000 Subject: [PATCH 1/2] feat(export): export each decoder layer as layerwise calibration finishes it Layerwise calibration can already resume, but only through a full-precision scratch checkpoint, and a completed run still owes a second whole-model export pass -- itself needing a GPU session. Setting layerwise.export_dir writes each decoder layer to a quantized HF shard as soon as that layer is calibrated, so finishing the last calibrated layer finishes the checkpoint and a run that outlives its session resumes owing only the remaining layers plus finalize(). One shard per layer, model-layer-{idx:05d}.safetensors, is the resume invariant: "shard exists" means "layer done" across a restart. finalize() then exports the tail, writes the config artifacts, and builds the index from the shards on disk, so an earlier run's layers are picked up as they are. Because export converts each layer in place, and a resumed run never recalibrates the layers it skipped, the in-memory model is not valid for inference afterwards; hf_ptq forces --skip_generate and says so. Supported: FP8, NVFP4, FP8_PB_REAL, and mixed layers, resident or under accelerate offload. Refused up front, each because a per-layer pass cannot reproduce what the whole-model path does globally: AWQ/SVDQuant (pre-quant-scale fusion), weight-tied quantized modules (sync_tied_input_amax), multi-process jobs, split rules, MTP, multimodal, and the second-exporter flags. Verified byte-identical against export_hf_checkpoint on five model/format pairings, including 123,513 tensors with 0 differing on an offloaded Qwen3.6-35B-A3B, plus kill-and-resume at scale and vLLM generation equality. Also includes a pre-existing main fix this depends on: _is_layerwise used getattr on an algorithm that YAML parses as a dict, so it answered False for every layerwise recipe in the repo and the batch-size probe it gates was never skipped. Behaviour change: --batch_size 0 now yields batch_size=1 for layerwise recipes, as its comment intends. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 1 + examples/hf_ptq/example_utils.py | 123 ++++- examples/hf_ptq/hf_ptq.py | 130 ++++- modelopt/torch/export/layerwise_export.py | 449 +++++++++++++++++ modelopt/torch/export/model_config.py | 5 + modelopt/torch/export/unified_export_hf.py | 37 +- .../export/unified_export_hf_streaming.py | 30 +- modelopt/torch/quantization/config.py | 27 +- modelopt/torch/quantization/mode.py | 2 + modelopt/torch/quantization/model_calib.py | 45 +- .../quantization/utils/layerwise_calib.py | 84 +++- ..._experts_only-kv_fp8_layerwise_export.yaml | 64 +++ modelopt_recipes/ptq.md | 7 +- tests/examples/hf_ptq/test_example_utils.py | 131 +++++ .../gpu/torch/export/test_layerwise_export.py | 453 ++++++++++++++++++ .../quantization/test_config_validation.py | 7 +- 16 files changed, 1504 insertions(+), 91 deletions(-) create mode 100644 modelopt/torch/export/layerwise_export.py create mode 100644 modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export.yaml create mode 100644 tests/gpu/torch/export/test_layerwise_export.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 95e3f12ef58..7be270b3edc 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,7 @@ Changelog - Add a calibration-free streaming Kimi-K3 converter and checkpoint-mirror recipe for NVFP4 routed experts with ``input_scale=1.0`` and 128x128 block-FP8 KDA/MLA attention weights. The converter operates shard-by-shard on the source checkpoint's packed MXFP4 experts instead of loading the 2.8T model through the in-memory ``hf_ptq.py`` path. - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Supports FP8 and NVFP4 on single-process models, resident or offloaded; other formats and placements raise ``NotImplementedError`` before calibration starts. *Megatron Framework (M-LM / M-Bridge)* diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 997b2eef125..862bc0f7fe5 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -1134,12 +1134,90 @@ def copy_custom_model_files( print("No checkpoint sidecar files found to copy") +def _layerwise_blocks(algorithm) -> list[dict]: + """Every ``layerwise`` block in the algorithm, which may be one entry or a list.""" + entries = algorithm if isinstance(algorithm, list) else [algorithm] + return [ + e["layerwise"] + for e in entries + if isinstance(e, dict) and isinstance(e.get("layerwise"), dict) + ] + + +def recipe_layerwise_blocks(recipe) -> list[dict]: + """Every ``layerwise`` block in a recipe's algorithm(s), in order, normalized to dicts. + + Reads the parsed *recipe*, where YAML gives plain dicts and the deprecated + ``--auto_quantize_*`` path gives config objects; :func:`_layerwise_blocks` reads the + resolved ``quant_cfg``, which is always dicts. + """ + quantize = getattr(recipe, "quantize", None) + algorithm = getattr(quantize, "algorithm", None) + entries = algorithm if isinstance(algorithm, list) else [algorithm] + blocks = [] + for entry in entries: + block = ( + entry.get("layerwise") if isinstance(entry, dict) else getattr(entry, "layerwise", None) + ) + if block is not None: + blocks.append(block if isinstance(block, dict) else block.model_dump()) + return blocks + + def _layerwise_checkpoint_dir(algorithm) -> str | None: - """Return the nested ``layerwise.checkpoint_dir``, or None.""" - if not isinstance(algorithm, dict): + """First ``layerwise.checkpoint_dir`` across the algorithm entries, or None.""" + return next( + (b["checkpoint_dir"] for b in _layerwise_blocks(algorithm) if b.get("checkpoint_dir")), + None, + ) + + +def layerwise_export_block(algorithm) -> dict | None: + """The one ``layerwise`` block that owns per-layer export, or None. + + Export finalizes each layer's shard during calibration, so a later pass would change + the model after its checkpoint was written: exactly one entry may set ``export_dir``, + and it must be the last. + """ + entries = algorithm if isinstance(algorithm, list) else [algorithm] + exporting = [ + (i, e["layerwise"]) + for i, e in enumerate(entries) + if isinstance(e, dict) + and isinstance(e.get("layerwise"), dict) + and e["layerwise"].get("export_dir") is not None + ] + if not exporting: return None - nested = algorithm.get("layerwise") or {} - return nested.get("checkpoint_dir") if isinstance(nested, dict) else None + if len(exporting) > 1: + raise ValueError( + f"{len(exporting)} algorithm entries set layerwise.export_dir; only one " + "calibration pass can own the exported checkpoint." + ) + index, block = exporting[0] + if index != len(entries) - 1: + raise ValueError( + f"layerwise.export_dir is set on algorithm entry {index} of {len(entries)}; it " + "must be the last, since a later pass would change the model after its shards " + "were written." + ) + return block + + +def default_layerwise_resume_dir(quant_cfg: dict, export_path: str) -> tuple[dict, bool]: + """Derive ``layerwise.checkpoint_dir`` from ``export_path`` when unset. + + A sibling, not a child: nothing deletes the resume state, so inside ``export_path`` it + would ship in the checkpoint. An explicit path is left alone. + """ + quant_cfg = copy.deepcopy(quant_cfg) + # The exporting block specifically: another pass's explicit checkpoint_dir says nothing + # about where this one resumes from. + block = layerwise_export_block(quant_cfg.get("algorithm")) + if block is None or block.get("checkpoint_dir") is not None: + return quant_cfg, False + block["checkpoint_dir"] = export_path.rstrip("/") + ".layerwise_resume" + return quant_cfg, True def needs_checkpoint_path_update(quant_cfg: dict) -> bool: @@ -1156,8 +1234,7 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> tuple[dict, str] Returns ``(updated_quant_cfg, resolved_path)`` so the caller can log or reference the resolved path without re-deriving the dict shape. """ - base_dir = _layerwise_checkpoint_dir(quant_cfg["algorithm"]) - assert base_dir is not None # guaranteed by needs_checkpoint_path_update + assert needs_checkpoint_path_update(quant_cfg), "no layerwise.checkpoint_dir to resolve" name = model_path.rstrip("/") if "/" in name and not os.path.isabs(name): @@ -1166,13 +1243,43 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> tuple[dict, str] name = Path(name).name config_hash = hashlib.sha256(json.dumps(quant_cfg, default=str).encode()).hexdigest()[:8] - resolved = os.path.join(base_dir, f"{name}_{config_hash}") + suffix = f"{name}_{config_hash}" quant_cfg = copy.deepcopy(quant_cfg) - quant_cfg["algorithm"]["layerwise"]["checkpoint_dir"] = resolved + # Each pass keeps its own base, so two layerwise passes cannot resolve onto one manifest. + exporting = layerwise_export_block(quant_cfg.get("algorithm")) + resolved = None + for block in _layerwise_blocks(quant_cfg.get("algorithm")): + if block.get("checkpoint_dir") is None: + continue + block["checkpoint_dir"] = os.path.join(block["checkpoint_dir"], suffix) + if resolved is None or block is exporting: + resolved = block["checkpoint_dir"] + assert resolved is not None # needs_checkpoint_path_update found one above 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 via ``layerwise.export_dir``; its value is a placeholder, since the + destination is per-run. Raises when nothing was retargeted: the caller decides to skip + the real export from a separately parsed recipe, so a silent no-op would leave + ``--export_path`` empty on a run reporting success. + """ + quant_cfg = copy.deepcopy(quant_cfg) + algorithm = quant_cfg.get("algorithm") + block = layerwise_export_block(algorithm) + if block is None: + raise ValueError( + "layerwise export is enabled but no layerwise.export_dir was found to retarget " + f"in algorithm={algorithm!r}. The exported shards would go to the recipe's " + "placeholder path instead of --export_path." + ) + block["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 a56a62b54b5..0e3c8035bf8 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -34,6 +34,7 @@ cleanup_distributed, copy_custom_model_files, create_vlm_calibration_loop, + default_layerwise_resume_dir, get_model, get_processor, get_tokenizer, @@ -43,9 +44,11 @@ mlflow_run, mtp_layer_prefixes_from_checkpoint, needs_checkpoint_path_update, + recipe_layerwise_blocks, resolve_checkpoint_dir, resolve_mlflow_args, run_nemotron_vl_preview, + set_layerwise_export_dir, setup_distributed_args, validate_fsdp2_supported, ) @@ -768,6 +771,65 @@ 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 export_quantized's branches: a second exporter would overwrite --export_path. + 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()"), + ( + # int8_sq is the export-format constant, int8_smoothquant the qformat preset. + "--qformat int8_smoothquant", + any(t in args.qformat for t in ("int8_sq", "int8_smoothquant")), + "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, @@ -870,11 +932,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( @@ -1128,17 +1201,22 @@ def quantize_main( aq_config = None fixed_quantize_config = None - def _is_layerwise(obj): - if isinstance(obj, ModelOptPTQRecipe): - return _is_layerwise(obj.quantize.algorithm) - if isinstance(obj, ModelOptAutoQuantizeRecipe): - return obj.quantize is not None and _is_layerwise(obj.quantize.algorithm) - 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) + layerwise_cfgs = recipe_layerwise_blocks(recipe) + is_layerwise = any(cfg.get("enable", False) for cfg in layerwise_cfgs) + + # The value is a placeholder, replaced with --export_path below; presence is the switch. + args.layerwise_export = any(cfg.get("export_dir") is not None for cfg in layerwise_cfgs) + if args.layerwise_export: + if isinstance(recipe, ModelOptAutoQuantizeRecipe): + # Only the mono-quantize path retargets export_dir and runs the refusals; + # auto_quantize would export to the placeholder and skip the real export. + raise NotImplementedError( + "layerwise.export_dir is not supported with an AutoQuantize recipe; " + "use a PTQ recipe, or drop export_dir and export afterwards." + ) + if not args.skip_generate: + print("Layerwise export: forcing --skip_generate, the model is left in export form.") + args.skip_generate = True if args.batch_size == 0: # For VL models with image-text calibration, skip automatic batch size detection @@ -1263,6 +1341,11 @@ def _is_layerwise(obj): # Complementary to recipe `*mtp*` wildcards (name-match); this catches MTP layers # identified by index. mtp_layer_prefixes = getattr(full_model, "_mtp_layer_prefixes", None) + if args.layerwise_export and not mtp_layer_prefixes: + # Only the FSDP2 loader flags these before quantization. Per-layer export has + # to refuse *before* calibration, or the run writes a complete-looking + # checkpoint and only then discovers it is missing the MTP weights. + mtp_layer_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path) if mtp_layer_prefixes: quant_cfg = copy.deepcopy(quant_cfg) for prefix in mtp_layer_prefixes: @@ -1270,6 +1353,21 @@ def _is_layerwise(obj): quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) print(f"Excluding MTP layer from quantization: {pattern}") + # Before resolve_checkpoint_dir, which hashes the config: with the placeholder + # still in it, two --export_path values would share one checkpoint dir. + if args.layerwise_export: + assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) + quant_cfg = set_layerwise_export_dir(quant_cfg, args.export_path) + print(f"Layerwise export enabled: writing quantized shards to {args.export_path}") + # The shards are only a resume artifact if the manifest that names the resume + # point survives alongside them; see default_layerwise_resume_dir. + quant_cfg, moved = default_layerwise_resume_dir(quant_cfg, args.export_path) + if moved: + print( + "Layerwise checkpoint_dir co-located with the export path so a resumed " + "run finds its manifest next to the shards it must not overwrite." + ) + 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..cce90da38d3 --- /dev/null +++ b/modelopt/torch/export/layerwise_export.py @@ -0,0 +1,449 @@ +# 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 pathlib import Path + +import torch +import torch.nn as nn +from safetensors import safe_open +from safetensors.torch import save_file + +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + requires_weight_materialization, +) +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils import distributed as dist + +from .layer_utils import is_moe, sync_moe_gate_up_amax +from .model_config import FUSION_FREE_FORMATS, QUANTIZATION_NVFP4 +from .model_utils import TiedWeightMap +from .quant_aware_conversion import build_reverse_name_mapper, revert_quant_config_names +from .quant_utils import _postprocess_single_tensor, get_quant_config, get_quantization_format +from .registry import ExportContext, PrepareMoEInputsRegistry +from .unified_export_hf import ( + _add_mtp_exclusions, + _dispatch_export_handler, + _fuse_shared_input_modules, + _prepare_moe_inputs, + _resolve_export_dtype, + _write_hf_export_config, + collect_shared_input_modules, + save_non_weight_artifacts, +) +from .unified_export_hf_streaming import _assert_no_split_rules + +# Fusable per layer because the groups (q/k/v, gate/up) never cross a decoder boundary. +# AWQ and SVDQuant also need pre-quant-scale steps, which are still whole-model. +_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, keyed by index so a re-export overwrites.""" + return f"model-layer-{layer_idx:05d}.safetensors" + + +def _is_quantized_module(module: nn.Module) -> bool: + """By type, not name: fused experts name theirs ``gate_up_proj_weight_quantizer``.""" + return any( + isinstance(child, (TensorQuantizer, SequentialQuantizer)) for child in module.children() + ) + + +def _module_formats(model: nn.Module) -> set: + """Every distinct format present. ``get_quantization_format`` stops at the first.""" + 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]: + """Quantized modules sharing a weight with another. + + Grouped by name, which survives offload: a ``data_ptr`` grouping sees nothing when the + weights are on meta and would pass vacuously. Falls back to ``data_ptr`` when the model + publishes no map (transformers < 5). + """ + tied_map = TiedWeightMap(model) + groups: dict[str, list[str]] = {} + 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): + continue + key = tied_map.group_key(f"{name}.weight") + if key is not None: + groups.setdefault(key, []).append(name) + elif not weight.is_meta and weight.data_ptr(): + # data_ptr() is 0 for meta tensors and DTensors, grouping unrelated modules. + by_ptr.setdefault(weight.data_ptr(), []).append(name) + tied = {n for names in groups.values() if len(names) > 1 for n in names} + tied |= {n for names in by_ptr.values() if len(names) > 1 for n in names} + return sorted(tied) + + +def assert_formats_supported(module: nn.Module, scope: str) -> None: + """Raise unless every format in ``module`` can be reproduced per layer. + + Called before calibration to fail early, and again per exported layer -- AWQ and + SVDQuant only become visible once the calibrator registers their discriminators. + """ + unsupported = sorted(str(f) for f in _module_formats(module) - SUPPORTED_FORMATS) + if unsupported: + raise NotImplementedError( + f"layerwise export does not support quantization format(s) {unsupported} " + f"({scope}): 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)}." + ) + + +def assert_layerwise_export_supported(model: nn.Module) -> None: + """Raise unless per-layer export is valid for this model.""" + assert_formats_supported(model, "before calibration") + + tied = _tied_quantized_modules(model) + 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. " + "Conversion quantizes every nn.Linear and nn.Embedding, so disabling their " + "quantizers does not lift this -- tie_word_embeddings models need " + "export_hf_checkpoint()." + ) + + 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." + ) + + +class LayerwiseExporter: + """Writes one decoder layer's quantized shard per call, then the tail and index. + + Built before calibration, driven per layer, finalized after the last. ``finalize`` + indexes the shards on disk, so an earlier run's layers are picked up as they are. + """ + + def __init__( + self, + model: nn.Module, + export_dir: Path | str, + dtype: torch.dtype | None = None, + ) -> None: + """Validate support and capture model-level state. + + Runs before calibration, so nothing amax-dependent exists yet. + """ + assert_layerwise_export_supported(model) + # Splits regroup tensors across the whole state dict; no per-layer pass reverses that. + _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 + + self._ctx = ExportContext(model=model, dtype=_resolve_export_dtype(model, dtype)) + + self._export_dir = Path(export_dir) + self._export_dir.mkdir(parents=True, exist_ok=True) + # Not get_kv_cache_dtype: it does not recurse, so on the root it answers None. + self._kv_cache_format = get_quant_config( + model, is_modelopt_qlora=self._ctx.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." + ) + + def export_layer( + self, + layer_idx: int, + layer_module: nn.Module, + layer_inputs: list | None = None, + ) -> None: + """Pack one calibrated layer into its shard, converting it in place. + + ``layer_inputs`` are the layer's cached calibration activations, replayed once so + a fusing format can rediscover which modules share an input; omit them only when + nothing fuses. + """ + # Local, as in every other export module: the plugin imports transformers. + from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + + assert not self._finalized, "export_layer() called after finalize()" + if layer_module is not self._layers[layer_idx]: + # Not an assert: -O would strip it, and the failure is silent -- layer N's + # tensors land in layer M's shard and the index looks perfectly well formed. + raise RuntimeError( + f"layer_idx {layer_idx} does not match the module passed; calibration and " + "export disagree on decoder layer order." + ) + + assert_formats_supported(layer_module, "once calibrated") + + layer_name = self._layer_names[layer_idx] + tensors: dict[str, torch.Tensor] = {} + + # Order matters at both seams: scales derive from amax, so they must be final + # before packing, and the restack consumes packed per-expert tensors. + _prepare_moe_inputs(layer_module, self._ctx.dtype, self._ctx.is_modelopt_qlora) + self._unify_shared_quantization_params(layer_module, layer_inputs) + + 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(tensors, str(self._export_dir / layer_shard_name(layer_idx))) + + def _unify_shared_quantization_params( + self, layer_module: nn.Module, layer_inputs: list | None + ) -> None: + """Unify the quantization parameters across the modules a fused kernel merges. + + The per-layer half of ``requantize_resmooth_fused_llm_layers``: one input scale per + shared-input group, one weight_scale_2 per expert gate/up pair. Its pre-quant-scale + steps are AWQ/SVDQuant-only and refused. + """ + # A set, not get_quantization_format: that stops at the first hit, so a mixed + # FP8-attention/NVFP4-expert layer would report fp8 and skip fusing entirely. + if _module_formats(layer_module) - FUSION_FREE_FORMATS: + self._fuse_shared_input_scales(layer_module, layer_inputs) + sync_moe_gate_up_amax(layer_module) + + def _fuse_shared_input_scales(self, layer_module: nn.Module, layer_inputs: list | None) -> None: + """Rediscover the groups that share an input, on real activations, and fuse them.""" + layer_format = get_quantization_format(layer_module) + if not layer_inputs: + raise RuntimeError( + f"layer format {layer_format!r} needs input-sharing groups to fuse its " + "scales, but no layer_inputs were supplied to rediscover them." + ) + + args, kwargs = layer_inputs[0] + input_to_linear, _ = collect_shared_input_modules( + layer_module, lambda: layer_module(*args, **kwargs) + ) + _fuse_shared_input_modules( + self._ctx.model, input_to_linear, quantization_format=layer_format + ) + + def finalize(self) -> dict: + """Export the tail, write the config artifacts, and index all shards. + + Leaves ``export_dir`` a complete checkpoint; no ``export_hf_checkpoint()`` needed. + """ + assert not self._finalized, "finalize() called twice" + self._finalized = True + + model = self._ctx.model + quant_config = get_quant_config(model, is_modelopt_qlora=self._ctx.is_modelopt_qlora) + _add_mtp_exclusions(model, quant_config) + # No gate/up sync here: export_layer did every layer, and the tail has no experts. + if getattr(model, "hf_quantizer", None) is not None: + model.hf_quantizer = None + # Names must match the tensors', or a loader reads 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) + + name_to_module = dict(model.named_modules()) + # Recomputed, not snapshotted in __init__: calibration adds modules inside the + # layers (SharedQuantState), and a stale set would leave them to the tail pass. + decoder_owned_ids = {id(m) for layer in self._layers for m in layer.modules()} + + 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) + + # Offloaded tail modules are on meta and _collect drops meta silently, so each + # needs its own materialization window. + for name, module in model.named_modules(): + if id(module) in decoder_owned_ids: + continue + if not requires_weight_materialization(module, model, name_to_module): + continue + with enable_weight_access_and_writeback(module, model, 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 decoder_owned_ids or id(module) in handled_ids: + continue + if _holds_meta_tensor(module): + # Packing would raise deep in the handler; skipping would drop it silently. + 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) + + save_file(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) + warnings.warn( + "The exported checkpoint is complete, but per-layer export leaves the model in " + "export form: it must not be used for inference." + ) + return quant_config + + def completed_layers(self) -> int: + """How many leading layers have a shard. Contiguous: a gap means the rest never ran.""" + n = 0 + while (self._export_dir / layer_shard_name(n)).exists(): + n += 1 + return n + + def assert_no_orphan_shards(self) -> None: + """Refuse to redo work when shards exist but no usable resume record does.""" + done = self.completed_layers() + if not done: + return + raise RuntimeError( + f"{self._export_dir} already holds shards for layers 0..{done - 1}, but the " + "layerwise checkpoint directory has no usable resume record, so calibration " + "would restart at layer 0 and overwrite them. Either restore the checkpoint " + f"directory that produced these shards, or delete {self._export_dir} to " + "re-export." + ) + + def assert_shards_present(self, upto: int) -> None: + """Require shards for layers ``[0, upto)``, which a resume intends to skip. + + Otherwise a mismatched checkpoint/export pair only surfaces after a full run. + """ + 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.""" + 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._ctx.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) + out[new_key] = new_value.detach().contiguous().cpu() + + def _write_index(self) -> None: + """Build ``model.safetensors.index.json`` from the shards on disk. + + From disk because a resumed run never saw the earlier shards in memory; by layer + count rather than a glob, so a longer previous run's leftovers cannot leak in. + """ + # Out of the index already; delete them so the directory *is* the checkpoint. + for stale in self._export_dir.glob("model-layer-*.safetensors"): + if int(stale.stem.rsplit("-", 1)[1]) >= len(self._layers): + stale.unlink() + + shards = [self._export_dir / layer_shard_name(i) for i in range(len(self._layers))] + shards.append(self._export_dir / _TAIL_SHARD) + + 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 _shard_data_bytes(path: Path) -> int: + """Payload size of a safetensors file: total minus the 8-byte prefix and header.""" + 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 4dae93b5aa7..d98211cc3fb 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", ""))] @@ -1449,6 +1448,34 @@ 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. + + For exporters that stream weights out themselves and never hand a state dict to + ``save_pretrained``, which is not an option here: MoE models (e.g. DSR1) share expert + storage across layers, so safetensors' shared-tensor check fires even on an empty dict. + The ``*.py`` files are what ``trust_remote_code`` checkpoints (e.g. NemotronH) need. + """ + _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 50b5d797c6a..a6493aac44c 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"] @@ -422,28 +418,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 460a914f9ef..3e3f20a5999 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -754,6 +754,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.", @@ -820,16 +837,6 @@ def _coerce_layerwise(cls, value): """Coerce ``layerwise=None``/``LayerwiseConfig`` to dict form.""" return _coerce_layerwise_input(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." - ) - return self - @model_validator(mode="after") def _validate_non_mutating_layerwise_supported(self): """Enforce the ``calib_mutates_weights=False`` whitelist.""" 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..d4266b289b9 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -34,6 +34,7 @@ from modelopt.torch.quantization.utils.layerwise_calib import ( LayerActivationCollector, _CheckpointState, + _reconcile_export_with_resume, ) from modelopt.torch.utils import print_rank_0, warn_rank_0 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master @@ -2060,16 +2061,12 @@ def layerwise_calibrate( skip / run / capture strategy so that inter-layer logic in parent modules (e.g. mask construction) executes naturally without model-specific hooks. - If ``checkpoint_dir`` is passed (via ``calib_kwargs``), per-layer checkpoints - are saved after each layer completes. On restart, calibration resumes from - the last completed layer. - - ``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). + Every knob arrives through ``calib_kwargs`` from :class:`LayerwiseConfig`, which + documents them; ``export_dir`` additionally leaves the model in export form, so it + must not be used for inference afterwards. """ 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,14 +2087,29 @@ def layerwise_calibrate( num_layers = len(transformer_layers) print_rank_0(f"Layerwise calibration: Found {num_layers} transformer layers") + # Before calibration, so unsupported models fail in seconds not hours. + 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 _reconcile_export_with_resume( + exporter, checkpoint_dir, start_layer, num_layers + ): + exporter.finalize() + print_rank_0(f"Layerwise export: finalized existing shards in {export_dir}") + return + layer_pbar = tqdm( total=num_layers, initial=start_layer, @@ -2168,9 +2180,17 @@ def _layer_forward_loop(m, _inputs=layer_inputs): next_inputs = input_getter.cache_outputs_for_next_layer_calib( layer, forward_loop ) + if exporter is not None: + # As above, for the fusion probe. Only when one runs: without an + # exporter nothing touches the layer before _set_layer_states does. + layer._layerwise_calib.mode = "original" elif is_last: next_inputs = None + # After the next-layer capture in both orderings: final state. + if exporter is not None: + exporter.export_layer(layer_idx, layer, layer_inputs) + if ckpt: ckpt.save(layer_idx, model, transformer_layers, next_inputs) @@ -2185,6 +2205,15 @@ 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; their shards come from the earlier run." + ) + 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..e04f83d380b 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")) @@ -561,6 +566,37 @@ def detect_resume_point(checkpoint_dir: str) -> tuple[int, dict] | None: return (last + 1, manifest) +def _reconcile_export_with_resume( + exporter, checkpoint_dir: str | None, start_layer: int, num_layers: int +) -> bool: + """Reconcile the shards on disk with the layer calibration will start from. + + Returns True when every layer already has a shard, so only ``finalize()`` is owed. + """ + manifest = _read_manifest(checkpoint_dir) if checkpoint_dir is not None else None + last = (manifest or {}).get("last_completed_layer") + total = (manifest or {}).get("num_layers") + if total is not None and total != num_layers: + raise ValueError( + f"Layerwise checkpoint at {checkpoint_dir} was written for {total} layers " + f"but this model has {num_layers}. Use a fresh checkpoint_dir." + ) + # detect_resume_point returns None once the manifest is complete, which puts start_layer + # back at 0 and would recalibrate everything the shards already hold. + if last is not None and last + 1 >= num_layers: + exporter.assert_shards_present(num_layers) + return True + + if start_layer > 0: + exporter.assert_shards_present(start_layer) + elif checkpoint_dir is not None: + # Starting at 0 with a checkpoint_dir means no usable resume record, so calibration + # would silently overwrite finished shards. Without one there is no resume to lose, + # and re-exporting is the documented behaviour. + exporter.assert_no_orphan_shards() + return False + + class _CheckpointState: """Manages checkpoint save and restore for layerwise calibration. @@ -580,6 +616,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 +630,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 hold each layer's result, + # so resume skips rather than restores. + 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 +645,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 +658,9 @@ def from_folder( ("num_layers", num_layers), ("save_every", save_every), ("calib_mutates_weights", calib_mutates_weights), + # Else a resume 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 +679,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 +720,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() @@ -723,6 +768,15 @@ def full_restore(self, layers: nn.ModuleList, model: nn.Module) -> None: print_rank_0(f"Checkpoint: restored {self.start_layer} previously calibrated layers") + def _prune_stale_next_inputs(self, keep: int) -> None: + """Drop every layer's cached activations but the committed boundary's.""" + for idx in range(self.num_layers): + if idx == keep: + continue + stale = os.path.join(_layer_dir(self.checkpoint_dir, idx), "next_inputs.pt") + if os.path.exists(stale): + os.remove(stale) + def save( self, layer_idx: int, @@ -747,14 +801,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,7 +841,13 @@ def save( self.num_layers, save_every=self.save_every, calib_mutates_weights=self.calib_mutates_weights, + save_layer_state=self.save_layer_state, ) + # Per-layer export only: its resume dir is auto-derived, so an activation set per + # layer would dwarf a checkpoint the user never opted into. After the manifest, so + # a crash mid-write still resumes from the previous boundary. + if not self.save_layer_state: + self._prune_stale_next_inputs(keep=layer_idx) window_start = self._last_saved_layer + 1 self._last_saved_layer = layer_idx window_size = layer_idx - window_start + 1 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..049922cc4fb --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export.yaml @@ -0,0 +1,64 @@ +# 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 its value to --export_path, since the destination is per-run. + + An interrupted run then resumes without recalibrating or re-exporting finished layers, + which is the point for a PTQ run that outlasts its GPU session. The resume state lives + beside the checkpoint at .layerwise_resume unless you set + layerwise.checkpoint_dir yourself. Export rediscovers the q/k/v and gate/up scale-fusion + groups per layer, so every tensor matches a whole-model export in key, dtype, shape and + value. Only the shard layout differs: one file per layer instead of packed shards. + + Single-process models, resident or accelerate-offloaded. 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 + # 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/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index edc5d926810..1735289ac3f 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -29,7 +29,7 @@ supported combinations. ### The shipped recipes
-All 24 general/ptq/ recipes (click to expand) +All 25 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -48,6 +48,7 @@ supported combinations. | `nvfp4_experts_only-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | max | | `nvfp4_experts_only-kv_fp8_layerwise` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise | | `nvfp4_experts_only-kv_fp8_layerwise_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (non-mutating, for disk offload) | +| `nvfp4_experts_only-kv_fp8_layerwise_export` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (exports each layer as it is calibrated) | | `nvfp4_experts_only_mse-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | MSE + FP8 sweep | | `nvfp4_experts_only_input_scale1-kv_fp8_cast` | NVFP4 W4A4, MoE experts only, expert `input_scale` pinned to 1.0 | FP8 (constant amax) | max (weights); expert activations uncalibrated | | `nvfp4_omlp_only-kv_fp8` | NVFP4 W4A4, o_proj + MLP/MoE | FP8 (calibrated) | max | @@ -195,6 +196,10 @@ How the quantization scales are searched. The default (no suffix) is `max`. - **`layerwise`** (`nvfp4_experts_only-kv_fp8_layerwise`) — max calibration done one decoder layer at a time to **lower peak memory**; same numerics as the non-layerwise variant. +- **`layerwise export`** (`nvfp4_experts_only-kv_fp8_layerwise_export`) — the same + calibration, additionally writing each decoder layer to the export checkpoint as + soon as it is calibrated. A run interrupted part-way **resumes without redoing + finished layers**, and no separate export pass is needed. Same numerics again. These can also be **stacked** when a single method isn't enough — e.g. `mse` + `gptq` combines an MSE-searched weight scale with GPTQ's layerwise update. diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index b904143dcd7..e532af09fed 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -560,3 +560,134 @@ def from_pretrained(*args, **kwargs): example_utils.get_model("checkpoint", device="cpu", trust_remote_code=trust_remote_code) assert used["path"] == ("bundled" if expect_bundled_code else "builtin") + + +def _layerwise(**kwargs): + return {"enable": True, **kwargs} + + +def _blocks(quant_cfg): + algorithm = quant_cfg["algorithm"] + entries = algorithm if isinstance(algorithm, list) else [algorithm] + return [e["layerwise"] for e in entries if "layerwise" in e] + + +@pytest.mark.parametrize( + ("algorithm", "expected"), + [ + pytest.param( + {"method": "max", "layerwise": _layerwise(export_dir="/placeholder")}, + ["/out.layerwise_resume"], + id="single-entry", + ), + pytest.param( + [ + {"method": "awq", "layerwise": _layerwise()}, + {"method": "max", "layerwise": _layerwise(export_dir="/placeholder")}, + ], + [None, "/out.layerwise_resume"], + id="only-the-exporting-entry", + ), + pytest.param( + [ + {"method": "awq", "layerwise": _layerwise(checkpoint_dir="/theirs")}, + {"method": "max", "layerwise": _layerwise(export_dir="/placeholder")}, + ], + ["/theirs", "/out.layerwise_resume"], + id="another-entrys-explicit-path-is-not-this-ones", + ), + ], +) +def test_default_layerwise_resume_dir_targets_the_exporting_entry(algorithm, expected): + """Only the pass that exports gets a derived resume dir, and only if it lacks one.""" + updated, changed = example_utils.default_layerwise_resume_dir({"algorithm": algorithm}, "/out") + + assert [b.get("checkpoint_dir") for b in _blocks(updated)] == expected + assert changed is True + + +def test_resolve_checkpoint_dir_keeps_each_entrys_base(): + """Two layerwise passes must not resolve onto one manifest.""" + algorithm = [ + {"method": "awq", "layerwise": _layerwise(checkpoint_dir="/theirs")}, + {"method": "max", "layerwise": _layerwise(checkpoint_dir="/ours", export_dir="/ph")}, + ] + + updated, resolved = example_utils.resolve_checkpoint_dir({"algorithm": algorithm}, "/m/Model") + + theirs, ours = (b["checkpoint_dir"] for b in _blocks(updated)) + assert theirs.startswith("/theirs/") and ours.startswith("/ours/") + assert theirs != ours + # The exporting pass owns the path the caller reports. + assert resolved == ours + + +@pytest.mark.parametrize( + ("algorithm", "match"), + [ + pytest.param( + [ + {"layerwise": _layerwise(export_dir="/a")}, + {"layerwise": _layerwise(export_dir="/b")}, + ], + "only one calibration pass", + id="two-exporting-entries", + ), + pytest.param( + [{"layerwise": _layerwise(export_dir="/a")}, {"method": "max"}], + "must be the last", + id="a-later-pass-would-change-the-model", + ), + ], +) +def test_set_layerwise_export_dir_refuses_ambiguous_ownership(algorithm, match): + with pytest.raises(ValueError, match=match): + example_utils.set_layerwise_export_dir({"algorithm": algorithm}, "/out") + + +class _Block: + """A config object, as the deprecated ``--auto_quantize_*`` path builds.""" + + def __init__(self, **fields): + self._fields = fields + + def model_dump(self): + return dict(self._fields) + + +class _Recipe: + def __init__(self, algorithm): + self.quantize = SimpleNamespace(algorithm=algorithm) + + +@pytest.mark.parametrize( + ("recipe", "expected"), + [ + pytest.param(None, [], id="no-recipe"), + pytest.param(_Recipe(None), [], id="no-algorithm"), + pytest.param(_Recipe({"method": "max"}), [], id="algorithm-without-layerwise"), + pytest.param( + _Recipe({"method": "max", "layerwise": {"enable": True}}), + [{"enable": True}], + id="dict-entry", + ), + pytest.param( + _Recipe( + [ + {"method": "awq", "layerwise": {"enable": True}}, + {"method": "max", "layerwise": {"enable": True, "export_dir": "/x"}}, + ] + ), + [{"enable": True}, {"enable": True, "export_dir": "/x"}], + id="list-keeps-algorithm-order", + ), + pytest.param( + _Recipe(SimpleNamespace(layerwise=_Block(enable=True, export_dir="/x"))), + [{"enable": True, "export_dir": "/x"}], + id="config-object-entry", + ), + ], +) +def test_recipe_layerwise_blocks(recipe, expected): + """Both recipe shapes normalize to dicts, so callers need no shape-aware access.""" + assert example_utils.recipe_layerwise_blocks(recipe) == expected 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..1bc820a05c2 --- /dev/null +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -0,0 +1,453 @@ +# 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 contextlib +import copy +import json +from unittest.mock import patch + +import pytest +import torch +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe +from safetensors.torch import load_file + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.layerwise_export import LayerwiseExporter, layer_shard_name +from modelopt.torch.export.unified_export_hf import export_hf_checkpoint + +NUM_LAYERS = 4 +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" + + +def _fp8_cfg(): + return copy.deepcopy(mtq.FP8_DEFAULT_CFG) + + +def _kv_cache_cfg(): + return mtq.update_quant_cfg_with_kv_cache_quant( + copy.deepcopy(mtq.FP8_DEFAULT_CFG), copy.deepcopy(mtq.FP8_KV_CFG["quant_cfg"]) + ) + + +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 _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 _int4_awq_cfg(): + return copy.deepcopy(mtq.INT4_AWQ_CFG) + + +def _nvfp4_awq_cfg(): + cfg = copy.deepcopy(mtq.NVFP4_AWQ_LITE_CFG) + cfg["quant_cfg"].append({"quantizer_name": "*o_proj*", "enable": False}) + return cfg + + +@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) + + +@contextlib.contextmanager +def _dies_at_layer(layer_idx: int): + """Lose the session mid-model, the way a GPU timeout would.""" + real = LayerwiseExporter.export_layer + + def die(self, idx, *args, **kwargs): + if idx == layer_idx: + raise RuntimeError("interrupted") + return real(self, idx, *args, **kwargs) + + with patch.object(LayerwiseExporter, "export_layer", die): + yield + + +@pytest.mark.parametrize("interrupt_at", [None, 2], ids=["fresh", "resumed"]) +@pytest.mark.parametrize( + ("make_cfg", "layerwise_extra", "expected_key_suffix"), + [ + # NVFP4 fuses q/k/v and gate/up scales, so per-layer rediscovery has to match. + pytest.param(_nvfp4_cfg, {}, ("weight_scale_2",), id="nvfp4"), + # The probe runs the layer directly, so the capture must leave it in "original". + pytest.param( + _nvfp4_cfg, + {"get_qdq_activations_from_prev_layer": True}, + ("weight_scale_2",), + id="nvfp4_qdq_from_prev_layer", + ), + # A layer holding two formats must still fuse the one that needs it. + pytest.param(_mixed_fp8_nvfp4_cfg, {}, None, id="mixed_fp8_nvfp4"), + # KV scales only survive if the format is read off the whole quant config: from + # the root module alone it is None, and the per-tensor pass then asserts on the + # first *_bmm_quantizer._amax it sees. + pytest.param(_kv_cache_cfg, {}, ("k_scale", "v_scale"), id="kv_cache"), + pytest.param(_fp8_cfg, {}, None, id="fp8"), + ], +) +def test_export_matches_whole_model_export( + tmp_path, make_cfg, layerwise_extra, expected_key_suffix, interrupt_at +): + """Exporting per layer during calibration must yield the same checkpoint. + + ``interrupt_at`` crosses every config with a lost-session resume, since the two + interact: a resumed run re-enters the export path part-way through the model. + """ + baseline_dir = tmp_path / "baseline" + base = make_cfg() + base["algorithm"] = {"method": "max", "layerwise": {"enable": True, **layerwise_extra}} + export_hf_checkpoint(mtq.quantize(_build_model(), base, _calib), export_dir=baseline_dir) + + export_dir = tmp_path / "fused" + cfg = _layerwise_cfg(export_dir, tmp_path / "ckpt", base=make_cfg()) + cfg["algorithm"]["layerwise"].update(layerwise_extra) + if interrupt_at is not None: + with _dies_at_layer(interrupt_at), pytest.raises(RuntimeError, match="interrupted"): + mtq.quantize(_build_model(), cfg, _calib) + mtq.quantize(_build_model(), cfg, _calib) + + exported = _load_checkpoint(export_dir) + if expected_key_suffix: + assert any(k.endswith(expected_key_suffix) for k in exported), ( + f"no {expected_key_suffix} keys in the exported checkpoint" + ) + _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) + # 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_index_resolves_every_key_to_the_shard_holding_it(tmp_path): + """A loader resolves keys through the index; tensor equality never exercises that. + + A weight_map entry naming the wrong shard compares equal to a whole-model export and + still fails in vLLM or transformers. + """ + export_dir = tmp_path / "fused" + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib) + + weight_map = json.loads((export_dir / "model.safetensors.index.json").read_text())["weight_map"] + on_disk = {} + for shard in sorted(set(weight_map.values())): + assert (export_dir / shard).is_file(), f"index names a missing shard {shard}" + on_disk.update(dict.fromkeys(load_file(export_dir / shard), shard)) + + assert set(weight_map) == set(on_disk), "index and shards disagree on which keys exist" + assert all(on_disk[k] == v for k, v in weight_map.items()), "key routed to the wrong shard" + + +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")) + # output_meta is not reconstructible from exported weights, so it stays. + assert list(checkpoint_dir.rglob("output_meta.pt")) + # The cached activations are the bulk of the resume dir, and only the committed + # boundary's are resumable -- keeping one per layer would dwarf the checkpoint. + assert not list(checkpoint_dir.rglob("next_inputs.pt")), ( + "a completed run has nothing to resume from, so no activation cache should remain" + ) + + +def test_resume_skips_exported_layers(tmp_path, baseline_checkpoint): + """A run resuming mid-model must still produce the full, correct checkpoint.""" + export_dir = tmp_path / "fused" + checkpoint_dir = tmp_path / "ckpt" + # Die partway, the way a lost GPU session would: only the committed boundary is + # resumable, so rewinding a *finished* run's manifest would not reproduce this state. + with _dies_at_layer(2), pytest.raises(RuntimeError, match="interrupted"): + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + + assert (export_dir / layer_shard_name(1)).is_file(), "layer 1 was never committed" + assert not (export_dir / layer_shard_name(2)).exists(), "layer 2 should not have landed" + + # Shards 0..1 are on disk and must be reused rather than recalculated. + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + + _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(export_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_complete_manifest_finalizes_without_recalibrating(tmp_path, baseline_checkpoint): + """A crash between the last shard and finalize must cost only the finalize. + + detect_resume_point returns None once the manifest is complete, so start_layer falls + back to 0 and every layer would be recalculated and overwritten. + """ + export_dir = tmp_path / "fused" + checkpoint_dir = tmp_path / "ckpt" + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + + # What a crash after the final ckpt.save looks like: every shard and a complete + # manifest on disk, but no tail, index or config yet. + (export_dir / "model-tail.safetensors").unlink() + (export_dir / "model.safetensors.index.json").unlink() + layer_mtimes = {p.name: p.stat().st_mtime for p in export_dir.glob("model-layer-*.safetensors")} + assert layer_mtimes + + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + + # The layer shards must be reused verbatim, not rewritten. + for name, mtime in layer_mtimes.items(): + assert (export_dir / name).stat().st_mtime == mtime, f"{name} was rewritten" + _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(export_dir)) + + +@pytest.mark.parametrize("damage", ["deleted", "no_resume_point"]) +def test_shards_without_resume_record_refuse(tmp_path, damage): + """A lost resume record must not silently overwrite finished shards. + + Either way start_layer falls back to 0, so assert_shards_present checks nothing. + """ + export_dir = tmp_path / "fused" + checkpoint_dir = tmp_path / "ckpt" + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + + manifest = checkpoint_dir / "manifest.json" + if damage == "deleted": + manifest.unlink() + else: + record = json.loads(manifest.read_text()) + record.pop("last_completed_layer") + manifest.write_text(json.dumps(record)) + + with pytest.raises(RuntimeError, match="no usable resume record"): + mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + + +def test_export_without_checkpoint_dir_may_overwrite(tmp_path): + """Used directly, without checkpoint_dir, there is no resume to lose. + + ``hf_ptq`` derives one from ``--export_path`` so its users get resume by default; a + library caller that omits it is opting out, and re-exporting from scratch is then the + documented behaviour rather than an error. + """ + export_dir = tmp_path / "fused" + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + cfg["algorithm"] = { + "method": "max", + "layerwise": {"enable": True, "export_dir": str(export_dir)}, + } + mtq.quantize(_build_model(), cfg, _calib) + mtq.quantize(_build_model(), copy.deepcopy(cfg), _calib) # must not raise + + +def _build_moe_model(): + torch.manual_seed(0) + model = get_tiny_qwen3_moe(num_experts=16, num_experts_per_tok=1).cuda().eval() + model.config.architectures = ["Qwen3MoeForCausalLM"] + return model + + +@pytest.mark.parametrize( + "make_cfg", + [ + pytest.param(_fp8_cfg, id="quantized"), + # `enable: False` leaves the quantizer module in place, so this refuses too -- + # excluding lm_head from quantization is not a way past the gate. + pytest.param( + lambda: {**_fp8_cfg(), "quant_cfg": [{"quantizer_name": "*", "enable": False}]}, + id="quantizers_disabled", + ), + ], +) +def test_tied_embeddings_are_refused(tmp_path, make_cfg): + """Every tie_word_embeddings model is refused, which is why export needs no tie handling.""" + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=NUM_LAYERS, tie_word_embeddings=True).cuda().eval() + model.config.architectures = ["LlamaForCausalLM"] + + cfg = _layerwise_cfg(tmp_path / "fused", tmp_path / "ckpt", base=make_cfg()) + with pytest.raises(NotImplementedError, match="weight-tied quantized modules"): + mtq.quantize(model, cfg, _calib) + + +def test_moe_export_matches(tmp_path): + """MoE layers take a different path: fused expert inputs and gate/up amax sync.""" + baseline_dir = tmp_path / "baseline" + base = _nvfp4_cfg() + base["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + export_hf_checkpoint(mtq.quantize(_build_moe_model(), base, _calib), export_dir=baseline_dir) + + export_dir = tmp_path / "fused" + mtq.quantize( + _build_moe_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()), _calib + ) + + _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) + + +def test_export_consumes_the_model_without_affecting_the_checkpoint(tmp_path): + """Per-layer export converts each layer in place; the shard is written first. + + The layer is dead state by then -- the next layer's inputs were captured before the + call, and resume skips finished layers in favour of their shards. What must not drift + is the checkpoint, so this pins that the in-place conversion is invisible to it. + """ + 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" + model = mtq.quantize( + _build_model(), + _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()), + _calib, + ) + + _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) + # The returned model is in export form, which is why hf_ptq forces --skip_generate. + packed = [ + n + for n, m in model.named_modules() + if hasattr(m, "weight_scale") or hasattr(m, "input_scale") + ] + assert packed, "expected the exported model to be left in export form" + + +@pytest.mark.parametrize( + ("make_cfg", "method", "match"), + [ + # Visible from the config: int4_awq is keyed on num_bits/SequentialQuantizer. + pytest.param(_int4_awq_cfg, "max", "awq", id="int4_awq_from_config"), + # Only visible afterwards: the NVFP4 discriminators (_pre_quant_scale, + # svdquant_lora_a) are registered by the calibrator, so the constructor's gate + # sees plain nvfp4 and the check has to run again on the first exported layer. + pytest.param(_nvfp4_awq_cfg, "awq_lite", "nvfp4_awq", id="nvfp4_awq_after_calibration"), + ], +) +def test_awq_is_refused(tmp_path, make_cfg, method, match): + """AWQ needs the pre-quant-scale steps, which are still whole-model.""" + cfg = _layerwise_cfg(tmp_path / "fused", tmp_path / "ckpt", base=make_cfg()) + cfg["algorithm"]["method"] = method + cfg["algorithm"]["layerwise"]["calib_mutates_weights"] = method != "max" + + with pytest.raises(NotImplementedError, match=match): + 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 546dd21f85f..cd08959a976 100644 --- a/tests/unit/torch/quantization/test_config_validation.py +++ b/tests/unit/torch/quantization/test_config_validation.py @@ -609,9 +609,9 @@ def test_dict_form_no_deprecation(self): warnings.simplefilter("error", DeprecationWarning) MaxCalibConfig(layerwise={"enable": True}) - def test_checkpoint_dir_requires_enable(self): - with pytest.raises(ValidationError, match=r"requires layerwise.enable=True"): - MaxCalibConfig(layerwise={"checkpoint_dir": "/x"}) + def test_directories_are_inert_without_enable(self): + """Not an error: with enable=False nothing reads them.""" + assert MaxCalibConfig(layerwise={"checkpoint_dir": "/x"}).layerwise.enable is False @pytest.mark.parametrize( ("cfg_cls", "expected_qdq"), @@ -648,6 +648,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 From 7eafa8cfa7409e63204eb5bf2eb30caa7b36e9b2 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Sun, 30 Aug 2026 05:49:08 +0000 Subject: [PATCH 2/2] fix(export): read the quant config before export_layer consumes the quantizers get_quant_config reports on the quantizer modules, and per-layer export replaces them as it goes, so reading it in finalize() described a model with no quantizers left. The exported checkpoint then advertised quant_algo=null with an empty quantized_layers while its weights were packed NVFP4 -- a loader would not apply the format. Under the shipped experts-only recipe, _write_hf_export_config saw neither a quant_algo nor a kv_cache_quant_algo and skipped hf_quant_config.json entirely. Snapshot it in __init__ instead, next to the kv-cache format already captured there -- which is why that one field came out right while the rest did not. The values are set by mtq.quantize before calibration, and exclude_modules is unaffected by it: the pre-calibration snapshot reproduces the whole-model path's post-calibration list exactly. Uniform FP8 and NVFP4 hid this because their configs survive the conversion; only a mixed model loses its algo, which is the shape every shipped layerwise-export recipe uses. Found by comparing configs, which nothing did: the equivalence tests asserted the artifacts existed but never that they said the same thing. _assert_same_quant_config now compares hf_quant_config.json and config.json's quantization_config, presence included. With the fix reverted it fails test_moe_export_matches and passes the ten uniform-format cases, matching what a Qwen3.6-35B-A3B export shows. Verified on that model: 123,513 tensors, 0 differing, all three config artifacts identical to a whole-model export built from main. Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/export/layerwise_export.py | 9 ++++--- .../gpu/torch/export/test_layerwise_export.py | 26 +++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index cce90da38d3..a196c9a9e52 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -197,10 +197,11 @@ def __init__( self._export_dir = Path(export_dir) self._export_dir.mkdir(parents=True, exist_ok=True) + # Read here, not in finalize(): it reports on the quantizer modules, which + # export_layer replaces as it goes, so by finalize() the model looks unquantized. + self._quant_config = get_quant_config(model, is_modelopt_qlora=self._ctx.is_modelopt_qlora) # Not get_kv_cache_dtype: it does not recurse, so on the root it answers None. - self._kv_cache_format = get_quant_config( - model, is_modelopt_qlora=self._ctx.is_modelopt_qlora - )["quantization"]["kv_cache_quant_algo"] + self._kv_cache_format = self._quant_config["quantization"]["kv_cache_quant_algo"] self._finalized = False self._name_mapper = None @@ -298,7 +299,7 @@ def finalize(self) -> dict: self._finalized = True model = self._ctx.model - quant_config = get_quant_config(model, is_modelopt_qlora=self._ctx.is_modelopt_qlora) + quant_config = self._quant_config _add_mtp_exclusions(model, quant_config) # No gate/up sync here: export_layer did every layer, and the tail has no experts. if getattr(model, "hf_quantizer", None) is not None: diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index 1bc820a05c2..ca40f72e875 100644 --- a/tests/gpu/torch/export/test_layerwise_export.py +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -85,6 +85,30 @@ def _assert_same_checkpoint(expected, actual): assert torch.equal(got.float(), want.float()), f"{key}: values differ" +def _assert_same_quant_config(baseline_dir, export_dir): + """The metadata a loader reads, which tensor equality never covers. + + ``get_quant_config`` reports on the quantizer modules, and per-layer export replaces + them as it goes, so a config read too late describes an unquantized model while the + weights are packed. + """ + for name, key in ( + ("hf_quant_config.json", "quantization"), + ("config.json", "quantization_config"), + ): + want, got = baseline_dir / name, export_dir / name + assert want.is_file() == got.is_file(), ( + f"{name}: present in baseline={want.is_file()} but exported={got.is_file()}" + ) + if not want.is_file(): + continue + expected = json.loads(want.read_text()).get(key) + actual = json.loads(got.read_text()).get(key) + assert actual == expected, ( + f"{name}[{key}] differs:\n baseline={expected}\n fused={actual}" + ) + + def _fp8_cfg(): return copy.deepcopy(mtq.FP8_DEFAULT_CFG) @@ -220,6 +244,7 @@ def test_export_matches_whole_model_export( f"no {expected_key_suffix} keys in the exported checkpoint" ) _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) + _assert_same_quant_config(baseline_dir, 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" @@ -401,6 +426,7 @@ def test_moe_export_matches(tmp_path): ) _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) + _assert_same_quant_config(baseline_dir, export_dir) def test_export_consumes_the_model_without_affecting_the_checkpoint(tmp_path):