diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index 7d71f9e106f..e2f5aa2425b 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -91,6 +91,8 @@ nemotron_h, minimax_m3_vl, llama4_text, + kimi_vl, + kimi_k25, nemotron_h_omni, qwen2_moe, qwen3, diff --git a/modelopt/torch/models/deepseek_v3/specs.py b/modelopt/torch/models/deepseek_v3/specs.py index dab92dab2eb..d1b20b1338b 100644 --- a/modelopt/torch/models/deepseek_v3/specs.py +++ b/modelopt/torch/models/deepseek_v3/specs.py @@ -23,6 +23,15 @@ __all__: list[str] = [] +# No gate entry on purpose, whatever transformers' own plan has: the router is logits-only and the +# MoE block computes GLOBAL top-k ids, so the experts already receive unmasked global ids -- what +# the all-to-all dispatch wants. +BASE_MODEL_EP_PLAN = { + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", +} + # DeepseekV3MoE is invisible to the generic MoE detection: its class name does not end # in "SparseMoeBlock" and it calls its router ``gate``, so neither the name test nor the # structural router+experts test in is_moe matches. block_names below is what puts it on @@ -37,16 +46,7 @@ model_type="deepseek_v3", min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), - ep_spec=EPSpec( - # Stock transformers ships no plan for deepseek_v3. No gate entry on purpose: the - # router is logits-only and the MoE block computes GLOBAL top-k ids, so the experts - # already receive unmasked global ids -- what the all-to-all dispatch wants. - base_model_ep_plan={ - "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", - }, - ), + ep_spec=EPSpec(base_model_ep_plan=BASE_MODEL_EP_PLAN), moe_spec=MoESpec( block_names=("DeepseekV3MoE",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/torch/models/kimi_k25/__init__.py b/modelopt/torch/models/kimi_k25/__init__.py new file mode 100644 index 00000000000..2cf9aee48e6 --- /dev/null +++ b/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 diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py new file mode 100644 index 00000000000..1d882417510 --- /dev/null +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -0,0 +1,156 @@ +# 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.7-Code (``kimi_k25``). + +a remote-code VLM whose text backbone is a DeepSeek-V3 variant with compressed-tensors-compressed experts. + +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 on-the-fly decompressor enforces: earlier releases' +``Concatenate`` trips over the sibling keys it has already consumed. The decompressor couples tightly to +transformers' still-evolving ``WeightConverter`` conversion-ops API. +""" + +import sys + +from modelopt.torch.utils.moe import apply_distributed_support + +from .. import _onfly_decompress +from .._remote_code import stage_full_remote_closure +from ..kimi_vl.ep import _drop_tie_weights_kwargs, _install_compat_import_shims + + +def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): + """Load Kimi-K2.7'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 expert decompressor. + """ + if type(hf_config).__module__.startswith("transformers."): + # From 5.13, a load without trust_remote_code builds transformers' own port, which the + # patch below never reaches. + raise ValueError( + "kimi_k25: expert parallelism runs on the checkpoint's own modeling code; " + "load it with trust_remote_code=True." + ) + if not hasattr(hf_config, "text_config"): + raise ValueError("kimi_k25.apply: hf_config has no .text_config (expected a VLM config)") + if ckpt_path is not None and not trust_remote_code: + # Patching below imports the checkpoint's own modeling code, and transformers' + # get_class_from_dynamic_module runs it without checking trust_remote_code itself. + raise ValueError( + "kimi_k25: expert parallelism imports the checkpoint's own modeling code; " + "load it with trust_remote_code=True." + ) + + # transformers stages only 1-hop relative imports; Kimi's remote code is 2 hops deep. Install the + # full-closure staging shim BEFORE any remote-code load (this plugin's get_class_from_dynamic_module + # below, and from_pretrained later in load_and_prepare_ep). + _install_compat_import_shims() + stage_full_remote_closure() + + # 1. 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) + + # 2. 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. + _onfly_decompress.install_onfly_decompress(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}`. + 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. If ckpt_path isn't available we skip the patch and + # the run will fall back to Kimi's per-expert nn.ModuleList layout (non-EP-capable) -- caller + # should always thread ckpt_path through apply_distributed_support. + if ckpt_path is not None: + _patch_kimi_modeling_to_native(ckpt_path, trust_remote_code) + return hf_config + + +def _patch_kimi_modeling_to_native(ckpt_path: str, trust_remote_code: bool) -> None: + """Point the remote module's ``DeepseekV3ForCausalLM`` at the transformers-native class. + + Load Kimi's remote-code ``modeling_kimi_k25`` module and replace its ``DeepseekV3ForCausalLM`` + symbol with the transformers-native class. After this, ``KimiK25ForConditionalGeneration.__init__`` + builds the EP-capable native deepseek_v3 instead of Kimi's per-expert Linears version. + """ + from transformers.dynamic_module_utils import get_class_from_dynamic_module + from transformers.models.deepseek_v3 import DeepseekV3ForCausalLM as NativeDeepseekV3ForCausalLM + + # Force the dynamic-module loader to stage + import modeling_kimi_k25 into sys.modules. + get_class_from_dynamic_module( + "modeling_kimi_k25.KimiK25ForConditionalGeneration", + ckpt_path, + trust_remote_code=trust_remote_code, + ) + + patched = [] + for mod_name, mod in list(sys.modules.items()): + if mod is None: + continue + # The dynamic-module loader names the module like + # ``transformers_modules...modeling_kimi_k25``. + if mod_name.endswith(".modeling_kimi_k25") and hasattr(mod, "DeepseekV3ForCausalLM"): + setattr(mod, "DeepseekV3ForCausalLM", NativeDeepseekV3ForCausalLM) + # Like Kimi-VL's: the remote tie_weights() takes no kwargs, but transformers' post_init and + # load call tie_weights(recompute_mapping=False) -> TypeError. + kimi_cls = getattr(mod, "KimiK25ForConditionalGeneration", 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_vl_kwargs_shim", False): + kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie) + patched.append(mod_name) + if not patched: + raise RuntimeError( + "kimi_k25: could not find a loaded modeling_kimi_k25 in sys.modules to monkey-patch" + ) + print( + f"[kimi_k25] monkey-patched DeepseekV3ForCausalLM -> transformers-native in: {patched}", + flush=True, + ) diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py new file mode 100644 index 00000000000..cc6757478a7 --- /dev/null +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -0,0 +1,37 @@ +# 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.7: 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. +""" + +from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN +from ..specs import EPSpec, ModelSpec, register +from .ep import prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_k25", + # transformers ships Kimi-K2.5 itself from 5.13. Expert parallelism still runs on the + # checkpoint's own code (``trust_remote_code=True``), which ``prepare`` patches. + min_transformers_version="5.13", + # 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), + ) +) diff --git a/modelopt/torch/models/kimi_vl/__init__.py b/modelopt/torch/models/kimi_vl/__init__.py new file mode 100644 index 00000000000..ec89db2722b --- /dev/null +++ b/modelopt/torch/models/kimi_vl/__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_vl specs (HF model type ``kimi_vl``).""" + +from . import specs diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py new file mode 100644 index 00000000000..99177d1c017 --- /dev/null +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -0,0 +1,171 @@ +# 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-VL-A3B (``kimi_vl``). + +a remote-code VLM whose text backbone is a Moonlight (DeepSeek-V3-lite: MLA, ``q_lora_rank=None``, 64 routed experts, 2 +shared) MoE and whose vision tower is MoonViT. The bf16 sibling of ``kimi_k25`` -- SAME shape, but NO compressed-tensors +decompressor (the checkpoint ships plain bf16 per-expert weights). + +Challenge (identical to kimi_k25): the bundled ``modeling_kimi_vl.py`` builds its language model from +an INLINE ``DeepseekV3ForCausalLM`` whose ``DeepseekV3MoE`` is an ``nn.ModuleList`` of per-expert +``gate_proj``/``up_proj``/``down_proj`` Linears -- incompatible with transformers' EP machinery, which +needs the FUSED 3-D ``gate_up_proj``/``down_proj`` Parameters the transformers-native ``deepseek_v3`` +modeling ships. + +Solution (mirrors kimi_k25): + 1. Shim the import symbols ``modeling_kimi_vl`` needs that newer transformers dropped (e.g. + ``is_torch_fx_available``) so the module IMPORTS -- they only guard torch.fx / legacy attn-mask + paths that our native-LM + MoonViT forward never execute. + 2. Swap ``hf_config.text_config`` to a transformers-native ``DeepseekV3Config`` (drops the remote + ``auto_map`` so the native class -- and its ``base_model_ep_plan`` -- is used for the LM). + 3. Apply the deepseek_v3 EP plan to that native ``text_config`` (post_init prefixes it to the VLM + root as ``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}``). + 4. Monkey-patch ``modeling_kimi_vl.DeepseekV3ForCausalLM`` -> the transformers-native class, so + ``KimiVLForConditionalGeneration.__init__`` (``self.language_model = DeepseekV3ForCausalLM( + config.text_config)``) builds the EP-capable native LM while MoonViT + the KimiVL wrapper stay + from the bundled code. Native deepseek_v3's per-expert WeightConverter fuses the checkpoint's + per-expert ``.weight`` tensors into the 3-D params on load (no decompressor -- weights are bf16). + +Version: matches Kimi-K2.7 -- validated on transformers 5.9.0 (has the EP framework + the legacy +symbols the remote code imports). The compat shims make it tolerant of a few later drops too. +""" + +import sys + +from modelopt.torch.utils.moe import apply_distributed_support + +from .. import _onfly_decompress +from .._remote_code import stage_full_remote_closure + + +def _install_compat_import_shims(): + """Add stubs for symbols ``modeling_kimi_vl`` imports that newer transformers removed, so the module imports. + + Each guards a torch.fx / legacy-attn-mask code path the native LM + MoonViT forward never hit, so a benign stub is + safe. Idempotent (only sets missing attrs). + """ + try: + import transformers.utils.import_utils as iu + + if not hasattr(iu, "is_torch_fx_available"): + iu.is_torch_fx_available = lambda: False + except Exception as e: + print(f"[kimi_vl] compat-shim warning (import_utils): {type(e).__name__}: {e}", flush=True) + try: + import transformers.pytorch_utils as pu + + if not hasattr(pu, "is_torch_greater_or_equal_than_1_13"): + pu.is_torch_greater_or_equal_than_1_13 = True + except Exception as e: + print(f"[kimi_vl] compat-shim warning (pytorch_utils): {type(e).__name__}: {e}", flush=True) + try: + import transformers.modeling_attn_mask_utils as amu + + if not hasattr(amu, "_prepare_4d_causal_attention_mask"): + + def _prepare_4d_causal_attention_mask(*args, **kwargs): # pragma: no cover + raise RuntimeError( + "kimi_vl: remote _prepare_4d_causal_attention_mask stub called -- the native LM " + "builds its own mask, so this path should be unreachable." + ) + + amu._prepare_4d_causal_attention_mask = _prepare_4d_causal_attention_mask + except Exception as e: + print( + f"[kimi_vl] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}", flush=True + ) + + +def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): + """Load Kimi-VL's language model as the transformers-native deepseek_v3. + + Swap text_config to native deepseek_v3 + apply the EP plan + monkey-patch Kimi-VL's modeling to + build the native EP-capable DeepseekV3ForCausalLM for the language sub-model. No decompressor: the + Kimi-VL checkpoint is plain bf16. + """ + if not hasattr(hf_config, "text_config"): + raise ValueError("kimi_vl.apply: hf_config has no .text_config (expected a VLM config)") + if ckpt_path is not None and not trust_remote_code: + # Patching below imports the checkpoint's own modeling code, and transformers' + # get_class_from_dynamic_module runs it without checking trust_remote_code itself. + raise ValueError( + "kimi_vl: expert parallelism imports the checkpoint's own modeling code; " + "load it with trust_remote_code=True." + ) + + _install_compat_import_shims() + stage_full_remote_closure() + + # 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) + # EP plan onto the native text_config (bubbles up to the VLM root under language_model.*). + apply_distributed_support("deepseek_v3", hf_config.text_config) + + if ckpt_path is not None: + _patch_kimi_vl_modeling_to_native(ckpt_path, trust_remote_code) + return hf_config + + +def _drop_tie_weights_kwargs(orig_tie): + """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it.""" + + def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc. + return orig_tie(self) + + setattr(_tie_weights, "_kimi_vl_kwargs_shim", True) + return _tie_weights + + +def _patch_kimi_vl_modeling_to_native(ckpt_path: str, trust_remote_code: bool) -> None: + """Point the remote module's inline ``DeepseekV3ForCausalLM`` at the native class. + + Load Kimi-VL's remote ``modeling_kimi_vl`` and replace its inline ``DeepseekV3ForCausalLM`` with + the transformers-native class, so ``KimiVLForConditionalGeneration.__init__`` builds the EP-capable + native deepseek_v3 for ``self.language_model`` (MoonViT + wrapper stay from the bundled code). + """ + from transformers.dynamic_module_utils import get_class_from_dynamic_module + from transformers.models.deepseek_v3 import DeepseekV3ForCausalLM as NativeDeepseekV3ForCausalLM + + get_class_from_dynamic_module( + "modeling_kimi_vl.KimiVLForConditionalGeneration", + ckpt_path, + trust_remote_code=trust_remote_code, + ) + + patched = [] + for mod_name, mod in list(sys.modules.items()): + if mod is None: + continue + if mod_name.endswith(".modeling_kimi_vl") and hasattr(mod, "DeepseekV3ForCausalLM"): + setattr(mod, "DeepseekV3ForCausalLM", NativeDeepseekV3ForCausalLM) + # transformers 5.9's init_weights calls self.tie_weights(recompute_mapping=False), but + # Kimi-VL's remote tie_weights() has the older no-kwarg signature -> TypeError. Wrap it to + # swallow the extra kwargs (it only ties the LM head/embeddings; the kwargs don't apply). + kimi_cls = getattr(mod, "KimiVLForConditionalGeneration", 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_vl_kwargs_shim", False): + kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie) + patched.append(mod_name) + if not patched: + raise RuntimeError( + "kimi_vl: could not find a loaded modeling_kimi_vl in sys.modules to monkey-patch" + ) + print( + f"[kimi_vl] monkey-patched DeepseekV3ForCausalLM -> transformers-native (+tie_weights shim) in: " + f"{patched}", + flush=True, + ) diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py new file mode 100644 index 00000000000..7742f4251b1 --- /dev/null +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -0,0 +1,36 @@ +# 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_vl specs (HF model type ``kimi_vl``). + +Kimi-VL-A3B: a remote-code VLM whose text backbone is a Moonlight MoE. Its bundled +modeling builds the language model from per-expert Linears, which the EP machinery cannot shard, so +``ep.py`` swaps in the transformers-native deepseek_v3 language model. +""" + +from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN +from ..specs import EPSpec, ModelSpec, register +from .ep import prepare + +__all__: list[str] = [] + +register( + ModelSpec( + 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), + ) +) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py new file mode 100644 index 00000000000..8185f838972 --- /dev/null +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -0,0 +1,143 @@ +# 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. + +"""Expert-parallel adapter for Kimi-K2.7 (``kimi_k25``), a remote-code VLM on DeepSeek-V3. + +The adapter rebuilds the language model as transformers' native deepseek_v3 and installs the +on-the-fly expert decompressor. The checkpoint's remote code is not available here, so the tests +stub the dynamic-module loader with a stand-in remote module. Every process-wide patch is undone +after each test. +""" + +import sys +import types +from types import SimpleNamespace + +import pytest + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.8"): + pytest.skip( + "Kimi-K2.5 expert parallelism needs transformers >= 5.8, for the decompression", + allow_module_level=True, + ) + +import transformers.dynamic_module_utils as dmu +from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig + +from modelopt.torch.models import _onfly_decompress, get_spec +from modelopt.torch.models.kimi_k25 import ep as kimi_k25 +from modelopt.torch.utils.moe import _routed_expert_keys + +QUANT = {"quant_method": "compressed-tensors", "format": "pack-quantized"} +REMOTE = "transformers_modules.Kimi_hyphen_K2_dot_7.0123abcd" # how the loader names remote code + + +class _RemoteTextConfig(PretrainedConfig): + """Stands in for the remote DeepSeek-V3 config Kimi-K2.7 bundles as its ``text_config``.""" + + model_type = "kimi_k2_remote_text" + + +def _vlm_config(): + text = _RemoteTextConfig( + hidden_size=64, + num_hidden_layers=2, + n_routed_experts=8, + moe_intermediate_size=16, + auto_map={"AutoConfig": "configuration_deepseek.DeepseekV3Config"}, + quantization_config=QUANT.copy(), + ) + return SimpleNamespace(text_config=text, quantization_config=QUANT.copy()) + + +@pytest.fixture +def remote(monkeypatch): + """Stand in for the checkpoint's remote code: importing it registers ``modeling_kimi_k25``, + whose ``DeepseekV3ForCausalLM`` is the per-expert remote language model.""" + module = types.ModuleType(f"{REMOTE}.modeling_kimi_k25") + module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + module.loads = [] + + def import_remote_class(class_ref, ckpt_path, **kwargs): + module.loads.append((class_ref, ckpt_path, kwargs)) + monkeypatch.setitem(sys.modules, module.__name__, module) + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + return module + + +def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkeypatch, remote): + """The text config becomes transformers' own ``DeepseekV3Config`` (no remote ``auto_map``, no + stale ``quantization_config``) with an EP plan for the fused experts, naming the experts the + spec states for a run that skips prepare; the decompressor is installed while the VLM config + still carries the compressed-tensors scheme; the remote-code staging shim is left installed; + without a checkpoint nothing remote is imported.""" + installs = [] + monkeypatch.setattr( + _onfly_decompress, + "install_onfly_decompress", + lambda cfg: installs.append(getattr(cfg, "quantization_config", None)), + ) + staging = dmu.get_cached_module_file + config = _vlm_config() + + assert kimi_k25.prepare(config) is config + + text = config.text_config + assert type(text) is DeepseekV3Config + assert getattr(text, "auto_map", None) is None + assert getattr(text, "quantization_config", None) is None + assert (text.hidden_size, text.n_routed_experts, text.moe_intermediate_size) == (64, 8, 16) + assert text.base_model_ep_plan["layers.*.mlp.experts.gate_up_proj"] == "grouped_gemm" + assert text.base_model_ep_plan["layers.*.mlp.experts.down_proj"] == "grouped_gemm" + assert _routed_expert_keys(text.base_model_ep_plan) == _routed_expert_keys( + get_spec("kimi_k25").ep_spec.prepared_ep_plan + ) + assert installs == [QUANT] + assert dmu.get_cached_module_file is not staging + assert remote.loads == [] + + +def test_prepare_with_a_checkpoint_points_the_remote_modeling_at_native_deepseek_v3( + monkeypatch, remote +): + """``KimiK25ForConditionalGeneration.__init__`` builds ``DeepseekV3ForCausalLM(text_config)`` + from its own module's namespace; that name must now be transformers' native class, so the + language model gets fused, EP-capable experts. Other remote modules keep their own class.""" + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + bystander = types.ModuleType(f"{REMOTE}.modeling_deepseek") + remote_lm = bystander.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + monkeypatch.setitem(sys.modules, bystander.__name__, bystander) + + kimi_k25.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + + assert remote.loads == [ + ("modeling_kimi_k25.KimiK25ForConditionalGeneration", "/ckpt", {"trust_remote_code": True}) + ] + assert remote.DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + assert bystander.DeepseekV3ForCausalLM is remote_lm + + +def test_prepare_refuses_transformers_own_port(): + """From 5.13 a load without ``trust_remote_code`` builds transformers' port, which the patch + never reaches; ``prepare`` must refuse it rather than run unpatched.""" + native = pytest.importorskip("transformers.models.kimi_k25.configuration_kimi_k25") + with pytest.raises(ValueError, match="trust_remote_code=True"): + kimi_k25.prepare(native.Kimi_K25Config(), ckpt_path="/ckpt") diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py new file mode 100644 index 00000000000..9ce96c68155 --- /dev/null +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -0,0 +1,237 @@ +# 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. + +"""Expert-parallel adapter for Kimi-VL-A3B (``kimi_vl``), a remote-code VLM on a Moonlight MoE. + +The adapter rebuilds the language model as transformers' native deepseek_v3; its experts are plain +bf16, so nothing is decompressed. The checkpoint's remote code is not available here, so the tests +stub the dynamic-module loader with a stand-in remote module. Every process-wide patch is undone +after each test. +""" + +import importlib +import sys +import types +from types import SimpleNamespace + +import pytest + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.6"): + pytest.skip("Kimi-VL expert parallelism needs transformers >= 5.6", allow_module_level=True) + +import transformers.dynamic_module_utils as dmu +import transformers.utils.import_utils as import_utils +from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig + +from modelopt.torch.models import _onfly_decompress, get_spec +from modelopt.torch.models.kimi_k25 import ep as kimi_k25 +from modelopt.torch.models.kimi_vl import ep as kimi_vl +from modelopt.torch.utils.moe import _routed_expert_keys + +# Symbols the remote modeling imports that newer transformers dropped. +COMPAT_SYMBOLS = ( + ("transformers.utils.import_utils", "is_torch_fx_available"), + ("transformers.pytorch_utils", "is_torch_greater_or_equal_than_1_13"), + ("transformers.modeling_attn_mask_utils", "_prepare_4d_causal_attention_mask"), +) + + +def _restore(monkeypatch, owner, name): + """Have ``monkeypatch`` put ``owner.`` back after the test, or remove it if it was absent.""" + if hasattr(owner, name): + monkeypatch.setattr(owner, name, getattr(owner, name)) + else: + monkeypatch.setattr(owner, name, None, raising=False) + monkeypatch.delattr(owner, name) + + +class _RemoteTextConfig(PretrainedConfig): + """Stands in for the remote DeepSeek-V3 config Kimi-VL bundles as its ``text_config``.""" + + model_type = "kimi_vl_remote_text" + + +def _vlm_config(): + text = _RemoteTextConfig( + hidden_size=64, + num_hidden_layers=2, + n_routed_experts=8, + moe_intermediate_size=16, + q_lora_rank=None, + auto_map={"AutoConfig": "configuration_kimi_vl.DeepseekV3Config"}, + ) + return SimpleNamespace(text_config=text) + + +@pytest.fixture +def remote(monkeypatch, request): + """Stand in for the checkpoint's remote code: importing it registers its modeling module + (Kimi-VL's ``modeling_kimi_vl`` unless a test names another), with the inline per-expert + ``DeepseekV3ForCausalLM`` and a VLM class whose ``tie_weights`` has the old no-kwarg signature, + as both Kimi checkpoints ship it. Also restores the compat symbols and loader hooks prepare + installs.""" + modeling, vlm_name = getattr( + request, "param", ("modeling_kimi_vl", "KimiVLForConditionalGeneration") + ) + for module_name, name in COMPAT_SYMBOLS: + try: + owner = importlib.import_module(module_name) + except ImportError: + continue + _restore(monkeypatch, owner, name) + module = types.ModuleType(f"transformers_modules.Kimi.0123abcd.{modeling}") + module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + module.loads, module.tied = [], [] + + def tie_weights(self): + module.tied.append(self) + return "tied" + + module.vlm = type(vlm_name, (), {"tie_weights": tie_weights}) + setattr(module, vlm_name, module.vlm) + + def import_remote_class(class_ref, ckpt_path, **kwargs): + module.loads.append((class_ref, ckpt_path, kwargs)) + monkeypatch.setitem(sys.modules, module.__name__, module) + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + return module + + +def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monkeypatch, remote): + """The text config becomes transformers' own ``DeepseekV3Config`` (no remote ``auto_map``) with + an EP plan for the fused experts, naming the experts the spec states for a run that skips + prepare; the experts are bf16, so no decompressor is installed; the symbols the remote modeling + imports exist; the remote-code staging shim is left installed.""" + monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) + monkeypatch.setattr( + _onfly_decompress, + "install_onfly_decompress", + lambda cfg: pytest.fail("Kimi-VL's experts are bf16: there is nothing to decompress"), + ) + staging = dmu.get_cached_module_file + config = _vlm_config() + + assert kimi_vl.prepare(config) is config + + text = config.text_config + assert type(text) is DeepseekV3Config + assert getattr(text, "auto_map", None) is None + assert (text.hidden_size, text.n_routed_experts, text.moe_intermediate_size) == (64, 8, 16) + assert text.base_model_ep_plan["layers.*.mlp.experts.gate_up_proj"] == "grouped_gemm" + assert text.base_model_ep_plan["layers.*.mlp.experts.down_proj"] == "grouped_gemm" + assert _routed_expert_keys(text.base_model_ep_plan) == _routed_expert_keys( + get_spec("kimi_vl").ep_spec.prepared_ep_plan + ) + assert import_utils.is_torch_fx_available() is False + assert dmu.get_cached_module_file is not staging + assert remote.loads == [] + + +@pytest.mark.parametrize( + ("adapter", "remote"), + [ + (kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration")), + (kimi_k25, ("modeling_kimi_k25", "KimiK25ForConditionalGeneration")), + ], + ids=["vl", "k25"], + indirect=["remote"], +) +def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_weights( + monkeypatch, adapter, remote +): + """The VLM class builds its language model from the module-level ``DeepseekV3ForCausalLM``, + which must now be transformers' native class. Both checkpoints' remote modeling imports a + symbol newer transformers dropped, and has an old no-kwarg ``tie_weights`` that must accept the + ``recompute_mapping`` newer transformers passes. Running prepare again does not wrap + ``tie_weights`` twice.""" + monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + shim = remote.vlm.tie_weights + adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + + vlm_class = f"{remote.__name__.rsplit('.', 1)[-1]}.{remote.vlm.__name__}" + assert remote.loads == [(vlm_class, "/ckpt", {"trust_remote_code": True})] * 2 + assert remote.DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + assert import_utils.is_torch_fx_available() is False + assert remote.vlm.tie_weights is shim + model = remote.vlm() + assert model.tie_weights(missing_keys=set(), recompute_mapping=False) == "tied" + assert remote.tied == [model] + + +# Kimi-K2.5's adapter is Kimi-VL's with on-the-fly decompression added; these hold for both. + + +@pytest.mark.parametrize( + ("swap", "modeling"), + [ + (lambda: kimi_vl._patch_kimi_vl_modeling_to_native("/ckpt", False), "modeling_kimi_vl"), + (lambda: kimi_k25._patch_kimi_modeling_to_native("/ckpt", True), "modeling_kimi_k25"), + ], + ids=["vl", "k25"], +) +def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch, swap, modeling): + """Skipping the swap silently would build the per-expert layout that EP cannot shard.""" + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", lambda *args, **kwargs: None) + + with pytest.raises(RuntimeError, match=modeling): + swap() + + +@pytest.mark.parametrize("adapter", [kimi_vl, kimi_k25], ids=["vl", "k25"]) +@pytest.mark.parametrize( + ("config", "kwargs", "match"), + [ + (SimpleNamespace, {}, "text_config"), + # transformers' get_class_from_dynamic_module runs the checkpoint's code without checking + # trust_remote_code, so prepare must refuse before reaching it. + (_vlm_config, {"ckpt_path": "/ckpt"}, "trust_remote_code=True"), + ], + ids=["no-text-config", "untrusted-checkpoint"], +) +def test_prepare_refuses_before_patching_anything(monkeypatch, adapter, config, kwargs, match): + """Before prepare patches anything process-wide, such as the loader's remote-code staging, or + imports the checkpoint's code.""" + + def unstaged_loader(*args, **kwargs): + """The loader before any staging shim wraps it.""" + + monkeypatch.setattr(dmu, "get_cached_module_file", unstaged_loader) + monkeypatch.setattr( + dmu, + "get_class_from_dynamic_module", + lambda *args, **kwargs: pytest.fail("the checkpoint's code was imported"), + ) + + with pytest.raises(ValueError, match=match): + adapter.prepare(config(), **kwargs) + + assert dmu.get_cached_module_file is unstaged_loader + + +@pytest.mark.parametrize( + ("model_type", "adapter"), [("kimi_vl", kimi_vl), ("kimi_k25", kimi_k25)], ids=["vl", "k25"] +) +def test_the_spec_routes_expert_parallelism_through_its_own_adapter(model_type, adapter): + """The checkpoint's code builds no EP-capable experts, so its whole EP setup is this adapter's + ``prepare``: another Kimi adapter's would load it without its own swap or decompression.""" + assert get_spec(model_type).ep_spec.prepare is adapter.prepare diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 7592dc8da7f..706af4c1eb4 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -480,6 +480,44 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( resolve_distribution_plan(str(tmp_path)) +def _remote_kimi(moe): + """A remote-code Kimi VLM as the checkpoint's own code builds it, unpatched by its spec's + ``prepare``: each MoE layer of its language model holds its routed experts as per-expert + modules in an ``nn.ModuleList``, at ``language_model.model.layers...experts``.""" + + def layer(): + block = nn.Module() + block.gate = nn.Linear(2, 4, bias=False) + block.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(4)) + return nn.ModuleDict({moe: block}) + + model = nn.Module() + model.language_model = nn.Module() + model.language_model.model = nn.Module() + model.language_model.model.layers = nn.ModuleList(layer() for _ in range(2)) + return model + + +@pytest.mark.parametrize(("model_type", "moe"), [("kimi_vl", "mlp"), ("kimi_k25", "mlp")]) +def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( + run_as, serve, monkeypatch, tmp_path, model_type, moe +): + """Remote-code Kimi checkpoints get their plan from their spec's ``prepare``, which an EP degree + of 1 skips: it patches their code. The plan their spec states still finds the experts their own + code builds, so whole on every rank they are refused like Qwen3-MoE.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "1") + monkeypatch.setattr( + get_spec(model_type).ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") + ) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi(moe)) + run_as(4) + serve(SimpleNamespace(model_type=model_type, text_config=SimpleNamespace(n_routed_experts=8))) + + with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + resolve_distribution_plan(str(tmp_path)) + + def test_the_plan_a_spec_would_supply(monkeypatch): """What that guard reads at an EP degree of 1: the spec's plan, or the plan it states its ``prepare`` attaches, read without running ``prepare``; nothing for a model with no spec or no