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

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)
 

modelopt/torch/quantization/plugins/huggingface.py

@@ -1796,17 +1796,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
 

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

@@ -29,7 +29,7 @@ supported combinations.
 ### The shipped recipes
 
 <details>
-<summary>All 25 <code>general/ptq/</code> recipes (click to expand)</summary>
+<summary>All 26 <code>general/ptq/</code> recipes (click to expand)</summary>
 
 | 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 |