[OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter

modelopt/torch/models/__init__.py

@@ -92,6 +92,7 @@
     minimax_m3_vl,
     llama4_text,
     kimi_vl,
+    kimi_k25,
     nemotron_h_omni,
     qwen2_moe,
     qwen3,

modelopt/torch/models/_kimi.py

@@ -23,7 +23,7 @@
 
 # modelopt.torch.utils and transformers are imported inside the functions: the model registry imports
 # this module, and it must import without transformers, the optional [hf] extra.
-from ._remote_code import keep_patched
+from ._remote_code import _apply, keep_patched, stage_full_remote_closure
 
 
 def checkpoint_to_patch(
@@ -114,6 +114,55 @@ def _prepare_4d_causal_attention_mask(*args, **kwargs):  # pragma: no cover
         print_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}")
 
 
+def make_importable() -> None:
+    """The names newer transformers dropped, and the code's two-deep relative imports staged."""
+    install_compat_import_shims()
+    # transformers stages only 1-hop relative imports; Kimi's remote code is 2 hops deep.
+    stage_full_remote_closure()
+
+
+def make_buildable(
+    modeling: str, vlm_class: str, hf_config, ckpt_path: str | None, trust_remote_code: bool
+) -> None:
+    """What both checkpoints' own code needs to import and build on this transformers, and no more.
+
+    The names newer transformers dropped, the code's two-deep relative imports staged, and the VLM
+    class's ``tie_weights`` taking the keyword arguments transformers' ``post_init`` passes it. None
+    of it changes what the model computes or loads, so the plan decision runs it before ``prepare``,
+    whatever plan comes out. A config of transformers' own port, or code this load may not run, has
+    nothing of the checkpoint's to build. Idempotent.
+    """
+    from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+    if type(hf_config).__module__.startswith("transformers.") or ckpt_path is None:
+        return
+    if not trust_remote_code:
+        return
+    make_importable()
+    get_class_from_dynamic_module(
+        f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code
+    )
+    # Remote code only: transformers' own port of a model can share the module's name.
+    _apply(f".{modeling}", lambda module: _shim_tie_weights(module, vlm_class))
+
+
+def _shim_tie_weights(module, vlm_class: str) -> bool:
+    """Have ``module``'s VLM class take the keyword arguments transformers passes ``tie_weights``.
+
+    The remote ``tie_weights()`` takes none, but transformers' ``post_init`` and load call
+    ``tie_weights(recompute_mapping=False)``: a TypeError at construction. Returns whether it changed
+    the class.
+    """
+    kimi_cls = getattr(module, vlm_class, None)
+    if kimi_cls is None:
+        return False
+    own_tie = vars(kimi_cls).get("tie_weights")
+    if own_tie is None or getattr(own_tie, "_kimi_tie_weights_shim", False):
+        return False
+    kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie)
+    return True
+
+
 def _drop_tie_weights_kwargs(orig_tie):
     """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it.
 
@@ -134,10 +183,10 @@ def _native_for_the_prepared_config(own_class, native_class):
     """The ``DeepseekV3ForCausalLM`` the patched module builds its language model with.
 
     transformers' native class for the native ``DeepseekV3Config`` ``prepare`` puts on
-    ``text_config``, and the checkpoint's own class for its own config. Deciding the plan prepares a
-    copy of the config before the plan is known, and the patch outlives that copy: a load the plan
-    keeps off expert parallelism builds from the config as read, and so gets the checkpoint's own
-    language model, as it would with no patch.
+    ``text_config``, and the checkpoint's own class for its own config. The patch outlives the load
+    it was made for: a later load in the process that its plan keeps off expert parallelism builds
+    from the config as read, and so gets the checkpoint's own language model, as it would with no
+    patch.
     """
     from transformers import DeepseekV3Config
 
@@ -183,13 +232,7 @@ def to_native(module):
         module.DeepseekV3ForCausalLM = _native_for_the_prepared_config(
             current, NativeDeepseekV3ForCausalLM
         )
-        # The remote tie_weights() takes no kwargs, but transformers' post_init and load call
-        # tie_weights(recompute_mapping=False) -> TypeError.
-        kimi_cls = getattr(module, vlm_class, None)
-        if kimi_cls is not None:
-            own_tie = vars(kimi_cls).get("tie_weights")
-            if own_tie is not None and not getattr(own_tie, "_kimi_tie_weights_shim", False):
-                kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie)
+        _shim_tie_weights(module, vlm_class)
         return True
 
     # The dynamic-module loader names the module like

modelopt/torch/models/kimi_k25/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+"""kimi_k25 specs (HF model type ``kimi_k25``)."""
+
+from . import specs

modelopt/torch/models/kimi_k25/ep.py

@@ -0,0 +1,128 @@
+# 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.
+
+"""DP=EP a2a support for Kimi-K2.5 (``kimi_k25``).
+
+a remote-code VLM whose text backbone is a DeepSeek-V3 variant with compressed-tensors-compressed experts.
+Kimi-K2.6 and Kimi-K2.7-Code ship the same model type and modeling classes.
+
+The challenge: Kimi's bundled ``modeling_deepseek.py`` defines ``DeepseekV3MoE`` as an
+``nn.ModuleList`` of per-expert ``DeepseekV3MLP`` (separate ``gate_proj`` / ``up_proj`` /
+``down_proj`` Linears per expert). That layout is incompatible with transformers' EP machinery,
+which expects the **fused** 3-D ``gate_up_proj`` / ``down_proj`` Parameters that the
+transformers-native ``deepseek_v3`` modeling ships. We need the native modeling to get EP-capable
+experts, but we also want the rest of the VLM (vision tower + mm projector) to come from Kimi's
+custom code.
+
+Solution -- monkey-patch ``modeling_kimi_k25``'s ``DeepseekV3ForCausalLM`` import to point at the
+transformers-native class, AND swap ``hf_config.text_config`` to a transformers-native
+``DeepseekV3Config`` (the native class won't accept Kimi's custom config subclass). Then
+``KimiK25ForConditionalGeneration.__init__`` (which does ``self.language_model =
+DeepseekV3ForCausalLM(config.text_config)``) builds the native EP-capable model for the language
+sub-model, while vision_tower + mm_projector are still built from Kimi's custom code.
+
+On-the-fly compressed-tensors expert decompress happens during from_pretrained's
+``WeightConverter`` machinery (installed by ``_onfly_decompress.install_onfly_decompress``); the
+fused per-expert source ``.weight`` tensors are decompressed + merged into the 3-D
+``gate_up_proj`` / ``down_proj`` parameters. Format-agnostic (compressed_tensors decides the
+dequant from the checkpoint's quantization_config).
+
+EP plan: attached to the (now native) ``text_config``. transformers' ``post_init`` aggregates
+child-module plans up to the root with name prefix, so the final root ``_ep_plan`` keys end up as
+``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj,...}`` -- matching the actual
+parameter paths in the loaded VLM.
+
+Version floor: transformers >= 5.8, which the spec declares. The on-the-fly decompressor, installed
+for a compressed-tensors checkpoint, needs it: earlier releases' ``Concatenate`` trips over the
+sibling keys it has already consumed. The decompressor couples tightly to
+transformers' still-evolving ``WeightConverter`` conversion-ops API.
+"""
+
+from .. import _kimi, _onfly_decompress
+
+# The checkpoint formats prepare has the experts decompress as they load; the spec states them for
+# the plan decision, which runs before prepare.
+DECOMPRESSED_FORMATS = ("compressed-tensors",)
+
+
+def make_buildable(
+    hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_
+) -> None:
+    """Have the checkpoint's own code import and build on this transformers.
+
+    The plan decision needs it before ``prepare``; it changes nothing the model computes or loads.
+    """
+    _kimi.make_buildable(
+        "modeling_kimi_k25",
+        "KimiK25ForConditionalGeneration",
+        hf_config,
+        ckpt_path,
+        trust_remote_code,
+    )
+
+
+def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_):
+    """Load Kimi-K2.5's language model as the transformers-native deepseek_v3.
+
+    Swap text_config to native deepseek_v3 + monkey-patch Kimi's modeling to use the native
+    DeepseekV3ForCausalLM + apply the EP plan + install the on-the-fly decompressor for
+    compressed-tensors experts.
+    """
+    # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe
+    # imports the registry.
+    from modelopt.torch.utils.moe import apply_distributed_support
+
+    ckpt_path = _kimi.checkpoint_to_patch("kimi_k25", hf_config, ckpt_path, trust_remote_code)
+
+    # 1. Install on-the-fly compressed-tensors expert decompressor. Operates on transformers-native
+    #    deepseek_v3's per-expert WeightConverter (which fuses per-expert .weight tensors into the
+    #    3-D gate_up_proj/down_proj Parameters). Format-agnostic; driven by checkpoint's
+    #    quantization_config, which Kimi-K2.5 carries only on the text_config step 2 replaces. It
+    #    goes first, so a checkpoint it refuses leaves everything as it was. Experts stored
+    #    otherwise are left to the expert-parallel load's format handling, which reads their scheme
+    #    from the copy of quantization_config Kimi's config class keeps on the top-level config;
+    #    step 2 leaves that copy alone. A checkpoint without that copy is refused.
+    quantization_config = _onfly_decompress._find_quantization_config(hf_config) or {}
+    if quantization_config.get("quant_method") in DECOMPRESSED_FORMATS:
+        _onfly_decompress.install_onfly_decompress(hf_config)
+    else:
+        _kimi.refuse_quantization_only_on_text_config("kimi_k25", hf_config)
+
+    # Before any remote-code load: this plugin's get_class_from_dynamic_module below, and
+    # from_pretrained later in load_and_prepare_ep.
+    _kimi.make_importable()  # the swap below loads the code and adapts its tie_weights
+
+    # 2. Replace text_config with a transformers-native DeepseekV3Config built from the same fields.
+    #    Kimi's __init__ passes config.text_config to DeepseekV3ForCausalLM; the native class rejects
+    #    Kimi's custom config subclass. to_native_backbone_config also drops auto_map +
+    #    quantization_config so transformers builds the native modeling instead of re-routing back
+    #    to Kimi's remote code via auto_map.
+    hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config)
+
+    # 3. Apply the EP plan to the now-native text_config. transformers' post_init walks children
+    #    and prefixes their _ep_plan onto the parent's: text_config's plan ends up at the VLM root
+    #    as `language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}`.
+    hf_config.text_config = apply_distributed_support("deepseek_v3", hf_config.text_config)
+
+    # 4. Force-load Kimi's modeling_kimi_k25 module (so we can patch its DeepseekV3ForCausalLM
+    #    symbol) and swap to the native class.
+    _kimi.patch_language_model_to_native(
+        "kimi_k25",
+        "modeling_kimi_k25",
+        "KimiK25ForConditionalGeneration",
+        ckpt_path,
+        trust_remote_code,
+    )
+    return hf_config

modelopt/torch/models/kimi_k25/specs.py

@@ -0,0 +1,55 @@
+# 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.
+
+"""kimi_k25 specs (HF model type ``kimi_k25``).
+
+Kimi-K2.5 (Kimi-K2.6 and Kimi-K2.7-Code share its model type): the compressed-expert sibling of
+``kimi_vl`` -- same swap to a native deepseek_v3 language model, plus on-the-fly decompression of
+its packed experts at load.
+
+Only the expert-parallel section is registered under ``kimi_k25``: the MoE and export sections stay
+``deepseek_v3``'s, so lookups by ``kimi_k25`` find neither yet. Nor do ``is_moe``'s name and structure
+fallbacks recognize the native ``DeepseekV3MoE`` blocks ``prepare`` installs, so the HF export's
+MoE-block handling skips them. Sharing the sections is left to a follow-up.
+"""
+
+from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN
+from ..specs import EPSpec, ModelSpec, register
+from .ep import DECOMPRESSED_FORMATS, make_buildable, prepare
+
+__all__: list[str] = []
+
+register(
+    ModelSpec(
+        model_type="kimi_k25",
+        # transformers ships Kimi-K2.5 itself from 5.13, so modeling_source stays "transformers": the
+        # registry pairs it with that release, and a "remote_code" claim would be stale from 5.13 on.
+        # Expert parallelism still runs on the checkpoint's own code (``trust_remote_code=True``),
+        # which ``prepare`` patches, refusing the port. So the version describes transformers' port,
+        # not the model this spec's load builds: a MoE or export section added here must name the
+        # native deepseek_v3 language model's modules.
+        min_transformers_version="5.13",
+        # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan,
+        # and decompresses the experts as they load, which needs transformers 5.8.
+        ep_spec=EPSpec(
+            prepare=prepare,
+            prepared_ep_plan=BASE_MODEL_EP_PLAN,
+            plan_config_attr="text_config",
+            make_buildable=make_buildable,
+            decompressed_formats=DECOMPRESSED_FORMATS,
+            requires_transformers="5.8",
+        ),
+    )
+)

modelopt/torch/models/kimi_vl/ep.py

@@ -21,7 +21,22 @@
 """
 
 from .. import _kimi, _onfly_decompress
-from .._remote_code import stage_full_remote_closure
+
+
+def make_buildable(
+    hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_
+) -> None:
+    """Have the checkpoint's own code import and build on this transformers.
+
+    The plan decision needs it before ``prepare``; it changes nothing the model computes or loads.
+    """
+    _kimi.make_buildable(
+        "modeling_kimi_vl",
+        "KimiVLForConditionalGeneration",
+        hf_config,
+        ckpt_path,
+        trust_remote_code,
+    )
 
 
 def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_):
@@ -38,8 +53,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool
     ckpt_path = _kimi.checkpoint_to_patch("kimi_vl", hf_config, ckpt_path, trust_remote_code)
     _kimi.refuse_quantization_only_on_text_config("kimi_vl", hf_config)
 
-    _kimi.install_compat_import_shims()
-    stage_full_remote_closure()
+    _kimi.make_importable()  # the swap below loads the code and adapts its tie_weights
 
     # Native deepseek_v3 config for the LM (drops the remote auto_map so the native class is used).
     hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config)

modelopt/torch/models/kimi_vl/specs.py

@@ -27,7 +27,7 @@
 
 from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN
 from ..specs import EPSpec, ModelSpec, register
-from .ep import prepare
+from .ep import make_buildable, prepare
 
 __all__: list[str] = []
 
@@ -36,6 +36,11 @@
         model_type="kimi_vl",
         modeling_source="remote_code",
         # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan.
-        ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN),
+        ep_spec=EPSpec(
+            prepare=prepare,
+            prepared_ep_plan=BASE_MODEL_EP_PLAN,
+            plan_config_attr="text_config",
+            make_buildable=make_buildable,
+        ),
     )
 )