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)
 

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}")