From 2a51f3883f2de06a8b48f77440d9313db1748cb9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:05:28 +0000 Subject: [PATCH 01/17] Add the Kimi-K2.5 expert-parallel adapter Kimi-K2.5 (and Kimi-K2.6 and Kimi-K2.7-Code, which share its model type) is Kimi-VL's shape with compressed-tensors INT4 experts. Its adapter swaps in transformers' native deepseek_v3 language model as Kimi-VL's does, through the patch they share, and installs the on-the-fly decompressor first of all, so a checkpoint the decompressor refuses leaves the config and the loader untouched. It installs it only for a checkpoint whose quantization_config declares compressed-tensors: bf16 experts load as they are, and any other scheme is left to the expert-parallel load's format handling, which reads the copy of quantization_config Kimi's config class keeps on the top level. A scheme on text_config alone, with no top-level copy, is refused before anything changes. The spec declares the transformers 5.8 the decompression needs, and records that transformers ships its own Kimi-K2.5 port from 5.13, which the adapter refuses: expert parallelism runs on the checkpoint's own code (trust_remote_code=True). The shared adapter tests and the remote-Kimi refusal test run on Kimi-K2.5 too. Moved out of the Kimi adapters change, whose history is archived. Signed-off-by: Shengliang Xu --- modelopt/torch/models/__init__.py | 1 + modelopt/torch/models/kimi_k25/__init__.py | 18 ++ modelopt/torch/models/kimi_k25/ep.py | 111 ++++++++++ modelopt/torch/models/kimi_k25/specs.py | 46 +++++ tests/unit/torch/models/kimi_vl/test_ep.py | 16 +- tests/unit/torch/models/test_kimi_k25_ep.py | 195 ++++++++++++++++++ .../unit/torch/utils/test_model_load_utils.py | 2 +- 7 files changed, 383 insertions(+), 6 deletions(-) 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 tests/unit/torch/models/test_kimi_k25_ep.py diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index b87f3872f73..bd748e773e0 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -92,6 +92,7 @@ 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..4a6fc6ad12f --- /dev/null +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DP=EP a2a support for Kimi-K2.5 (``kimi_k25``). + +a remote-code VLM whose text backbone is a DeepSeek-V3 variant with compressed-tensors-compressed experts. +Kimi-K2.6 and Kimi-K2.7-Code ship the same model type and modeling classes. + +The challenge: Kimi's bundled ``modeling_deepseek.py`` defines ``DeepseekV3MoE`` as an +``nn.ModuleList`` of per-expert ``DeepseekV3MLP`` (separate ``gate_proj`` / ``up_proj`` / +``down_proj`` Linears per expert). That layout is incompatible with transformers' EP machinery, +which expects the **fused** 3-D ``gate_up_proj`` / ``down_proj`` Parameters that the +transformers-native ``deepseek_v3`` modeling ships. We need the native modeling to get EP-capable +experts, but we also want the rest of the VLM (vision tower + mm projector) to come from Kimi's +custom code. + +Solution -- monkey-patch ``modeling_kimi_k25``'s ``DeepseekV3ForCausalLM`` import to point at the +transformers-native class, AND swap ``hf_config.text_config`` to a transformers-native +``DeepseekV3Config`` (the native class won't accept Kimi's custom config subclass). Then +``KimiK25ForConditionalGeneration.__init__`` (which does ``self.language_model = +DeepseekV3ForCausalLM(config.text_config)``) builds the native EP-capable model for the language +sub-model, while vision_tower + mm_projector are still built from Kimi's custom code. + +On-the-fly compressed-tensors expert decompress happens during from_pretrained's +``WeightConverter`` machinery (installed by ``_onfly_decompress.install_onfly_decompress``); the +fused per-expert source ``.weight`` tensors are decompressed + merged into the 3-D +``gate_up_proj`` / ``down_proj`` parameters. Format-agnostic (compressed_tensors decides the +dequant from the checkpoint's quantization_config). + +EP plan: attached to the (now native) ``text_config``. transformers' ``post_init`` aggregates +child-module plans up to the root with name prefix, so the final root ``_ep_plan`` keys end up as +``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj,...}`` -- matching the actual +parameter paths in the loaded VLM. + +Version floor: transformers >= 5.8, which the spec declares. The on-the-fly decompressor, installed +for a compressed-tensors checkpoint, needs it: earlier releases' ``Concatenate`` trips over the +sibling keys it has already consumed. The decompressor couples tightly to +transformers' still-evolving ``WeightConverter`` conversion-ops API. +""" + +from .. import _kimi, _onfly_decompress +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.5's language model as the transformers-native deepseek_v3. + + Swap text_config to native deepseek_v3 + monkey-patch Kimi's modeling to use the native + DeepseekV3ForCausalLM + apply the EP plan + install the on-the-fly decompressor for + compressed-tensors experts. + """ + # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe + # imports the registry. + from modelopt.torch.utils.moe import apply_distributed_support + + ckpt_path = _kimi.checkpoint_to_patch("kimi_k25", hf_config, ckpt_path, trust_remote_code) + + # 1. Install on-the-fly compressed-tensors expert decompressor. Operates on transformers-native + # deepseek_v3's per-expert WeightConverter (which fuses per-expert .weight tensors into the + # 3-D gate_up_proj/down_proj Parameters). Format-agnostic; driven by checkpoint's + # quantization_config, which Kimi-K2.5 carries only on the text_config step 2 replaces. It + # goes first, so a checkpoint it refuses leaves everything as it was. Experts stored + # otherwise are left to the expert-parallel load's format handling, which reads their scheme + # from the copy of quantization_config Kimi's config class keeps on the top-level config; + # step 2 leaves that copy alone. A checkpoint without that copy is refused. + quantization_config = _onfly_decompress._find_quantization_config(hf_config) or {} + if quantization_config.get("quant_method") == "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 + # below, and from_pretrained later in load_and_prepare_ep). + _kimi.install_compat_import_shims() + stage_full_remote_closure() + + # 2. Replace text_config with a transformers-native DeepseekV3Config built from the same fields. + # Kimi's __init__ passes config.text_config to DeepseekV3ForCausalLM; the native class rejects + # Kimi's custom config subclass. to_native_backbone_config also drops auto_map + + # quantization_config so transformers builds the native modeling instead of re-routing back + # to Kimi's remote code via auto_map. + hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config) + + # 3. Apply the EP plan to the now-native text_config. transformers' post_init walks children + # and prefixes their _ep_plan onto the parent's: text_config's plan ends up at the VLM root + # as `language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}`. + hf_config.text_config = apply_distributed_support("deepseek_v3", hf_config.text_config) + + # 4. Force-load Kimi's modeling_kimi_k25 module (so we can patch its DeepseekV3ForCausalLM + # symbol) and swap to the native class. + _kimi.patch_language_model_to_native( + "kimi_k25", + "modeling_kimi_k25", + "KimiK25ForConditionalGeneration", + ckpt_path, + trust_remote_code, + ) + return hf_config diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py new file mode 100644 index 00000000000..41168532000 --- /dev/null +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""kimi_k25 specs (HF model type ``kimi_k25``). + +Kimi-K2.5 (Kimi-K2.6 and Kimi-K2.7-Code share its model type): the compressed-expert sibling of +``kimi_vl`` -- same swap to a native deepseek_v3 language model, plus on-the-fly decompression of +its packed experts at load. + +Only the expert-parallel section is registered under ``kimi_k25``: the MoE and export sections stay +``deepseek_v3``'s, so lookups by ``kimi_k25`` find neither yet. Sharing them is left to a follow-up. +""" + +from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN +from ..specs import EPSpec, ModelSpec, register +from .ep import prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_k25", + # transformers ships Kimi-K2.5 itself from 5.13. Expert parallelism still runs on the + # checkpoint's own code (``trust_remote_code=True``), which ``prepare`` patches. So the + # version describes transformers' port, not the model this spec's load builds: a MoE or + # export section added here must name the native deepseek_v3 language model's modules. + min_transformers_version="5.13", + # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan, + # and decompresses the experts as they load, which needs transformers 5.8. + ep_spec=EPSpec( + prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, requires_transformers="5.8" + ), + ) +) diff --git a/tests/unit/torch/models/kimi_vl/test_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py index bd5158077a9..25c56133689 100644 --- a/tests/unit/torch/models/kimi_vl/test_ep.py +++ b/tests/unit/torch/models/kimi_vl/test_ep.py @@ -42,6 +42,7 @@ from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig 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 @@ -159,11 +160,14 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke assert dmu.get_cached_module_file is not staging -# What the Kimi adapters share, through each adapter's prepare. +# 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"))], - ids=["vl"], + [ + (kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration")), + (kimi_k25, ("modeling_kimi_k25", "KimiK25ForConditionalGeneration")), + ], + ids=["vl", "k25"], indirect=["remote"], ) @@ -233,7 +237,7 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_has_nothing_to_swap( adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) -@pytest.mark.parametrize("adapter", [kimi_vl], ids=["vl"]) +@pytest.mark.parametrize("adapter", [kimi_vl, kimi_k25], ids=["vl", "k25"]) @pytest.mark.parametrize( ("config", "kwargs", "match"), [ @@ -282,7 +286,9 @@ def unstaged_loader(*args, **kwargs): assert dmu.get_cached_module_file is unstaged_loader -@pytest.mark.parametrize(("model_type", "adapter"), [("kimi_vl", kimi_vl)], ids=["vl"]) +@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.""" 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..b873abc910d --- /dev/null +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -0,0 +1,195 @@ +# 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.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 +stub the dynamic-module loader with a stand-in remote module. Every process-wide patch is undone +after each test. +""" + +import sys +import types +from types import SimpleNamespace + +import pytest + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.8"): + pytest.skip( + "Kimi-K2.5 expert parallelism needs transformers >= 5.8, for the decompression", + allow_module_level=True, + ) + +import transformers.dynamic_module_utils as dmu +from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig + +from modelopt.torch.models import _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 + +QUANT = {"quant_method": "compressed-tensors", "format": "pack-quantized"} +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.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, + n_routed_experts=8, + moe_intermediate_size=16, + auto_map={"AutoConfig": "configuration_deepseek.DeepseekV3Config"}, + quantization_config=QUANT.copy(), + ) + return SimpleNamespace(text_config=text) + + +@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. 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 = [] + + 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) + 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 + + +@pytest.mark.parametrize( + "quantization_config", + [ + QUANT, + # Left to the expert-parallel load's format handling: bf16 experts load as they are, and + # any other scheme is read from 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 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, + "install_onfly_decompress", + lambda cfg: installs.append(_onfly_decompress._find_quantization_config(cfg)), + ) + 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 + + text = config.text_config + assert type(text) is DeepseekV3Config + assert getattr(text, "auto_map", None) is None + assert getattr(text, "quantization_config", None) is None + assert (text.hidden_size, text.n_routed_experts, text.moe_intermediate_size) == (64, 8, 16) + assert text.base_model_ep_plan["layers.*.mlp.experts.gate_up_proj"] == "grouped_gemm" + assert text.base_model_ep_plan["layers.*.mlp.experts.down_proj"] == "grouped_gemm" + assert _routed_expert_keys(text.base_model_ep_plan) == _routed_expert_keys( + get_spec("kimi_k25").ep_spec.prepared_ep_plan + ) + assert installs == ([QUANT] if quantization_config is QUANT else []) + assert dmu.get_cached_module_file is not staging + + +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_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.""" + 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.""" + native = pytest.importorskip("transformers.models.kimi_k25.configuration_kimi_k25") + 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) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 817da74a8c6..b942c392446 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -844,7 +844,7 @@ def layer(): return model -@pytest.mark.parametrize("model_type", ["kimi_vl"]) +@pytest.mark.parametrize("model_type", ["kimi_vl", "kimi_k25"]) def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( run_as, serve, monkeypatch, tmp_path, model_type ): From 21d43db8868397d0f0ccab20260ad2af3abda810 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:10:39 +0000 Subject: [PATCH 02/17] Say why the Kimi-K2.5 spec keeps modeling_source at transformers, and what its missing MoE section skips transformers ships Kimi-K2.5 from 5.13, and the registry pairs a transformers modeling_source with that release, so a remote_code claim would be stale from 5.13 on, though prepare runs the checkpoint's own code. As for Kimi-VL, is_moe does not recognize the native DeepseekV3MoE blocks prepare installs until the sections are shared. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/specs.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index 41168532000..4b8dd922332 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -20,7 +20,9 @@ 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. +``deepseek_v3``'s, so lookups by ``kimi_k25`` find neither yet. Nor do ``is_moe``'s name and structure +fallbacks recognize the native ``DeepseekV3MoE`` blocks ``prepare`` installs, so the HF export's +MoE-block handling skips them. Sharing the sections is left to a follow-up. """ from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN @@ -32,10 +34,12 @@ register( ModelSpec( model_type="kimi_k25", - # transformers ships Kimi-K2.5 itself from 5.13. Expert parallelism still runs on the - # checkpoint's own code (``trust_remote_code=True``), which ``prepare`` patches. So the - # version describes transformers' port, not the model this spec's load builds: a MoE or - # export section added here must name the native deepseek_v3 language model's modules. + # transformers ships Kimi-K2.5 itself from 5.13, so modeling_source stays "transformers": the + # registry pairs it with that release, and a "remote_code" claim would be stale from 5.13 on. + # Expert parallelism still runs on the checkpoint's own code (``trust_remote_code=True``), + # which ``prepare`` patches, refusing the port. So the version describes transformers' port, + # not the model this spec's load builds: a MoE or export section added here must name the + # native deepseek_v3 language model's modules. min_transformers_version="5.13", # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan, # and decompresses the experts as they load, which needs transformers 5.8. From 3c0336bd15cfd20326d72bb52902609f58df77a5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:10:48 +0000 Subject: [PATCH 03/17] Keep only the sibling-module claim in the Kimi-K2.5 swap test The shared adapter tests, run for Kimi-K2.5 too, already check the load and the swap. What is left is Kimi-K2.5's own: the swap, scoped to modeling_kimi_k25, leaves modeling_deepseek's DeepseekV3ForCausalLM alone. The test fails when the scope is dropped. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k25_ep.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/test_kimi_k25_ep.py index b873abc910d..d0dd29ebad3 100644 --- a/tests/unit/torch/models/test_kimi_k25_ep.py +++ b/tests/unit/torch/models/test_kimi_k25_ep.py @@ -38,7 +38,7 @@ ) import transformers.dynamic_module_utils as dmu -from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig +from transformers import DeepseekV3Config, PretrainedConfig from modelopt.torch.models import _kimi, _onfly_decompress, _remote_code, get_spec from modelopt.torch.models.kimi_k25 import ep as kimi_k25 @@ -137,12 +137,10 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3( assert dmu.get_cached_module_file is not staging -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.""" +def test_the_swap_leaves_a_sibling_remote_module_alone(monkeypatch, remote): + """Kimi-K2.5's ``modeling_kimi_k25`` imports ``modeling_deepseek``, which defines a + ``DeepseekV3ForCausalLM`` of its own; the swap, scoped to ``modeling_kimi_k25``, must not reach + it. The swap itself is the shared adapter tests'.""" monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) bystander = types.ModuleType(f"{REMOTE}.modeling_deepseek") remote_lm = bystander.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) @@ -150,10 +148,6 @@ def test_prepare_with_a_checkpoint_points_the_remote_modeling_at_native_deepseek 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 From 15c11ebd8c983333b9176d3816b315fa6a8a74da Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:26:57 +0000 Subject: [PATCH 04/17] Prepare a remote MoE's spec only for an expert-parallel plan The plan decision ran the model spec's prepare on a copy of the config before the plan was known, wherever it could still be expert-parallel. For the Kimi models prepare patches the checkpoint's own code, the loader and the config for the whole process, so a run whose plan came out FSDP2-sharded, or whole and refused, kept those patches. The decision now reads the plan the spec states (prepared_ep_plan), on the config Kimi's experts are counted from, and finds the experts through it in the model the checkpoint's own code builds; only an expert-parallel plan then runs prepare, on every rank. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/specs.py | 5 +- modelopt/torch/models/kimi_vl/specs.py | 4 +- modelopt/torch/utils/moe.py | 12 +++-- .../torch/utils/plugins/model_load_utils.py | 35 +++++++++---- .../unit/torch/utils/test_model_load_utils.py | 49 ++++++++++++++++++- 5 files changed, 89 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index 4b8dd922332..5fedf8e0b31 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -44,7 +44,10 @@ # 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" + prepare=prepare, + prepared_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="text_config", + requires_transformers="5.8", ), ) ) diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py index 73ba5e78b2c..8a1e5d1a1cb 100644 --- a/modelopt/torch/models/kimi_vl/specs.py +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -36,6 +36,8 @@ model_type="kimi_vl", modeling_source="remote_code", # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. - ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN), + ep_spec=EPSpec( + prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config" + ), ) ) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index b9ed3863c63..e7e2c9ef348 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -226,6 +226,7 @@ def apply_distributed_support( *, ckpt_path: "str | None" = None, trust_remote_code: bool = False, + prepare: bool = True, ): """Prepare ``hf_config`` so ``model_type`` can be loaded expert-parallel, from its model spec. @@ -241,6 +242,10 @@ def apply_distributed_support( one (:func:`_ep_config`), not only on the config the spec would take: on the root, or on another nested config. + With ``prepare=False``, the spec's ``prepare`` does not run, and the plan it would attach + (``prepared_ep_plan``) goes on the config instead: what a plan decision reads, which changes + nothing outside the config, as a load that is not expert-parallel must find the process. + Returns the config to load, unchanged for a model type with no spec or no EP section. The plan is written into the config it is given, so a caller that keeps its own config passes a copy, as every caller here does. The generic all-to-all transport is not wired here -- it self-registers @@ -277,10 +282,11 @@ def apply_distributed_support( f"config.{' or config.'.join(attrs)}, which this config does not have, nor a plan " "of its own." ) - if ep.base_model_ep_plan and not existing: - target.base_model_ep_plan = dict(ep.base_model_ep_plan) + plan = ep.base_model_ep_plan or (None if prepare else ep.prepared_ep_plan) + if plan and not existing: + target.base_model_ep_plan = dict(plan) - if ep.prepare is not None: + if prepare and ep.prepare is not None: hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) return hf_config diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 00e2b4e9379..33841e84cc4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1896,7 +1896,7 @@ def checkpoint_read() -> bool: # class the sharded loads refuse rules it out too: the decision then plans the whole model on # every rank, or refuses, without the spec, whose prepare could fail, or patch the process, # for a load that never shards. - prepared = [hf_config] + stated = [hf_config] if dist.broadcast_from( "spec", lambda: ( @@ -1905,28 +1905,45 @@ def checkpoint_read() -> bool: and _ep_degree(ep_size, hf_config.model_type) != 1 ), ): - # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. A - # rank that cannot apply it fails every rank, as a checkpoint read does above. - def spec_applied() -> bool: - prepared[0] = apply_distributed_support( + # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. Its + # plan only: its prepare, which may patch the process, runs for an expert-parallel plan + # alone. A rank that cannot apply it fails every rank, as a checkpoint read does above. + def spec_stated() -> bool: + stated[0] = apply_distributed_support( hf_config.model_type, copy.deepcopy(hf_config), ckpt_path=ckpt_path, trust_remote_code=trust_remote_code, + prepare=False, ) return True - dist.all_reduce("model spec", spec_applied, all) + dist.all_reduce("model spec", spec_stated, all) plan = dist.broadcast_from( "plan", lambda: _decide_distribution( - ckpt_path, prepared[0], world_size, ep_size, dense, trust_remote_code, auto_class + ckpt_path, stated[0], world_size, ep_size, dense, trust_remote_code, auto_class ), ) if plan != DistributionPlan(): print_rank_0(f"Distribution: {plan!r}") - # Only the expert-parallel load shards by the plan the spec attached. - return plan, prepared[0] if plan.is_expert_parallel else hf_config + if not plan.is_expert_parallel: + return plan, hf_config + # Only the expert-parallel load shards by the plan the spec attached, and only it gets the spec + # prepared, on every rank. + prepared = [hf_config] + + def spec_applied() -> bool: + prepared[0] = apply_distributed_support( + hf_config.model_type, + copy.deepcopy(hf_config), + ckpt_path=ckpt_path, + trust_remote_code=trust_remote_code, + ) + return True + + dist.all_reduce("model spec prepared", spec_applied, all) + return plan, prepared[0] def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b942c392446..85ed001ee05 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -868,6 +868,50 @@ def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( resolve_distribution_plan(str(tmp_path)) +@needs_ep_transformers +@pytest.mark.parametrize("model_type", ["kimi_vl", "kimi_k25"]) +@pytest.mark.parametrize( + ("ranks", "dense", "expert_parallel"), + [(3, "fsdp", False), (3, "replicate", None), (4, "replicate", True)], + ids=["fsdp2", "whole", "expert-parallel"], +) +def test_a_remote_moe_is_prepared_for_an_expert_parallel_plan_alone( + run_as, serve, monkeypatch, tmp_path, model_type, ranks, dense, expert_parallel +): + """A Kimi spec's ``prepare`` patches the checkpoint's own code, the loader's hooks and the + config for the whole process, so only an expert-parallel plan runs it, once decided. The + decision finds the experts through the plan the spec states, in the model the checkpoint's own + code builds. Sharded by FSDP2, or left whole on every rank (refused), the run never prepares, + and gets the config as read.""" + calls = [] + plan = get_spec(model_type).ep_spec.prepared_ep_plan + + def prepare(hf_config, **kwargs): + calls.append(kwargs) + hf_config.text_config = SimpleNamespace(n_routed_experts=4, base_model_ep_plan=dict(plan)) + return hf_config + + monkeypatch.setattr(get_spec(model_type).ep_spec, "prepare", prepare) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi()) + run_as(ranks) # 4 experts shard over 4 ranks, and over no EP degree of 3 + as_read = serve( + SimpleNamespace(model_type=model_type, text_config=SimpleNamespace(n_routed_experts=4)) + ) + + if expert_parallel is None: + with pytest.raises(ValueError, match="can be expert-parallel, but the plan resolved"): + resolve_distribution_plan(str(tmp_path), dense=dense) + assert calls == [] + return + decided, config = resolve_distribution_plan(str(tmp_path), dense=dense) + assert decided.is_expert_parallel == expert_parallel + assert len(calls) == expert_parallel + if expert_parallel: + assert config.text_config.base_model_ep_plan == plan + else: + assert config is as_read and vars(config.text_config) == {"n_routed_experts": 4} + + def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): """Asked for no expert parallelism but for dense sharding, the MoE is sharded whole by FSDP2, which is not the waste of a replica on every rank.""" @@ -1735,8 +1779,9 @@ def test_a_load_prepares_its_config_once( ) prepared = [] - def apply_distributed_support(model_type, hf_config, **kwargs): - prepared.append(model_type) + def apply_distributed_support(model_type, hf_config, prepare=True, **kwargs): + if prepare: # the plan alone, which deciding reads, prepares nothing + prepared.append(model_type) hf_config.base_model_ep_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan return hf_config From 80a8a262def4aee5b08b468a678144fc9efc3e90 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:33:11 +0000 Subject: [PATCH 05/17] Make a remote Kimi's code importable for the plan decision, before prepare The plan decision now builds the model the checkpoint's own code builds, before the spec's prepare runs, and Kimi-VL's and Kimi-K2.5's code only imports on current transformers with the names it lost shimmed and its two-deep imports staged. EPSpec.make_importable runs just that, on every rank in the decision's spec step, whatever plan comes out: it patches no class, config, converter or loader hook. Both Kimi specs name the shared step, which their prepare runs too. EPSpec.prepare, prepared_ep_plan and resolve_distribution_plan now say that prepare runs only for an expert-parallel plan, and the end-to-end test runs the real Kimi prepare through the decision: FSDP2-sharded, or left whole and refused, the checkpoint's own language model class, the loader and the config as read are untouched; an expert-parallel plan swaps in the native language model. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 22 ++++-- modelopt/torch/models/kimi_k25/ep.py | 9 +-- modelopt/torch/models/kimi_k25/specs.py | 2 + modelopt/torch/models/kimi_vl/ep.py | 4 +- modelopt/torch/models/kimi_vl/specs.py | 6 +- modelopt/torch/models/specs.py | 24 ++++-- modelopt/torch/utils/moe.py | 10 ++- .../torch/utils/plugins/model_load_utils.py | 12 +-- tests/unit/torch/models/kimi_vl/test_ep.py | 78 ++++++++++++++++++- .../unit/torch/utils/test_model_load_utils.py | 44 ----------- 10 files changed, 131 insertions(+), 80 deletions(-) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index f48972306af..c1cd531c05f 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -23,7 +23,7 @@ # modelopt.torch.utils and transformers are imported inside the functions: the model registry imports # this module, and it must import without transformers, the optional [hf] extra. -from ._remote_code import keep_patched +from ._remote_code import keep_patched, stage_full_remote_closure def checkpoint_to_patch( @@ -114,6 +114,18 @@ def _prepare_4d_causal_attention_mask(*args, **kwargs): # pragma: no cover print_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}") +def make_importable() -> None: + """What both checkpoints' own code needs to import on this transformers, and nothing more. + + The names newer transformers dropped, and the code's two-deep relative imports staged. It patches + no class, config, converter or loader hook, so the plan decision runs it before ``prepare``, + whatever plan comes out. Idempotent. + """ + install_compat_import_shims() + # transformers stages only 1-hop relative imports; Kimi's remote code is 2 hops deep. + stage_full_remote_closure() + + def _drop_tie_weights_kwargs(orig_tie): """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it. @@ -134,10 +146,10 @@ def _native_for_the_prepared_config(own_class, native_class): """The ``DeepseekV3ForCausalLM`` the patched module builds its language model with. transformers' native class for the native ``DeepseekV3Config`` ``prepare`` puts on - ``text_config``, and the checkpoint's own class for its own config. Deciding the plan prepares a - copy of the config before the plan is known, and the patch outlives that copy: a load the plan - keeps off expert parallelism builds from the config as read, and so gets the checkpoint's own - language model, as it would with no patch. + ``text_config``, and the checkpoint's own class for its own config. The patch outlives the load + it was made for: a later load in the process that its plan keeps off expert parallelism builds + from the config as read, and so gets the checkpoint's own language model, as it would with no + patch. """ from transformers import DeepseekV3Config diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 4a6fc6ad12f..78e268d3c18 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -51,7 +51,6 @@ """ 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, **_): @@ -81,11 +80,9 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool 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 - # below, and from_pretrained later in load_and_prepare_ep). - _kimi.install_compat_import_shims() - stage_full_remote_closure() + # Before any remote-code load: this plugin's get_class_from_dynamic_module below, and + # from_pretrained later in load_and_prepare_ep. + _kimi.make_importable() # 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 diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index 5fedf8e0b31..183d60a766d 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -25,6 +25,7 @@ MoE-block handling skips them. Sharing the sections is left to a follow-up. """ +from .._kimi import make_importable from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN from ..specs import EPSpec, ModelSpec, register from .ep import prepare @@ -47,6 +48,7 @@ prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", + make_importable=make_importable, requires_transformers="5.8", ), ) diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 9d6899f0210..69c5993fdd8 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -21,7 +21,6 @@ """ 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, **_): @@ -38,8 +37,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool ckpt_path = _kimi.checkpoint_to_patch("kimi_vl", hf_config, ckpt_path, trust_remote_code) _kimi.refuse_quantization_only_on_text_config("kimi_vl", hf_config) - _kimi.install_compat_import_shims() - stage_full_remote_closure() + _kimi.make_importable() # 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) diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py index 8a1e5d1a1cb..da95cc1e02e 100644 --- a/modelopt/torch/models/kimi_vl/specs.py +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -25,6 +25,7 @@ MoE-block handling skips them. Sharing the sections is left to a follow-up. """ +from .._kimi import make_importable from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN from ..specs import EPSpec, ModelSpec, register from .ep import prepare @@ -37,7 +38,10 @@ modeling_source="remote_code", # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. ep_spec=EPSpec( - prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config" + prepare=prepare, + prepared_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="text_config", + make_importable=make_importable, ), ) ) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index fba161de33d..752f8959d0a 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -213,19 +213,27 @@ class EPSpec(SpecSection): fields above have been applied, so a ``prepare`` that returns another config than the one it was handed must carry over the ``base_model_ep_plan`` they attached. - It runs before the run's plan is decided, wherever the plan could still be expert-parallel, so - a run whose EP degree comes out at 1 has run it too: it must leave a load of the whole model, - on every rank or FSDP2-sharded, as it would be without it, and may run more than once.""" + It runs only once a run's plan is decided expert-parallel, on every rank, so a load of the whole + model, on every rank or FSDP2-sharded, never meets what it patches. It may run more than once in + a process.""" + + make_importable: "Callable[[], Any] | None" = None + """For a checkpoint carrying its own modeling code: what that code needs to import on the + installed transformers, so the plan decision can build the model before ``prepare`` runs. + + Called with no arguments, on every rank, before the plan is decided, whatever plan comes out. So + it may only make the code importable (shim names transformers dropped, stage the modules it + imports): it patches no class, config, converter or loader hook. It may run more than once.""" prepared_ep_plan: dict[str, str] | None = None """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare`` supplies. - Stated, not applied: ``prepare`` still attaches it. A run that rules expert parallelism out - skips ``prepare``, which patches the model's code, yet must still tell an expert-parallel MoE - it would hold whole on every rank from a dense model. It finds the routed experts through this - plan, in the model as the checkpoint's own code builds it: ``prepare`` replaces the experts in - place, so the plan names their containers there too.""" + Stated, not applied: ``prepare`` still attaches it. The plan decision runs before ``prepare``, + which only an expert-parallel plan runs, and a run that rules expert parallelism out must still + tell an expert-parallel MoE it would hold whole on every rank from a dense model. Both find the + routed experts through this plan, in the model as the checkpoint's own code builds it: + ``prepare`` replaces the experts in place, so the plan names their containers there too.""" unsupported: str | None = None """Why this model cannot be quantized expert-parallel yet, when its config would route it there. diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index e7e2c9ef348..5b40154e4e4 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -242,9 +242,11 @@ def apply_distributed_support( one (:func:`_ep_config`), not only on the config the spec would take: on the root, or on another nested config. - With ``prepare=False``, the spec's ``prepare`` does not run, and the plan it would attach - (``prepared_ep_plan``) goes on the config instead: what a plan decision reads, which changes - nothing outside the config, as a load that is not expert-parallel must find the process. + With ``prepare=False``, the spec's ``prepare`` does not run: the plan it would attach + (``prepared_ep_plan``) goes on the config instead, and only what makes the checkpoint's own code + importable runs (``make_importable``). That is what a plan decision reads, and it patches no + class, config, converter or loader hook, as a load that is not expert-parallel must find the + process. Returns the config to load, unchanged for a model type with no spec or no EP section. The plan is written into the config it is given, so a caller that keeps its own config passes a copy, as @@ -288,6 +290,8 @@ def apply_distributed_support( if prepare and ep.prepare is not None: hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) + elif not prepare and ep.make_importable is not None: + ep.make_importable() return hf_config diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 33841e84cc4..b671f74e755 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1848,12 +1848,12 @@ def resolve_distribution_plan( Unless the call or rank 0's knobs rule expert parallelism out (a single process, an ``auto_class`` the sharded loads refuse, or an EP degree of 1, which auto is for a model whose - spec refuses expert parallelism on the installed transformers), every rank applies the model's - spec to a copy of the config first - (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what - supplies the expert-parallel plan, and an expert-parallel plan returns that copy. Any other - plan returns the config without it, which the export would otherwise write out; the spec's - ``prepare`` has run for it all the same, which its contract allows + spec refuses expert parallelism on the installed transformers), every rank reads the model's + spec into a copy of the config first, its plan without its ``prepare`` + (:func:`~modelopt.torch.utils.moe.apply_distributed_support` with ``prepare=False``): for + several models it is what supplies the expert-parallel plan. An expert-parallel plan then has + every rank prepare a copy of the config as read, which it returns. Any other plan returns the + config as read, and its load never meets what ``prepare`` patches (:attr:`~modelopt.torch.models.EPSpec.prepare`). A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and diff --git a/tests/unit/torch/models/kimi_vl/test_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py index 25c56133689..6b033f9b926 100644 --- a/tests/unit/torch/models/kimi_vl/test_ep.py +++ b/tests/unit/torch/models/kimi_vl/test_ep.py @@ -37,14 +37,19 @@ if Version(transformers.__version__) < Version("5.7"): pytest.skip("Kimi-VL expert parallelism needs transformers >= 5.7", allow_module_level=True) +import transformers.core_model_loading as cml import transformers.dynamic_module_utils as dmu +import transformers.modeling_utils as mu import transformers.utils.import_utils as import_utils +from torch import nn from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig 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 import distributed as dist from modelopt.torch.utils.moe import _routed_expert_keys +from modelopt.torch.utils.plugins import model_load_utils # Symbols the remote modeling imports that newer transformers dropped. COMPAT_SYMBOLS = ( @@ -200,10 +205,10 @@ def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_wei def test_a_load_off_expert_parallelism_builds_the_checkpoint_s_own_language_model( monkeypatch, adapter, remote ): - """Deciding the plan prepares a copy of the config on every rank before the plan is known, and - the patch outlives that copy. A load the plan keeps off expert parallelism (replicated, or - FSDP2) builds from the config as read, so its VLM gets the checkpoint's own language model, as - with no patch; only the native config ``prepare`` makes builds transformers' native one.""" + """The patch outlives the load it was made for. A later load in the process that its plan keeps + off expert parallelism (replicated, or FSDP2) builds from the config as read, so its VLM gets + the checkpoint's own language model, as with no patch; only the native config ``prepare`` makes + builds transformers' native one.""" monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) as_read = _vlm_config() prepared = adapter.prepare(copy.deepcopy(as_read), ckpt_path="/ckpt", trust_remote_code=True) @@ -213,6 +218,71 @@ def test_a_load_off_expert_parallelism_builds_the_checkpoint_s_own_language_mode assert type(remote.vlm(as_read).language_model) is remote.remote_lm +def _built_by_its_own_code(): + """The VLM the checkpoint's own code builds, on the meta device: each MoE layer of its language + model holds its routed experts as per-expert modules, at ``language_model.model.layers..mlp.experts``.""" + + def layer(): + mlp = nn.Module() + mlp.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(8)) + return nn.ModuleDict({"mlp": mlp}) + + 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)) + model.tie_weights = lambda: None + return model + + +@_BOTH_ADAPTERS +@pytest.mark.parametrize( + ("ranks", "dense", "outcome"), + [(3, "fsdp", "fsdp2"), (3, "replicate", "refused"), (4, "replicate", "expert-parallel")], + ids=["fsdp2", "whole", "expert-parallel"], +) +def test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code( + monkeypatch, tmp_path, adapter, remote, ranks, dense, outcome +): + """Deciding the plan builds the model the checkpoint's own code builds, importable but unpatched, + and finds its 8 experts through the plan the spec states. Only an expert-parallel plan (4 ranks) + then runs ``prepare``, which swaps in the native language model. FSDP2-sharded, or left whole on + 3 ranks, which no EP degree of the 8 experts divides (refused), the run keeps the checkpoint's + own class, the loader and the config as read.""" + model_type = "kimi_vl" if adapter is kimi_vl else "kimi_k25" + as_read = _vlm_config() + as_read.model_type = model_type + monkeypatch.setattr(dist, "size", lambda group=None: ranks) + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: copy.deepcopy(as_read) + ) + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: _built_by_its_own_code() + ) + hooks = (mu.get_model_conversion_mapping, cml.spawn_materialize) + + if outcome == "refused": + with pytest.raises(ValueError, match="can be expert-parallel, but the plan resolved"): + model_load_utils.resolve_distribution_plan( + str(tmp_path), dense=dense, trust_remote_code=True + ) + else: + plan, config = model_load_utils.resolve_distribution_plan( + str(tmp_path), dense=dense, trust_remote_code=True + ) + assert plan.is_expert_parallel == (outcome == "expert-parallel") + assert import_utils.is_torch_fx_available() is False # importable on every plan + assert (remote.DeepseekV3ForCausalLM is remote.remote_lm) == (outcome != "expert-parallel") + if outcome == "expert-parallel": + assert type(config.text_config) is DeepseekV3Config + assert config.text_config.base_model_ep_plan + else: + assert (mu.get_model_conversion_mapping, cml.spawn_materialize) == hooks + if outcome == "fsdp2": + assert type(config.text_config) is _RemoteTextConfig + assert config.text_config.base_model_ep_plan is None + + @_BOTH_ADAPTERS @pytest.mark.parametrize( ("missing", "match"), diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 85ed001ee05..aed67aaba0a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -868,50 +868,6 @@ def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( resolve_distribution_plan(str(tmp_path)) -@needs_ep_transformers -@pytest.mark.parametrize("model_type", ["kimi_vl", "kimi_k25"]) -@pytest.mark.parametrize( - ("ranks", "dense", "expert_parallel"), - [(3, "fsdp", False), (3, "replicate", None), (4, "replicate", True)], - ids=["fsdp2", "whole", "expert-parallel"], -) -def test_a_remote_moe_is_prepared_for_an_expert_parallel_plan_alone( - run_as, serve, monkeypatch, tmp_path, model_type, ranks, dense, expert_parallel -): - """A Kimi spec's ``prepare`` patches the checkpoint's own code, the loader's hooks and the - config for the whole process, so only an expert-parallel plan runs it, once decided. The - decision finds the experts through the plan the spec states, in the model the checkpoint's own - code builds. Sharded by FSDP2, or left whole on every rank (refused), the run never prepares, - and gets the config as read.""" - calls = [] - plan = get_spec(model_type).ep_spec.prepared_ep_plan - - def prepare(hf_config, **kwargs): - calls.append(kwargs) - hf_config.text_config = SimpleNamespace(n_routed_experts=4, base_model_ep_plan=dict(plan)) - return hf_config - - monkeypatch.setattr(get_spec(model_type).ep_spec, "prepare", prepare) - monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi()) - run_as(ranks) # 4 experts shard over 4 ranks, and over no EP degree of 3 - as_read = serve( - SimpleNamespace(model_type=model_type, text_config=SimpleNamespace(n_routed_experts=4)) - ) - - if expert_parallel is None: - with pytest.raises(ValueError, match="can be expert-parallel, but the plan resolved"): - resolve_distribution_plan(str(tmp_path), dense=dense) - assert calls == [] - return - decided, config = resolve_distribution_plan(str(tmp_path), dense=dense) - assert decided.is_expert_parallel == expert_parallel - assert len(calls) == expert_parallel - if expert_parallel: - assert config.text_config.base_model_ep_plan == plan - else: - assert config is as_read and vars(config.text_config) == {"n_routed_experts": 4} - - def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): """Asked for no expert parallelism but for dense sharding, the MoE is sharded whole by FSDP2, which is not the waste of a replica on every rank.""" From bced888bf45cad216b016663cad8b305cd2c5143 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:35:08 +0000 Subject: [PATCH 06/17] Say plainly when a spec's prepare, make_importable and stated plan run From ep-routing's review of the move: prepare's doc names the collective step it runs in, on what config, and the load that prepares a handed plan; make_importable's names its gate and why it must be idempotent; prepared_ep_plan's says deciding attaches it with no prepare; and resolve_distribution_plan names both collective steps. Signed-off-by: Shengliang Xu --- modelopt/torch/models/specs.py | 24 ++++++++++--------- .../torch/utils/plugins/model_load_utils.py | 11 +++++---- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 752f8959d0a..04a01c7d344 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -213,27 +213,29 @@ class EPSpec(SpecSection): fields above have been applied, so a ``prepare`` that returns another config than the one it was handed must carry over the ``base_model_ep_plan`` they attached. - It runs only once a run's plan is decided expert-parallel, on every rank, so a load of the whole - model, on every rank or FSDP2-sharded, never meets what it patches. It may run more than once in - a process.""" + It runs only once a plan is decided expert-parallel: on every rank, in the plan decision's + collective "model spec prepared" step, on a copy of the config as read, so a refusal there + raises on every rank; or, for an expert-parallel plan handed in, in the load, after the load's + own checks, the GPU refusal included. It may patch the model's code, which is why nothing else + runs it. It may run more than once in a process.""" make_importable: "Callable[[], Any] | None" = None """For a checkpoint carrying its own modeling code: what that code needs to import on the installed transformers, so the plan decision can build the model before ``prepare`` runs. - Called with no arguments, on every rank, before the plan is decided, whatever plan comes out. So - it may only make the code importable (shim names transformers dropped, stage the modules it - imports): it patches no class, config, converter or loader hook. It may run more than once.""" + Called with no arguments, on every rank, before the plan is decided, for every plan that could + still be expert-parallel, whatever plan comes out. So it must be idempotent, and may only import + or stage the code (shim names transformers dropped, stage the modules it imports): it patches no + class, config, converter or loader hook.""" prepared_ep_plan: dict[str, str] | None = None """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare`` supplies. - Stated, not applied: ``prepare`` still attaches it. The plan decision runs before ``prepare``, - which only an expert-parallel plan runs, and a run that rules expert parallelism out must still - tell an expert-parallel MoE it would hold whole on every rank from a dense model. Both find the - routed experts through this plan, in the model as the checkpoint's own code builds it: - ``prepare`` replaces the experts in place, so the plan names their containers there too.""" + Deciding the plan attaches it with no ``prepare`` run, so the decision counts the routed experts + in the model the checkpoint's own code builds. A run that rules expert parallelism out finds + them the same way, to tell an expert-parallel MoE it would hold whole on every rank from a dense + model. ``prepare`` replaces the experts in place, so the plan names their containers there too.""" unsupported: str | None = None """Why this model cannot be quantized expert-parallel yet, when its config would route it there. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index b671f74e755..f1a0dcfdb80 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1848,12 +1848,13 @@ def resolve_distribution_plan( Unless the call or rank 0's knobs rule expert parallelism out (a single process, an ``auto_class`` the sharded loads refuse, or an EP degree of 1, which auto is for a model whose - spec refuses expert parallelism on the installed transformers), every rank reads the model's - spec into a copy of the config first, its plan without its ``prepare`` - (:func:`~modelopt.torch.utils.moe.apply_distributed_support` with ``prepare=False``): for + spec refuses expert parallelism on the installed transformers), every rank applies the model's + spec to a copy of the config first, in the collective "model spec" step: its plan, and what + makes the checkpoint's own code importable, with no ``prepare`` run + (:func:`~modelopt.torch.utils.moe.apply_distributed_support` with ``prepare=False``). For several models it is what supplies the expert-parallel plan. An expert-parallel plan then has - every rank prepare a copy of the config as read, which it returns. Any other plan returns the - config as read, and its load never meets what ``prepare`` patches + every rank prepare a copy of the config as read, in the collective "model spec prepared" step, + and returns that copy. Any other plan returns the config as read, with no ``prepare`` run (:attr:`~modelopt.torch.models.EPSpec.prepare`). A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and From 3aade0d09405f072d2b7c6ba890252d767a3c1da Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:56:42 +0000 Subject: [PATCH 07/17] Count a checkpoint its spec decompresses as one the expert-parallel load takes decision and the expert-parallel load would refuse its checkpoint's format. Kimi-K2.5's compressed-tensors checkpoint reads as such a format on the config the decision sees, before prepare installs the decompressor, so a decided Kimi-K2.5 run planned without expert parallelism. EPSpec.decompressed_formats states the formats a spec's prepare decompresses as the experts load; the decision counts them among those the load takes, and Kimi-K2.5's spec names compressed- tensors. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/specs.py | 1 + modelopt/torch/models/specs.py | 6 +++++ .../torch/utils/plugins/model_load_utils.py | 6 ++++- .../unit/torch/utils/test_model_load_utils.py | 27 +++++++++++++++++++ 4 files changed, 39 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index 183d60a766d..4347091666c 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -49,6 +49,7 @@ prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", make_importable=make_importable, + decompressed_formats=("compressed-tensors",), requires_transformers="5.8", ), ) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 04a01c7d344..2c09b3a9ab4 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -228,6 +228,12 @@ class EPSpec(SpecSection): or stage the code (shim names transformers dropped, stage the modules it imports): it patches no class, config, converter or loader hook.""" + decompressed_formats: tuple[str, ...] = () + """The checkpoint formats (a ``quantization_config``'s ``quant_method``) whose experts + ``prepare`` has decompress as they load, which the expert-parallel load then takes beside + unquantized checkpoints. The plan decision runs before ``prepare``, and reads this to count such + a checkpoint one the expert-parallel load takes.""" + prepared_ep_plan: dict[str, str] | None = None """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare`` supplies. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f1a0dcfdb80..866e76b1d20 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2026,7 +2026,11 @@ def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> Exception | None: parallelism out, or the transformers release it needs, is not asked here: the plan decision and the load ask the spec first (:func:`~modelopt.torch.utils.moe.check_distributed_support`). """ - if (method := _ckpt_quant_method(hf_config)) not in _EP_SOURCE_FORMATS: + # The formats it takes include those the spec's prepare, which runs only once the plan is + # decided, has the experts decompress as they load. + ep = _ep_spec(hf_config.model_type) + takes = _EP_SOURCE_FORMATS | frozenset(ep.decompressed_formats if ep is not None else ()) + if (method := _ckpt_quant_method(hf_config)) not in takes: return NotImplementedError( f"The expert-parallel load cannot dequantize a {method!r} checkpoint." ) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index aed67aaba0a..893cbf26fd5 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -825,6 +825,33 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) +@pytest.mark.skipif( + Version(transformers.__version__) < Version(_ep_floor("kimi_k25")), + reason="Kimi-K2.5's expert parallelism needs a later transformers", +) +@pytest.mark.parametrize( + ("model_type", "expert_parallel"), [("kimi_k25", True), ("kimi_vl", False)] +) +def test_a_checkpoint_its_spec_decompresses_is_decided_expert_parallel( + run_as, serve, monkeypatch, tmp_path, model_type, expert_parallel +): + """The decision runs before ``prepare``, which installs the decompressor. A compressed-tensors + checkpoint whose spec decompresses its experts as they load (Kimi-K2.5's) is planned + expert-parallel all the same; one whose spec does not (Kimi-VL's) is planned without it, as the + expert-parallel load would refuse it.""" + spec = get_spec(model_type).ep_spec + monkeypatch.setattr(spec, "prepare", lambda hf_config, **kwargs: hf_config) + monkeypatch.setattr(spec, "make_importable", lambda: None) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi()) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + run_as(4) + scheme = {"quant_method": "compressed-tensors"} + text = SimpleNamespace(n_routed_experts=4, quantization_config=scheme) + serve(SimpleNamespace(model_type=model_type, quantization_config=scheme, text_config=text)) + + assert resolve_distribution_plan(str(tmp_path))[0].is_expert_parallel == expert_parallel + + def _remote_kimi(): """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 From 8af85cfa9820a16ed66d237a45643bc5013870d7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:15:38 +0000 Subject: [PATCH 08/17] Mirror the Kimi-K2.5 adapter's module in its tests As CONTRIBUTING asks, and as the Kimi-VL adapter's tests now do: the Kimi-K2.5 adapter's tests move to tests/unit/torch/models/kimi_k25/ test_ep.py, mirroring kimi_k25/ep.py. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/kimi_k25/__init__.py | 15 +++++++++++++++ .../{test_kimi_k25_ep.py => kimi_k25/test_ep.py} | 0 2 files changed, 15 insertions(+) create mode 100644 tests/unit/torch/models/kimi_k25/__init__.py rename tests/unit/torch/models/{test_kimi_k25_ep.py => kimi_k25/test_ep.py} (100%) diff --git a/tests/unit/torch/models/kimi_k25/__init__.py b/tests/unit/torch/models/kimi_k25/__init__.py new file mode 100644 index 00000000000..1275d78dff5 --- /dev/null +++ b/tests/unit/torch/models/kimi_k25/__init__.py @@ -0,0 +1,15 @@ +# 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. + diff --git a/tests/unit/torch/models/test_kimi_k25_ep.py b/tests/unit/torch/models/kimi_k25/test_ep.py similarity index 100% rename from tests/unit/torch/models/test_kimi_k25_ep.py rename to tests/unit/torch/models/kimi_k25/test_ep.py From 1eed0d04e0f38450bf297bee38465cefc51edff5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:29:43 +0000 Subject: [PATCH 09/17] Have a remote Kimi's code build, not only import, before the plan decision Measured on Kimi-K2.5's own code: since the plan decision builds the checkpoint's own model on the meta device before prepare, the build fails on transformers 5.18, as post_init calls the VLM class's tie_weights with recompute_mapping, which the remote tie_weights does not take. prepare's swap used to adapt it first. EPSpec.make_importable becomes make_buildable(hf_config, ckpt_path=..., trust_remote_code=...): what the checkpoint's own code needs to import and build on the installed transformers, which may adapt a method to the arguments transformers now passes but changes nothing the model computes or loads. Kimi-VL's and Kimi-K2.5's step also loads the code and adapts its VLM class's tie_weights; the contract now names what it may not patch, as Claude's review asked: an MoE or expert class, a config, a weight converter or a weight-loading hook. New test test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_ transformers builds the tiny Kimi-VL remote code through transformers' own loader after the step, and fails with that TypeError without it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 60 +++++++++++++++---- modelopt/torch/models/kimi_k25/ep.py | 18 +++++- modelopt/torch/models/kimi_k25/specs.py | 5 +- modelopt/torch/models/kimi_vl/ep.py | 18 +++++- modelopt/torch/models/kimi_vl/specs.py | 5 +- modelopt/torch/models/specs.py | 18 +++--- modelopt/torch/utils/moe.py | 9 ++- tests/unit/torch/models/kimi_vl/test_ep.py | 41 ++++++++++++- .../unit/torch/utils/test_model_load_utils.py | 2 +- 9 files changed, 140 insertions(+), 36 deletions(-) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index c1cd531c05f..fbb33a658f6 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -115,17 +115,57 @@ def _prepare_4d_causal_attention_mask(*args, **kwargs): # pragma: no cover def make_importable() -> None: - """What both checkpoints' own code needs to import on this transformers, and nothing more. - - The names newer transformers dropped, and the code's two-deep relative imports staged. It patches - no class, config, converter or loader hook, so the plan decision runs it before ``prepare``, - whatever plan comes out. Idempotent. - """ + """The names newer transformers dropped, and the code's two-deep relative imports staged.""" install_compat_import_shims() # transformers stages only 1-hop relative imports; Kimi's remote code is 2 hops deep. stage_full_remote_closure() +def make_buildable( + modeling: str, vlm_class: str, hf_config, ckpt_path: str | None, trust_remote_code: bool +) -> None: + """What both checkpoints' own code needs to import and build on this transformers, and no more. + + The names newer transformers dropped, the code's two-deep relative imports staged, and the VLM + class's ``tie_weights`` taking the keyword arguments transformers' ``post_init`` passes it. None + of it changes what the model computes or loads, so the plan decision runs it before ``prepare``, + whatever plan comes out. A config of transformers' own port, or code this load may not run, has + nothing of the checkpoint's to build. Idempotent. + """ + import sys + + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + make_importable() + if type(hf_config).__module__.startswith("transformers.") or ckpt_path is None: + return + if not trust_remote_code: + return + get_class_from_dynamic_module( + f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code + ) + for name, module in list(sys.modules.items()): + if name.endswith(f".{modeling}") and module is not None: + _shim_tie_weights(module, vlm_class) + + +def _shim_tie_weights(module, vlm_class: str) -> bool: + """Have ``module``'s VLM class take the keyword arguments transformers passes ``tie_weights``. + + The remote ``tie_weights()`` takes none, but transformers' ``post_init`` and load call + ``tie_weights(recompute_mapping=False)``: a TypeError at construction. Returns whether it changed + the class. + """ + kimi_cls = getattr(module, vlm_class, None) + if kimi_cls is None: + return False + own_tie = vars(kimi_cls).get("tie_weights") + if own_tie is None or getattr(own_tie, "_kimi_tie_weights_shim", False): + return False + kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie) + return True + + def _drop_tie_weights_kwargs(orig_tie): """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it. @@ -195,13 +235,7 @@ def to_native(module): module.DeepseekV3ForCausalLM = _native_for_the_prepared_config( current, NativeDeepseekV3ForCausalLM ) - # The remote tie_weights() takes no kwargs, but transformers' post_init and load call - # tie_weights(recompute_mapping=False) -> TypeError. - kimi_cls = getattr(module, vlm_class, None) - if kimi_cls is not None: - own_tie = vars(kimi_cls).get("tie_weights") - if own_tie is not None and not getattr(own_tie, "_kimi_tie_weights_shim", False): - kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie) + _shim_tie_weights(module, vlm_class) return True # The dynamic-module loader names the module like diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 78e268d3c18..0895215af78 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -53,6 +53,22 @@ from .. import _kimi, _onfly_decompress +def make_buildable( + hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_ +) -> None: + """Have the checkpoint's own code import and build on this transformers. + + The plan decision needs it before ``prepare``; it changes nothing the model computes or loads. + """ + _kimi.make_buildable( + "modeling_kimi_k25", + "KimiK25ForConditionalGeneration", + hf_config, + ckpt_path, + trust_remote_code, + ) + + def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): """Load Kimi-K2.5's language model as the transformers-native deepseek_v3. @@ -82,7 +98,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # Before any remote-code load: this plugin's get_class_from_dynamic_module below, and # from_pretrained later in load_and_prepare_ep. - _kimi.make_importable() + _kimi.make_importable() # the swap below loads the code and adapts its tie_weights # 2. Replace text_config with a transformers-native DeepseekV3Config built from the same fields. # Kimi's __init__ passes config.text_config to DeepseekV3ForCausalLM; the native class rejects diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index 4347091666c..182691f0263 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -25,10 +25,9 @@ MoE-block handling skips them. Sharing the sections is left to a follow-up. """ -from .._kimi import make_importable from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN from ..specs import EPSpec, ModelSpec, register -from .ep import prepare +from .ep import make_buildable, prepare __all__: list[str] = [] @@ -48,7 +47,7 @@ prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", - make_importable=make_importable, + make_buildable=make_buildable, decompressed_formats=("compressed-tensors",), requires_transformers="5.8", ), diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 69c5993fdd8..1b80f6ce066 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -23,6 +23,22 @@ from .. import _kimi, _onfly_decompress +def make_buildable( + hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_ +) -> None: + """Have the checkpoint's own code import and build on this transformers. + + The plan decision needs it before ``prepare``; it changes nothing the model computes or loads. + """ + _kimi.make_buildable( + "modeling_kimi_vl", + "KimiVLForConditionalGeneration", + hf_config, + ckpt_path, + trust_remote_code, + ) + + def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): """Load Kimi-VL's language model as the transformers-native deepseek_v3. @@ -37,7 +53,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool ckpt_path = _kimi.checkpoint_to_patch("kimi_vl", hf_config, ckpt_path, trust_remote_code) _kimi.refuse_quantization_only_on_text_config("kimi_vl", hf_config) - _kimi.make_importable() + _kimi.make_importable() # the swap below loads the code and adapts its tie_weights # Native deepseek_v3 config for the LM (drops the remote auto_map so the native class is used). hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config) diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py index da95cc1e02e..bc511f806c9 100644 --- a/modelopt/torch/models/kimi_vl/specs.py +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -25,10 +25,9 @@ MoE-block handling skips them. Sharing the sections is left to a follow-up. """ -from .._kimi import make_importable from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN from ..specs import EPSpec, ModelSpec, register -from .ep import prepare +from .ep import make_buildable, prepare __all__: list[str] = [] @@ -41,7 +40,7 @@ prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", - make_importable=make_importable, + make_buildable=make_buildable, ), ) ) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 2c09b3a9ab4..566bebd4922 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -219,14 +219,16 @@ class EPSpec(SpecSection): own checks, the GPU refusal included. It may patch the model's code, which is why nothing else runs it. It may run more than once in a process.""" - make_importable: "Callable[[], Any] | None" = None - """For a checkpoint carrying its own modeling code: what that code needs to import on the - installed transformers, so the plan decision can build the model before ``prepare`` runs. - - Called with no arguments, on every rank, before the plan is decided, for every plan that could - still be expert-parallel, whatever plan comes out. So it must be idempotent, and may only import - or stage the code (shim names transformers dropped, stage the modules it imports): it patches no - class, config, converter or loader hook.""" + make_buildable: "Callable[..., Any] | None" = None + """For a checkpoint carrying its own modeling code: what that code needs to import and build on + the installed transformers, so the plan decision can build the model before ``prepare`` runs. + + Called as ``make_buildable(hf_config, ckpt_path=..., trust_remote_code=...)``, on every rank, + before the plan is decided, for every plan that could still be expert-parallel, whatever plan + comes out. So it must be idempotent and change nothing the model computes or loads: it may shim + names transformers dropped, wrap the remote-code staging to reach the modules the code imports, + and adapt a method to the arguments transformers now passes it, but patches no MoE or expert + class, config, weight converter or weight-loading hook.""" decompressed_formats: tuple[str, ...] = () """The checkpoint formats (a ``quantization_config``'s ``quant_method``) whose experts diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 5b40154e4e4..08d40dc9ac4 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -244,9 +244,8 @@ def apply_distributed_support( With ``prepare=False``, the spec's ``prepare`` does not run: the plan it would attach (``prepared_ep_plan``) goes on the config instead, and only what makes the checkpoint's own code - importable runs (``make_importable``). That is what a plan decision reads, and it patches no - class, config, converter or loader hook, as a load that is not expert-parallel must find the - process. + import and build runs (``make_buildable``). That is what a plan decision reads, and it changes + nothing the model computes or loads, as a load that is not expert-parallel must find the process. Returns the config to load, unchanged for a model type with no spec or no EP section. The plan is written into the config it is given, so a caller that keeps its own config passes a copy, as @@ -290,8 +289,8 @@ def apply_distributed_support( if prepare and ep.prepare is not None: hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) - elif not prepare and ep.make_importable is not None: - ep.make_importable() + elif not prepare and ep.make_buildable is not None: + ep.make_buildable(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) return hf_config diff --git a/tests/unit/torch/models/kimi_vl/test_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py index 6b033f9b926..5e93c0fbde2 100644 --- a/tests/unit/torch/models/kimi_vl/test_ep.py +++ b/tests/unit/torch/models/kimi_vl/test_ep.py @@ -21,8 +21,11 @@ after each test. """ +import contextlib import copy import importlib +import json +import shutil import sys import types from types import SimpleNamespace @@ -41,8 +44,15 @@ import transformers.dynamic_module_utils as dmu import transformers.modeling_utils as mu import transformers.utils.import_utils as import_utils +from _test_utils.torch import kimi_vl_tiny_modeling from torch import nn -from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + DeepseekV3Config, + DeepseekV3ForCausalLM, + PretrainedConfig, +) from modelopt.torch.models import _onfly_decompress, _remote_code, get_spec from modelopt.torch.models.kimi_k25 import ep as kimi_k25 @@ -272,6 +282,7 @@ def test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code( ) assert plan.is_expert_parallel == (outcome == "expert-parallel") assert import_utils.is_torch_fx_available() is False # importable on every plan + assert getattr(vars(remote.vlm)["tie_weights"], "_kimi_tie_weights_shim", False) # buildable assert (remote.DeepseekV3ForCausalLM is remote.remote_lm) == (outcome != "expert-parallel") if outcome == "expert-parallel": assert type(config.text_config) is DeepseekV3Config @@ -283,6 +294,34 @@ def test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code( assert config.text_config.base_model_ep_plan is None +def test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_transformers( + monkeypatch, tmp_path +): + """The plan decision builds the checkpoint's own code on the meta device before ``prepare`` + runs, and Kimi-VL's VLM class has a ``tie_weights`` that takes none of the keyword arguments + transformers' ``post_init`` passes it. ``make_buildable`` adapts it, and swaps nothing: the + module still builds the language model from its own ``DeepseekV3ForCausalLM``.""" + for module_name, name in COMPAT_SYMBOLS: + with contextlib.suppress(ImportError): + _restore(monkeypatch, importlib.import_module(module_name), name) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + shutil.copy(kimi_vl_tiny_modeling.__file__, tmp_path / "modeling_kimi_vl.py") + remote = "modeling_kimi_vl.KimiVLForConditionalGeneration" + auto_map = {"AutoConfig": "modeling_kimi_vl.KimiVLConfig", "AutoModelForCausalLM": remote} + text = {"model_type": "deepseek_v3", "hidden_size": 16, "num_hidden_layers": 1} + stored = {"model_type": "kimi_vl", "auto_map": auto_map, "text_config": text} + (tmp_path / "config.json").write_text(json.dumps(stored)) + config = AutoConfig.from_pretrained(tmp_path, trust_remote_code=True) + + kimi_vl.make_buildable(config, ckpt_path=str(tmp_path), trust_remote_code=True) + model = model_load_utils.build_meta_causal_lm( + str(tmp_path), True, None, copy.deepcopy(config), auto_class=AutoModelForCausalLM + ) + + assert type(model).tie_weights._kimi_tie_weights_shim + assert sys.modules[type(model).__module__].DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + + @_BOTH_ADAPTERS @pytest.mark.parametrize( ("missing", "match"), diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 893cbf26fd5..97da0d337c4 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -841,7 +841,7 @@ def test_a_checkpoint_its_spec_decompresses_is_decided_expert_parallel( expert-parallel load would refuse it.""" spec = get_spec(model_type).ep_spec monkeypatch.setattr(spec, "prepare", lambda hf_config, **kwargs: hf_config) - monkeypatch.setattr(spec, "make_importable", lambda: None) + monkeypatch.setattr(spec, "make_buildable", lambda *args, **kwargs: None) monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi()) monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) run_as(4) From ea86e8f46bac96a5a102bdeca6d2bfaa1ef28c96 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:29:54 +0000 Subject: [PATCH 10/17] Derive the Kimi-K2.5 decompressor test's expectation from its spec From Claude's review: the test expected the decompressor for compressed-tensors by a literal, beside the spec's decompressed_formats that the plan decision reads. It now expects an install exactly for the formats the spec states, so the spec and prepare cannot drift apart without failing it. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/kimi_k25/test_ep.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/models/kimi_k25/test_ep.py b/tests/unit/torch/models/kimi_k25/test_ep.py index d0dd29ebad3..69105898412 100644 --- a/tests/unit/torch/models/kimi_k25/test_ep.py +++ b/tests/unit/torch/models/kimi_k25/test_ep.py @@ -107,8 +107,9 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3( """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 for compressed-tensors - experts only, while the text config it replaces still carries their scheme; the remote-code - staging shim is left installed.""" + experts only, while the text config it replaces still carries their scheme, and exactly for the + formats the spec states it decompresses, which the plan decision counts loadable before prepare + runs; the remote-code staging shim is left installed.""" installs = [] monkeypatch.setattr( _onfly_decompress, @@ -133,7 +134,9 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3( assert _routed_expert_keys(text.base_model_ep_plan) == _routed_expert_keys( get_spec("kimi_k25").ep_spec.prepared_ep_plan ) - assert installs == ([QUANT] if quantization_config is QUANT else []) + decompressed = get_spec("kimi_k25").ep_spec.decompressed_formats + expected = (quantization_config or {}).get("quant_method") in decompressed + assert installs == ([QUANT] if expected else []) assert dmu.get_cached_module_file is not staging From fe075f40010219a0a2ab9f952559bd0c585d5fc2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:29:57 +0000 Subject: [PATCH 11/17] Pin a rank failing alone where the decided expert-parallel plan is prepared From Claude's review: the step that prepares a decided expert-parallel plan is a collective of its own, which carries every prepare refusal, and no test reached it with a rank failing alone. The rank-failure test gains a prepared case: rank 1 states the spec as rank 0 does and fails only where prepare runs, and every rank must raise its error; it fails when that step is not a collective. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 97da0d337c4..95a7007db71 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1989,9 +1989,12 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, def _one_rank_cannot_apply_the_spec_job(rank, size, ckpt, step): apply = model_load_utils.apply_distributed_support + prepares = [] def apply_on_rank0_only(*args, **kwargs): - if rank == 1: + prepares.append(kwargs.get("prepare", True)) + # "prepared": rank 1 states the spec as rank 0 does, and fails only where prepare runs. + if rank == 1 and (step != "prepared" or prepares[-1]): raise OSError("this rank cannot import the checkpoint's code") return apply(*args, **kwargs) @@ -2011,7 +2014,7 @@ def apply_on_rank0_only(*args, **kwargs): mock.patch.object(model_load_utils, "load_and_prepare_ep", _loaded_in_step) ) try: - if step == "decision": + if step in ("decision", "prepared"): resolve_distribution_plan(ckpt) else: load_pretrained(ckpt, distribution=DistributionPlan(2, ep_size=2)) @@ -2020,14 +2023,17 @@ def apply_on_rank0_only(*args, **kwargs): else: raised = "no error" assert "cannot import the checkpoint's code" in raised, raised + # Deciding an expert-parallel plan, every rank reached the step that prepares it. + assert (True in prepares) == (step != "decision"), prepares @needs_ep_transformers -@pytest.mark.parametrize("step", ["decision", "load"]) +@pytest.mark.parametrize("step", ["decision", "prepared", "load"]) def test_a_rank_that_cannot_apply_the_spec_fails_every_rank(cpu_dist_workers, tmp_path, step): """A spec's preparation reads the checkpoint and imports its code on each rank, which can fail - on one alone; the others must not wait for it, in the plan's broadcast when deciding, nor in the - load a handed expert-parallel plan starts.""" + on one alone; the others must not wait for it, in the plan's broadcast when deciding, in the + step that prepares the expert-parallel plan decided, nor in the load a handed expert-parallel + plan starts.""" _deepseek_v3(n_routed_experts=8).save_pretrained(tmp_path) cpu_dist_workers.run(_one_rank_cannot_apply_the_spec_job, ckpt=str(tmp_path), step=step) From 978b2704b6cb6d069aef752ce7348f24bde1ea8b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:33:32 +0000 Subject: [PATCH 12/17] Unload the build-step test's remote Kimi code after it The test imports the tiny Kimi-VL checkpoint's code through transformers' loader, which keeps it in sys.modules; a later test then found it as a loaded modeling_kimi_vl. It now drops the remote modules it loaded. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/kimi_vl/test_ep.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/tests/unit/torch/models/kimi_vl/test_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py index 5e93c0fbde2..3f96a03e030 100644 --- a/tests/unit/torch/models/kimi_vl/test_ep.py +++ b/tests/unit/torch/models/kimi_vl/test_ep.py @@ -311,15 +311,22 @@ def test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_transformers text = {"model_type": "deepseek_v3", "hidden_size": 16, "num_hidden_layers": 1} stored = {"model_type": "kimi_vl", "auto_map": auto_map, "text_config": text} (tmp_path / "config.json").write_text(json.dumps(stored)) - config = AutoConfig.from_pretrained(tmp_path, trust_remote_code=True) + loaded = set(sys.modules) + try: + config = AutoConfig.from_pretrained(tmp_path, trust_remote_code=True) - kimi_vl.make_buildable(config, ckpt_path=str(tmp_path), trust_remote_code=True) - model = model_load_utils.build_meta_causal_lm( - str(tmp_path), True, None, copy.deepcopy(config), auto_class=AutoModelForCausalLM - ) + kimi_vl.make_buildable(config, ckpt_path=str(tmp_path), trust_remote_code=True) + model = model_load_utils.build_meta_causal_lm( + str(tmp_path), True, None, copy.deepcopy(config), auto_class=AutoModelForCausalLM + ) - assert type(model).tie_weights._kimi_tie_weights_shim - assert sys.modules[type(model).__module__].DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + assert type(model).tie_weights._kimi_tie_weights_shim + assert sys.modules[type(model).__module__].DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + finally: + # The checkpoint's code stays imported otherwise, for later tests to find as a Kimi module. + for name in set(sys.modules) - loaded: + if name.startswith("transformers_modules."): + del sys.modules[name] @_BOTH_ADAPTERS From c83024c4381f950fead60dce5283f87bf3229b20 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 01:10:50 +0000 Subject: [PATCH 13/17] Run the Kimi build step only when there is code of the checkpoint's to build From Claude's review: make_buildable installed the import shims and the staging before the checks that say nothing of the checkpoint's will be built, a config of transformers' own port, no checkpoint path, or code the load may not run. It now returns before them in those cases, so the step leaves the process as it found it. New test test_the_build_step_installs_nothing_when_there_is_nothing_to_build fails with the shims first. The build-step test also checks that the VLM's language model stays the checkpoint's own, per-expert class. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 2 +- tests/unit/torch/models/kimi_vl/test_ep.py | 26 +++++++++++++++++++++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index fbb33a658f6..6f40c4750f2 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -136,11 +136,11 @@ def make_buildable( from transformers.dynamic_module_utils import get_class_from_dynamic_module - make_importable() if type(hf_config).__module__.startswith("transformers.") or ckpt_path is None: return if not trust_remote_code: return + make_importable() get_class_from_dynamic_module( f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code ) diff --git a/tests/unit/torch/models/kimi_vl/test_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py index 3f96a03e030..c2d7c820991 100644 --- a/tests/unit/torch/models/kimi_vl/test_ep.py +++ b/tests/unit/torch/models/kimi_vl/test_ep.py @@ -294,6 +294,27 @@ def test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code( assert config.text_config.base_model_ep_plan is None +@pytest.mark.parametrize( + "kwargs", + [{"ckpt_path": None, "trust_remote_code": True}, {"ckpt_path": "/ckpt"}], + ids=["no-checkpoint", "untrusted"], +) +def test_the_build_step_installs_nothing_when_there_is_nothing_to_build(monkeypatch, kwargs): + """With no checkpoint, or code this load may not run, nothing of the checkpoint's is imported or + built, so the build step installs no import shim and no staging either.""" + for module_name, name in COMPAT_SYMBOLS: + with contextlib.suppress(ImportError): + _restore(monkeypatch, importlib.import_module(module_name), name) + monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) + staging = dmu.get_cached_module_file + monkeypatch.setattr(dmu, "get_cached_module_file", staging) + + kimi_vl.make_buildable(_vlm_config(), **kwargs) + + assert dmu.get_cached_module_file is staging + assert not hasattr(import_utils, "is_torch_fx_available") + + def test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_transformers( monkeypatch, tmp_path ): @@ -321,7 +342,10 @@ def test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_transformers ) assert type(model).tie_weights._kimi_tie_weights_shim - assert sys.modules[type(model).__module__].DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + # Nothing swapped: the language model is the checkpoint's own, its experts per expert. + own = sys.modules[type(model).__module__].DeepseekV3ForCausalLM + assert not getattr(own, "_kimi_native_swap", False) + assert type(model.language_model) is own finally: # The checkpoint's code stays imported otherwise, for later tests to find as a Kimi module. for name in set(sys.modules) - loaded: From 84c06e0855e306afa9d8c92f935750a112aeb99c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 01:11:00 +0000 Subject: [PATCH 14/17] State the formats Kimi-K2.5's prepare decompresses once From Claude's review: compressed-tensors was written both in prepare's condition for installing the decompressor and in the spec's decompressed_formats, which the plan decision reads before prepare runs. Both now read kimi_k25.ep.DECOMPRESSED_FORMATS. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k25/ep.py | 6 +++++- modelopt/torch/models/kimi_k25/specs.py | 4 ++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py index 0895215af78..b9b6451360b 100644 --- a/modelopt/torch/models/kimi_k25/ep.py +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -52,6 +52,10 @@ from .. import _kimi, _onfly_decompress +# The checkpoint formats prepare has the experts decompress as they load; the spec states them for +# the plan decision, which runs before prepare. +DECOMPRESSED_FORMATS = ("compressed-tensors",) + def make_buildable( hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_ @@ -91,7 +95,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool # from the copy of quantization_config Kimi's config class keeps on the top-level config; # step 2 leaves that copy alone. A checkpoint without that copy is refused. quantization_config = _onfly_decompress._find_quantization_config(hf_config) or {} - if quantization_config.get("quant_method") == "compressed-tensors": + if quantization_config.get("quant_method") in DECOMPRESSED_FORMATS: _onfly_decompress.install_onfly_decompress(hf_config) else: _kimi.refuse_quantization_only_on_text_config("kimi_k25", hf_config) diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py index 182691f0263..93035a9d3b6 100644 --- a/modelopt/torch/models/kimi_k25/specs.py +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -27,7 +27,7 @@ from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN from ..specs import EPSpec, ModelSpec, register -from .ep import make_buildable, prepare +from .ep import DECOMPRESSED_FORMATS, make_buildable, prepare __all__: list[str] = [] @@ -48,7 +48,7 @@ prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", make_buildable=make_buildable, - decompressed_formats=("compressed-tensors",), + decompressed_formats=DECOMPRESSED_FORMATS, requires_transformers="5.8", ), ) From 5ae670582e55d3ddbb7162b71ec9285f41214e71 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 03:44:55 +0000 Subject: [PATCH 15/17] Keep the Kimi build step's tie_weights shim to remote code From Claude's review: make_buildable adapted tie_weights on every loaded module whose name ends with the modeling module's, which from transformers 5.13 includes transformers' own Kimi-K2.5 port, transformers.models.kimi_k25.modeling_kimi_k25, whose class a later load of the port would then use with its keyword arguments dropped. It now goes through _remote_code._apply, which reaches remote modules only, as the swap does. test_the_build_step_leaves_transformers_own_port_alone fails with the unscoped scan. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 9 +++------ tests/unit/torch/models/kimi_k25/test_ep.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index 6f40c4750f2..3914104dfdd 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -23,7 +23,7 @@ # modelopt.torch.utils and transformers are imported inside the functions: the model registry imports # this module, and it must import without transformers, the optional [hf] extra. -from ._remote_code import keep_patched, stage_full_remote_closure +from ._remote_code import _apply, keep_patched, stage_full_remote_closure def checkpoint_to_patch( @@ -132,8 +132,6 @@ def make_buildable( whatever plan comes out. A config of transformers' own port, or code this load may not run, has nothing of the checkpoint's to build. Idempotent. """ - import sys - from transformers.dynamic_module_utils import get_class_from_dynamic_module if type(hf_config).__module__.startswith("transformers.") or ckpt_path is None: @@ -144,9 +142,8 @@ def make_buildable( get_class_from_dynamic_module( f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code ) - for name, module in list(sys.modules.items()): - if name.endswith(f".{modeling}") and module is not None: - _shim_tie_weights(module, vlm_class) + # Remote code only: transformers' own port of a model can share the module's name. + _apply(f".{modeling}", lambda module: _shim_tie_weights(module, vlm_class)) def _shim_tie_weights(module, vlm_class: str) -> bool: diff --git a/tests/unit/torch/models/kimi_k25/test_ep.py b/tests/unit/torch/models/kimi_k25/test_ep.py index 69105898412..9fc68baaa87 100644 --- a/tests/unit/torch/models/kimi_k25/test_ep.py +++ b/tests/unit/torch/models/kimi_k25/test_ep.py @@ -154,6 +154,25 @@ def test_the_swap_leaves_a_sibling_remote_module_alone(monkeypatch, remote): assert bystander.DeepseekV3ForCausalLM is remote_lm +def test_the_build_step_leaves_transformers_own_port_alone(monkeypatch, remote): + """From 5.13 transformers ships its own Kimi-K2.5, whose module shares the remote one's name + (``transformers.models.kimi_k25.modeling_kimi_k25``). The build step adapts the remote VLM + class's ``tie_weights`` and must not reach the port's, which a later load of it would use.""" + + def own_tie(self): + return "tied" + + port = types.ModuleType("transformers.models.kimi_k25.modeling_kimi_k25") + port.KimiK25ForConditionalGeneration = type("Port", (), {"tie_weights": own_tie}) + monkeypatch.setitem(sys.modules, port.__name__, port) + remote.KimiK25ForConditionalGeneration = type("Remote", (), {"tie_weights": own_tie}) + + kimi_k25.make_buildable(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + + assert vars(port.KimiK25ForConditionalGeneration)["tie_weights"] is own_tie + assert vars(remote.KimiK25ForConditionalGeneration)["tie_weights"]._kimi_tie_weights_shim + + def test_a_checkpoint_the_decompressor_refuses_is_refused_before_anything_changes( monkeypatch, remote ): From 7a1f1cff28c6f1d8b3900c3a41e6526880059d1f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:09:50 +0000 Subject: [PATCH 16/17] Test that a spec with a prepare states the plan the decision reads Deciding the plan no longer runs prepare, so a spec whose plan only prepare attaches would have its MoE read as dense and planned without expert parallelism, with no refusal. Both Kimi specs state theirs; the test pins it for every EP spec. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_ep_specs.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 7a4bd9ef182..76b23b2a26b 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -87,6 +87,17 @@ def test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal(model_type): ) +@pytest.mark.parametrize("model_type", EP_MODEL_TYPES) +def test_a_spec_with_a_prepare_states_its_plan(model_type): + """Deciding the plan runs no prepare, so a plan only prepare attaches goes unseen there: the + decision reads the MoE as dense and plans it without expert parallelism, with no refusal.""" + ep = get_spec(model_type).ep_spec + assert not ep.prepare or ep.base_model_ep_plan or ep.prepared_ep_plan, ( + f"{model_type} has a prepare and states no plan, which the plan decision needs, as it runs " + "no prepare: state the plan prepare attaches as prepared_ep_plan" + ) + + @pytest.mark.parametrize("model_type", EP_MODEL_TYPES) def test_a_plan_goes_where_the_load_reads_one(model_type): """A spelling the load does not read is read as no plan at all, and the MoE planned as dense.""" From b88d349be4d7830883093ca5440cb0fa95675ea7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:33:40 +0000 Subject: [PATCH 17/17] Say how the plan decision treats the formats Kimi-K2.5's prepare decompresses Signed-off-by: Shengliang Xu #2746 no longer keeps a loadable flag in the decision: it holds the expert-parallel load's refusal against an MoE, and the formats a spec states it decompresses are not among the reasons. --- tests/unit/torch/models/kimi_k25/test_ep.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/models/kimi_k25/test_ep.py b/tests/unit/torch/models/kimi_k25/test_ep.py index 9fc68baaa87..6e51d9aba44 100644 --- a/tests/unit/torch/models/kimi_k25/test_ep.py +++ b/tests/unit/torch/models/kimi_k25/test_ep.py @@ -108,8 +108,8 @@ def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3( 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 for compressed-tensors experts only, while the text config it replaces still carries their scheme, and exactly for the - formats the spec states it decompresses, which the plan decision counts loadable before prepare - runs; the remote-code staging shim is left installed.""" + formats the spec states it decompresses, which the plan decision, made before prepare runs, does + not hold against the expert-parallel load; the remote-code staging shim is left installed.""" installs = [] monkeypatch.setattr( _onfly_decompress,