diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 862bc0f7fe5..0db85c574be 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -41,6 +41,7 @@ AutoModelForCausalLM, AutoProcessor, AutoTokenizer, + PretrainedConfig, PreTrainedTokenizerBase, ProcessorMixin, ) @@ -703,6 +704,18 @@ 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: + """Re-apply the caller's ``attn_implementation`` after ``__init__``. + + Kimi-K3's remote code overwrites it unconditionally, so an uninstalled backend fails at + the export forward rather than at load. Drop once the checkpoint is fixed upstream. + """ + sub_configs = (v for v in vars(model.config).values() if isinstance(v, PretrainedConfig)) + for cfg in (model.config, *sub_configs): + if getattr(cfg, "_attn_implementation", None) not in (None, attn_implementation): + cfg._attn_implementation = attn_implementation + + def _get_config_dtype(config): config_dtype = ( getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.bfloat16 @@ -967,6 +980,10 @@ def has_pack_quantized_config(config): **model_kwargs2, ) model.eval() + + 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/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 0e3c8035bf8..b60f5b420b5 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -78,6 +78,7 @@ has_spec_opt, save_expert_token_count_table, ) +from modelopt.torch.export.layerwise_export import export_parent 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 @@ -753,12 +754,18 @@ def mono_quantize( else None, ) - if calibration_only: - language_model = mtq.calibrate( - language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop - ) - else: - language_model = mtq.quantize(language_model, quant_cfg, forward_loop=calibrate_loop) + # A VLM calibrates its language model but must export the whole thing. Both branches + # reach the same layerwise machinery, so the link is set over both. + parent = full_model if args.layerwise_export else language_model + with export_parent(parent): + if calibration_only: + language_model = mtq.calibrate( + language_model, quant_cfg["algorithm"], forward_loop=calibrate_loop + ) + else: + language_model = mtq.quantize( + language_model, quant_cfg, forward_loop=calibrate_loop + ) # For VL models, update full_model to use the quantized language model if is_nemotron_vl_model: @@ -779,12 +786,14 @@ 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." - ) + 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 the exported tensor " + "names cannot be resolved against the VLM." + ) if mtp_layer_prefixes: raise NotImplementedError( @@ -859,15 +868,17 @@ def export_quantized( 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 - ) + # Not under per-layer export: it already wrote a config with the + # quantization_config, and this source config is unquantized. + if not args.layerwise_export: + 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: diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index a196c9a9e52..a1ddf58f3f5 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -16,7 +16,9 @@ """Write each decoder layer's quantized checkpoint shard as soon as it is calibrated.""" import contextlib +import contextvars import json +import re import warnings from pathlib import Path @@ -82,8 +84,8 @@ def _module_formats(model: nn.Module) -> set: } -def _tied_quantized_modules(model: nn.Module) -> list[str]: - """Quantized modules sharing a weight with another. +def _tied_weight_modules(model: nn.Module) -> list[str]: + """Modules sharing a weight with another, quantized or not. Grouped by name, which survives offload: a ``data_ptr`` grouping sees nothing when the weights are on meta and would pass vacuously. Falls back to ``data_ptr`` when the model @@ -94,7 +96,7 @@ def _tied_quantized_modules(model: nn.Module) -> list[str]: by_ptr: dict[int, list[str]] = {} for name, module in model.named_modules(): weight = getattr(module, "weight", None) - if weight is None or not _is_quantized_module(module): + if weight is None: continue key = tied_map.group_key(f"{name}.weight") if key is not None: @@ -127,16 +129,15 @@ def assert_layerwise_export_supported(model: nn.Module) -> None: """Raise unless per-layer export is valid for this model.""" assert_formats_supported(model, "before calibration") - tied = _tied_quantized_modules(model) + tied = _tied_weight_modules(model) if tied: raise NotImplementedError( - f"layerwise export does not support weight-tied quantized modules {tied[:6]}: " - "the whole-model path merges their input_quantizer amaxes via " - "sync_tied_input_amax so both sides share one input_scale, which a per-layer " - "pass cannot do because a tie partner may be uncalibrated or already written. " - "Conversion quantizes every nn.Linear and nn.Embedding, so disabling their " - "quantizers does not lift this -- tie_word_embeddings models need " - "export_hf_checkpoint()." + f"layerwise export does not support weight-tied modules {tied[:6]}: quantized, " + "the whole-model path merges their input_quantizer amaxes via sync_tied_input_amax " + "so both sides share one input_scale, which a per-layer pass cannot do because a " + "tie partner may be uncalibrated or already written; unquantized, save_pretrained " + "drops the duplicate key and writing shards directly does not. " + "tie_word_embeddings models need export_hf_checkpoint()." ) if dist.is_initialized() and dist.size() > 1: @@ -146,6 +147,75 @@ def assert_layerwise_export_supported(model: nn.Module) -> None: ) +_export_parent: contextvars.ContextVar[nn.Module | None] = contextvars.ContextVar( + "layerwise_export_parent", default=None +) + + +@contextlib.contextmanager +def export_parent(parent: nn.Module): + """Export the checkpoint for ``parent`` while calibration runs on one of its submodules. + + The decoder layers are the same objects either way, so walking the parent yields + parent-namespace names, the full config and the untouched towers with no prefixing. + """ + token = _export_parent.set(parent) + try: + yield + finally: + _export_parent.reset(token) + + +def _resolve_export_parent(model: nn.Module) -> nn.Module: + """Return the model the checkpoint should describe. Membership is by identity, not name.""" + parent = _export_parent.get() + if parent is None or parent is model: + return model + if all(m is not model for m in parent.modules()): + raise ValueError( + f"export_parent() was given a {type(parent).__name__} that does not contain the " + "calibrated model." + ) + return parent + + +def build_legacy_name_mapper(model: nn.Module): + r"""Hub-name mapper for transformers < 5, or ``None``. + + ``save_pretrained`` is what reverses ``_checkpoint_conversion_mapping``, and per-layer + export writes shards directly without it. ``build_reverse_name_mapper`` is no help + either: it reads transformers 5's ``conversion_mapping`` and raises on 4.x. + + Rules are inverted (the mapping is stored hub -> in-memory) and applied longest-prefix + first, or a short one shadows a longer (``lm_head`` inside ``model.language_model...``). + The hub side is a regex: its groups are stripped exactly as save_pretrained strips them + when reversing the same mapping, and the remainder has its backslashes neutralised so a + pattern like ``layers\.(\d+)`` substitutes literally instead of raising on ``\d``. + """ + mapping = getattr(model, "_checkpoint_conversion_mapping", None) + if not mapping: + return None + rules = sorted( + ( + ( + re.compile("^" + re.escape(mem)), + re.sub(r"\(.*\)", "", hub.lstrip("^")).replace("\\", "\\\\"), + ) + 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 + + class LayerwiseExporter: """Writes one decoder layer's quantized shard per call, then the tail and index. @@ -163,6 +233,7 @@ def __init__( Runs before calibration, so nothing amax-dependent exists yet. """ + model = _resolve_export_parent(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) @@ -208,10 +279,12 @@ def __init__( try: self._name_mapper = build_reverse_name_mapper(model) except Exception as exc: - warnings.warn( - f"Reverse name mapper unavailable ({exc}); exported tensor names may not " - "match the original HF hub checkpoint." - ) + self._name_mapper = build_legacy_name_mapper(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." + ) def export_layer( self, diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 4acb4d30dfa..f0a2f0733ff 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1800,13 +1800,20 @@ def get_homogeneous_hf_decoder_layers(model: nn.Module) -> nn.ModuleList | None: if not _is_supported_hf_model(model): return None - 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 + # Take the shallowest ``layers`` but keep descending while wrappers remain: the nesting + # order varies, e.g. Kimi-K3 keeps its layers at ``language_model.model.layers``. + decoder, seen = model, set() + while id(decoder) not in seen: + seen.add(id(decoder)) + if isinstance(getattr(decoder, "layers", None), nn.ModuleList): + 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/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml b/modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml new file mode 100644 index 00000000000..ea1e26d897e --- /dev/null +++ b/modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export.yaml @@ -0,0 +1,66 @@ +# 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. + + Kimi-K3's counterpart to general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export, from + which it differs only in the narrower expert scope below -- the reason it does not live + in the general tier. It was validated on a model too large to hold resident, but nothing + here configures offload: that comes from --offload_folder and the 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 1735289ac3f..e201f6a0d6f 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -360,6 +360,15 @@ checkpoint's** quant config verbatim: cache remain BF16. Because the 2.8T source uses packed MXFP4 expert tensors, use the calibration-free streaming converter in `examples/kimi/` rather than the in-memory `hf_ptq.py` flow. +- **`models/moonshotai/Kimi-K3/ptq/nvfp4_experts-kv_fp8_layerwise_export`** is the + in-memory counterpart to the above: NVFP4 W4A4 on the routed experts with an FP8 + KV cache, calibrated and **exported one decoder layer at a time**, so a run that + outlasts its GPU session resumes without recalibrating or re-exporting finished + layers. It scopes experts as `*.experts.*` rather than the general recipes' + `*block_sparse_moe*`, which on K3 also matches `shared_experts.*` and + `routed_expert_*` -- 552 modules the vendor left unquantized. Pair it with + `--offload_folder` and per-device memory budgets; nothing in the recipe itself + configures offload. - **`models/mistralai/Mistral-Medium-3.5-128B/ptq/nvfp4-max-calib`** mirrors `nvidia/Mistral-Medium-3.5-128B-NVFP4`: decoder MLP layers 4–86 use NVFP4 W4A4, edge MLP layers 0–3 and 87 use FP8 W8A8, and all attention projections diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index ca40f72e875..0e6e4cc4ca5 100644 --- a/tests/gpu/torch/export/test_layerwise_export.py +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -22,11 +22,21 @@ import pytest import torch -from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe +import torch.nn as nn +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 ( + LayerwiseExporter, + export_parent, + 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 @@ -409,7 +419,7 @@ def test_tied_embeddings_are_refused(tmp_path, make_cfg): model.config.architectures = ["LlamaForCausalLM"] cfg = _layerwise_cfg(tmp_path / "fused", tmp_path / "ckpt", base=make_cfg()) - with pytest.raises(NotImplementedError, match="weight-tied quantized modules"): + with pytest.raises(NotImplementedError, match="weight-tied modules"): mtq.quantize(model, cfg, _calib) @@ -477,3 +487,73 @@ def test_awq_is_refused(tmp_path, make_cfg, method, match): with pytest.raises(NotImplementedError, match=match): mtq.quantize(_build_model(), cfg, _calib) + + +def _disable_quant_on_towers(vlm): + """Mirror ``extract_and_prepare_language_model_from_vl``: towers get disabled quantizers. + + hf_ptq runs this before calibration, so by the time the exporter walks the parent the + vision tower and projector carry ``TensorQuantizer`` children. That changes what + ``_module_formats`` and the tail dispatch see, so the test has to reproduce it. + """ + lineage = get_language_model_from_vl(vlm) + language_model, ancestors = lineage[-1], lineage[:-1] + disabled = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} + memo = set(ancestors) | {language_model} + for ancestor in ancestors: + for _, module in ancestor.named_children(): + if module not in memo: + mtq.quantize(module, copy.deepcopy(disabled), forward_loop=None) + memo.add(module) + return language_model + + +def _build_vlm(): + torch.manual_seed(0) + model = get_tiny_gemma3vl(tie_word_embeddings=False).cuda().eval() + # The kwarg only reaches text_config; the outer config still ties lm_head to the + # embedding, and tied weights are refused. Untie for real so this test stays on the + # namespace/towers/config behaviour it exists for. + model.config.tie_word_embeddings = False + model._tied_weights_keys = {} + model.all_tied_weights_keys = {} + model.lm_head.weight = nn.Parameter(model.lm_head.weight.detach().clone()) + # is_multimodal_model reads this, and the tiny fixtures leave it unset. + model.config.architectures = ["Gemma3ForConditionalGeneration"] + return model + + +def _calib_vlm(language_model): + # use_cache=False: this calls the text model directly, so a retained cache would make + # batch 2 build a mask for 31 positions against 16 keys. + for batch in CALIB_BATCHES: + language_model(batch.cuda(), use_cache=False) + + +def test_vlm_export_matches_whole_model_export(tmp_path): + """A VLM calibrates its language model but must export the whole model.""" + baseline_vlm = _build_vlm() + baseline_lm = _disable_quant_on_towers(baseline_vlm) + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True}} + mtq.quantize(baseline_lm, cfg, _calib_vlm) + baseline_dir = tmp_path / "baseline" + export_hf_checkpoint(baseline_vlm, export_dir=baseline_dir) + + vlm = _build_vlm() + language_model = _disable_quant_on_towers(vlm) + export_dir = tmp_path / "fused" + with export_parent(vlm): + mtq.quantize(language_model, _layerwise_cfg(export_dir, tmp_path / "ckpt"), _calib_vlm) + + exported = _load_checkpoint(export_dir) + _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) + _assert_same_quant_config(baseline_dir, export_dir) + + assert any(k.startswith("language_model.") for k in exported), "lost the VLM namespace" + assert any(k.startswith(("vision_tower", "multi_modal_projector")) for k in exported), ( + "the towers outside the language model were not exported" + ) + assert "vision_config" in json.loads((export_dir / "config.json").read_text()), ( + "config.json describes the submodel, not the VLM" + )