Kimi-K3 on layerwise fused export: 3 defect fixes, multimodal support, single-B200 NVFP4

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)
 

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:

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,

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
 

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 <export_path>.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

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