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))