diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 89a88f76458..74ec9c0df45 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -703,6 +703,99 @@ 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}") + 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:") + for line in changed: + 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 @@ -967,6 +1060,18 @@ 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) + + # 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) 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/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 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..b041123cfa7 --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml @@ -0,0 +1,63 @@ +# 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). + + 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. +quantize: + algorithm: + method: max + layerwise: + enable: true + # max only updates _amax, so the exported shard stays valid for its layer. + calib_mutates_weights: false + # Presence enables per-layer export; the value is replaced with --export_path. + export_dir: /tmp/modelopt_layerwise_export + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers 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 | 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.