From 86c6d91caddc3a0ac8f7c9d19f71f3fdfda92131 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 23:13:08 +0000 Subject: [PATCH] feat(export): support multimodal and MTP models in layerwise export Per-layer export called finalize() from inside layerwise_calibrate, which is the wrong scope for it, and both refused model families were refused for that reason. Calibration only sees the module it was handed. A VLM calibrates its language model, so the shards, the exclusions and config.json all described that submodel rather than the whole VLM. And calibration runs before orphaned MTP weights are loaded, by which point every shard was already written, so they could not be passed at all. The exporter is now created by whoever owns the export and announced on the model. It publishes itself on the export root and, for a VLM, on the language model too, so whichever of the two mtq.quantize is handed finds it; calibration binds it and drives it per layer, and export_hf_checkpoint dispatches to it. A layerwise VLM run is the same mtq.quantize(...) / export_hf_checkpoint(...) pair as a plain LLM, and orphaned MTP tensors are an ordinary finalize() argument. mtq.quantize and mtq.calibrate are unchanged: a layerwise-only feature does not belong in the public quantization API. Construction is inert, since the caller builds the exporter before there are quantizers to validate or read a config from; bind() does that, from calibration after quantizer insertion and before any layer is converted, so unsupported models still fail in seconds rather than hours. Only the pass that sets export_dir drives the exporter, since a list-form algorithm runs one per entry and an earlier pass must not convert layers a later one still has to calibrate. Calibration now writes only the layer shards; finalize() adds the tail shard, the index and the config artifacts. It is announced rather than held so a config-only caller can reach it, and the message says what is still owed. Tested end to end through hf_ptq against a baseline exported by main: Qwen3-VL-8B (1254 keys, 0 differing) and GLM-4.7-Flash (28119 keys, 0 differing, all 212 orphaned MTP tensors in the tail shard and the index). Co-Authored-By: Claude Opus 5 Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- CHANGELOG.rst | 5 +- examples/hf_ptq/hf_ptq.py | 130 ++++++------ modelopt/torch/export/layerwise_export.py | 100 +++++++-- modelopt/torch/export/unified_export_hf.py | 9 + modelopt/torch/quantization/config.py | 13 +- modelopt/torch/quantization/model_calib.py | 30 ++- .../gpu/torch/export/test_layerwise_export.py | 197 ++++++++++++++++-- 7 files changed, 374 insertions(+), 110 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aa45af04f7f..69ad253374b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,10 @@ Changelog **New Features** +*Quantization* + +- 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. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. + **Backward Breaking Changes** **Deprecations** @@ -28,7 +32,6 @@ Changelog - Add opt-in FP8 Vision Encoder recipes under the ``qwen3_vl`` and ``qwen3_5`` model types. The vision-only recipe keeps the language model and KV cache in high precision; the joint recipe quantizes Vision Encoder and language-model Linears and uses FP8 KV-cache cast. Both quantize primary and deepstack merger Linears where present, while leaving patch embedding and vision-attention BMMs in high precision. Exported checkpoints require an inference runtime that supports quantized Vision Encoder Linears. - 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/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index ede1ce706a0..821d82b93ed 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -40,6 +40,7 @@ get_tokenizer, is_enc_dec, is_nemotron_vl, + layerwise_export_block, load_mtp_weights, mlflow_run, mtp_layer_prefixes_from_checkpoint, @@ -78,6 +79,7 @@ has_spec_opt, save_expert_token_count_table, ) +from modelopt.torch.export.layerwise_export import LayerwiseExporter from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights @@ -778,6 +780,10 @@ def mono_quantize( else None, ) + if args.layerwise_export: + # Announces itself on the model; calibration and the export both pick it up there. + LayerwiseExporter(full_model, args.export_path) + if calibration_only: language_model = mtq.calibrate( language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop @@ -796,27 +802,23 @@ 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. +def assert_layerwise_export_compatible(args, full_model, algorithm) -> None: + """Refuse layerwise export before calibration starts, not after the run is paid for. - 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. + Layerwise export writes each layer's shard during calibration and finishes the checkpoint + in finalize() afterwards, so anything that would rewrite or contradict that checkpoint 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." - ) + block = layerwise_export_block(algorithm) + if block is not None: + entries = algorithm if isinstance(algorithm, list) else [algorithm] + owner = next(e for e in entries if isinstance(e, dict) and e.get("layerwise") is block) + if not owner.get("method"): + raise NotImplementedError( + "layerwise.export_dir needs a calibration method: without one there is no " + "per-layer pass to write the shards, so the export would find nothing. Set " + "algorithm.method, or export without layerwise.export_dir." + ) if has_spec_opt(full_model): raise NotImplementedError( @@ -855,6 +857,27 @@ def assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) -> ) +def save_source_config(args, export_path) -> None: + """Copy the source model's config to the export path, for VLMs the exporters skip.""" + print(f"Saving original model config to {export_path}") + config_kwargs = {"trust_remote_code": args.trust_remote_code} + if args.attn_implementation is not None: + config_kwargs["attn_implementation"] = args.attn_implementation + AutoConfig.from_pretrained(args.pyt_ckpt_path, **config_kwargs).save_pretrained(export_path) + + +def save_processor_config(args, export_path) -> None: + """Copy the processor config, without which a VLM checkpoint cannot preprocess images.""" + try: + print(f"Saving processor config to {export_path}") + AutoProcessor.from_pretrained( + args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code + ).save_pretrained(export_path) + except Exception as e: + print(f"Warning: Could not save processor config: {e}") + print("This is normal for some VLM architectures that don't use AutoProcessor") + + def export_quantized( args: argparse.Namespace, full_model: torch.nn.Module, @@ -880,29 +903,12 @@ def export_quantized( print(f"Quantized speculative decoding checkpoint exported to: {export_path}") return - # Check if the model is a multimodal/VLM model - is_vlm = is_multimodal_model(full_model) - - if is_vlm: - # Save original model config and the processor config to the export path for VLMs. - print(f"Saving original model config to {export_path}") - - config_kwargs = {"trust_remote_code": args.trust_remote_code} - if args.attn_implementation is not None: - config_kwargs["attn_implementation"] = args.attn_implementation - AutoConfig.from_pretrained(args.pyt_ckpt_path, **config_kwargs).save_pretrained( - export_path - ) - - # Try to save processor config if available - try: - print(f"Saving processor config to {export_path}") - AutoProcessor.from_pretrained( - args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code - ).save_pretrained(export_path) - except Exception as e: - print(f"Warning: Could not save processor config: {e}") - print("This is normal for some VLM architectures that don't use AutoProcessor") + if is_multimodal_model(full_model): + # Per-layer export writes its own config.json with quantization_config, which + # the source config would replace; it never writes a processor config. + if not args.layerwise_export: + save_source_config(args, export_path) + save_processor_config(args, export_path) start_time = time.time() is_tensorrt_llm_export = ( @@ -957,22 +963,11 @@ def export_quantized( if mtp_layer_prefixes: full_model._mtp_layer_prefixes = mtp_layer_prefixes - 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, - ) + export_hf_checkpoint( + full_model, + export_dir=export_path, + extra_state_dict=mtp_state_dict, + ) if args.qformat == "w4a16_nvfp4": warnings.warn( @@ -1231,7 +1226,17 @@ def quantize_main( 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) + args.layerwise_export = any( + cfg.get("export_dir") is not None and cfg.get("enable", False) for cfg in layerwise_cfgs + ) + if not args.layerwise_export and any( + cfg.get("export_dir") is not None for cfg in layerwise_cfgs + ): + warnings.warn( + "layerwise.export_dir is set but layerwise.enable is not, so there is no " + "per-layer pass to write the shards: the whole-model export runs instead, which " + "holds the full state dict in host memory." + ) if args.layerwise_export: if isinstance(recipe, ModelOptAutoQuantizeRecipe): # Only the mono-quantize path retargets export_dir and runs the refusals; @@ -1368,9 +1373,8 @@ def quantize_main( # 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. + # Only the FSDP2 loader flags these before quantization, and the exclusions must + # be in quant_cfg before mtq.quantize converts the first layer. mtp_layer_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path) if mtp_layer_prefixes: quant_cfg = copy.deepcopy(quant_cfg) @@ -1382,7 +1386,7 @@ def quantize_main( # 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) + assert_layerwise_export_compatible(args, full_model, quant_cfg.get("algorithm")) 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 diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index a196c9a9e52..55f576e2ffb 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -35,7 +35,7 @@ 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 .model_utils import TiedWeightMap, get_language_model_from_vl 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 @@ -57,6 +57,10 @@ SUPPORTED_FORMATS = FUSION_FREE_FORMATS | _PER_LAYER_FUSABLE_FORMATS +#: Set on the model handed to ``mtq.quantize``, so calibration and the export that follows +#: it reach the same exporter. +LAYERWISE_EXPORTER_ATTR = "_layerwise_exporter" + _TAIL_SHARD = "model-tail.safetensors" _INDEX_FILE = "model.safetensors.index.json" @@ -159,10 +163,50 @@ def __init__( export_dir: Path | str, dtype: torch.dtype | None = None, ) -> None: - """Validate support and capture model-level state. + """Name the model the checkpoint describes and where it goes. + + Nothing is inspected: the caller builds this before ``mtq.quantize``, when there is + no quantizer yet to validate or read a config from. :meth:`bind` does that. + """ + self._model = model + self._export_dir = Path(export_dir) + self._export_dir.mkdir(parents=True, exist_ok=True) + self._dtype = dtype + self._bound = False + self._finalized = False + self._announced_on: list[nn.Module] = [] + # A VLM calibrates its language model but exports the whole thing, so announce on + # both: whichever of the two mtq.quantize is handed will find this exporter. + self.announce(model) + lineage = get_language_model_from_vl(model) + if lineage: + self.announce(lineage[-1]) + + @property + def export_dir(self) -> Path: + """Where the shards go. The exporter owns this, not the caller's config.""" + return self._export_dir + + def announce(self, module: nn.Module) -> None: + """Publish this exporter on ``module`` for a later pass to pick up. + + Calibration and export are handed different models -- a VLM calibrates its language + model but exports the whole thing -- so each end is told separately. + """ + setattr(module, LAYERWISE_EXPORTER_ATTR, self) + if not any(m is module for m in self._announced_on): + self._announced_on.append(module) + + def bind(self, calibrated_layers: list[nn.Module]) -> None: + """Validate the model and snapshot what the tail pass needs. - Runs before calibration, so nothing amax-dependent exists yet. + Called from calibration, after quantizer insertion and before any layer is converted + -- the only window where both hold. """ + if self._bound: + # Only the first pass has a model with no layer converted yet. + return + model = self._model assert_layerwise_export_supported(model) # Splits regroup tensors across the whole state dict; no per-layer pass reverses that. _assert_no_split_rules(model) @@ -184,7 +228,14 @@ def __init__( "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. + # A length difference is the one mismatch export_layer's per-layer identity check + # cannot catch: every call would pass and _write_index would then miss a shard. + if len(layers) != len(calibrated_layers): + raise RuntimeError( + f"the exporter found {len(layers)} decoder layers but calibration will drive " + f"{len(calibrated_layers)}, so layer_idx would not agree. The exporter's " + "model must contain exactly the layers being calibrated." + ) self._layers = layers layer_ids = {id(m): i for i, m in enumerate(layers)} self._layer_names: dict[int, str] = {} @@ -193,16 +244,12 @@ def __init__( 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) - # 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._ctx = ExportContext(model=model, dtype=_resolve_export_dtype(model, self._dtype)) + # get_quant_config reports on the quantizer modules, which export_layer replaces as + # it goes, so by finalize() the model would look 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 = self._quant_config["quantization"]["kv_cache_quant_algo"] - self._finalized = False self._name_mapper = None try: @@ -213,6 +260,9 @@ def __init__( "match the original HF hub checkpoint." ) + # Last, so a bind() that raised is not mistaken for a completed one on retry. + self._bound = True + def export_layer( self, layer_idx: int, @@ -228,6 +278,7 @@ def export_layer( # Local, as in every other export module: the plugin imports transformers. from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + assert self._bound, "export_layer() before bind()" 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 @@ -290,12 +341,23 @@ def _fuse_shared_input_scales(self, layer_module: nn.Module, layer_inputs: list self._ctx.model, input_to_linear, quantization_format=layer_format ) - def finalize(self) -> dict: + def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> dict: """Export the tail, write the config artifacts, and index all shards. - Leaves ``export_dir`` a complete checkpoint; no ``export_hf_checkpoint()`` needed. + ``extra_state_dict`` carries tensors with no slot in ``model.state_dict()`` -- MTP + weights, whichever convention the checkpoint uses. They are already in export form, + so only the hub-name reversal applies, and they win on a name clash exactly as they + do in ``export_hf_checkpoint``. """ - assert not self._finalized, "finalize() called twice" + # Not asserts: finalize() is called by whoever owns the export, and -O would strip + # the contract they are stating. + if not self._bound: + raise RuntimeError( + "finalize() before calibration bound the exporter: layerwise calibration " + "never ran, so there are no layer shards to finish." + ) + if self._finalized: + raise RuntimeError("finalize() called twice; the checkpoint is already written.") self._finalized = True model = self._ctx.model @@ -354,10 +416,20 @@ def finalize(self) -> dict: continue self._collect(tail, name, tensor) + for name, tensor in (extra_state_dict or {}).items(): + key = self._name_mapper(name) if self._name_mapper is not None else name + tail[key] = tensor.detach().contiguous().cpu() + 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) + # Left attached, the exporter follows the model into any save or deepcopy. + for module in self._announced_on: + if getattr(module, LAYERWISE_EXPORTER_ATTR, None) is self: + delattr(module, LAYERWISE_EXPORTER_ATTR) + self._announced_on.clear() + warnings.warn( "The exported checkpoint is complete, but per-layer export leaves the model in " "export form: it must not be used for inference." diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index d98211cc3fb..b75e20cb2fb 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1554,6 +1554,15 @@ def export_hf_checkpoint( :func:`_postprocess_safetensors` for diffusion model exports. See its docstring for supported keys. """ + # Local: layerwise_export imports this module. + from .layerwise_export import LAYERWISE_EXPORTER_ATTR + + exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) + if exporter is not None: + # Per-layer export wrote the shards during calibration; this writes the rest. + exporter.finalize(extra_state_dict=extra_state_dict) + return + export_dir = Path(export_dir) export_dir.mkdir(parents=True, exist_ok=True) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 43920b5b5dd..541cf5cefdf 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -759,15 +759,16 @@ class LayerwiseConfig(ModeloptBaseConfig): 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 " + "this directory the moment its calibration finishes, replacing the separate " "``export_hf_checkpoint()`` pass and its full-precision intermediate. " + "Calibration writes only the layer shards; the checkpoint does not load until " + "``finalize()`` is called on the exporter attached to the model, which adds " + "the tail shard, the index and the config artifacts. " "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." + "models, resident or accelerate-offloaded; AWQ, SVDQuant, multi-process jobs " + "and weight-tied quantized modules raise NotImplementedError. Per-layer export " + "leaves the model in memory in export form, never valid for inference." ), ) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index d4266b289b9..8a66951d5a3 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2087,12 +2087,23 @@ 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. + from modelopt.torch.export.layerwise_export import LAYERWISE_EXPORTER_ATTR, LayerwiseExporter + exporter = None + finalize_hint = "" if export_dir is not None: - from modelopt.torch.export.layerwise_export import LayerwiseExporter - - exporter = LayerwiseExporter(model, export_dir) + # Only the pass that sets export_dir drives it: an earlier layerwise pass must not + # convert layers a later one still has to calibrate. + exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) + if exporter is None: + exporter = LayerwiseExporter(model, export_dir) + # Only for one built here: a caller that announced its own finishes it too. + finalize_hint = ( + f" Call finalize() on model.{LAYERWISE_EXPORTER_ATTR} to write the tail " + "shard, the index and the config artifacts; it does not load until then." + ) + # Before calibration, so unsupported models fail in seconds not hours. + exporter.bind(calibrated_layers=list(transformer_layers)) ckpt = _CheckpointState.from_folder( checkpoint_dir, @@ -2106,8 +2117,10 @@ def layerwise_calibrate( 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}") + warn_rank_0( + f"Layerwise export: every layer shard in {exporter.export_dir} is already " + f"written.{finalize_hint}" + ) return layer_pbar = tqdm( @@ -2206,8 +2219,9 @@ def _layer_forward_loop(m, _inputs=layer_inputs): 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}") + warn_rank_0( + f"Layerwise export: wrote every layer shard to {exporter.export_dir}.{finalize_hint}" + ) if start_layer > 0: warn_rank_0( f"This run resumed at layer {start_layer}, so layers 0..{start_layer - 1} " diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index ca40f72e875..de30c2848af 100644 --- a/tests/gpu/torch/export/test_layerwise_export.py +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -22,11 +22,20 @@ import pytest import torch -from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe +from _test_utils.torch.transformers_models import ( + get_tiny_gemma3vl, + 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.layerwise_export import ( + LAYERWISE_EXPORTER_ATTR, + LayerwiseExporter, + layer_shard_name, +) +from modelopt.torch.export.model_utils import get_language_model_from_vl from modelopt.torch.export.unified_export_hf import export_hf_checkpoint NUM_LAYERS = 4 @@ -46,6 +55,23 @@ def _build_model(): return model +def _layerwise_quantize(model, cfg, extra_state_dict=None, export_model=None): + """Attach the exporter, calibrate, then finish the checkpoint. + + ``export_model`` names a wider root than the calibrated model, as a VLM pipeline does; + omitted, calibration builds and announces its own. + """ + if export_model is not None: + setattr( + model, + LAYERWISE_EXPORTER_ATTR, + LayerwiseExporter(export_model, cfg["algorithm"]["layerwise"]["export_dir"]), + ) + mtq.quantize(model, cfg, _calib) + getattr(model, LAYERWISE_EXPORTER_ATTR).finalize(extra_state_dict=extra_state_dict) + return model + + def _layerwise_cfg(export_dir, checkpoint_dir, base=None): cfg = copy.deepcopy(base or mtq.FP8_DEFAULT_CFG) cfg["algorithm"] = { @@ -236,7 +262,7 @@ def test_export_matches_whole_model_export( 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) + _layerwise_quantize(_build_model(), cfg) exported = _load_checkpoint(export_dir) if expected_key_suffix: @@ -250,6 +276,142 @@ def test_export_matches_whole_model_export( assert (export_dir / artifact).is_file(), f"{artifact} missing" +def test_config_only_export_is_finishable(tmp_path): + """No caller-supplied exporter: calibration builds one and announces it to be finished. + + Held in a local instead, the run would leave layer shards with no tail, index or config. + Finalize clears the announcement, or the exporter follows the model into save/deepcopy. + """ + export_dir = tmp_path / "fused" + model = _build_model() + mtq.quantize(model, _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib) + + exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) + assert exporter is not None, "calibration did not attach the exporter it built" + exporter.finalize() + for artifact in ("model-tail.safetensors", "model.safetensors.index.json", "config.json"): + assert (export_dir / artifact).is_file(), f"{artifact} missing" + assert getattr(model, LAYERWISE_EXPORTER_ATTR, None) is None, ( + "the exporter is still attached after the export finished" + ) + + +def test_only_the_pass_owning_export_dir_exports(tmp_path, baseline_checkpoint): + """Only the entry that sets export_dir may drive the exporter. + + An earlier layerwise pass reaching it would convert every layer into export form before + this pass has calibrated it, capturing an intermediate state in the shards. + """ + export_dir = tmp_path / "fused" + cfg = _layerwise_cfg(export_dir, tmp_path / "ckpt") + layerwise = cfg["algorithm"]["layerwise"] + # Its own checkpoint_dir: a manifest completed by pass 1 would make pass 2 resume past + # every layer. + first = copy.deepcopy(layerwise) + del first["export_dir"] + first["checkpoint_dir"] = str(tmp_path / "ckpt_first") + cfg["algorithm"] = [ + {"method": "max", "layerwise": first}, + {"method": "max", "layerwise": copy.deepcopy(layerwise)}, + ] + model = _build_model() + setattr(model, LAYERWISE_EXPORTER_ATTR, LayerwiseExporter(model, export_dir)) + exported_by = [] + real = LayerwiseExporter.export_layer + + def record(self, idx, *args, **kwargs): + exported_by.append(idx) + return real(self, idx, *args, **kwargs) + + with patch.object(LayerwiseExporter, "export_layer", record): + mtq.quantize(model, cfg, _calib) + getattr(model, LAYERWISE_EXPORTER_ATTR).finalize() + + assert exported_by == list(range(NUM_LAYERS)), ( + f"each layer must be exported exactly once, by the owning pass; got {exported_by}" + ) + _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(export_dir)) + + +def test_exporter_rooted_off_the_calibrated_layers_is_refused(tmp_path): + """layer_idx only means the same thing on both sides while the layer lists match.""" + torch.manual_seed(0) + model = get_tiny_llama(num_hidden_layers=NUM_LAYERS).cuda().eval() + model.config.architectures = ["LlamaForCausalLM"] + other = get_tiny_llama(num_hidden_layers=NUM_LAYERS + 1).cuda().eval() + other.config.architectures = ["LlamaForCausalLM"] + + export_dir = tmp_path / "fused" + setattr(model, LAYERWISE_EXPORTER_ATTR, LayerwiseExporter(other, export_dir)) + with pytest.raises(RuntimeError, match="layer_idx would not agree"): + mtq.quantize(model, _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib) + + +def test_orphaned_tensors_reach_the_tail_shard(tmp_path): + """MTP weights load after calibration, so they can only be passed once it is over.""" + export_dir = tmp_path / "fused" + orphans = { + "mtp.layers.0.weight": torch.ones(4, 4, dtype=torch.bfloat16), + "mtp.norm.weight": torch.ones(4, dtype=torch.bfloat16), + } + _layerwise_quantize( + _build_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt"), extra_state_dict=orphans + ) + + exported = _load_checkpoint(export_dir) + for key, value in orphans.items(): + assert key in exported, f"{key} missing from the exported checkpoint" + assert torch.equal(exported[key].cpu(), value) + weight_map = json.loads((export_dir / "model.safetensors.index.json").read_text())["weight_map"] + assert set(orphans) <= set(weight_map), "orphans written but left out of the index" + + +def test_vlm_export_follows_the_documented_flow(tmp_path): + """quantize(language_model) then export_hf_checkpoint(vlm), with no layerwise plumbing. + + The exporter finds the language model itself, and export_hf_checkpoint dispatches to it, + so a VLM caller writes the same two lines a plain LLM does. + """ + torch.manual_seed(0) + # sliding_window past the calibration length: the tiny default (16) masks against a + # shorter window than the batch and the text forward fails before export is reached. + vlm = get_tiny_gemma3vl(tie_word_embeddings=False, sliding_window=1024).cuda().eval() + # The plain attribute is what TiedWeightMap reads; the config kwarg only reaches text_config. + vlm.all_tied_weights_keys = {} + language_model = get_language_model_from_vl(vlm)[-1] + + export_dir = tmp_path / "fused" + cfg = _layerwise_cfg(export_dir, tmp_path / "ckpt") + LayerwiseExporter(vlm, export_dir) + mtq.quantize(language_model, cfg, _calib) + export_hf_checkpoint(vlm, export_dir=export_dir) + + exported = _load_checkpoint(export_dir) + assert any(k.startswith("language_model.model.layers.") for k in exported), ( + "decoder keys lost the VLM namespace" + ) + assert any(k.startswith("vision_tower.") for k in exported), "the vision tower was not exported" + assert getattr(vlm, LAYERWISE_EXPORTER_ATTR, None) is None + assert getattr(language_model, LAYERWISE_EXPORTER_ATTR, None) is None + + +def test_exporter_root_widens_the_checkpoint_to_the_parent(tmp_path): + """Rooting the exporter at the parent is what puts the shards in its namespace.""" + torch.manual_seed(0) + parent = get_tiny_llama(num_hidden_layers=NUM_LAYERS).cuda().eval() + parent.config.architectures = ["LlamaForCausalLM"] + inner = parent.model + + export_dir = tmp_path / "fused" + _layerwise_quantize(inner, _layerwise_cfg(export_dir, tmp_path / "ckpt"), export_model=parent) + + exported = _load_checkpoint(export_dir) + assert any(k.startswith("model.layers.") for k in exported), ( + "layer keys lost the parent namespace" + ) + assert any(k.startswith("lm_head") for k in exported), "the parent's tail was not exported" + + 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. @@ -257,7 +419,7 @@ def test_index_resolves_every_key_to_the_shard_holding_it(tmp_path): still fails in vLLM or transformers. """ export_dir = tmp_path / "fused" - mtq.quantize(_build_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib) + _layerwise_quantize(_build_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt")) weight_map = json.loads((export_dir / "model.safetensors.index.json").read_text())["weight_map"] on_disk = {} @@ -272,7 +434,7 @@ def test_index_resolves_every_key_to_the_shard_holding_it(tmp_path): 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) + _layerwise_quantize(_build_model(), _layerwise_cfg(tmp_path / "fused", checkpoint_dir)) assert not list(checkpoint_dir.rglob("weights.pt")) assert not list(checkpoint_dir.rglob("quantizer_buffers.pt")) @@ -292,13 +454,13 @@ def test_resume_skips_exported_layers(tmp_path, baseline_checkpoint): # 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) + _layerwise_quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir)) 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) + _layerwise_quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir)) _assert_same_checkpoint(baseline_checkpoint, _load_checkpoint(export_dir)) @@ -306,7 +468,7 @@ def test_resume_skips_exported_layers(tmp_path, baseline_checkpoint): 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) + _layerwise_quantize(_build_model(), _layerwise_cfg(tmp_path / "fused", checkpoint_dir)) manifest_path = checkpoint_dir / "manifest.json" manifest = json.loads(manifest_path.read_text()) @@ -314,8 +476,8 @@ def test_resume_without_matching_shards_fails_fast(tmp_path): 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 + _layerwise_quantize( + _build_model(), _layerwise_cfg(tmp_path / "empty_export", checkpoint_dir) ) @@ -327,7 +489,7 @@ def test_complete_manifest_finalizes_without_recalibrating(tmp_path, baseline_ch """ export_dir = tmp_path / "fused" checkpoint_dir = tmp_path / "ckpt" - mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + _layerwise_quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir)) # 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. @@ -336,7 +498,7 @@ def test_complete_manifest_finalizes_without_recalibrating(tmp_path, baseline_ch 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) + _layerwise_quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir)) # The layer shards must be reused verbatim, not rewritten. for name, mtime in layer_mtimes.items(): @@ -352,7 +514,7 @@ def test_shards_without_resume_record_refuse(tmp_path, damage): """ export_dir = tmp_path / "fused" checkpoint_dir = tmp_path / "ckpt" - mtq.quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir), _calib) + _layerwise_quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir)) manifest = checkpoint_dir / "manifest.json" if damage == "deleted": @@ -363,7 +525,7 @@ def test_shards_without_resume_record_refuse(tmp_path, damage): 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) + _layerwise_quantize(_build_model(), _layerwise_cfg(export_dir, checkpoint_dir)) def test_export_without_checkpoint_dir_may_overwrite(tmp_path): @@ -421,8 +583,8 @@ def test_moe_export_matches(tmp_path): 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 + _layerwise_quantize( + _build_moe_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()) ) _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) @@ -442,10 +604,9 @@ def test_export_consumes_the_model_without_affecting_the_checkpoint(tmp_path): export_hf_checkpoint(mtq.quantize(_build_model(), base, _calib), export_dir=baseline_dir) export_dir = tmp_path / "fused" - model = mtq.quantize( + model = _layerwise_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))