From 1899b986a20cf722cc17a25d55446add2758a72e Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:17:19 -0700 Subject: [PATCH 1/6] fix(export): find decoder layers through arbitrarily nested wrappers get_homogeneous_hf_decoder_layers unwrapped .model then .language_model once each, so it only reached layers exactly two wrappers deep and in that order. Kimi-K3 keeps its decoder at language_model.model.layers, so the walk stopped on the intermediate wrapper, returned None, and the architecture was reported unsupported -- which takes layerwise calibration and per-layer export out of reach for the model that needs them most. Unwrap iteratively instead, bounded so a cycle cannot hang. Also re-apply --attn_implementation after model construction. Kimi-K3's remote code overwrites _attn_implementation to flash_attention_2 unconditionally in __init__, ignoring the flag; export runs a trace forward, so an unavailable backend fails there rather than at load. Sub-configs are walked because remote code typically rewrites the nested text_config, and layer modules hold a reference to the same object. Only applied when the caller passed the flag explicitly, so nothing changes by default. Both were found in the previous Kimi-K3 run (PR #2008 workflow) but were never upstreamed; the branch carrying them was deleted after that PR merged. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 9f86b2857cfdbf13aed2c1f75c2d97c46d0395fb) --- examples/hf_ptq/example_utils.py | 42 +++++++++++++++++++ .../torch/quantization/plugins/huggingface.py | 25 ++++++++--- 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 89a88f76458..50e28d07456 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -703,6 +703,42 @@ def _resolve_init_config(hf_config, auto_model_module, ckpt_path, config_kwargs) return hf_config +def _force_attn_implementation(model, attn_implementation: str) -> None: + """Set ``_attn_implementation`` on the model config and every nested sub-config. + + Some remote modeling code overrides the requested attention implementation inside + ``__init__`` -- Kimi-K3 rewrites it to ``flash_attention_2`` unconditionally, ignoring + ``--attn_implementation``. Export runs a trace forward, so a backend whose compiled + extension is unavailable in this environment fails there rather than at load. + + Sub-configs are walked because multimodal models keep separate ones per tower, and + remote code typically rewrites the *nested* config (Kimi-K3 rewrites ``text_config`` + from ``KimiLinearModel.__init__``). Layer modules hold a reference to the same config + object, so overriding it here takes effect on the next forward. Only applied when the + caller asked for an implementation explicitly. + """ + pending, seen, changed = [model.config], set(), [] + while pending: + cfg = pending.pop() + if cfg is None or id(cfg) in seen: + continue + seen.add(id(cfg)) + current = getattr(cfg, "_attn_implementation", None) + if current is not None and current != attn_implementation: + try: + cfg._attn_implementation = attn_implementation + changed.append(f"{type(cfg).__name__}: {current} -> {attn_implementation}") + except Exception as e: # pragma: no cover - depends on remote config class + warnings.warn(f"Could not apply attn_implementation on {type(cfg).__name__}: {e}") + for sub in ("text_config", "vision_config", "audio_config", "decoder", "encoder"): + pending.append(getattr(cfg, sub, None)) + + if changed: + print("Re-applied the requested attention implementation after model init:") + for line in changed: + print(f" {line}") + + def _get_config_dtype(config): config_dtype = ( getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.bfloat16 @@ -967,6 +1003,12 @@ def has_pack_quantized_config(config): **model_kwargs2, ) model.eval() + + # Honour the caller's explicit choice even when remote modeling code overwrote it + # during __init__ (see _force_attn_implementation). + if attn_implementation is not None: + _force_attn_implementation(model, attn_implementation) + if has_pack_quantized_config(hf_config): _unpack_compressed_linear_weights(model, ckpt_path) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index dde24513086..b108b1a0dc2 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1794,17 +1794,30 @@ def is_homogeneous_hf_model(model: nn.Module) -> bool: return len(layer_classes) == 1 +#: How deep to unwrap before giving up. Each level is one of the wrappers below, and no +#: released architecture nests more than a handful; the bound only stops a cycle. +_MAX_DECODER_UNWRAP_DEPTH = 8 + + def get_homogeneous_hf_decoder_layers(model: nn.Module) -> nn.ModuleList | None: if not _is_supported_hf_model(model): return None + # Unwrap iteratively rather than testing each wrapper once: multimodal models nest them + # in either order and to arbitrary depth. Kimi-K3 keeps its layers at + # ``language_model.model.layers``, so a single ``.model`` then ``.language_model`` walk + # stops on the intermediate wrapper and reports the architecture unsupported. decoder = model - if hasattr(decoder, "model"): - decoder = decoder.model - if hasattr(decoder, "language_model"): - decoder = decoder.language_model - if hasattr(decoder, "layers"): - return decoder.layers + for _ in range(_MAX_DECODER_UNWRAP_DEPTH): + if hasattr(decoder, "layers"): + return decoder.layers + for attr in ("model", "language_model"): + inner = getattr(decoder, attr, None) + if isinstance(inner, nn.Module): + decoder = inner + break + else: + return None return None From 969c0a89cc020d9ec64583559082e78a3905a1cb Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:27:18 -0700 Subject: [PATCH 2/6] feat(recipes): experts-only NVFP4 layerwise fused export for offloaded models Pairs layerwise.export_dir with checkpoint_dir for a PTQ run too large to hold resident, so an interrupted run resumes without recalibrating or re-exporting finished layers. That combination is the point for a run that outlasts its GPU session -- a kill loses at most the in-flight layer. Scopes experts as '*.experts.*' rather than '*block_sparse_moe*'. On fine-grained MoE the broader glob also matches shared_experts.*, routed_expert_{up,down}_proj and routed_expert_norm; on Kimi-K3 that is 552 additional modules the vendor left unquantized, one of which is an RMSNorm. shared_experts is missed by the narrow glob because its path contains '_experts.' rather than '.experts.', per fnmatch semantics. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 02ed6b545ad9f96130e85a4a6296ead883022c2c) --- ..._only-kv_fp8_layerwise_export_offload.yaml | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml new file mode 100644 index 00000000000..2be18d6522c --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: > + NVFP4 W4A4 on routed experts only, FP8 KV cache, max layerwise calibration, with each + decoder layer exported to its own shard as soon as it is calibrated. Same intent as + nvfp4_experts_only-kv_fp8_layerwise_export, but scoped and paired for a model too large + to hold resident: use with --offload_folder and per-device memory budgets. + + Expert scoping is '*.experts.*' rather than '*block_sparse_moe*'. On fine-grained MoE + models the broader glob also matches shared_experts.*, routed_expert_{up,down}_proj and + routed_expert_norm -- on Kimi-K3 that is 552 extra modules the vendor deliberately left + unquantized, one of which is an RMSNorm. '*.experts.*' matches only the routed expert + projections; shared_experts is missed because its path contains '_experts.' rather than + '.experts.' (fnmatch semantics, conversion.py). + + Pairs checkpoint_dir with export_dir so an interrupted run resumes without recalibrating + or re-exporting finished layers. That is the point of the combination for a run that + outlasts its GPU session: a killed run loses at most the in-flight layer. + + A resumed run never recalibrates the layers it skipped, so the exported checkpoint is + complete but the in-memory model is not and must not be used for inference. +quantize: + algorithm: + method: max + layerwise: + enable: true + # max only updates _amax, so the exported shard stays valid for its layer. + calib_mutates_weights: false + checkpoint_dir: /tmp/modelopt_layerwise_ckpt + # Presence enables per-layer export; the value is replaced with --export_path. + export_dir: /tmp/modelopt_layerwise_export + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers From 0b369dde92afd5aa9c97c3b92ecf1f8863310e27 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 13 Aug 2026 14:13:02 -0700 Subject: [PATCH 3/6] fix(hf_ptq): pin offloaded weights that are read outside their own forward accelerate materializes an offloaded weight in that module's pre-forward hook and returns it to meta in the matching post-forward. A weight read from a *sibling's* forward is therefore on meta at the moment it is used. Kimi-K3 does exactly this: _apply_attn_res computes norm.weight.float() * proj.weight.squeeze(0).float() from the decoder layer's forward, reaching into six children (three call sites in modeling_kimi_linear.py). The result is RuntimeError: Tensor on device meta is not on the expected device cuda:0! which is why a disk-offloaded K3 could not run a forward at all. Setting the tensor resident is not enough on its own -- post_forward walks the module's tensors and pushes every one back to meta, so the hook has to go. Detaching alone is not enough either: AlignDevicesHook.detach_hook restores each tensor to original_devices[name] and skips meta, which is precisely what a disk-offloaded param has, so it would be left on meta. Retargeting original_devices at the execution device first makes detach materialize the values itself, through accelerate's own code path rather than a hand-rolled copy. Safe because these modules' forward is never called -- only their raw .weight is read -- so removing the hook removes nothing that was doing work. The tensors are one row each, (1, hidden) and (hidden,), so pinning all six on a 93-layer model costs single-digit MB. Verified on a disk-offloaded tiny Kimi-K3: 10 externally-read params on meta before, 0 after, and the forward completes with finite logits where it previously raised. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 4c50e132ab1995d09f59e45d0ebe768facaa4911) --- examples/hf_ptq/example_utils.py | 67 +++++++++++++++++++++++++++++++- 1 file changed, 65 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 50e28d07456..74ec9c0df45 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -730,8 +730,10 @@ def _force_attn_implementation(model, attn_implementation: str) -> None: changed.append(f"{type(cfg).__name__}: {current} -> {attn_implementation}") except Exception as e: # pragma: no cover - depends on remote config class warnings.warn(f"Could not apply attn_implementation on {type(cfg).__name__}: {e}") - for sub in ("text_config", "vision_config", "audio_config", "decoder", "encoder"): - pending.append(getattr(cfg, sub, None)) + pending.extend( + getattr(cfg, sub, None) + for sub in ("text_config", "vision_config", "audio_config", "decoder", "encoder") + ) if changed: print("Re-applied the requested attention implementation after model init:") @@ -739,6 +741,61 @@ def _force_attn_implementation(model, attn_implementation: str) -> None: print(f" {line}") +#: Modules whose weights are read from *another* module's forward, so accelerate never +#: materializes them. Kimi-K3's ``_apply_attn_res`` does +#: ``norm.weight.float() * proj.weight.squeeze(0).float()`` from the decoder layer's +#: forward, reaching into these six children (``modeling_kimi_linear.py``, three call +#: sites). Each is one row -- ``(1, hidden)`` and ``(hidden,)`` -- so pinning all of them +#: on a 93-layer model costs single-digit MB. +_EXTERNALLY_READ_PARAM_SUFFIXES = ( + "self_attention_res_proj", + "self_attention_res_norm", + "mlp_res_proj", + "mlp_res_norm", + "output_attn_res_proj", + "output_attn_res_norm", +) + + +def _pin_externally_read_params( + model, suffixes: tuple[str, ...] = _EXTERNALLY_READ_PARAM_SUFFIXES +) -> int: + """Make offloaded weights that are read outside their own forward permanently resident. + + accelerate materializes an offloaded weight in *that module's* pre-forward hook and + returns it to meta in the matching post-forward. A weight read from a sibling's forward + is therefore on meta at the moment it is used, which surfaces as + ``Tensor on device meta is not on the expected device cuda:0``. + + Setting the tensor is not enough on its own: ``post_forward`` walks the module's tensors + and pushes every one back to meta, so the hook has to go. Detaching alone is not enough + either -- ``AlignDevicesHook.detach_hook`` restores each tensor to + ``original_devices[name]`` and *skips* meta, which is precisely what a disk-offloaded + param has, so it would be left on meta. Retargeting ``original_devices`` at the + execution device first makes detach do the materialization itself, using accelerate's + own code path rather than a hand-rolled copy. + + Safe because these modules' ``forward`` is never called -- only their raw ``.weight`` is + read -- so removing the hook removes nothing that was doing work. + + Returns the number of modules pinned. + """ + from accelerate.hooks import remove_hook_from_module + + pinned = 0 + for name, module in model.named_modules(): + if not name.endswith(suffixes): + continue + hook = getattr(module, "_hf_hook", None) + if hook is None or not getattr(hook, "offload", False): + continue + device = hook.execution_device + hook.original_devices = dict.fromkeys(getattr(hook, "original_devices", {}), device) + remove_hook_from_module(module) + pinned += 1 + return pinned + + def _get_config_dtype(config): config_dtype = ( getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.bfloat16 @@ -1009,6 +1066,12 @@ def has_pack_quantized_config(config): if attn_implementation is not None: _force_attn_implementation(model, attn_implementation) + # Offloaded weights that a sibling's forward reads would otherwise be on meta when used. + if _disk_offload: + n_pinned = _pin_externally_read_params(model) + if n_pinned: + print(f"Pinned {n_pinned} externally-read modules so offload cannot meta them.") + if has_pack_quantized_config(hf_config): _unpack_compressed_linear_weights(model, ckpt_path) From d4f0d50d5dc76face10d263b42c0147a92959f0d Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:09:19 +0000 Subject: [PATCH 4/6] docs(recipes): document the offloaded layerwise export recipe The recipe shipped undocumented; the parent branch's docs test now requires a row and the count kept in step. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt_recipes/ptq.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index e6956f3d847..9b14c3adae9 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -29,7 +29,7 @@ supported combinations. ### The shipped recipes
-All 25 general/ptq/ recipes (click to expand) +All 26 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -49,6 +49,7 @@ supported combinations. | `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-kv_fp8_layerwise_export_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise, tuned for accelerate-offloaded models | | `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 | From f9c44f744cb4d523e537ee004832fcd769db549c Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:26:57 +0000 Subject: [PATCH 5/6] feat(export): support multimodal models in per-layer export Calibration runs on the extracted language model, so the exporter has to be told which model the checkpoint describes. EXPORT_PARENT_ATTR carries that link and resolve_export_parent() turns it into a key prefix, found by identity so a module that merely looks like the language model cannot be mistaken for it. Three things then move into the parent's namespace: tensor keys gain the prefix, the vision tower and projector are collected in finalize() -- calibration never saw them and every other pass walks the submodel -- and the quant config's module references are rewritten, with those towers added to exclude_modules so a loader does not read plain BF16 as quantized. The config artifacts are written from the parent, and the VLM path no longer overwrites config.json afterwards, which would have stripped quantization_config off a finished checkpoint. The refusal narrows rather than disappears: a VLM whose language model is not reachable from the full model still cannot be exported, because the prefix would be undefined and the shards would silently describe the submodel alone. Reimplemented against the current exporter rather than cherry-picked -- the name mapper is now built from the export model, since Gemma3-VL stores the decoder at model.language_model.layers but publishes it as language_model.model.layers, and a submodel-scoped mapper cannot produce the published name. build_legacy_name_mapper covers transformers < 5, where save_pretrained rather than the exporter reverses _checkpoint_conversion_mapping. The test fails without the parent link: the vision tower is absent entirely. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- examples/hf_ptq/hf_ptq.py | 41 ++++- modelopt/torch/export/layerwise_export.py | 172 +++++++++++++++++- .../gpu/torch/export/test_layerwise_export.py | 70 ++++++- 3 files changed, 264 insertions(+), 19 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 66bc31f2cfa..fda06a2480f 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -77,6 +77,7 @@ has_spec_opt, save_expert_token_count_table, ) +from modelopt.torch.export.layerwise_export import EXPORT_PARENT_ATTR 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 @@ -121,6 +122,17 @@ def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool: mto.enable_huggingface_checkpointing() +def _link_export_parent(language_model, full_model) -> None: + """Point per-layer export at the model the checkpoint must describe. + + Straight into ``__dict__``: ``nn.Module.__setattr__`` would register the parent as a + submodule of its own child, and that cycle makes ``named_modules()`` recurse forever. + Plain attribute lookup still finds it. + """ + if language_model is not full_model: + language_model.__dict__[EXPORT_PARENT_ATTR] = full_model + + def extract_and_prepare_language_model_from_vl(full_model): """Extract language model from VL model and disable quantization for non-language components. @@ -630,6 +642,8 @@ def load_model(args: argparse.Namespace): if extracted_lm is not None: language_model = extracted_lm model_type = extracted_model_type + if args.layerwise_export: + _link_export_parent(language_model, full_model) else: if args.specdec_offline_dataset is not None: language_model = full_model @@ -656,6 +670,8 @@ def load_model(args: argparse.Namespace): if extracted_lm is not None: language_model = extracted_lm model_type = extracted_model_type + if args.layerwise_export: + _link_export_parent(language_model, full_model) tokenizer = get_tokenizer(args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code) @@ -778,12 +794,20 @@ def assert_layerwise_export_compatible(args, full_model, mtp_layer_prefixes) -> 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." - ) + # Calibration runs on the extracted language model, so the exporter is told which + # model the checkpoint describes; it then prefixes tensor keys, writes the full VLM + # config and picks up the untouched towers. That needs the submodel to be reachable + # from the full model -- otherwise the prefix is undefined and the shards would + # silently describe the submodel alone. + lineage = get_language_model_from_vl(full_model) + language_model = lineage[-1] if lineage else None + if language_model is None or all(m is not language_model for m in full_model.modules()): + raise NotImplementedError( + "layerwise.export_dir does not support this multimodal model: its language " + "model could not be located inside the full model, so exported tensor names " + "cannot be resolved against the VLM namespace. Export without " + "layerwise.export_dir." + ) if mtp_layer_prefixes: raise NotImplementedError( @@ -857,7 +881,10 @@ def export_quantized( # Check if the model is a multimodal/VLM model is_vlm = is_multimodal_model(full_model) - if is_vlm: + # Skipped under per-layer export: calibration already wrote a config carrying the + # quantization_config, and the source config is unquantized, so writing it over the + # top would strip that metadata off a finished checkpoint. + if is_vlm and not args.layerwise_export: # Save original model config and the processor config to the export path for VLMs. print(f"Saving original model config to {export_path}") diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 62ba5a64776..6fc95a63b18 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -18,6 +18,7 @@ import contextlib import hashlib import json +import re import warnings from collections.abc import Callable from pathlib import Path @@ -65,6 +66,81 @@ def _is_quantized_module(module: nn.Module) -> bool: ) +#: Set by the caller on the calibrated submodel to name the model the checkpoint should +#: describe. Multimodal pipelines calibrate the extracted language model, but the exported +#: checkpoint has to describe the whole VLM or its config and tensor names disagree. +EXPORT_PARENT_ATTR = "_layerwise_export_parent" + + +def resolve_export_parent(model: nn.Module) -> tuple[nn.Module, str]: + """Return ``(model the checkpoint describes, key prefix for its tensors)``. + + Without :data:`EXPORT_PARENT_ATTR` this is ``(model, "")``. With it, the prefix is the + dotted path from parent down to the calibrated submodel, found by identity so a module + that merely looks like the language model cannot be mistaken for it. + """ + parent = getattr(model, EXPORT_PARENT_ATTR, None) + if parent is None or parent is model: + return model, "" + for name, module in parent.named_modules(): + if module is model: + return parent, f"{name}." if name else "" + raise ValueError( + f"{EXPORT_PARENT_ATTR} was set to a {type(parent).__name__} that does not contain " + "the calibrated model, so exported tensor names cannot be resolved against it." + ) + + +def build_legacy_name_mapper(model: nn.Module): + """Hub-name mapper for transformers < 5, or None if the model declares no mapping. + + ``save_pretrained`` is what reverses ``_checkpoint_conversion_mapping`` on 4.x. Per-layer + export writes shards with ``save_file`` and never passes through it, so without this it + would publish in-memory names where the whole-model path publishes hub names. + """ + mapping = getattr(model, "_checkpoint_conversion_mapping", None) + if not mapping: + return None + # Stored hub-pattern -> in-memory-prefix, so invert. Longest in-memory prefix first, or + # a shorter one shadows a longer one (``lm_head`` inside ``model.language_model...``). + rules = sorted( + ((re.compile("^" + re.escape(mem)), hub.lstrip("^")) for hub, mem in mapping.items()), + key=lambda r: -len(r[0].pattern), + ) + + def _map(name: str) -> str: + for pattern, replacement in rules: + new, n = pattern.subn(replacement, name, count=1) + if n: + return new + return name + + return _map + + +def _prefix_quant_config_names(quant_config: dict, prefix: str, siblings: list[str]) -> None: + """Move a quant config into the parent's namespace, in place. + + Module references gain ``prefix`` so they name the tensors the shards were written + under. Every subtree outside the calibrated submodel is then added to + ``exclude_modules``: calibration never saw them, so a loader that does not find them + excluded would read plain BF16 weights as quantized. + """ + if not prefix: + return + quantization = quant_config.get("quantization") + if not isinstance(quantization, dict): + return + excluded = quantization.get("exclude_modules") + if isinstance(excluded, list): + quantization["exclude_modules"] = sorted( + [prefix + m for m in excluded] + [f"{s}*" for s in siblings] + ) + quantized = quantization.get("quantized_layers") + if isinstance(quantized, dict): + quantization["quantized_layers"] = {prefix + k: v for k, v in quantized.items()} + + def _module_formats(model: nn.Module) -> set: """Every distinct format present. ``get_quantization_format`` stops at the first.""" return { @@ -259,18 +335,38 @@ def __init__( self._kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] self._finalized = False + # These differ for a multimodal pipeline: calibration runs on the extracted language + # model, but the checkpoint must describe the whole VLM. Everything that walks + # modules stays on `model`; only tensor keys and the config artifacts move up. + self._export_model, self._key_prefix = resolve_export_parent(model) + self._sibling_prefixes = ( + [ + name + for name, _ in self._export_model.named_children() + if name != self._key_prefix.split(".", 1)[0] + ] + if self._key_prefix + else [] + ) + self._name_mapper = None try: - self._name_mapper = build_reverse_name_mapper(model) + # From the export model: the mapper reverses transformers' internal layout into + # hub names across the whole checkpoint, and the two namespaces differ. Gemma3-VL + # stores the decoder at model.language_model.layers but publishes it as + # language_model.model.layers, which a submodel-scoped mapper cannot produce. + self._name_mapper = build_reverse_name_mapper(self._export_model) except Exception as exc: - warnings.warn( - f"Reverse name mapper unavailable ({exc}); exported tensor names may not " - "match the original HF hub checkpoint." - ) + self._name_mapper = build_legacy_name_mapper(self._export_model) + if self._name_mapper is None: + warnings.warn( + f"Reverse name mapper unavailable ({exc}); exported tensor names may not " + "match the original HF hub checkpoint." + ) # By name, not data_ptr: layers and tail are separate passes. raw_tied_keys: set[str] = ( - set(getattr(model, "_tied_weights_keys", None) or []) - if getattr(model.config, "tie_word_embeddings", False) + set(getattr(self._export_model, "_tied_weights_keys", None) or []) + if getattr(self._export_model.config, "tie_word_embeddings", False) else set() ) self._tied_alias_keys: set[str] = ( @@ -456,10 +552,17 @@ def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> d mapped = self._name_mapper(name) if self._name_mapper is not None else name tail.setdefault(mapped, tensor.detach().contiguous().cpu()) + # Subtrees of the parent outside the calibrated submodel -- a VLM's vision tower and + # projector. Calibration never saw them and every pass above walks `model`, so + # without this they are simply absent. Unquantized, so no export handler is needed. + self._collect_sibling_subtrees(tail) + save_file(_copy_storage_aliases(tail), str(self._export_dir / _TAIL_SHARD)) self._write_index() - save_non_weight_artifacts(model, self._export_dir) - _write_hf_export_config(model, quant_config, self._export_dir) + _prefix_quant_config_names(quant_config, self._key_prefix, self._sibling_prefixes) + # The parent's config: the checkpoint has to describe the whole model, towers included. + save_non_weight_artifacts(self._export_model, self._export_dir) + _write_hf_export_config(self._export_model, quant_config, self._export_dir) return quant_config def _bind_identity(self, quant_config: dict) -> None: @@ -549,6 +652,51 @@ def assert_shards_present(self, upto: int) -> None: "export directories are from different runs; delete one and restart." ) + def _collect_sibling_subtrees(self, tail: dict[str, torch.Tensor]) -> None: + """Add the parent's tensors that live outside the calibrated submodel. + + No-op unless :func:`resolve_export_parent` found a parent. These keys are already in + the parent's namespace, so the prefix must not be applied again -- only the hub-name + reversal, which spans the whole checkpoint. + """ + if not self._key_prefix: + return + from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + requires_weight_materialization, + ) + + parent = self._export_model + parent_modules = dict(parent.named_modules()) + inner_ids = {id(m) for m in self._ctx.model.modules()} + + def _store(key: str, tensor: torch.Tensor) -> None: + mapped = self._name_mapper(key) if self._name_mapper is not None else key + if tensor is not None and not tensor.is_meta and mapped not in self._tied_alias_keys: + tail.setdefault(mapped, tensor.detach().contiguous().cpu()) + + for name, module in parent.named_modules(): + if not name or id(module) in inner_ids or name.startswith(self._key_prefix): + continue + if not requires_weight_materialization(module, parent, parent_modules): + continue + with enable_weight_access_and_writeback( + module, parent, parent_modules, writeback=False + ): + for key, tensor in module.state_dict().items(): + _store(f"{name}.{key}", tensor) + + for key, tensor in parent.state_dict().items(): + if key.startswith(self._key_prefix): + continue + if tensor is not None and tensor.is_meta: + raise RuntimeError( + f"{key!r} is on meta and its module was offered no materialization " + "window, so it cannot be exported. Export without export_dir and use " + "export_hf_checkpoint() for this model." + ) + _store(key, tensor) + def _collect(self, out: dict[str, torch.Tensor], full_key: str, tensor: torch.Tensor) -> None: """Apply per-tensor export postprocessing and hub-name reversal, or drop the tensor.""" from .quant_utils import _postprocess_single_tensor @@ -556,7 +704,11 @@ def _collect(self, out: dict[str, torch.Tensor], full_key: str, tensor: torch.Te 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 + self._key_prefix + full_key, + tensor, + 448, + self._kv_cache_format, + self._ctx.is_modelopt_qlora, ) if new_key is None or new_value is None: return diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index 34f70f5502f..c2da74cff26 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 ( + EXPORT_PARENT_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 @@ -415,6 +424,63 @@ def test_moe_export_matches(tmp_path): _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) +def _build_vlm(): + """A tiny VLM -- the shape per-layer export used to refuse outright.""" + torch.manual_seed(0) + model = get_tiny_gemma3vl().cuda().eval() + # is_multimodal_model reads this, and the tiny fixtures leave it unset. + model.config.architectures = ["Gemma3ForConditionalGeneration"] + return model + + +def _vlm_language_model(vlm): + lineage = get_language_model_from_vl(vlm) + assert lineage, "fixture is expected to expose a language model lineage" + return lineage[-1] + + +def _calib_vlm(language_model): + # use_cache=False: this calls the text model directly rather than the CausalLM wrapper, + # so a KV cache would carry across batches and the second batch would build a mask for + # more positions than there are keys. + for batch in CALIB_BATCHES: + language_model(batch.cuda(), use_cache=False) + + +def _vlm_cfg(): + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + return cfg + + +def test_vlm_export_matches_whole_model_export(tmp_path): + """Calibrating the submodel must still export a checkpoint describing the whole VLM. + + Without the parent link the shards and config would describe the language model alone: + tensor keys missing their prefix, the vision tower absent, the config the text model's. + """ + baseline_vlm = _build_vlm() + mtq.quantize(_vlm_language_model(baseline_vlm), _vlm_cfg(), _calib_vlm) + baseline_dir = tmp_path / "baseline" + export_hf_checkpoint(baseline_vlm, export_dir=baseline_dir) + + export_dir = tmp_path / "fused" + vlm = _build_vlm() + language_model = _vlm_language_model(vlm) + language_model.__dict__[EXPORT_PARENT_ATTR] = vlm + cfg = _vlm_cfg() + cfg["algorithm"]["layerwise"] |= { + "export_dir": str(export_dir), + "checkpoint_dir": str(tmp_path / "ckpt"), + "calib_mutates_weights": False, + } + mtq.quantize(language_model, cfg, _calib_vlm) + + exported = _load_checkpoint(export_dir) + assert any("vision" in k for k in exported), "vision tower missing from the checkpoint" + _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) + + def test_export_does_not_mutate_the_model(tmp_path): """Exporting a layer must leave the model exactly as calibration left it. From 9066997fe063e97411415dc610425c3bbd29e3cc Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:32:59 +0000 Subject: [PATCH 6/6] fix(recipes): stop the offload recipe pinning resume state to /tmp The recipe shipped checkpoint_dir: /tmp/modelopt_layerwise_ckpt, which is container-local -- a run that outlasts its GPU session came back to a wiped manifest and restarted at layer 0, the exact failure this recipe exists to avoid. Leaving it unset lets hf_ptq derive .layerwise_resume, which lives next to the shards it describes. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- ...fp4_experts_only-kv_fp8_layerwise_export_offload.yaml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml index 2be18d6522c..b041123cfa7 100644 --- a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml @@ -34,9 +34,11 @@ metadata: projections; shared_experts is missed because its path contains '_experts.' rather than '.experts.' (fnmatch semantics, conversion.py). - Pairs checkpoint_dir with export_dir so an interrupted run resumes without recalibrating - or re-exporting finished layers. That is the point of the combination for a run that - outlasts its GPU session: a killed run loses at most the in-flight layer. + An interrupted run resumes without recalibrating or re-exporting finished layers, losing + at most the in-flight one -- the point of the combination for a run that outlasts its GPU + session. The resume state lives beside the checkpoint at .layerwise_resume + unless you set layerwise.checkpoint_dir yourself; do not point it at container-local + storage, or a run that survives its session comes back to a wiped manifest. 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. @@ -47,7 +49,6 @@ quantize: enable: true # max only updates _amax, so the exported shard stays valid for its layer. calib_mutates_weights: false - checkpoint_dir: /tmp/modelopt_layerwise_ckpt # Presence enables per-layer export; the value is replaced with --export_path. export_dir: /tmp/modelopt_layerwise_export quant_cfg: