From db64043239dd1ca7e6fcf65f68053abb3374c5a9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:53:38 +0000 Subject: [PATCH 01/20] Add the Kimi-K2.5 and Kimi-VL expert-parallel adapters Both checkpoints ship their own modeling, which builds the language model's MoE from an nn.ModuleList of per-expert Linears that transformers' expert-parallel machinery cannot shard. Their text backbone is DeepSeek-V3, so the adapters swap in transformers' native deepseek_v3 language model -- fused gate_up_proj / down_proj experts and the plan to shard them -- and keep the bundled vision tower and projector from the checkpoint's own code. Kimi-K2.5's experts are compressed-tensors INT4, decompressed as they load. Signed-off-by: Shengliang Xu --- modelopt/torch/models/__init__.py | 2 + modelopt/torch/models/kimi_k25/__init__.py | 18 ++ modelopt/torch/models/kimi_k25/ep.py | 133 +++++++++++++++ modelopt/torch/models/kimi_k25/specs.py | 33 ++++ modelopt/torch/models/kimi_vl/__init__.py | 18 ++ modelopt/torch/models/kimi_vl/ep.py | 164 +++++++++++++++++++ modelopt/torch/models/kimi_vl/specs.py | 34 ++++ tests/unit/torch/models/test_kimi_k25_ep.py | 147 +++++++++++++++++ tests/unit/torch/models/test_kimi_vl_ep.py | 173 ++++++++++++++++++++ 9 files changed, 722 insertions(+) create mode 100644 modelopt/torch/models/kimi_k25/__init__.py create mode 100644 modelopt/torch/models/kimi_k25/ep.py create mode 100644 modelopt/torch/models/kimi_k25/specs.py create mode 100644 modelopt/torch/models/kimi_vl/__init__.py create mode 100644 modelopt/torch/models/kimi_vl/ep.py create mode 100644 modelopt/torch/models/kimi_vl/specs.py create mode 100644 tests/unit/torch/models/test_kimi_k25_ep.py create mode 100644 tests/unit/torch/models/test_kimi_vl_ep.py 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/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..a30bb93314f --- /dev/null +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -0,0 +1,133 @@ +# 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 + + +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 not hasattr(hf_config, "text_config"): + raise ValueError("kimi_k25.apply: hf_config has no .text_config (expected a VLM config)") + + # 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). + 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) + 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..3bc901ed3aa --- /dev/null +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -0,0 +1,33 @@ +# 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 ..specs import EPSpec, ModelSpec, register +from .ep import prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_k25", + modeling_source="remote_code", + ep_spec=EPSpec(prepare=prepare), + ) +) 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..055c6a0d89f --- /dev/null +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -0,0 +1,164 @@ +# 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)") + + _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..142771685f4 --- /dev/null +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -0,0 +1,34 @@ +# 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 ..specs import EPSpec, ModelSpec, register +from .ep import prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_vl", + modeling_source="remote_code", + ep_spec=EPSpec(prepare=prepare), + ) +) 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..742c3ee91af --- /dev/null +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -0,0 +1,147 @@ +# 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.6"): + pytest.skip("Kimi-K2.7 expert parallelism needs transformers >= 5.6", 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 + +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; 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 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_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch): + """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_kimi_k25"): + kimi_k25._patch_kimi_modeling_to_native("/ckpt", True) + + +def test_prepare_rejects_a_config_without_a_text_config(): + with pytest.raises(ValueError, match="text_config"): + kimi_k25.prepare(SimpleNamespace()) + + +def test_the_kimi_k25_spec_routes_expert_parallelism_through_this_adapter(): + """kimi_k25 is remote code, so its whole EP setup is this adapter's ``prepare``.""" + spec = get_spec("kimi_k25") + assert spec.modeling_source == "remote_code" + assert spec.ep_spec.prepare is kimi_k25.prepare 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..f0e7a637834 --- /dev/null +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -0,0 +1,173 @@ +# 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 +from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig + +from modelopt.torch.models import _onfly_decompress, get_spec +from modelopt.torch.models.kimi_vl import ep as kimi_vl + +# 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): + """Stand in for the checkpoint's remote code: importing it registers ``modeling_kimi_vl``, with + the inline per-expert ``DeepseekV3ForCausalLM`` and a VLM class whose ``tie_weights`` has the + old no-kwarg signature. Also restores the compat symbols and loader hooks prepare installs.""" + for module_name, name in COMPAT_SYMBOLS: + try: + owner = importlib.import_module(module_name) + except ImportError: + continue + _restore(monkeypatch, owner, name) + module = types.ModuleType("transformers_modules.Kimi_hyphen_VL.0123abcd.modeling_kimi_vl") + module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + module.loads, module.tied = [], [] + + class KimiVLForConditionalGeneration: + def tie_weights(self): + module.tied.append(self) + return "tied" + + module.KimiVLForConditionalGeneration = KimiVLForConditionalGeneration + + 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; the experts are bf16, so no decompressor is installed; the + symbols the remote modeling imports exist; the remote-code staging shim is left installed.""" + import transformers.utils.import_utils as import_utils + + 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 import_utils.is_torch_fx_available() is False + assert dmu.get_cached_module_file is not staging + assert remote.loads == [] + + +def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_weights(remote): + """``KimiVLForConditionalGeneration`` builds its language model from the module-level + ``DeepseekV3ForCausalLM``, which must now be transformers' native class; its old no-kwarg + ``tie_weights`` must accept the ``recompute_mapping`` newer transformers passes. Running + prepare again does not wrap ``tie_weights`` twice.""" + kimi_vl.prepare(_vlm_config(), ckpt_path="/ckpt") + shim = remote.KimiVLForConditionalGeneration.tie_weights + kimi_vl.prepare(_vlm_config(), ckpt_path="/ckpt") + + vlm_class = "modeling_kimi_vl.KimiVLForConditionalGeneration" + assert remote.loads == [(vlm_class, "/ckpt", {"trust_remote_code": False})] * 2 + assert remote.DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + assert remote.KimiVLForConditionalGeneration.tie_weights is shim + model = remote.KimiVLForConditionalGeneration() + assert model.tie_weights(missing_keys=set(), recompute_mapping=False) == "tied" + assert remote.tied == [model] + + +def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch): + """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_kimi_vl"): + kimi_vl._patch_kimi_vl_modeling_to_native("/ckpt", False) + + +def test_prepare_rejects_a_config_without_a_text_config(): + with pytest.raises(ValueError, match="text_config"): + kimi_vl.prepare(SimpleNamespace()) + + +def test_the_kimi_vl_spec_routes_expert_parallelism_through_this_adapter(): + """kimi_vl is remote code, so its whole EP setup is this adapter's ``prepare``.""" + spec = get_spec("kimi_vl") + assert spec.modeling_source == "remote_code" + assert spec.ep_spec.prepare is kimi_vl.prepare From a6f56695fd88b50021f3c253d290cc214e6210b5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 11:26:26 +0000 Subject: [PATCH 02/20] Skip the tests below the transformers release the code requires The module skipped below 5.6, but the code under test requires 5.8 and raises below it, so a 5.6 or 5.7 run failed instead of skipping. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k25_ep.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 742c3ee91af..24720699b69 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -31,8 +31,11 @@ from packaging.version import Version -if Version(transformers.__version__) < Version("5.6"): - pytest.skip("Kimi-K2.7 expert parallelism needs transformers >= 5.6", allow_module_level=True) +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 3acbb33c5dc6553a2ee83fb964ea5f54146dbfc6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Mon, 12 Oct 2026 01:31:14 +0000 Subject: [PATCH 03/20] Record that transformers ships Kimi-K2.5 from 5.13, and keep EP on its remote code transformers 5.13 added kimi_k25, so its spec can no longer claim remote code; test_specs_vs_transformers fails on the latest transformers until it names the release. A load without trust_remote_code now builds transformers' port, which prepare's patch never reaches, so prepare refuses that config instead of running unpatched. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 7 +++++++ modelopt/torch/models/kimi_k25/specs.py | 4 +++- tests/unit/torch/models/test_kimi_k25_ep.py | 11 +++++++++-- 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index a30bb93314f..dc3bec05d89 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -62,6 +62,13 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool 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)") diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index 3bc901ed3aa..a65f0d57f83 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -27,7 +27,9 @@ register( ModelSpec( model_type="kimi_k25", - modeling_source="remote_code", + # 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", ep_spec=EPSpec(prepare=prepare), ) ) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 24720699b69..17b45599386 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -143,8 +143,15 @@ def test_prepare_rejects_a_config_without_a_text_config(): kimi_k25.prepare(SimpleNamespace()) +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") + + def test_the_kimi_k25_spec_routes_expert_parallelism_through_this_adapter(): - """kimi_k25 is remote code, so its whole EP setup is this adapter's ``prepare``.""" + """The checkpoint's code builds no EP-capable experts, so its whole EP setup is ``prepare``.""" spec = get_spec("kimi_k25") - assert spec.modeling_source == "remote_code" assert spec.ep_spec.prepare is kimi_k25.prepare From f906500acf15f8a256c570f2f16e075c1d9c5811 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 02:12:37 +0000 Subject: [PATCH 04/20] State the expert plan Kimi-K2.5's and Kimi-VL's prepare attaches Both specs' prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. An EP degree of 1 skips prepare, so their specs supplied no plan there, and the remote-code checkpoints, whose configs carry none, loaded whole on every rank with a replicated dense part instead of being refused. Their specs now state deepseek_v3's plan, hoisted to a constant both import. It finds the experts the checkpoint's own code builds, at the same paths. Where transformers ships a deepseek_v3 plan of its own, prepare keeps that one, which names the same experts. The comment that moves with the plan no longer says transformers ships none, which stopped being true in 5.18. Signed-off-by: Shengliang Xu --- modelopt/torch/models/deepseek_v3/specs.py | 20 +++++----- modelopt/torch/models/kimi_k25/specs.py | 4 +- modelopt/torch/models/kimi_vl/specs.py | 4 +- tests/unit/torch/models/test_kimi_k25_ep.py | 11 ++++-- tests/unit/torch/models/test_kimi_vl_ep.py | 9 ++++- .../utils/test_resolve_distribution_plan.py | 38 +++++++++++++++++++ 6 files changed, 69 insertions(+), 17 deletions(-) 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/specs.py b/modelopt/torch/models/kimi_k25/specs.py index a65f0d57f83..cc6757478a7 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -19,6 +19,7 @@ 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 @@ -30,6 +31,7 @@ # 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", - ep_spec=EPSpec(prepare=prepare), + # 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/specs.py b/modelopt/torch/models/kimi_vl/specs.py index 142771685f4..7742f4251b1 100644 --- a/modelopt/torch/models/kimi_vl/specs.py +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -20,6 +20,7 @@ ``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 @@ -29,6 +30,7 @@ ModelSpec( model_type="kimi_vl", modeling_source="remote_code", - ep_spec=EPSpec(prepare=prepare), + # 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 index 17b45599386..b56d4e46238 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -42,6 +42,7 @@ 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 @@ -84,9 +85,10 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): 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; 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.""" + 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, @@ -105,6 +107,9 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkey 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 == [] diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index f0e7a637834..a6b28695fc5 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -40,6 +40,7 @@ from modelopt.torch.models import _onfly_decompress, get_spec 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 = ( @@ -109,8 +110,9 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): 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; the experts are bf16, so no decompressor is installed; the - symbols the remote modeling imports exist; the remote-code staging shim is left installed.""" + 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.""" import transformers.utils.import_utils as import_utils monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) @@ -130,6 +132,9 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke 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 == [] diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 2d3f01d8ff2..39e7979a344 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -598,6 +598,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 From 02e3af0528a210e71398993c102a4fc13c5fa2d7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:05:46 +0000 Subject: [PATCH 05/20] Show that the Kimi adapters refuse a config before patching anything Without prepare's own text_config check, to_native_backbone_config raises the same ValueError, so the refusal tests passed either way. The check exists to refuse before prepare installs the remote-code staging and, for Kimi-VL, the import shims, which are process-wide. The tests now hand the loader an unwrapped stand-in and require it to be left as it was, which fails once the check is gone. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k25_ep.py | 11 ++++++++++- tests/unit/torch/models/test_kimi_vl_ep.py | 11 ++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index b56d4e46238..77ae139156e 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -143,10 +143,19 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch kimi_k25._patch_kimi_modeling_to_native("/ckpt", True) -def test_prepare_rejects_a_config_without_a_text_config(): +def test_prepare_rejects_a_config_without_a_text_config(monkeypatch): + """Before prepare patches anything process-wide, such as the loader's remote-code staging.""" + + def unstaged_loader(*args, **kwargs): + """The loader before any staging shim wraps it.""" + + monkeypatch.setattr(dmu, "get_cached_module_file", unstaged_loader) + with pytest.raises(ValueError, match="text_config"): kimi_k25.prepare(SimpleNamespace()) + assert dmu.get_cached_module_file is unstaged_loader + def test_prepare_refuses_transformers_own_port(): """From 5.13 a load without ``trust_remote_code`` builds transformers' port, which the patch diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index a6b28695fc5..4db5132e189 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -166,10 +166,19 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch kimi_vl._patch_kimi_vl_modeling_to_native("/ckpt", False) -def test_prepare_rejects_a_config_without_a_text_config(): +def test_prepare_rejects_a_config_without_a_text_config(monkeypatch): + """Before prepare patches anything process-wide, such as the loader's remote-code staging.""" + + def unstaged_loader(*args, **kwargs): + """The loader before any staging shim wraps it.""" + + monkeypatch.setattr(dmu, "get_cached_module_file", unstaged_loader) + with pytest.raises(ValueError, match="text_config"): kimi_vl.prepare(SimpleNamespace()) + assert dmu.get_cached_module_file is unstaged_loader + def test_the_kimi_vl_spec_routes_expert_parallelism_through_this_adapter(): """kimi_vl is remote code, so its whole EP setup is this adapter's ``prepare``.""" From baad29372140ee389066ab8cccb6721b600a1d7a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:51:35 +0000 Subject: [PATCH 06/20] Trim the Kimi adapter tests: one parametrized test per shared behavior The swap-fails-loudly, missing-text_config refusal and spec-routing tests were written once per adapter. Each is now one test in test_kimi_vl_ep.py, parametrized over Kimi-VL and Kimi-K2.5, and each case still fails on its own adapter's regression: a silent swap, staging before the text_config check, a spec wired to the other adapter's prepare. The VL routing test's modeling_source assert goes, since test_specs_vs_transformers checks that claim. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k25_ep.py | 28 ----------------- tests/unit/torch/models/test_kimi_vl_ep.py | 35 +++++++++++++++------ 2 files changed, 25 insertions(+), 38 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 77ae139156e..8185f838972 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -135,37 +135,9 @@ def test_prepare_with_a_checkpoint_points_the_remote_modeling_at_native_deepseek assert bystander.DeepseekV3ForCausalLM is remote_lm -def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch): - """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_kimi_k25"): - kimi_k25._patch_kimi_modeling_to_native("/ckpt", True) - - -def test_prepare_rejects_a_config_without_a_text_config(monkeypatch): - """Before prepare patches anything process-wide, such as the loader's remote-code staging.""" - - def unstaged_loader(*args, **kwargs): - """The loader before any staging shim wraps it.""" - - monkeypatch.setattr(dmu, "get_cached_module_file", unstaged_loader) - - with pytest.raises(ValueError, match="text_config"): - kimi_k25.prepare(SimpleNamespace()) - - assert dmu.get_cached_module_file is unstaged_loader - - 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") - - -def test_the_kimi_k25_spec_routes_expert_parallelism_through_this_adapter(): - """The checkpoint's code builds no EP-capable experts, so its whole EP setup is ``prepare``.""" - spec = get_spec("kimi_k25") - assert spec.ep_spec.prepare is kimi_k25.prepare diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 4db5132e189..6e82f81e569 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -39,6 +39,7 @@ 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 @@ -158,15 +159,27 @@ def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_wei assert remote.tied == [model] -def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch): +# 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_kimi_vl"): - kimi_vl._patch_kimi_vl_modeling_to_native("/ckpt", False) + with pytest.raises(RuntimeError, match=modeling): + swap() -def test_prepare_rejects_a_config_without_a_text_config(monkeypatch): +@pytest.mark.parametrize("adapter", [kimi_vl, kimi_k25], ids=["vl", "k25"]) +def test_prepare_rejects_a_config_without_a_text_config(monkeypatch, adapter): """Before prepare patches anything process-wide, such as the loader's remote-code staging.""" def unstaged_loader(*args, **kwargs): @@ -175,13 +188,15 @@ def unstaged_loader(*args, **kwargs): monkeypatch.setattr(dmu, "get_cached_module_file", unstaged_loader) with pytest.raises(ValueError, match="text_config"): - kimi_vl.prepare(SimpleNamespace()) + adapter.prepare(SimpleNamespace()) assert dmu.get_cached_module_file is unstaged_loader -def test_the_kimi_vl_spec_routes_expert_parallelism_through_this_adapter(): - """kimi_vl is remote code, so its whole EP setup is this adapter's ``prepare``.""" - spec = get_spec("kimi_vl") - assert spec.modeling_source == "remote_code" - assert spec.ep_spec.prepare is kimi_vl.prepare +@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 From 6e0afe7abf83e3ff5edfe73dc6cbbe0cbee4afaf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:59:36 +0000 Subject: [PATCH 07/20] Import transformers' import_utils at the Kimi-VL tests' module scope test_prepare_rebuilds_the_language_model_config_without_a_decompressor imported transformers.utils.import_utils inside the test. It is core transformers, which the module already imports after its importorskip and version guards. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_vl_ep.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 6e82f81e569..1048d782e6e 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -36,6 +36,7 @@ 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 @@ -114,8 +115,6 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke 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.""" - import transformers.utils.import_utils as import_utils - monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) monkeypatch.setattr( _onfly_decompress, From 2feef7b9941959b8780c4cd0a6ca9b8f2189ef67 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 22:01:46 +0000 Subject: [PATCH 08/20] Load Kimi-K2.5's own remote code through its expert-parallel adapter The real Kimi-K2.5 remote code (moonshotai/Kimi-K2.5) failed to load through the adapter on every supported transformers release. Its modeling_deepseek.py imports is_torch_fx_available, which transformers 5.x dropped, and its KimiK25ForConditionalGeneration.tie_weights takes no arguments while transformers' post_init and load call it with recompute_mapping=False. Kimi-VL's adapter already handles both; give Kimi-K2.5's the same import shims and tie_weights wrapper. The Kimi-VL tests' stand-in remote code now takes the modeling module and VLM class to stand in for, and the swap-and-shim test runs for both adapters; the Kimi-K2.5 case fails without this change. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 9 +++ tests/unit/torch/models/test_kimi_vl_ep.py | 64 ++++++++++++++-------- 2 files changed, 51 insertions(+), 22 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index dc3bec05d89..a3e84b61e0d 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -54,6 +54,7 @@ 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, **_): @@ -75,6 +76,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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. @@ -129,6 +131,13 @@ def _patch_kimi_modeling_to_native(ckpt_path: str, trust_remote_code: bool) -> N # ``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( diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 1048d782e6e..9a8208a9c80 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -80,26 +80,31 @@ def _vlm_config(): @pytest.fixture -def remote(monkeypatch): - """Stand in for the checkpoint's remote code: importing it registers ``modeling_kimi_vl``, with - the inline per-expert ``DeepseekV3ForCausalLM`` and a VLM class whose ``tie_weights`` has the - old no-kwarg signature. Also restores the compat symbols and loader hooks prepare installs.""" +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("transformers_modules.Kimi_hyphen_VL.0123abcd.modeling_kimi_vl") + module = types.ModuleType(f"transformers_modules.Kimi.0123abcd.{modeling}") module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) module.loads, module.tied = [], [] - class KimiVLForConditionalGeneration: - def tie_weights(self): - module.tied.append(self) - return "tied" + def tie_weights(self): + module.tied.append(self) + return "tied" - module.KimiVLForConditionalGeneration = KimiVLForConditionalGeneration + 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)) @@ -140,20 +145,35 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke assert remote.loads == [] -def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_weights(remote): - """``KimiVLForConditionalGeneration`` builds its language model from the module-level - ``DeepseekV3ForCausalLM``, which must now be transformers' native class; its old no-kwarg - ``tie_weights`` must accept the ``recompute_mapping`` newer transformers passes. Running - prepare again does not wrap ``tie_weights`` twice.""" - kimi_vl.prepare(_vlm_config(), ckpt_path="/ckpt") - shim = remote.KimiVLForConditionalGeneration.tie_weights - kimi_vl.prepare(_vlm_config(), ckpt_path="/ckpt") +@pytest.mark.parametrize( + ("adapter", "remote", "trust_remote_code"), + [ + (kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration"), False), + (kimi_k25, ("modeling_kimi_k25", "KimiK25ForConditionalGeneration"), True), + ], + ids=["vl", "k25"], + indirect=["remote"], +) +def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_weights( + monkeypatch, adapter, remote, trust_remote_code +): + """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=trust_remote_code) + shim = remote.vlm.tie_weights + adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=trust_remote_code) - vlm_class = "modeling_kimi_vl.KimiVLForConditionalGeneration" - assert remote.loads == [(vlm_class, "/ckpt", {"trust_remote_code": False})] * 2 + vlm_class = f"{remote.__name__.rsplit('.', 1)[-1]}.{remote.vlm.__name__}" + assert remote.loads == [(vlm_class, "/ckpt", {"trust_remote_code": trust_remote_code})] * 2 assert remote.DeepseekV3ForCausalLM is DeepseekV3ForCausalLM - assert remote.KimiVLForConditionalGeneration.tie_weights is shim - model = remote.KimiVLForConditionalGeneration() + 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] From 3f0082b9321a9ef05253e369177305cd736399f1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:19:38 +0000 Subject: [PATCH 09/20] Refuse to import an untrusted checkpoint's code in the Kimi adapters With a checkpoint path, the Kimi-K2.5 and Kimi-VL adapters load the checkpoint's own modeling through transformers' get_class_from_dynamic_module, which runs it without checking trust_remote_code: the flag lands in its kwargs unused. Refuse in prepare unless the caller passed trust_remote_code=True, before anything is staged or patched. load_pretrained, resolve_distribution_plan and hf_ptq forward the caller's flag, so a trusted load is unchanged. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 7 ++++ modelopt/torch/models/kimi_vl/ep.py | 7 ++++ tests/unit/torch/models/test_kimi_vl_ep.py | 38 +++++++++++++++------- 3 files changed, 41 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index a3e84b61e0d..1d882417510 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -72,6 +72,13 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool ) 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 diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 055c6a0d89f..99177d1c017 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -98,6 +98,13 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool """ 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() diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 9a8208a9c80..9ce96c68155 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -146,16 +146,16 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke @pytest.mark.parametrize( - ("adapter", "remote", "trust_remote_code"), + ("adapter", "remote"), [ - (kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration"), False), - (kimi_k25, ("modeling_kimi_k25", "KimiK25ForConditionalGeneration"), True), + (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, trust_remote_code + 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 @@ -164,12 +164,12 @@ def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_wei ``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=trust_remote_code) + 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=trust_remote_code) + 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": trust_remote_code})] * 2 + 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 @@ -198,16 +198,32 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch @pytest.mark.parametrize("adapter", [kimi_vl, kimi_k25], ids=["vl", "k25"]) -def test_prepare_rejects_a_config_without_a_text_config(monkeypatch, adapter): - """Before prepare patches anything process-wide, such as the loader's remote-code staging.""" +@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="text_config"): - adapter.prepare(SimpleNamespace()) + with pytest.raises(ValueError, match=match): + adapter.prepare(config(), **kwargs) assert dmu.get_cached_module_file is unstaged_loader From c12401ca4dbbdf27663b2486695a768e890ddfc1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 21:07:12 +0000 Subject: [PATCH 10/20] Refuse a Kimi prepare without the checkpoint's path, and declare Kimi-K2.5's transformers floor Both Kimi adapters refuse, before they change anything, a prepare with no checkpoint path: they patch the checkpoint's own modeling code, and without it the config was left half-prepared, a native language-model config and a plan for fused experts under Kimi's own unpatched modeling. The trust_remote_code check now applies to every prepare. Kimi-K2.5's spec declares the transformers 5.8 its on-the-fly decompression needs, so an older release is refused by check_distributed_support, with its way around, before prepare runs. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 14 ++++++++------ modelopt/torch/models/kimi_k25/specs.py | 7 +++++-- modelopt/torch/models/kimi_vl/ep.py | 10 +++++++--- tests/unit/torch/models/test_kimi_k25_ep.py | 17 ++++++++++++----- tests/unit/torch/models/test_kimi_vl_ep.py | 8 +++++--- 5 files changed, 37 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 1d882417510..fc277ce3836 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -72,7 +72,12 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool ) 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: + if ckpt_path is None: + raise ValueError( + "kimi_k25: expert parallelism patches the checkpoint's own modeling code, so it needs " + "the checkpoint's path (apply_distributed_support's ckpt_path)." + ) + if 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( @@ -105,11 +110,8 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool 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) + # symbol) and swap to the native class. + _patch_kimi_modeling_to_native(ckpt_path, trust_remote_code) return hf_config diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index cc6757478a7..b026101f0c3 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -31,7 +31,10 @@ # 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), + # 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, requires_transformers="5.8" + ), ) ) diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 99177d1c017..edb9b99fbe7 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -98,7 +98,12 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool """ 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: + if ckpt_path is None: + raise ValueError( + "kimi_vl: expert parallelism patches the checkpoint's own modeling code, so it needs " + "the checkpoint's path (apply_distributed_support's ckpt_path)." + ) + if 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( @@ -114,8 +119,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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) + _patch_kimi_vl_modeling_to_native(ckpt_path, trust_remote_code) return hf_config diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 8185f838972..568959256b1 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -42,7 +42,7 @@ 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 +from modelopt.torch.utils.moe import _routed_expert_keys, check_distributed_support QUANT = {"quant_method": "compressed-tensors", "format": "pack-quantized"} REMOTE = "transformers_modules.Kimi_hyphen_K2_dot_7.0123abcd" # how the loader names remote code @@ -87,8 +87,7 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkey """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.""" + still carries the compressed-tensors scheme; the remote-code staging shim is left installed.""" installs = [] monkeypatch.setattr( _onfly_decompress, @@ -98,7 +97,7 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkey staging = dmu.get_cached_module_file config = _vlm_config() - assert kimi_k25.prepare(config) is config + assert kimi_k25.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) is config text = config.text_config assert type(text) is DeepseekV3Config @@ -112,7 +111,6 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkey ) 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( @@ -135,6 +133,15 @@ def test_prepare_with_a_checkpoint_points_the_remote_modeling_at_native_deepseek assert bystander.DeepseekV3ForCausalLM is remote_lm +def test_a_transformers_too_old_to_decompress_is_refused_before_prepare(monkeypatch): + """The experts decompress as they load from transformers 5.8: the spec refuses an older + release before prepare swaps the language model's config.""" + monkeypatch.setattr(transformers, "__version__", "5.7.0") + + with pytest.raises(RuntimeError, match=r"transformers >= 5\.8"): + check_distributed_support("kimi_k25") + + 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.""" diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 9ce96c68155..a627def43aa 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -129,7 +129,7 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke staging = dmu.get_cached_module_file config = _vlm_config() - assert kimi_vl.prepare(config) is config + assert kimi_vl.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) is config text = config.text_config assert type(text) is DeepseekV3Config @@ -142,7 +142,6 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke ) assert import_utils.is_torch_fx_available() is False assert dmu.get_cached_module_file is not staging - assert remote.loads == [] @pytest.mark.parametrize( @@ -202,11 +201,14 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch ("config", "kwargs", "match"), [ (SimpleNamespace, {}, "text_config"), + # Without the checkpoint there is no modeling code to patch, and a half-prepared config + # would load Kimi's own language model under a plan naming none of its experts. + (_vlm_config, {"trust_remote_code": True}, "needs the checkpoint's path"), # 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"], + ids=["no-text-config", "no-checkpoint", "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 From 06eb5bc49ba3032ccb5a11f52e8c83d924164d6b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:20:43 +0000 Subject: [PATCH 11/20] Install Kimi-K2.5's decompressor before its text config is replaced Kimi-K2.5, like Kimi-K2.6 and Kimi-K2.7-Code, carries its compressed-tensors quantization_config only on text_config. prepare replaced text_config with a native DeepseekV3Config, which drops it, before it installed the on-the-fly decompressor, so the decompressor found no scheme and refused the compressed checkpoint ("config has no compressed-tensors quantization_config"). prepare now installs it first, ahead of the compat shims and the remote-code staging too, so a checkpoint the decompressor refuses leaves the config and the loader as they were. The adapter test's stand-in config carries the scheme only on its text config, as the checkpoint does. The kimi_k25 and kimi_vl specs say that only their expert-parallel section is registered: the MoE and export sections stay deepseek_v3's. The kimi_k25 docstrings name Kimi-K2.5, whose model type Kimi-K2.6 and Kimi-K2.7-Code share, and Kimi-VL's version note states the 5.6 floor it is held to. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 20 ++++++----- modelopt/torch/models/kimi_k25/specs.py | 8 +++-- modelopt/torch/models/kimi_vl/ep.py | 4 +-- modelopt/torch/models/kimi_vl/specs.py | 3 ++ tests/unit/torch/models/test_kimi_k25_ep.py | 37 +++++++++++++++++---- 5 files changed, 52 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index fc277ce3836..41edee2a966 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -13,9 +13,10 @@ # 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``). +"""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`` / @@ -58,7 +59,7 @@ 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. + """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 expert decompressor. @@ -85,25 +86,26 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool "load it with trust_remote_code=True." ) + # 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. + _onfly_decompress.install_onfly_decompress(hf_config) + # 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. + # 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) - # 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}`. diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index b026101f0c3..ee7f2451736 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -15,8 +15,12 @@ """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. +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. Sharing them is left to a follow-up. """ from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index edb9b99fbe7..c2212f70c6f 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -39,8 +39,8 @@ 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. +Version: no floor beyond the EP framework's 5.6; validated on transformers 5.9.0. The compat shims +stand in for the symbols the remote code imports that later releases dropped. """ import sys diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py index 7742f4251b1..a5d9eacec42 100644 --- a/modelopt/torch/models/kimi_vl/specs.py +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -18,6 +18,9 @@ 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. + +Only the expert-parallel section is registered under ``kimi_vl``: the MoE and export sections stay +``deepseek_v3``'s, so lookups by ``kimi_vl`` find neither yet. Sharing them is left to a follow-up. """ from ..deepseek_v3.specs import 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 index 568959256b1..79b95012e6f 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -13,7 +13,7 @@ # 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. +"""Expert-parallel adapter for Kimi-K2.5 (``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 @@ -45,16 +45,17 @@ from modelopt.torch.utils.moe import _routed_expert_keys, check_distributed_support QUANT = {"quant_method": "compressed-tensors", "format": "pack-quantized"} -REMOTE = "transformers_modules.Kimi_hyphen_K2_dot_7.0123abcd" # how the loader names remote code +REMOTE = "transformers_modules.Kimi_hyphen_K2_dot_5.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``.""" + """Stands in for the remote DeepSeek-V3 config Kimi-K2.5 bundles as its ``text_config``.""" model_type = "kimi_k2_remote_text" def _vlm_config(): + """Kimi-K2.5's config: the compressed-tensors scheme sits only on the text config.""" text = _RemoteTextConfig( hidden_size=64, num_hidden_layers=2, @@ -63,7 +64,7 @@ def _vlm_config(): auto_map={"AutoConfig": "configuration_deepseek.DeepseekV3Config"}, quantization_config=QUANT.copy(), ) - return SimpleNamespace(text_config=text, quantization_config=QUANT.copy()) + return SimpleNamespace(text_config=text) @pytest.fixture @@ -86,13 +87,14 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): 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.""" + spec states for a run that skips prepare; the decompressor is installed while the text config + it replaces still carries the compressed-tensors scheme; the remote-code staging shim is left + installed.""" installs = [] monkeypatch.setattr( _onfly_decompress, "install_onfly_decompress", - lambda cfg: installs.append(getattr(cfg, "quantization_config", None)), + lambda cfg: installs.append(_onfly_decompress._find_quantization_config(cfg)), ) staging = dmu.get_cached_module_file config = _vlm_config() @@ -133,6 +135,27 @@ def test_prepare_with_a_checkpoint_points_the_remote_modeling_at_native_deepseek assert bystander.DeepseekV3ForCausalLM is remote_lm +def test_a_checkpoint_the_decompressor_refuses_is_refused_before_anything_changes( + monkeypatch, remote +): + """The decompressor refuses a checkpoint it cannot decompress as it loads before prepare + replaces the language model's config or wraps the loader's remote-code staging.""" + + def refuse(cfg): + raise NotImplementedError("the checkpoint has 2 compressed-tensors config groups") + + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", refuse) + staging = dmu.get_cached_module_file + config = _vlm_config() + text = config.text_config + + with pytest.raises(NotImplementedError, match="config groups"): + kimi_k25.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) + + assert config.text_config is text + assert dmu.get_cached_module_file is staging + + def test_a_transformers_too_old_to_decompress_is_refused_before_prepare(monkeypatch): """The experts decompress as they load from transformers 5.8: the spec refuses an older release before prepare swaps the language model's config.""" From 6ed604cb7b50a74fd55163708a6ca6b8723f899a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:14:18 +0000 Subject: [PATCH 12/20] Print the Kimi adapters' notes on rank 0 only, under a shared label prepare runs on every rank (resolve_distribution_plan applies the model spec on each), so the Kimi-K2.5 and Kimi-VL patch notes and the compat-shim warnings printed once per rank; they now print on rank 0, as the decompressor's install note does. Both adapters use the shims, so their warnings and the attention-mask stub's error say kimi rather than kimi_vl. The deepseek_v3 plan's comment says which releases it serves: transformers ships no deepseek_v3 plan before 5.13, and from 5.13 its own wins, with the same experts and an ep_router gate entry the load swaps for torch_a2a_router. The native-port refusal test passes trust_remote_code=True and matches the part of the message only that refusal has, so it fails without the guard instead of matching the trust_remote_code refusal. Signed-off-by: Shengliang Xu --- modelopt/torch/models/deepseek_v3/specs.py | 8 +++++--- modelopt/torch/models/kimi_k25/ep.py | 6 +++--- modelopt/torch/models/kimi_vl/ep.py | 16 +++++++--------- tests/unit/torch/models/test_kimi_k25_ep.py | 4 ++-- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/models/deepseek_v3/specs.py b/modelopt/torch/models/deepseek_v3/specs.py index d1b20b1338b..b7074c80fbd 100644 --- a/modelopt/torch/models/deepseek_v3/specs.py +++ b/modelopt/torch/models/deepseek_v3/specs.py @@ -23,9 +23,11 @@ __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. +# The plan for transformers releases that ship none for deepseek_v3 (before 5.13). From 5.13 its own +# plan wins (an existing plan always does, see apply_distributed_support): the same experts, plus an +# ep_router entry for the gate, which the load swaps for torch_a2a_router (global ids, unmasked). No +# gate entry here 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", diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 41edee2a966..7f0b8bf7d9a 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -51,6 +51,7 @@ import sys +from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import apply_distributed_support from .. import _onfly_decompress @@ -154,7 +155,6 @@ def _patch_kimi_modeling_to_native(ckpt_path: str, trust_remote_code: bool) -> N 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, + print_rank_0( + f"[kimi_k25] monkey-patched DeepseekV3ForCausalLM -> transformers-native in: {patched}" ) diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index c2212f70c6f..5a032a92160 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -45,6 +45,7 @@ import sys +from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import apply_distributed_support from .. import _onfly_decompress @@ -63,14 +64,14 @@ def _install_compat_import_shims(): 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) + print_rank_0(f"[kimi] compat-shim warning (import_utils): {type(e).__name__}: {e}") 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) + print_rank_0(f"[kimi] compat-shim warning (pytorch_utils): {type(e).__name__}: {e}") try: import transformers.modeling_attn_mask_utils as amu @@ -78,15 +79,13 @@ def _install_compat_import_shims(): 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 " + "kimi: 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 - ) + print_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}") def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): @@ -168,8 +167,7 @@ def _patch_kimi_vl_modeling_to_native(ckpt_path: str, trust_remote_code: bool) - raise RuntimeError( "kimi_vl: could not find a loaded modeling_kimi_vl in sys.modules to monkey-patch" ) - print( + print_rank_0( f"[kimi_vl] monkey-patched DeepseekV3ForCausalLM -> transformers-native (+tie_weights shim) in: " - f"{patched}", - flush=True, + f"{patched}" ) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 79b95012e6f..12b66bcc4a6 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -169,5 +169,5 @@ 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") + with pytest.raises(ValueError, match="runs on the checkpoint's own modeling code"): + kimi_k25.prepare(native.Kimi_K25Config(), ckpt_path="/ckpt", trust_remote_code=True) From 583097be9b40d746d68725d89352b82bd10b47c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:27:32 +0000 Subject: [PATCH 13/20] Keep the Kimi adapters' native language model patched on a module reload Each Kimi adapter repointed its remote module's DeepseekV3ForCausalLM at transformers' native class once, in prepare. transformers re-executes a remote module in place when the files it hashes change, and loads it afresh for another path to the same checkpoint; either rebinds the module's globals, so from_pretrained built the checkpoint's own per-expert model on a config already swapped to the native one. An expert-parallel Kimi-K2.7-Code load over 32 GPUs failed so, on some ranks, reading rope_theta, which the native config keeps under rope_parameters. keep_patched applies a patch to every loaded module with a name suffix, and again whenever transformers resolves a class from one, so the module from_pretrained builds from carries it. Both Kimi adapters use it for the native class and the tie_weights shim. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_remote_code.py | 40 +++++++++++++++++++++ modelopt/torch/models/kimi_k25/ep.py | 36 +++++++++---------- modelopt/torch/models/kimi_vl/ep.py | 34 +++++++++--------- tests/unit/torch/models/test_remote_code.py | 35 ++++++++++++++++-- 4 files changed, 106 insertions(+), 39 deletions(-) diff --git a/modelopt/torch/models/_remote_code.py b/modelopt/torch/models/_remote_code.py index 7db39497bab..ac6c2f6e6f7 100644 --- a/modelopt/torch/models/_remote_code.py +++ b/modelopt/torch/models/_remote_code.py @@ -18,8 +18,11 @@ import glob import os import shutil +import sys import tempfile import warnings +from collections.abc import Callable +from types import ModuleType def stage_full_remote_closure() -> None: @@ -65,3 +68,40 @@ def _staged(pretrained_model_name_or_path, module_file, *args, **kwargs): setattr(_staged, "_modelopt_full_closure", True) dmu.get_cached_module_file = _staged + + +# Patches keep_patched applies to remote modules, by module-name suffix. +_KEPT_PATCHES: dict[str, Callable[[ModuleType], None]] = {} + + +def keep_patched(suffix: str, patch: Callable[[ModuleType], None]) -> list[str]: + """Apply ``patch`` to every loaded remote module named ``*``, and again whenever + transformers resolves a class from one. + + transformers re-executes a remote module in place when the files it hashes change, and loads it + afresh under another name for another path to the same checkpoint; either way the module's + globals are rebound and a patch applied once is gone. Wrapping ``get_class_in_module`` re-applies + the patch to the module a class is resolved from, so ``from_pretrained`` builds from the patched + module. Returns the names of the modules patched now. Idempotent per suffix. + """ + import transformers.dynamic_module_utils as dmu + + _KEPT_PATCHES[suffix] = patch + if not getattr(dmu.get_class_in_module, "_modelopt_keep_patched", False): + orig = dmu.get_class_in_module + + def _resolved(class_name, module_path, *args, **kwargs): + cls = orig(class_name, module_path, *args, **kwargs) + for kept_suffix, kept_patch in _KEPT_PATCHES.items(): + if cls.__module__.endswith(kept_suffix): + kept_patch(sys.modules[cls.__module__]) + return cls + + setattr(_resolved, "_modelopt_keep_patched", True) + dmu.get_class_in_module = _resolved + patched = [] + for name, module in list(sys.modules.items()): + if module is not None and name.endswith(suffix): + patch(module) + patched.append(name) + return patched diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 7f0b8bf7d9a..e16be66f7e0 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -49,13 +49,11 @@ transformers' still-evolving ``WeightConverter`` conversion-ops API. """ -import sys - from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import apply_distributed_support from .. import _onfly_decompress -from .._remote_code import stage_full_remote_closure +from .._remote_code import keep_patched, stage_full_remote_closure from ..kimi_vl.ep import _drop_tie_weights_kwargs, _install_compat_import_shims @@ -135,22 +133,22 @@ def _patch_kimi_modeling_to_native(ckpt_path: str, trust_remote_code: bool) -> N 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) + def to_native(module): + if not hasattr(module, "DeepseekV3ForCausalLM"): + return + module.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(module, "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) + + # The dynamic-module loader names the module like + # ``transformers_modules...modeling_kimi_k25``. It stays patched when + # from_pretrained resolves the class again, even from a re-executed module. + patched = keep_patched(".modeling_kimi_k25", to_native) if not patched: raise RuntimeError( "kimi_k25: could not find a loaded modeling_kimi_k25 in sys.modules to monkey-patch" diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 5a032a92160..1378d61fd94 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -43,13 +43,11 @@ stand in for the symbols the remote code imports that later releases dropped. """ -import sys - from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import apply_distributed_support from .. import _onfly_decompress -from .._remote_code import stage_full_remote_closure +from .._remote_code import keep_patched, stage_full_remote_closure def _install_compat_import_shims(): @@ -148,21 +146,21 @@ def _patch_kimi_vl_modeling_to_native(ckpt_path: str, trust_remote_code: bool) - 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) + def to_native(module): + if not hasattr(module, "DeepseekV3ForCausalLM"): + return + module.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(module, "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) + + # Kept patched when from_pretrained resolves the class again, even from a re-executed module. + patched = keep_patched(".modeling_kimi_vl", to_native) if not patched: raise RuntimeError( "kimi_vl: could not find a loaded modeling_kimi_vl in sys.modules to monkey-patch" diff --git a/tests/unit/torch/models/test_remote_code.py b/tests/unit/torch/models/test_remote_code.py index d324773e47b..4b2f9726e30 100644 --- a/tests/unit/torch/models/test_remote_code.py +++ b/tests/unit/torch/models/test_remote_code.py @@ -13,9 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Staging every module of a local checkpoint's remote code.""" +"""Staging every module of a local checkpoint's remote code, and keeping a patch on it.""" import shutil +import sys from types import SimpleNamespace import pytest @@ -23,7 +24,7 @@ dmu = pytest.importorskip("transformers.dynamic_module_utils") from modelopt.torch.models import _remote_code -from modelopt.torch.models._remote_code import stage_full_remote_closure +from modelopt.torch.models._remote_code import keep_patched, stage_full_remote_closure def test_every_module_of_a_local_checkpoint_is_staged(monkeypatch, tmp_path): @@ -84,3 +85,33 @@ def copy_cut_short(src, dst): dmu.get_cached_module_file(str(ckpt), "a.py") assert (stage / "b.py").read_text() == (ckpt / "b.py").read_text() + + +def test_a_patch_stays_on_a_module_transformers_loads_again(monkeypatch, tmp_path): + """transformers re-executes a remote module it loads again (when the files it hashes change, or + forced as here), rebinding the module's globals. The patch must be back on the module the class + is resolved from, which is the one from_pretrained builds from.""" + monkeypatch.setattr(dmu, "HF_MODULES_CACHE", str(tmp_path / "modules")) + monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) + monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) + loaded = set(sys.modules) + ckpt = tmp_path / "ckpt" + ckpt.mkdir() + (ckpt / "modeling_kimi_k25.py").write_text( + "class DeepseekV3ForCausalLM:\n pass\n\n\nclass KimiK25ForConditionalGeneration:\n pass\n" + ) + native = type("DeepseekV3ForCausalLM", (), {}) + ref = "modeling_kimi_k25.KimiK25ForConditionalGeneration" + try: + first = dmu.get_class_from_dynamic_module(ref, str(ckpt)) + patched = keep_patched( + ".modeling_kimi_k25", lambda module: setattr(module, "DeepseekV3ForCausalLM", native) + ) + again = dmu.get_class_from_dynamic_module(ref, str(ckpt), force_download=True) + + assert patched == [first.__module__] + assert again is not first # the module was executed again + assert sys.modules[again.__module__].DeepseekV3ForCausalLM is native + finally: + for name in set(sys.modules) - loaded: + del sys.modules[name] From f40e7039edddb5722712cd2a8ae028b856849269 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:32:37 +0000 Subject: [PATCH 14/20] Report a remote-code patch re-applied after a reload, on its rank keep_patched's patch now returns whether it changed the module. A re-application that did means transformers loaded the module again after it was patched, so the rank prints which module, and a run shows whether the reload happened rather than only that the load succeeded. The test also covers a second path to the same checkpoint, which transformers loads under another module name. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_remote_code.py | 17 ++++++++---- modelopt/torch/models/kimi_k25/ep.py | 6 +++-- modelopt/torch/models/kimi_vl/ep.py | 6 +++-- tests/unit/torch/models/test_remote_code.py | 29 +++++++++++++++------ 4 files changed, 41 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/models/_remote_code.py b/modelopt/torch/models/_remote_code.py index ac6c2f6e6f7..6214cb9f4da 100644 --- a/modelopt/torch/models/_remote_code.py +++ b/modelopt/torch/models/_remote_code.py @@ -71,10 +71,10 @@ def _staged(pretrained_model_name_or_path, module_file, *args, **kwargs): # Patches keep_patched applies to remote modules, by module-name suffix. -_KEPT_PATCHES: dict[str, Callable[[ModuleType], None]] = {} +_KEPT_PATCHES: dict[str, Callable[[ModuleType], bool]] = {} -def keep_patched(suffix: str, patch: Callable[[ModuleType], None]) -> list[str]: +def keep_patched(suffix: str, patch: Callable[[ModuleType], bool]) -> list[str]: """Apply ``patch`` to every loaded remote module named ``*``, and again whenever transformers resolves a class from one. @@ -82,7 +82,8 @@ def keep_patched(suffix: str, patch: Callable[[ModuleType], None]) -> list[str]: afresh under another name for another path to the same checkpoint; either way the module's globals are rebound and a patch applied once is gone. Wrapping ``get_class_in_module`` re-applies the patch to the module a class is resolved from, so ``from_pretrained`` builds from the patched - module. Returns the names of the modules patched now. Idempotent per suffix. + module. ``patch`` returns whether it changed the module; a re-application that did is reported on + the rank it happened on. Returns the names of the modules patched now. Idempotent per suffix. """ import transformers.dynamic_module_utils as dmu @@ -93,8 +94,14 @@ def keep_patched(suffix: str, patch: Callable[[ModuleType], None]) -> list[str]: def _resolved(class_name, module_path, *args, **kwargs): cls = orig(class_name, module_path, *args, **kwargs) for kept_suffix, kept_patch in _KEPT_PATCHES.items(): - if cls.__module__.endswith(kept_suffix): - kept_patch(sys.modules[cls.__module__]) + if cls.__module__.endswith(kept_suffix) and kept_patch(sys.modules[cls.__module__]): + from modelopt.torch.utils import distributed as dist + + print( + f"[rank {dist.rank()}] re-patched {cls.__module__}, which transformers loaded " + "again after it was patched", + flush=True, + ) return cls setattr(_resolved, "_modelopt_keep_patched", True) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index e16be66f7e0..c23aa9141f2 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -134,8 +134,9 @@ def _patch_kimi_modeling_to_native(ckpt_path: str, trust_remote_code: bool) -> N ) def to_native(module): - if not hasattr(module, "DeepseekV3ForCausalLM"): - return + current = getattr(module, "DeepseekV3ForCausalLM", NativeDeepseekV3ForCausalLM) + if current is NativeDeepseekV3ForCausalLM: + return False module.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. @@ -144,6 +145,7 @@ def to_native(module): 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) + return True # The dynamic-module loader names the module like # ``transformers_modules...modeling_kimi_k25``. It stays patched when diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 1378d61fd94..04c4b8bc1a2 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -147,8 +147,9 @@ def _patch_kimi_vl_modeling_to_native(ckpt_path: str, trust_remote_code: bool) - ) def to_native(module): - if not hasattr(module, "DeepseekV3ForCausalLM"): - return + current = getattr(module, "DeepseekV3ForCausalLM", NativeDeepseekV3ForCausalLM) + if current is NativeDeepseekV3ForCausalLM: + return False module.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 @@ -158,6 +159,7 @@ def to_native(module): 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) + return True # Kept patched when from_pretrained resolves the class again, even from a re-executed module. patched = keep_patched(".modeling_kimi_vl", to_native) diff --git a/tests/unit/torch/models/test_remote_code.py b/tests/unit/torch/models/test_remote_code.py index 4b2f9726e30..05c2c8cd99e 100644 --- a/tests/unit/torch/models/test_remote_code.py +++ b/tests/unit/torch/models/test_remote_code.py @@ -87,10 +87,14 @@ def copy_cut_short(src, dst): assert (stage / "b.py").read_text() == (ckpt / "b.py").read_text() -def test_a_patch_stays_on_a_module_transformers_loads_again(monkeypatch, tmp_path): +@pytest.mark.parametrize("load_again", ["forced", "another-path"]) +def test_a_patch_stays_on_a_module_transformers_loads_again( + monkeypatch, tmp_path, capsys, load_again +): """transformers re-executes a remote module it loads again (when the files it hashes change, or - forced as here), rebinding the module's globals. The patch must be back on the module the class - is resolved from, which is the one from_pretrained builds from.""" + forced), and loads a second copy under another name for another path to the same checkpoint; + either way the module's globals are fresh. The patch must be back on the module the class is + resolved from, which is the one from_pretrained builds from, and the re-patch is reported.""" monkeypatch.setattr(dmu, "HF_MODULES_CACHE", str(tmp_path / "modules")) monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) @@ -100,18 +104,27 @@ def test_a_patch_stays_on_a_module_transformers_loads_again(monkeypatch, tmp_pat (ckpt / "modeling_kimi_k25.py").write_text( "class DeepseekV3ForCausalLM:\n pass\n\n\nclass KimiK25ForConditionalGeneration:\n pass\n" ) + (tmp_path / "same_ckpt").symlink_to(ckpt) native = type("DeepseekV3ForCausalLM", (), {}) + + def to_native(module): + changed = module.DeepseekV3ForCausalLM is not native + module.DeepseekV3ForCausalLM = native + return changed + ref = "modeling_kimi_k25.KimiK25ForConditionalGeneration" try: first = dmu.get_class_from_dynamic_module(ref, str(ckpt)) - patched = keep_patched( - ".modeling_kimi_k25", lambda module: setattr(module, "DeepseekV3ForCausalLM", native) - ) - again = dmu.get_class_from_dynamic_module(ref, str(ckpt), force_download=True) + patched = keep_patched(".modeling_kimi_k25", to_native) + if load_again == "forced": + again = dmu.get_class_from_dynamic_module(ref, str(ckpt), force_download=True) + else: + again = dmu.get_class_from_dynamic_module(ref, str(tmp_path / "same_ckpt")) assert patched == [first.__module__] - assert again is not first # the module was executed again + assert again is not first # a fresh module, not the one patched above assert sys.modules[again.__module__].DeepseekV3ForCausalLM is native + assert f"re-patched {again.__module__}" in capsys.readouterr().out finally: for name in set(sys.modules) - loaded: del sys.modules[name] From aad520702db8c3aeaa2ab5264abd10f9f1535fb9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:57:03 +0000 Subject: [PATCH 15/20] Decompress kimi_k25 experts only for a compressed-tensors checkpoint prepare installed the on-the-fly decompressor for every kimi_k25 checkpoint, and the decompressor refuses one without a compressed-tensors quantization_config: a bf16 copy stopped on "config has no compressed-tensors quantization_config", a block-fp8 one on pydantic's ValidationError. The expert-parallel load reads both itself: bf16 experts as stored, block-fp8 ones dequantized through the quantization_config Kimi's config class copies to the top level. prepare now installs the decompressor only for compressed-tensors experts, which Kimi-K2.5, Kimi-K2.6 and Kimi-K2.7-Code all ship. Kimi-VL's tie_weights shim also says why dropping missing_keys is safe: the remote tie_weights only calls its language model's, and the Kimi checkpoints set tie_word_embeddings=False. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 11 +++++++--- modelopt/torch/models/kimi_vl/ep.py | 7 +++++- tests/unit/torch/models/test_kimi_k25_ep.py | 24 ++++++++++++++++----- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index c23aa9141f2..65d3840beac 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -61,7 +61,8 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool """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 expert decompressor. + DeepseekV3ForCausalLM + apply the EP plan + install the on-the-fly decompressor for + compressed-tensors experts. """ if type(hf_config).__module__.startswith("transformers."): # From 5.13, a load without trust_remote_code builds transformers' own port, which the @@ -89,8 +90,12 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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. - _onfly_decompress.install_onfly_decompress(hf_config) + # goes first, so a checkpoint it refuses leaves everything as it was. Experts stored + # otherwise are the expert-parallel load's to read: bf16 as they are, block-fp8 by the copy + # of quantization_config Kimi's config class keeps on the top-level config. + quantization_config = _onfly_decompress._find_quantization_config(hf_config) or {} + if quantization_config.get("quant_method") == "compressed-tensors": + _onfly_decompress.install_onfly_decompress(hf_config) # 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 diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 04c4b8bc1a2..2bdda1ea789 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -121,7 +121,12 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool def _drop_tie_weights_kwargs(orig_tie): - """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it.""" + """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it. + + Dropping ``missing_keys`` as well as ``recompute_mapping`` is safe because the remote + ``tie_weights`` only calls its language model's, and the Kimi checkpoints set + ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. + """ def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc. return orig_tie(self) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 12b66bcc4a6..575f95b2269 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -84,12 +84,25 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): return module -def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkeypatch, remote): +@pytest.mark.parametrize( + "quantization_config", + [ + QUANT, + # Left to the expert-parallel load: bf16 experts load as they are, and block-fp8 ones are + # dequantized by the copy of quantization_config Kimi's config class keeps on the top level. + None, + {"quant_method": "fp8", "weight_block_size": [128, 128]}, + ], + ids=["compressed-tensors", "bf16", "block-fp8"], +) +def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3( + monkeypatch, remote, quantization_config +): """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 text config - it replaces still carries the compressed-tensors scheme; the remote-code staging shim is left - installed.""" + spec states for a run that skips prepare; the decompressor is installed for compressed-tensors + experts only, while the text config it replaces still carries their scheme; the remote-code + staging shim is left installed.""" installs = [] monkeypatch.setattr( _onfly_decompress, @@ -98,6 +111,7 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkey ) staging = dmu.get_cached_module_file config = _vlm_config() + config.text_config.quantization_config = quantization_config assert kimi_k25.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) is config @@ -111,7 +125,7 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3(monkey 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 installs == ([QUANT] if quantization_config is QUANT else []) assert dmu.get_cached_module_file is not staging From c0e62f9d3bb6d9e2cb390ed6f7ca509ca59d9990 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 02:59:22 +0000 Subject: [PATCH 16/20] Say the Kimi-K2.5 adapter's version floor is the spec's Since this slice installs the decompressor only for a compressed-tensors checkpoint, the 5.8 floor is the one the spec declares; the module docstring said the decompressor enforces it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 65d3840beac..4be2a6d27de 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -44,8 +44,9 @@ ``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 +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 f8386d8b2dc1cafbcba30a702f8e777ad20ad112 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:47:11 +0000 Subject: [PATCH 17/20] Stop a Kimi load whose remote modeling has no DeepseekV3ForCausalLM Both Kimi adapters read the remote module's DeepseekV3ForCausalLM with the native class as the default, so a module that no longer defines that name (renamed, or imported inside the VLM class's __init__) looked already patched: the swap and the tie_weights wrapper were skipped, the adapter reported success, and the load went on to build the per-expert MoE under a plan naming fused experts. The swap now raises for such a module. keep_patched scans remote-code modules only, so transformers' own port, which shares the module name from 5.13 (transformers.models.kimi_k25.modeling_kimi_k25) and defines no such class, is left alone. The two adapters' patchers were one function three names apart, and Kimi-K2.5's imported Kimi-VL's private shims. Both now live in modelopt/torch/models/_kimi.py with the refusals the adapters share, so Kimi-VL also refuses transformers' own config, and that refusal names the port's module and the auto_map the checkpoint needs. The Kimi test fixtures restore the get_class_in_module wrapper and the kept patches; the Kimi-K2.5 file leaves out the compat shims, which the Kimi-VL file covers for both adapters. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 161 ++++++++++++++++++++ modelopt/torch/models/_remote_code.py | 5 +- modelopt/torch/models/kimi_k25/ep.py | 80 ++-------- modelopt/torch/models/kimi_vl/ep.py | 122 ++------------- tests/unit/torch/models/test_kimi_k25_ep.py | 9 +- tests/unit/torch/models/test_kimi_vl_ep.py | 47 +++--- tests/unit/torch/models/test_remote_code.py | 18 ++- 7 files changed, 239 insertions(+), 203 deletions(-) create mode 100644 modelopt/torch/models/_kimi.py diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py new file mode 100644 index 00000000000..ab42956aba8 --- /dev/null +++ b/modelopt/torch/models/_kimi.py @@ -0,0 +1,161 @@ +# 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. + +"""What the Kimi-K2.5 and Kimi-VL adapters change in the modeling code their checkpoints ship. + +Both checkpoints' own VLM class builds its language model as ``DeepseekV3ForCausalLM( +config.text_config)``, a DeepSeek-V3 whose MoE is an ``nn.ModuleList`` of per-expert Linears that +expert parallelism cannot shard. The adapters point that name at transformers' native class, and +keep the vision tower and projector from the checkpoint's code. +""" + +from modelopt.torch.utils import print_rank_0 + +from ._remote_code import keep_patched + + +def checkpoint_to_patch( + label: str, hf_config, ckpt_path: str | None, trust_remote_code: bool +) -> str: + """Return the checkpoint whose own modeling code the adapter ``label`` patches. + + Refuses first, before its ``prepare`` changes anything, a config the adapter cannot patch. + """ + if type(hf_config).__module__.startswith("transformers."): + # transformers' own port (Kimi-K2.5's, from 5.13), which the patch never reaches: a load + # without trust_remote_code builds it, and so does one of a checkpoint naming no config class. + raise ValueError( + f"{label}: expert parallelism runs on the checkpoint's own modeling code, but this " + f"config is transformers' own port ({type(hf_config).__module__}). Load it with " + "trust_remote_code=True, from a checkpoint whose auto_map names its own AutoConfig." + ) + if not hasattr(hf_config, "text_config"): + raise ValueError(f"{label}: hf_config has no .text_config (expected a VLM config)") + if ckpt_path is None: + raise ValueError( + f"{label}: expert parallelism patches the checkpoint's own modeling code, so it needs " + "the checkpoint's path (apply_distributed_support's ckpt_path)." + ) + if not trust_remote_code: + # Patching imports the checkpoint's own modeling code, and transformers' + # get_class_from_dynamic_module runs it without checking trust_remote_code itself. + raise ValueError( + f"{label}: expert parallelism imports the checkpoint's own modeling code; " + "load it with trust_remote_code=True." + ) + return ckpt_path + + +def install_compat_import_shims() -> None: + """Add stubs for symbols the Kimi modeling imports that newer transformers removed, so it 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_rank_0(f"[kimi] compat-shim warning (import_utils): {type(e).__name__}: {e}") + 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_rank_0(f"[kimi] compat-shim warning (pytorch_utils): {type(e).__name__}: {e}") + 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: 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_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}") + + +def _drop_tie_weights_kwargs(orig_tie): + """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it. + + Dropping ``missing_keys`` as well as ``recompute_mapping`` is safe because the remote + ``tie_weights`` only calls its language model's, and the Kimi checkpoints set + ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. + """ + + def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc. + return orig_tie(self) + + setattr(_tie_weights, "_kimi_tie_weights_shim", True) + return _tie_weights + + +def patch_language_model_to_native( + label: str, modeling: str, vlm_class: str, ckpt_path: str, trust_remote_code: bool +) -> None: + """Point the checkpoint's ``modeling`` module's ``DeepseekV3ForCausalLM`` at the native class. + + Loads the checkpoint's remote ``modeling`` module and replaces its ``DeepseekV3ForCausalLM`` with + transformers' native class, so ``vlm_class.__init__`` builds the EP-capable native deepseek_v3 + for ``self.language_model`` instead of the 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 the module into sys.modules. + get_class_from_dynamic_module( + f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code + ) + + def to_native(module): + current = getattr(module, "DeepseekV3ForCausalLM", None) + if current is None: + # Not "already native": the VLM class then builds its language model from a class this + # swap does not reach (renamed, or imported inside __init__), with per-expert Linears. + raise RuntimeError( + f"{label}: {module.__name__} defines no module-level DeepseekV3ForCausalLM to swap " + "for the transformers-native class; the checkpoint's modeling code changed." + ) + if current is NativeDeepseekV3ForCausalLM: + return False + module.DeepseekV3ForCausalLM = 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) + return True + + # The dynamic-module loader names the module like + # ``transformers_modules...``. It stays patched when from_pretrained + # resolves the class again, even from a re-executed module. + patched = keep_patched(f".{modeling}", to_native) + if not patched: + raise RuntimeError( + f"{label}: could not find a loaded {modeling} in sys.modules to monkey-patch" + ) + print_rank_0( + f"[{label}] monkey-patched DeepseekV3ForCausalLM -> transformers-native (+tie_weights shim) " + f"in: {patched}" + ) diff --git a/modelopt/torch/models/_remote_code.py b/modelopt/torch/models/_remote_code.py index 6214cb9f4da..b2755779978 100644 --- a/modelopt/torch/models/_remote_code.py +++ b/modelopt/torch/models/_remote_code.py @@ -107,8 +107,11 @@ def _resolved(class_name, module_path, *args, **kwargs): setattr(_resolved, "_modelopt_keep_patched", True) dmu.get_class_in_module = _resolved patched = [] + # Remote modules only: transformers' own port of a model can share the module's name, as + # transformers.models.kimi_k25.modeling_kimi_k25 does from 5.13. + remote = f"{dmu.TRANSFORMERS_DYNAMIC_MODULE_NAME}." for name, module in list(sys.modules.items()): - if module is not None and name.endswith(suffix): + if module is not None and name.startswith(remote) and name.endswith(suffix): patch(module) patched.append(name) return patched diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 4be2a6d27de..e298923887e 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -50,12 +50,10 @@ transformers' still-evolving ``WeightConverter`` conversion-ops API. """ -from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import apply_distributed_support -from .. import _onfly_decompress -from .._remote_code import keep_patched, stage_full_remote_closure -from ..kimi_vl.ep import _drop_tie_weights_kwargs, _install_compat_import_shims +from .. import _kimi, _onfly_decompress +from .._remote_code import stage_full_remote_closure def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): @@ -65,27 +63,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool DeepseekV3ForCausalLM + apply the EP plan + install the on-the-fly decompressor for compressed-tensors experts. """ - 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 None: - raise ValueError( - "kimi_k25: expert parallelism patches the checkpoint's own modeling code, so it needs " - "the checkpoint's path (apply_distributed_support's ckpt_path)." - ) - if 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." - ) + 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 @@ -101,7 +79,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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() + _kimi.install_compat_import_shims() stage_full_remote_closure() # 2. Replace text_config with a transformers-native DeepseekV3Config built from the same fields. @@ -118,49 +96,11 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 4. Force-load Kimi's modeling_kimi_k25 module (so we can patch its DeepseekV3ForCausalLM # symbol) and swap to the native class. - _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", + _kimi.patch_language_model_to_native( + "kimi_k25", + "modeling_kimi_k25", + "KimiK25ForConditionalGeneration", ckpt_path, - trust_remote_code=trust_remote_code, - ) - - def to_native(module): - current = getattr(module, "DeepseekV3ForCausalLM", NativeDeepseekV3ForCausalLM) - if current is NativeDeepseekV3ForCausalLM: - return False - module.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(module, "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) - return True - - # The dynamic-module loader names the module like - # ``transformers_modules...modeling_kimi_k25``. It stays patched when - # from_pretrained resolves the class again, even from a re-executed module. - patched = keep_patched(".modeling_kimi_k25", to_native) - if not patched: - raise RuntimeError( - "kimi_k25: could not find a loaded modeling_kimi_k25 in sys.modules to monkey-patch" - ) - print_rank_0( - f"[kimi_k25] monkey-patched DeepseekV3ForCausalLM -> transformers-native in: {patched}" + trust_remote_code, ) + return hf_config diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 2bdda1ea789..87975e18eaf 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -43,47 +43,10 @@ stand in for the symbols the remote code imports that later releases dropped. """ -from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import apply_distributed_support -from .. import _onfly_decompress -from .._remote_code import keep_patched, 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_rank_0(f"[kimi] compat-shim warning (import_utils): {type(e).__name__}: {e}") - 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_rank_0(f"[kimi] compat-shim warning (pytorch_utils): {type(e).__name__}: {e}") - 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: 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_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}") +from .. import _kimi, _onfly_decompress +from .._remote_code import stage_full_remote_closure def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): @@ -93,22 +56,9 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool 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 None: - raise ValueError( - "kimi_vl: expert parallelism patches the checkpoint's own modeling code, so it needs " - "the checkpoint's path (apply_distributed_support's ckpt_path)." - ) - if 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." - ) + ckpt_path = _kimi.checkpoint_to_patch("kimi_vl", hf_config, ckpt_path, trust_remote_code) - _install_compat_import_shims() + _kimi.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). @@ -116,63 +66,11 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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) - _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. - - Dropping ``missing_keys`` as well as ``recompute_mapping`` is safe because the remote - ``tie_weights`` only calls its language model's, and the Kimi checkpoints set - ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. - """ - - 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", + _kimi.patch_language_model_to_native( + "kimi_vl", + "modeling_kimi_vl", + "KimiVLForConditionalGeneration", ckpt_path, - trust_remote_code=trust_remote_code, - ) - - def to_native(module): - current = getattr(module, "DeepseekV3ForCausalLM", NativeDeepseekV3ForCausalLM) - if current is NativeDeepseekV3ForCausalLM: - return False - module.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(module, "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) - return True - - # Kept patched when from_pretrained resolves the class again, even from a re-executed module. - patched = keep_patched(".modeling_kimi_vl", to_native) - if not patched: - raise RuntimeError( - "kimi_vl: could not find a loaded modeling_kimi_vl in sys.modules to monkey-patch" - ) - print_rank_0( - f"[kimi_vl] monkey-patched DeepseekV3ForCausalLM -> transformers-native (+tie_weights shim) in: " - f"{patched}" + trust_remote_code, ) + return hf_config diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 575f95b2269..40cb1cb00b6 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -40,7 +40,7 @@ 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 import _kimi, _onfly_decompress, _remote_code, get_spec from modelopt.torch.models.kimi_k25 import ep as kimi_k25 from modelopt.torch.utils.moe import _routed_expert_keys, check_distributed_support @@ -70,7 +70,9 @@ def _vlm_config(): @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.""" + whose ``DeepseekV3ForCausalLM`` is the per-expert remote language model. Also restores the + loader hooks and kept patches prepare installs, and leaves out its compat import shims, which + the Kimi-VL tests cover for both adapters.""" module = types.ModuleType(f"{REMOTE}.modeling_kimi_k25") module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) module.loads = [] @@ -81,6 +83,9 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) + monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) + monkeypatch.setattr(_kimi, "install_compat_import_shims", lambda: None) return module diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index a627def43aa..86c2e0af373 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -39,7 +39,7 @@ 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 import _onfly_decompress, _remote_code, 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 @@ -84,8 +84,8 @@ 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.""" + as both Kimi checkpoints ship it. Also restores the compat symbols, loader hooks and kept + patches prepare installs.""" modeling, vlm_name = getattr( request, "param", ("modeling_kimi_vl", "KimiVLForConditionalGeneration") ) @@ -112,6 +112,8 @@ def import_remote_class(class_ref, ckpt_path, **kwargs): monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) + monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) return module @@ -144,7 +146,8 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke assert dmu.get_cached_module_file is not staging -@pytest.mark.parametrize( +# Kimi-K2.5's adapter is Kimi-VL's with on-the-fly decompression added; these hold for both. +_BOTH_ADAPTERS = pytest.mark.parametrize( ("adapter", "remote"), [ (kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration")), @@ -153,6 +156,9 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke ids=["vl", "k25"], indirect=["remote"], ) + + +@_BOTH_ADAPTERS def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_weights( monkeypatch, adapter, remote ): @@ -177,29 +183,36 @@ def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_wei assert remote.tied == [model] -# Kimi-K2.5's adapter is Kimi-VL's with on-the-fly decompression added; these hold for both. - - +@_BOTH_ADAPTERS @pytest.mark.parametrize( - ("swap", "modeling"), + ("missing", "match"), [ - (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"), + ("module", "could not find a loaded modeling_kimi"), + ("class", "defines no module-level DeepseekV3ForCausalLM"), ], - ids=["vl", "k25"], + ids=["no-module", "no-class"], ) -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) +def test_the_swap_fails_loudly_when_the_remote_modeling_has_nothing_to_swap( + monkeypatch, adapter, remote, missing, match +): + """Skipping the swap silently would build the per-expert layout that EP cannot shard: when the + remote modeling did not load, or names its language model otherwise (a class renamed, or + imported where the VLM class builds it), which is not the swap already done.""" + if missing == "module": + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", lambda *args, **kwargs: None) + else: + del remote.DeepseekV3ForCausalLM - with pytest.raises(RuntimeError, match=modeling): - swap() + with pytest.raises(RuntimeError, match=match): + adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) @pytest.mark.parametrize("adapter", [kimi_vl, kimi_k25], ids=["vl", "k25"]) @pytest.mark.parametrize( ("config", "kwargs", "match"), [ + # transformers' own port (Kimi-K2.5's, from 5.13), which the patch never reaches. + (PretrainedConfig, {"ckpt_path": "/ckpt", "trust_remote_code": True}, "own port"), (SimpleNamespace, {}, "text_config"), # Without the checkpoint there is no modeling code to patch, and a half-prepared config # would load Kimi's own language model under a plan naming none of its experts. @@ -208,7 +221,7 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_did_not_load(monkeypatch # trust_remote_code, so prepare must refuse before reaching it. (_vlm_config, {"ckpt_path": "/ckpt"}, "trust_remote_code=True"), ], - ids=["no-text-config", "no-checkpoint", "untrusted-checkpoint"], + ids=["transformers-port", "no-text-config", "no-checkpoint", "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 diff --git a/tests/unit/torch/models/test_remote_code.py b/tests/unit/torch/models/test_remote_code.py index 05c2c8cd99e..2370079fc86 100644 --- a/tests/unit/torch/models/test_remote_code.py +++ b/tests/unit/torch/models/test_remote_code.py @@ -17,7 +17,7 @@ import shutil import sys -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import pytest @@ -128,3 +128,19 @@ def to_native(module): finally: for name in set(sys.modules) - loaded: del sys.modules[name] + + +def test_only_remote_modules_are_patched(monkeypatch): + """transformers' own port of a model can share its remote module's name (from 5.13, + ``transformers.models.kimi_k25.modeling_kimi_k25``); a patch for the checkpoint's code must + leave it alone.""" + monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) + monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) + remote = ModuleType("transformers_modules.Kimi.0123abcd.modeling_kimi_k25") + port = ModuleType("transformers.models.kimi_k25.modeling_kimi_k25") + for module in (remote, port): + monkeypatch.setitem(sys.modules, module.__name__, module) + seen = [] + + assert keep_patched(".modeling_kimi_k25", seen.append) == [remote.__name__] + assert seen == [remote] From a26448700d1208923afc689a9c7ea922bb5fb0fa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 05:30:02 +0000 Subject: [PATCH 18/20] Refuse a Kimi scheme on text_config alone; re-patch remote modules only keep_patched's get_class_in_module wrapper re-applied a kept patch to the module defining the resolved class, matched by name suffix alone. That module need not be the remote one transformers loaded: a remote modeling_kimi_k25.py importing its class from transformers' own 5.13 port made the wrapper patch transformers.models.kimi_k25.modeling_kimi_k25 and fail the load with "the checkpoint's modeling code changed", and a class whose module is not in sys.modules raised KeyError. The wrapper now looks the module up with sys.modules.get and patches remote-code modules only, through the check the scan of loaded modules uses (_is_remote). Both Kimi adapters replace text_config with a native DeepseekV3Config that drops its quantization_config, and the expert-parallel load then reads the checkpoint's quantization off the top-level copy Kimi-K2.5's config class keeps. A scheme declared on text_config alone was dropped silently: block-fp8 experts would load as stored, never scaled by weight_scale_inv. prepare now refuses such a config before changing anything. Kimi-K2.5's compressed-tensors experts, which the on-the-fly decompressor takes over, are unaffected. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 20 +++++++++++++++++ modelopt/torch/models/_remote_code.py | 25 ++++++++++++++++----- modelopt/torch/models/kimi_k25/ep.py | 5 ++++- modelopt/torch/models/kimi_vl/ep.py | 1 + tests/unit/torch/models/test_kimi_k25_ep.py | 2 ++ tests/unit/torch/models/test_kimi_vl_ep.py | 22 +++++++++++++++++- tests/unit/torch/models/test_remote_code.py | 21 +++++++++++++---- 7 files changed, 85 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index ab42956aba8..e5b7bed8847 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -58,6 +58,26 @@ def checkpoint_to_patch( return ckpt_path +def refuse_quantization_only_on_text_config(label: str, hf_config) -> None: + """Refuse a checkpoint whose ``quantization_config`` only its ``text_config`` declares. + + The adapter replaces ``text_config`` with a native config without it. Unless the experts + decompress as they load, the expert-parallel load then reads the checkpoint's quantization off + the top-level config, which Kimi-K2.5's config class copies it to; without that copy, the load + would read quantized experts as stored. + """ + if ( + getattr(hf_config.text_config, "quantization_config", None) is not None + and getattr(hf_config, "quantization_config", None) is None + ): + raise ValueError( + f"{label}: the checkpoint declares quantization_config on text_config only, which the " + "native deepseek_v3 config replacing it drops, so the expert-parallel load would read " + "its experts as stored. Declare it on the top-level config too, or start from an " + "unquantized copy of the model." + ) + + def install_compat_import_shims() -> None: """Add stubs for symbols the Kimi modeling imports that newer transformers removed, so it imports. diff --git a/modelopt/torch/models/_remote_code.py b/modelopt/torch/models/_remote_code.py index b2755779978..1fe4bc0f01c 100644 --- a/modelopt/torch/models/_remote_code.py +++ b/modelopt/torch/models/_remote_code.py @@ -74,6 +74,17 @@ def _staged(pretrained_model_name_or_path, module_file, *args, **kwargs): _KEPT_PATCHES: dict[str, Callable[[ModuleType], bool]] = {} +def _is_remote(name: str, suffix: str) -> bool: + """Whether the module ``name`` is remote code named ``*``. + + transformers' own port of a model can share the remote module's name, as + transformers.models.kimi_k25.modeling_kimi_k25 does from 5.13. + """ + import transformers.dynamic_module_utils as dmu + + return name.startswith(f"{dmu.TRANSFORMERS_DYNAMIC_MODULE_NAME}.") and name.endswith(suffix) + + def keep_patched(suffix: str, patch: Callable[[ModuleType], bool]) -> list[str]: """Apply ``patch`` to every loaded remote module named ``*``, and again whenever transformers resolves a class from one. @@ -93,8 +104,15 @@ def keep_patched(suffix: str, patch: Callable[[ModuleType], bool]) -> list[str]: def _resolved(class_name, module_path, *args, **kwargs): cls = orig(class_name, module_path, *args, **kwargs) + # The class's own module, which need not be the remote one transformers loaded (that + # can import the class from transformers' own port) and need not be in sys.modules. + module = sys.modules.get(cls.__module__) for kept_suffix, kept_patch in _KEPT_PATCHES.items(): - if cls.__module__.endswith(kept_suffix) and kept_patch(sys.modules[cls.__module__]): + if ( + module is not None + and _is_remote(cls.__module__, kept_suffix) + and kept_patch(module) + ): from modelopt.torch.utils import distributed as dist print( @@ -107,11 +125,8 @@ def _resolved(class_name, module_path, *args, **kwargs): setattr(_resolved, "_modelopt_keep_patched", True) dmu.get_class_in_module = _resolved patched = [] - # Remote modules only: transformers' own port of a model can share the module's name, as - # transformers.models.kimi_k25.modeling_kimi_k25 does from 5.13. - remote = f"{dmu.TRANSFORMERS_DYNAMIC_MODULE_NAME}." for name, module in list(sys.modules.items()): - if module is not None and name.startswith(remote) and name.endswith(suffix): + if module is not None and _is_remote(name, suffix): patch(module) patched.append(name) return patched diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index e298923887e..93192ba1cee 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -71,10 +71,13 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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 the expert-parallel load's to read: bf16 as they are, block-fp8 by the copy - # of quantization_config Kimi's config class keeps on the top-level config. + # of quantization_config Kimi's config class keeps on the top-level config, which step 2 + # leaves alone. A checkpoint without that copy is refused. quantization_config = _onfly_decompress._find_quantization_config(hf_config) or {} if quantization_config.get("quant_method") == "compressed-tensors": _onfly_decompress.install_onfly_decompress(hf_config) + else: + _kimi.refuse_quantization_only_on_text_config("kimi_k25", hf_config) # 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 diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 87975e18eaf..b9a064afa96 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -57,6 +57,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool Kimi-VL checkpoint is plain bf16. """ 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() diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index 40cb1cb00b6..9ba4749c97a 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -117,6 +117,8 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3( staging = dmu.get_cached_module_file config = _vlm_config() config.text_config.quantization_config = quantization_config + if quantization_config is not QUANT: + config.quantization_config = quantization_config # the top-level copy the load reads assert kimi_k25.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) is config diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 86c2e0af373..7b8fe0317c2 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -79,6 +79,13 @@ def _vlm_config(): return SimpleNamespace(text_config=text) +def _block_fp8_on_the_text_config_only(): + config = _vlm_config() + block_fp8 = {"quant_method": "fp8", "weight_block_size": [128, 128]} + config.text_config.quantization_config = block_fp8 + return config + + @pytest.fixture def remote(monkeypatch, request): """Stand in for the checkpoint's remote code: importing it registers its modeling module @@ -220,8 +227,21 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_has_nothing_to_swap( # 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"), + # The native text config drops the scheme, and with no top-level copy for the load to + # read, the experts would load as stored. + ( + _block_fp8_on_the_text_config_only, + {"ckpt_path": "/ckpt", "trust_remote_code": True}, + "on text_config only", + ), + ], + ids=[ + "transformers-port", + "no-text-config", + "no-checkpoint", + "untrusted-checkpoint", + "quantized-text-config-only", ], - ids=["transformers-port", "no-text-config", "no-checkpoint", "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 diff --git a/tests/unit/torch/models/test_remote_code.py b/tests/unit/torch/models/test_remote_code.py index 2370079fc86..f57dd62f434 100644 --- a/tests/unit/torch/models/test_remote_code.py +++ b/tests/unit/torch/models/test_remote_code.py @@ -130,17 +130,30 @@ def to_native(module): del sys.modules[name] -def test_only_remote_modules_are_patched(monkeypatch): +@pytest.mark.parametrize( + "defined_in", + [ + "transformers.models.kimi_k25.modeling_kimi_k25", + "transformers_modules.Kimi.4567cdef.modeling_kimi_k25", # not in sys.modules + ], + ids=["port", "not-loaded"], +) +def test_only_remote_modules_are_patched(monkeypatch, defined_in): """transformers' own port of a model can share its remote module's name (from 5.13, ``transformers.models.kimi_k25.modeling_kimi_k25``); a patch for the checkpoint's code must - leave it alone.""" - monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) - monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) + leave it alone. So must the re-patch when transformers resolves a class, which reaches the + module defining it: a remote module can import its class from the port, or from a module + sys.modules does not hold.""" remote = ModuleType("transformers_modules.Kimi.0123abcd.modeling_kimi_k25") port = ModuleType("transformers.models.kimi_k25.modeling_kimi_k25") for module in (remote, port): monkeypatch.setitem(sys.modules, module.__name__, module) + cls = type("KimiK25ForConditionalGeneration", (), {"__module__": defined_in}) + monkeypatch.setattr(dmu, "get_class_in_module", lambda *args, **kwargs: cls) + monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) seen = [] assert keep_patched(".modeling_kimi_k25", seen.append) == [remote.__name__] + module_path = f"{remote.__name__.replace('.', '/')}.py" # the remote module, as loaded + assert dmu.get_class_in_module(cls.__name__, module_path) is cls assert seen == [remote] From b5a6ff6c1a97aad9e42631481666b3c4790475c7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 07:24:28 +0000 Subject: [PATCH 19/20] Name 5.7 as the floor Kimi-VL's expert parallelism inherits Kimi-VL has no release requirement of its own, so it inherits the expert-parallel load's, now transformers 5.7. Its adapter docstring and the tests' skip said 5.6. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_vl/ep.py | 2 +- tests/unit/torch/models/test_kimi_vl_ep.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index b9a064afa96..09e96fc2719 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -39,7 +39,7 @@ 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: no floor beyond the EP framework's 5.6; validated on transformers 5.9.0. The compat shims +Version: no floor beyond the EP framework's 5.7; validated on transformers 5.9.0. The compat shims stand in for the symbols the remote code imports that later releases dropped. """ diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 7b8fe0317c2..7f587f3f678 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -32,8 +32,8 @@ 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) +if Version(transformers.__version__) < Version("5.7"): + pytest.skip("Kimi-VL expert parallelism needs transformers >= 5.7", allow_module_level=True) import transformers.dynamic_module_utils as dmu import transformers.utils.import_utils as import_utils From 56f533002bb196df46f1b2f975cfc405cce92689 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:25:30 +0000 Subject: [PATCH 20/20] Load the text config apply_distributed_support returns in the Kimi adapters Both prepares discarded the config apply_distributed_support returns, which is a prepare's result when the spec has one. deepseek_v3's has none today, so the plan was attached in place all the same; if it gains one that rebuilds the config, the adapters now load that config rather than the unprepared one. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 2 +- modelopt/torch/models/kimi_vl/ep.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 93192ba1cee..f5d68f83b8c 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -95,7 +95,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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) + 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. diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 09e96fc2719..ba3120d4b10 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -65,7 +65,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # 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) + hf_config.text_config = apply_distributed_support("deepseek_v3", hf_config.text_config) _kimi.patch_language_model_to_native( "kimi_vl",