From 199f59f4759247506af6863993d1c725f196b1d8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:03:44 +0000 Subject: [PATCH 1/9] Add the Kimi-VL expert-parallel adapter Kimi-VL ships its own modeling, which builds the language model's MoE from an nn.ModuleList of per-expert Linears that transformers' expert- parallel machinery cannot shard. Its text backbone is DeepSeek-V3, so the adapter swaps in transformers' native deepseek_v3 language model (fused gate_up_proj/down_proj experts, and the plan to shard them) and keeps the bundled vision tower and projector from the checkpoint's code. The patch Kimi-VL and Kimi-K2.5 share lives in models/_kimi.py: the refusals prepare runs before it changes anything (transformers' own port, no text_config, no checkpoint path, untrusted code, a scheme on text_config alone), the import shims for symbols transformers 5 dropped, and the swap of the module-level DeepseekV3ForCausalLM, kept on a module reload through keep_patched, with a tie_weights wrapper for the kwargs newer transformers passes. deepseek_v3's expert-parallel plan becomes a shared BASE_MODEL_EP_PLAN, which the spec states as the plan prepare attaches, so a run at an EP degree of 1, which skips prepare, still finds the experts the checkpoint's own code builds and refuses to hold them whole on every rank. 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.py | 184 ++++++++++++ modelopt/torch/models/deepseek_v3/specs.py | 23 +- modelopt/torch/models/kimi_vl/__init__.py | 18 ++ modelopt/torch/models/kimi_vl/ep.py | 79 ++++++ modelopt/torch/models/kimi_vl/specs.py | 39 +++ tests/unit/torch/models/test_kimi_vl_ep.py | 266 ++++++++++++++++++ .../unit/torch/utils/test_model_load_utils.py | 38 +++ 8 files changed, 637 insertions(+), 11 deletions(-) create mode 100644 modelopt/torch/models/_kimi.py create mode 100644 modelopt/torch/models/kimi_vl/__init__.py create mode 100644 modelopt/torch/models/kimi_vl/ep.py create mode 100644 modelopt/torch/models/kimi_vl/specs.py create mode 100644 tests/unit/torch/models/test_kimi_vl_ep.py diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index 3d107b7e12d..b87f3872f73 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -91,6 +91,7 @@ nemotron_h, minimax_m3_vl, llama4_text, + kimi_vl, nemotron_h_omni, qwen2_moe, qwen3, diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py new file mode 100644 index 00000000000..f18e598bfd2 --- /dev/null +++ b/modelopt/torch/models/_kimi.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""What the Kimi-K2.5 and Kimi-VL adapters change in the modeling code their checkpoints ship. + +Both checkpoints' own VLM class builds its language model as ``DeepseekV3ForCausalLM( +config.text_config)``, a DeepSeek-V3 whose MoE is an ``nn.ModuleList`` of per-expert Linears that +expert parallelism cannot shard. The adapters point that name at transformers' native class, and +keep the vision tower and projector from the checkpoint's code. +""" + +# modelopt.torch.utils is imported inside the functions: the model registry imports this module. +from ._remote_code import keep_patched + + +def checkpoint_to_patch( + label: str, hf_config, ckpt_path: str | None, trust_remote_code: bool +) -> str: + """Return the checkpoint whose own modeling code the adapter ``label`` patches. + + Refuses first, before its ``prepare`` changes anything, a config the adapter cannot patch. + """ + if type(hf_config).__module__.startswith("transformers."): + # transformers' own port (Kimi-K2.5's, from 5.13), which the patch never reaches: a load + # without trust_remote_code builds it, and so does one of a checkpoint naming no config class. + raise ValueError( + f"{label}: expert parallelism runs on the checkpoint's own modeling code, but this " + f"config is transformers' own port ({type(hf_config).__module__}). Load it with " + "trust_remote_code=True, from a checkpoint whose auto_map names its own AutoConfig." + ) + if not hasattr(hf_config, "text_config"): + raise ValueError(f"{label}: hf_config has no .text_config (expected a VLM config)") + if ckpt_path is None: + raise ValueError( + f"{label}: expert parallelism patches the checkpoint's own modeling code, so it needs " + "the checkpoint's path (apply_distributed_support's ckpt_path)." + ) + if not trust_remote_code: + # Patching imports the checkpoint's own modeling code, and transformers' + # get_class_from_dynamic_module runs it without checking trust_remote_code itself. + raise ValueError( + f"{label}: expert parallelism imports the checkpoint's own modeling code; " + "load it with trust_remote_code=True." + ) + return ckpt_path + + +def refuse_quantization_only_on_text_config(label: str, hf_config) -> None: + """Refuse a checkpoint whose ``quantization_config`` only its ``text_config`` declares. + + The adapter replaces ``text_config`` with a native config without it. Unless the experts + decompress as they load, the expert-parallel load then reads the checkpoint's quantization off + the top-level config, which Kimi-K2.5's config class copies it to; without that copy, the load + would read quantized experts as stored. + """ + if ( + getattr(hf_config.text_config, "quantization_config", None) is not None + and getattr(hf_config, "quantization_config", None) is None + ): + raise ValueError( + f"{label}: the checkpoint declares quantization_config on text_config only, which the " + "native deepseek_v3 config replacing it drops, so the expert-parallel load would read " + "its experts as stored. Declare it on the top-level config too, or start from an " + "unquantized copy of the model." + ) + + +def install_compat_import_shims() -> None: + """Add stubs for symbols the Kimi modeling imports that newer transformers removed, so it imports. + + Each guards a torch.fx / legacy-attn-mask code path the native LM + MoonViT forward never hit, so a benign stub is + safe. Idempotent (only sets missing attrs). + """ + from modelopt.torch.utils import print_rank_0 + + try: + import transformers.utils.import_utils as iu + + if not hasattr(iu, "is_torch_fx_available"): + iu.is_torch_fx_available = lambda: False + except Exception as e: + print_rank_0(f"[kimi] compat-shim warning (import_utils): {type(e).__name__}: {e}") + try: + import transformers.pytorch_utils as pu + + if not hasattr(pu, "is_torch_greater_or_equal_than_1_13"): + pu.is_torch_greater_or_equal_than_1_13 = True + except Exception as e: + print_rank_0(f"[kimi] compat-shim warning (pytorch_utils): {type(e).__name__}: {e}") + try: + import transformers.modeling_attn_mask_utils as amu + + if not hasattr(amu, "_prepare_4d_causal_attention_mask"): + + def _prepare_4d_causal_attention_mask(*args, **kwargs): # pragma: no cover + raise RuntimeError( + "kimi: remote _prepare_4d_causal_attention_mask stub called -- the native LM " + "builds its own mask, so this path should be unreachable." + ) + + amu._prepare_4d_causal_attention_mask = _prepare_4d_causal_attention_mask + except Exception as e: + print_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}") + + +def _drop_tie_weights_kwargs(orig_tie): + """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it. + + Dropping ``missing_keys`` as well as ``recompute_mapping`` is safe because the remote + ``tie_weights`` only calls its language model's, and the Kimi checkpoints set + ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. + """ + + def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc. + return orig_tie(self) + + setattr(_tie_weights, "_kimi_tie_weights_shim", True) + return _tie_weights + + +def patch_language_model_to_native( + label: str, modeling: str, vlm_class: str, ckpt_path: str, trust_remote_code: bool +) -> None: + """Point the checkpoint's ``modeling`` module's ``DeepseekV3ForCausalLM`` at the native class. + + Loads the checkpoint's remote ``modeling`` module and replaces its ``DeepseekV3ForCausalLM`` with + transformers' native class, so ``vlm_class.__init__`` builds the EP-capable native deepseek_v3 + for ``self.language_model`` instead of the per-expert Linears version. + """ + from transformers.dynamic_module_utils import get_class_from_dynamic_module + from transformers.models.deepseek_v3 import DeepseekV3ForCausalLM as NativeDeepseekV3ForCausalLM + + from modelopt.torch.utils import print_rank_0 + + # Force the dynamic-module loader to stage + import the module into sys.modules. + get_class_from_dynamic_module( + f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code + ) + + def to_native(module): + current = getattr(module, "DeepseekV3ForCausalLM", None) + if current is None: + # Not "already native": the VLM class then builds its language model from a class this + # swap does not reach (renamed, or imported inside __init__), with per-expert Linears. + raise RuntimeError( + f"{label}: {module.__name__} defines no module-level DeepseekV3ForCausalLM to swap " + "for the transformers-native class; the checkpoint's modeling code changed." + ) + if current is NativeDeepseekV3ForCausalLM: + return False + module.DeepseekV3ForCausalLM = NativeDeepseekV3ForCausalLM + # The remote tie_weights() takes no kwargs, but transformers' post_init and load call + # tie_weights(recompute_mapping=False) -> TypeError. + kimi_cls = getattr(module, vlm_class, None) + if kimi_cls is not None: + own_tie = vars(kimi_cls).get("tie_weights") + if own_tie is not None and not getattr(own_tie, "_kimi_tie_weights_shim", False): + kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie) + return True + + # The dynamic-module loader names the module like + # ``transformers_modules...``. It stays patched when from_pretrained + # resolves the class again, even from a re-executed module. + patched = keep_patched(f".{modeling}", to_native) + if not patched: + raise RuntimeError( + f"{label}: could not find a loaded {modeling} in sys.modules to monkey-patch" + ) + print_rank_0( + f"[{label}] monkey-patched DeepseekV3ForCausalLM -> transformers-native (+tie_weights shim) " + f"in: {patched}" + ) diff --git a/modelopt/torch/models/deepseek_v3/specs.py b/modelopt/torch/models/deepseek_v3/specs.py index d2f170cfc06..b7074c80fbd 100644 --- a/modelopt/torch/models/deepseek_v3/specs.py +++ b/modelopt/torch/models/deepseek_v3/specs.py @@ -23,6 +23,17 @@ __all__: list[str] = [] +# The plan for transformers releases that ship none for deepseek_v3 (before 5.13). From 5.13 its own +# plan wins (an existing plan always does, see apply_distributed_support): the same experts, plus an +# ep_router entry for the gate, which the load swaps for torch_a2a_router (global ids, unmasked). No +# gate entry here on purpose: the router is logits-only and the MoE block computes GLOBAL top-k ids, +# so the experts already receive unmasked global ids -- what the all-to-all dispatch wants. +BASE_MODEL_EP_PLAN = { + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", +} + # DeepseekV3MoE is invisible to the generic MoE detection: its class name does not end # in "SparseMoeBlock" and it calls its router ``gate``, so neither the name test nor the # structural router+experts test in is_moe matches. block_names below is what puts it on @@ -37,17 +48,7 @@ model_type="deepseek_v3", min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), - ep_spec=EPSpec( - # Stock transformers ships no plan for deepseek_v3 before 5.13; from 5.13 the config's - # own plan wins (see apply_distributed_support) and this one goes unused. No gate entry - # on purpose: the router is logits-only and the MoE block computes GLOBAL top-k ids, so - # the experts already receive unmasked global ids -- what the all-to-all dispatch wants. - base_model_ep_plan={ - "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", - }, - ), + ep_spec=EPSpec(base_model_ep_plan=BASE_MODEL_EP_PLAN), moe_spec=MoESpec( block_names=("DeepseekV3MoE",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/torch/models/kimi_vl/__init__.py b/modelopt/torch/models/kimi_vl/__init__.py new file mode 100644 index 00000000000..ec89db2722b --- /dev/null +++ b/modelopt/torch/models/kimi_vl/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""kimi_vl specs (HF model type ``kimi_vl``).""" + +from . import specs diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py new file mode 100644 index 00000000000..d1891f5c781 --- /dev/null +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DP=EP a2a support for Kimi-VL-A3B (``kimi_vl``). + +a remote-code VLM whose text backbone is a Moonlight (DeepSeek-V3-lite: MLA, ``q_lora_rank=None``, 64 routed experts, 2 +shared) MoE and whose vision tower is MoonViT. The bf16 sibling of ``kimi_k25`` -- SAME shape, but NO compressed-tensors +decompressor (the checkpoint ships plain bf16 per-expert weights). + +Challenge (identical to kimi_k25): the bundled ``modeling_kimi_vl.py`` builds its language model from +an INLINE ``DeepseekV3ForCausalLM`` whose ``DeepseekV3MoE`` is an ``nn.ModuleList`` of per-expert +``gate_proj``/``up_proj``/``down_proj`` Linears -- incompatible with transformers' EP machinery, which +needs the FUSED 3-D ``gate_up_proj``/``down_proj`` Parameters the transformers-native ``deepseek_v3`` +modeling ships. + +Solution (mirrors kimi_k25): + 1. Shim the import symbols ``modeling_kimi_vl`` needs that newer transformers dropped (e.g. + ``is_torch_fx_available``) so the module IMPORTS -- they only guard torch.fx / legacy attn-mask + paths that our native-LM + MoonViT forward never execute. + 2. Swap ``hf_config.text_config`` to a transformers-native ``DeepseekV3Config`` (drops the remote + ``auto_map`` so the native class -- and its ``base_model_ep_plan`` -- is used for the LM). + 3. Apply the deepseek_v3 EP plan to that native ``text_config`` (post_init prefixes it to the VLM + root as ``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}``). + 4. Monkey-patch ``modeling_kimi_vl.DeepseekV3ForCausalLM`` -> the transformers-native class, so + ``KimiVLForConditionalGeneration.__init__`` (``self.language_model = DeepseekV3ForCausalLM( + config.text_config)``) builds the EP-capable native LM while MoonViT + the KimiVL wrapper stay + from the bundled code. Native deepseek_v3's per-expert WeightConverter fuses the checkpoint's + per-expert ``.weight`` tensors into the 3-D params on load (no decompressor -- weights are bf16). + +Version: no floor beyond the EP framework's 5.7; validated on transformers 5.9.0. The compat shims +stand in for the symbols the remote code imports that later releases dropped. +""" + +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-VL's language model as the transformers-native deepseek_v3. + + Swap text_config to native deepseek_v3 + apply the EP plan + monkey-patch Kimi-VL's modeling to + build the native EP-capable DeepseekV3ForCausalLM for the language sub-model. No decompressor: the + Kimi-VL checkpoint is plain bf16. + """ + # 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_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() + + # Native deepseek_v3 config for the LM (drops the remote auto_map so the native class is used). + hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config) + # EP plan onto the native text_config (bubbles up to the VLM root under language_model.*). + hf_config.text_config = apply_distributed_support("deepseek_v3", hf_config.text_config) + + _kimi.patch_language_model_to_native( + "kimi_vl", + "modeling_kimi_vl", + "KimiVLForConditionalGeneration", + ckpt_path, + trust_remote_code, + ) + return hf_config diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py new file mode 100644 index 00000000000..a5d9eacec42 --- /dev/null +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""kimi_vl specs (HF model type ``kimi_vl``). + +Kimi-VL-A3B: a remote-code VLM whose text backbone is a Moonlight MoE. Its bundled +modeling builds the language model from per-expert Linears, which the EP machinery cannot shard, so +``ep.py`` swaps in the transformers-native deepseek_v3 language model. + +Only the expert-parallel section is registered under ``kimi_vl``: the MoE and export sections stay +``deepseek_v3``'s, so lookups by ``kimi_vl`` find neither yet. Sharing them is left to a follow-up. +""" + +from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN +from ..specs import EPSpec, ModelSpec, register +from .ep import prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_vl", + modeling_source="remote_code", + # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. + ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN), + ) +) diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py new file mode 100644 index 00000000000..6d82fb11543 --- /dev/null +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Expert-parallel adapter for Kimi-VL-A3B (``kimi_vl``), a remote-code VLM on a Moonlight MoE. + +The adapter rebuilds the language model as transformers' native deepseek_v3; its experts are plain +bf16, so nothing is decompressed. The checkpoint's remote code is not available here, so the tests +stub the dynamic-module loader with a stand-in remote module. Every process-wide patch is undone +after each test. +""" + +import importlib +import sys +import types +from types import SimpleNamespace + +import pytest + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.7"): + pytest.skip("Kimi-VL expert parallelism needs transformers >= 5.7", allow_module_level=True) + +import transformers.dynamic_module_utils as dmu +import transformers.utils.import_utils as import_utils +from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig + +from modelopt.torch.models import _onfly_decompress, _remote_code, get_spec +from modelopt.torch.models.kimi_vl import ep as kimi_vl +from modelopt.torch.utils.moe import _routed_expert_keys + +# Symbols the remote modeling imports that newer transformers dropped. +COMPAT_SYMBOLS = ( + ("transformers.utils.import_utils", "is_torch_fx_available"), + ("transformers.pytorch_utils", "is_torch_greater_or_equal_than_1_13"), + ("transformers.modeling_attn_mask_utils", "_prepare_4d_causal_attention_mask"), +) + + +def _restore(monkeypatch, owner, name): + """Have ``monkeypatch`` put ``owner.`` back after the test, or remove it if it was absent.""" + if hasattr(owner, name): + monkeypatch.setattr(owner, name, getattr(owner, name)) + else: + monkeypatch.setattr(owner, name, None, raising=False) + monkeypatch.delattr(owner, name) + + +class _RemoteTextConfig(PretrainedConfig): + """Stands in for the remote DeepSeek-V3 config Kimi-VL bundles as its ``text_config``.""" + + model_type = "kimi_vl_remote_text" + + +def _vlm_config(): + text = _RemoteTextConfig( + hidden_size=64, + num_hidden_layers=2, + n_routed_experts=8, + moe_intermediate_size=16, + q_lora_rank=None, + auto_map={"AutoConfig": "configuration_kimi_vl.DeepseekV3Config"}, + ) + return SimpleNamespace(text_config=text) + + +def _block_fp8_on_the_text_config_only(): + config = _vlm_config() + block_fp8 = {"quant_method": "fp8", "weight_block_size": [128, 128]} + config.text_config.quantization_config = block_fp8 + return config + + +@pytest.fixture +def remote(monkeypatch, request): + """Stand in for the checkpoint's remote code: importing it registers its modeling module + (Kimi-VL's ``modeling_kimi_vl`` unless a test names another), with the inline per-expert + ``DeepseekV3ForCausalLM`` and a VLM class whose ``tie_weights`` has the old no-kwarg signature, + as both Kimi checkpoints ship it. Also restores the compat symbols, loader hooks and kept + patches prepare installs.""" + modeling, vlm_name = getattr( + request, "param", ("modeling_kimi_vl", "KimiVLForConditionalGeneration") + ) + for module_name, name in COMPAT_SYMBOLS: + try: + owner = importlib.import_module(module_name) + except ImportError: + continue + _restore(monkeypatch, owner, name) + module = types.ModuleType(f"transformers_modules.Kimi.0123abcd.{modeling}") + module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + module.loads, module.tied = [], [] + + def tie_weights(self): + module.tied.append(self) + return "tied" + + module.vlm = type(vlm_name, (), {"tie_weights": tie_weights}) + setattr(module, vlm_name, module.vlm) + + def import_remote_class(class_ref, ckpt_path, **kwargs): + module.loads.append((class_ref, ckpt_path, kwargs)) + monkeypatch.setitem(sys.modules, module.__name__, module) + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) + monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) + return module + + +def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monkeypatch, remote): + """The text config becomes transformers' own ``DeepseekV3Config`` (no remote ``auto_map``) with + an EP plan for the fused experts, naming the experts the spec states for a run that skips + prepare; the experts are bf16, so no decompressor is installed; the symbols the remote modeling + imports exist; the remote-code staging shim is left installed.""" + monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) + monkeypatch.setattr( + _onfly_decompress, + "install_onfly_decompress", + lambda cfg: pytest.fail("Kimi-VL's experts are bf16: there is nothing to decompress"), + ) + staging = dmu.get_cached_module_file + config = _vlm_config() + + assert kimi_vl.prepare(config, 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 (text.hidden_size, text.n_routed_experts, text.moe_intermediate_size) == (64, 8, 16) + assert text.base_model_ep_plan["layers.*.mlp.experts.gate_up_proj"] == "grouped_gemm" + assert text.base_model_ep_plan["layers.*.mlp.experts.down_proj"] == "grouped_gemm" + assert _routed_expert_keys(text.base_model_ep_plan) == _routed_expert_keys( + get_spec("kimi_vl").ep_spec.prepared_ep_plan + ) + assert import_utils.is_torch_fx_available() is False + assert dmu.get_cached_module_file is not staging + + +# What the Kimi adapters share, through each adapter's prepare. +_BOTH_ADAPTERS = pytest.mark.parametrize( + ("adapter", "remote"), + [(kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration"))], + ids=["vl"], + indirect=["remote"], +) + + +@_BOTH_ADAPTERS +def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_weights( + monkeypatch, adapter, remote +): + """The VLM class builds its language model from the module-level ``DeepseekV3ForCausalLM``, + which must now be transformers' native class. Both checkpoints' remote modeling imports a + symbol newer transformers dropped, and has an old no-kwarg ``tie_weights`` that must accept the + ``recompute_mapping`` newer transformers passes. Running prepare again does not wrap + ``tie_weights`` twice.""" + monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + shim = remote.vlm.tie_weights + adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + + vlm_class = f"{remote.__name__.rsplit('.', 1)[-1]}.{remote.vlm.__name__}" + assert remote.loads == [(vlm_class, "/ckpt", {"trust_remote_code": True})] * 2 + assert remote.DeepseekV3ForCausalLM is DeepseekV3ForCausalLM + assert import_utils.is_torch_fx_available() is False + assert remote.vlm.tie_weights is shim + model = remote.vlm() + assert model.tie_weights(missing_keys=set(), recompute_mapping=False) == "tied" + assert remote.tied == [model] + + +@_BOTH_ADAPTERS +@pytest.mark.parametrize( + ("missing", "match"), + [ + ("module", "could not find a loaded modeling_kimi"), + ("class", "defines no module-level DeepseekV3ForCausalLM"), + ], + ids=["no-module", "no-class"], +) +def test_the_swap_fails_loudly_when_the_remote_modeling_has_nothing_to_swap( + monkeypatch, adapter, remote, missing, match +): + """Skipping the swap silently would build the per-expert layout that EP cannot shard: when the + remote modeling did not load, or names its language model otherwise (a class renamed, or + imported where the VLM class builds it), which is not the swap already done.""" + if missing == "module": + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", lambda *args, **kwargs: None) + else: + del remote.DeepseekV3ForCausalLM + + with pytest.raises(RuntimeError, match=match): + adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + + +@pytest.mark.parametrize("adapter", [kimi_vl], ids=["vl"]) +@pytest.mark.parametrize( + ("config", "kwargs", "match"), + [ + # transformers' own port (Kimi-K2.5's, from 5.13), which the patch never reaches. + (PretrainedConfig, {"ckpt_path": "/ckpt", "trust_remote_code": True}, "own port"), + (SimpleNamespace, {}, "text_config"), + # Without the checkpoint there is no modeling code to patch, and a half-prepared config + # would load Kimi's own language model under a plan naming none of its experts. + (_vlm_config, {"trust_remote_code": True}, "needs the checkpoint's path"), + # transformers' get_class_from_dynamic_module runs the checkpoint's code without checking + # trust_remote_code, so prepare must refuse before reaching it. + (_vlm_config, {"ckpt_path": "/ckpt"}, "trust_remote_code=True"), + # The native text config drops the scheme, and with no top-level copy for the load to + # read, the experts would load as stored. + ( + _block_fp8_on_the_text_config_only, + {"ckpt_path": "/ckpt", "trust_remote_code": True}, + "on text_config only", + ), + ], + ids=[ + "transformers-port", + "no-text-config", + "no-checkpoint", + "untrusted-checkpoint", + "quantized-text-config-only", + ], +) +def test_prepare_refuses_before_patching_anything(monkeypatch, adapter, config, kwargs, match): + """Before prepare patches anything process-wide, such as the loader's remote-code staging, or + imports the checkpoint's code.""" + + def unstaged_loader(*args, **kwargs): + """The loader before any staging shim wraps it.""" + + monkeypatch.setattr(dmu, "get_cached_module_file", unstaged_loader) + monkeypatch.setattr( + dmu, + "get_class_from_dynamic_module", + lambda *args, **kwargs: pytest.fail("the checkpoint's code was imported"), + ) + + with pytest.raises(ValueError, match=match): + adapter.prepare(config(), **kwargs) + + assert dmu.get_cached_module_file is unstaged_loader + + +@pytest.mark.parametrize(("model_type", "adapter"), [("kimi_vl", kimi_vl)], ids=["vl"]) +def test_the_spec_routes_expert_parallelism_through_its_own_adapter(model_type, adapter): + """The checkpoint's code builds no EP-capable experts, so its whole EP setup is this adapter's + ``prepare``: another Kimi adapter's would load it without its own swap or decompression.""" + assert get_spec(model_type).ep_spec.prepare is adapter.prepare diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index a6328a6b2e0..e3d024312ff 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -825,6 +825,44 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) +def _remote_kimi(moe): + """A remote-code Kimi VLM as the checkpoint's own code builds it, unpatched by its spec's + ``prepare``: each MoE layer of its language model holds its routed experts as per-expert + modules in an ``nn.ModuleList``, at ``language_model.model.layers...experts``.""" + + def layer(): + block = nn.Module() + block.gate = nn.Linear(2, 4, bias=False) + block.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(4)) + return nn.ModuleDict({moe: block}) + + model = nn.Module() + model.language_model = nn.Module() + model.language_model.model = nn.Module() + model.language_model.model.layers = nn.ModuleList(layer() for _ in range(2)) + model.tie_weights = lambda: None # a PreTrainedModel's; nothing here is tied + return model + + +@pytest.mark.parametrize(("model_type", "moe"), [("kimi_vl", "mlp")]) +def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( + run_as, serve, monkeypatch, tmp_path, model_type, moe +): + """Remote-code Kimi checkpoints get their plan from their spec's ``prepare``, which an EP degree + of 1 skips: it patches their code. The plan their spec states still finds the experts their own + code builds, so whole on every rank they are refused like Qwen3-MoE.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + monkeypatch.setattr( + get_spec(model_type).ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") + ) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi(moe)) + run_as(4) + serve(SimpleNamespace(model_type=model_type, text_config=SimpleNamespace(n_routed_experts=8))) + + with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + resolve_distribution_plan(str(tmp_path)) + + @pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"]) def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead( run_as, serve, monkeypatch, tmp_path, ep_degree From b2e68bbd7eadc5d72fdd49e7f227b0b25bd976e1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:04:05 +0000 Subject: [PATCH 2/9] Parametrize the remote-Kimi refusal test over the model type alone Both Kimi checkpoints' code builds its routed experts at .mlp.experts, so the moe parameter varied nothing. The stand-in builds them there, and the test takes the model type only. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index e3d024312ff..b46975d3696 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -825,16 +825,16 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) -def _remote_kimi(moe): +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 - modules in an ``nn.ModuleList``, at ``language_model.model.layers...experts``.""" + modules in an ``nn.ModuleList``, at ``language_model.model.layers..mlp.experts``.""" def layer(): block = nn.Module() block.gate = nn.Linear(2, 4, bias=False) block.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(4)) - return nn.ModuleDict({moe: block}) + return nn.ModuleDict({"mlp": block}) model = nn.Module() model.language_model = nn.Module() @@ -844,9 +844,9 @@ def layer(): return model -@pytest.mark.parametrize(("model_type", "moe"), [("kimi_vl", "mlp")]) +@pytest.mark.parametrize("model_type", ["kimi_vl"]) def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( - run_as, serve, monkeypatch, tmp_path, model_type, moe + run_as, serve, monkeypatch, tmp_path, model_type ): """Remote-code Kimi checkpoints get their plan from their spec's ``prepare``, which an EP degree of 1 skips: it patches their code. The plan their spec states still finds the experts their own @@ -855,7 +855,7 @@ def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( monkeypatch.setattr( get_spec(model_type).ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") ) - monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi(moe)) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi()) run_as(4) serve(SimpleNamespace(model_type=model_type, text_config=SimpleNamespace(n_routed_experts=8))) From 06c4e44e87111eca3b2b8790de1137e8164dad47 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:09:38 +0000 Subject: [PATCH 3/9] Say what the Kimi-VL spec's missing MoE section and the tie shim leave out The spec registers only its expert-parallel section, and is_moe's fallbacks do not recognize the native DeepseekV3MoE blocks prepare installs, so the HF export's MoE-block handling skips them until the sections are shared. The tie_weights shim drops missing_keys: a tied checkpoint would still load its embedding into both, with lm_head.weight reported missing (checked on transformers 5.18 and 5.5.4). Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 3 ++- modelopt/torch/models/kimi_vl/specs.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index f18e598bfd2..6c23e01a00c 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -120,7 +120,8 @@ def _drop_tie_weights_kwargs(orig_tie): Dropping ``missing_keys`` as well as ``recompute_mapping`` is safe because the remote ``tie_weights`` only calls its language model's, and the Kimi checkpoints set - ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. + ``tie_word_embeddings=False``: nothing is tied, so no tied key is left to report to the load. A + tied one would still load its embedding into both, with ``lm_head.weight`` reported missing. """ def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc. diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py index a5d9eacec42..73ba5e78b2c 100644 --- a/modelopt/torch/models/kimi_vl/specs.py +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -20,7 +20,9 @@ ``ep.py`` swaps in the transformers-native deepseek_v3 language model. Only the expert-parallel section is registered under ``kimi_vl``: the MoE and export sections stay -``deepseek_v3``'s, so lookups by ``kimi_vl`` find neither yet. Sharing them is left to a follow-up. +``deepseek_v3``'s, so lookups by ``kimi_vl`` find neither yet. Nor do ``is_moe``'s name and structure +fallbacks recognize the native ``DeepseekV3MoE`` blocks ``prepare`` installs, so the HF export's +MoE-block handling skips them. Sharing the sections is left to a follow-up. """ from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN From 1d6e937986b5e6693872575c409786ce82437581 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 02:26:22 +0000 Subject: [PATCH 4/9] Build the native language model only from the config prepare made Deciding the plan applies the model spec to a copy of the config on every rank before the plan is known, and the Kimi adapters' prepare patches the checkpoint's modeling module for the process. A plan that then keeps the load off expert parallelism (replicated, or FSDP2) builds from the config as read, and the patched VLM class built transformers' native DeepseekV3ForCausalLM from the checkpoint's own text config, which it cannot take (on a stand-in, an AttributeError on pad_token_id). The swapped name now builds the native class only for the native DeepseekV3Config prepare puts on text_config, and the checkpoint's own class for its own config, so that load gets the language model it would with no patch. Both checkpoints' VLM classes use the name only to build their language model. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 32 ++++++++++++++++++---- tests/unit/torch/models/test_kimi_vl_ep.py | 29 ++++++++++++++++++-- 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index 6c23e01a00c..45de627e6f6 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -131,14 +131,34 @@ def _tie_weights(self, *args, **kwargs): # drop recompute_mapping etc. return _tie_weights +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. + """ + from transformers import DeepseekV3Config + + def build(config, *args, **kwargs): + built = native_class if isinstance(config, DeepseekV3Config) else own_class + return built(config, *args, **kwargs) + + setattr(build, "_kimi_native_swap", True) + return build + + def patch_language_model_to_native( label: str, modeling: str, vlm_class: str, ckpt_path: str, trust_remote_code: bool ) -> None: """Point the checkpoint's ``modeling`` module's ``DeepseekV3ForCausalLM`` at the native class. - Loads the checkpoint's remote ``modeling`` module and replaces its ``DeepseekV3ForCausalLM`` with - transformers' native class, so ``vlm_class.__init__`` builds the EP-capable native deepseek_v3 - for ``self.language_model`` instead of the per-expert Linears version. + Loads the checkpoint's remote ``modeling`` module and replaces its ``DeepseekV3ForCausalLM``, so + ``vlm_class.__init__`` builds the EP-capable native deepseek_v3 for ``self.language_model`` + instead of the per-expert Linears version, from the config ``prepare`` made + (:func:`_native_for_the_prepared_config`). """ from transformers.dynamic_module_utils import get_class_from_dynamic_module from transformers.models.deepseek_v3 import DeepseekV3ForCausalLM as NativeDeepseekV3ForCausalLM @@ -159,9 +179,11 @@ def to_native(module): f"{label}: {module.__name__} defines no module-level DeepseekV3ForCausalLM to swap " "for the transformers-native class; the checkpoint's modeling code changed." ) - if current is NativeDeepseekV3ForCausalLM: + if getattr(current, "_kimi_native_swap", False): return False - module.DeepseekV3ForCausalLM = NativeDeepseekV3ForCausalLM + 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) diff --git a/tests/unit/torch/models/test_kimi_vl_ep.py b/tests/unit/torch/models/test_kimi_vl_ep.py index 6d82fb11543..bd5158077a9 100644 --- a/tests/unit/torch/models/test_kimi_vl_ep.py +++ b/tests/unit/torch/models/test_kimi_vl_ep.py @@ -21,12 +21,14 @@ after each test. """ +import copy import importlib import sys import types from types import SimpleNamespace import pytest +import torch transformers = pytest.importorskip("transformers") @@ -102,14 +104,19 @@ def remote(monkeypatch, request): continue _restore(monkeypatch, owner, name) module = types.ModuleType(f"transformers_modules.Kimi.0123abcd.{modeling}") - module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + module.remote_lm = type("DeepseekV3ForCausalLM", (), {"__init__": lambda self, config: None}) + module.DeepseekV3ForCausalLM = module.remote_lm module.loads, module.tied = [], [] + def build(self, config=None): + if config is not None: # as both checkpoints' VLM classes build their language model + self.language_model = module.DeepseekV3ForCausalLM(config.text_config) + def tie_weights(self): module.tied.append(self) return "tied" - module.vlm = type(vlm_name, (), {"tie_weights": tie_weights}) + module.vlm = type(vlm_name, (), {"__init__": build, "tie_weights": tie_weights}) setattr(module, vlm_name, module.vlm) def import_remote_class(class_ref, ckpt_path, **kwargs): @@ -178,7 +185,6 @@ def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_wei vlm_class = f"{remote.__name__.rsplit('.', 1)[-1]}.{remote.vlm.__name__}" assert remote.loads == [(vlm_class, "/ckpt", {"trust_remote_code": True})] * 2 - assert remote.DeepseekV3ForCausalLM is DeepseekV3ForCausalLM assert import_utils.is_torch_fx_available() is False assert remote.vlm.tie_weights is shim model = remote.vlm() @@ -186,6 +192,23 @@ def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_wei assert remote.tied == [model] +@_BOTH_ADAPTERS +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.""" + 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) + + with torch.device("meta"): + assert type(remote.vlm(prepared).language_model) is DeepseekV3ForCausalLM + assert type(remote.vlm(as_read).language_model) is remote.remote_lm + + @_BOTH_ADAPTERS @pytest.mark.parametrize( ("missing", "match"), From 5921799e230baef3451b5391418a8dd4f598eca3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:27:50 +0000 Subject: [PATCH 5/9] Expect the left-whole refusal's arm this transformers takes for remote Kimi keeps the RuntimeError for a transformers below what the model's spec needs, which names the release. The remote-Kimi case expects whichever arm its transformers takes. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b46975d3696..22d209dc95d 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -80,7 +80,7 @@ from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan -from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS +from modelopt.torch.utils.moe import _EP_MIN_TRANSFORMERS, _ep_floor from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( _auto_class, @@ -850,7 +850,8 @@ def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( ): """Remote-code Kimi checkpoints get their plan from their spec's ``prepare``, which an EP degree of 1 skips: it patches their code. The plan their spec states still finds the experts their own - code builds, so whole on every rank they are refused like Qwen3-MoE.""" + code builds, so whole on every rank they are refused like Qwen3-MoE: a ValueError, or the + RuntimeError naming the release when this transformers is below what their spec needs.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") monkeypatch.setattr( get_spec(model_type).ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") @@ -858,8 +859,12 @@ def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi()) run_as(4) serve(SimpleNamespace(model_type=model_type, text_config=SimpleNamespace(n_routed_experts=8))) + below_floor = Version(transformers.__version__) < Version(_ep_floor(model_type)) - with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + with pytest.raises( + RuntimeError if below_floor else ValueError, + match="can be expert-parallel, but the plan resolved", + ): resolve_distribution_plan(str(tmp_path)) From b83a786f66beca5bfdb7699b15a923c5438275ca Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:11:21 +0000 Subject: [PATCH 6/9] Point a Kimi checkpoint quantized on its text config alone at an unquantized copy From meenchen-nv's review: the refusal of a quantization_config that only text_config declares advised declaring it on the top-level config too, which gets past this refusal only to meet the expert-parallel load's own refusal of a quantized checkpoint. It now names the way in, an unquantized copy. The module-scope comment also says why transformers is imported inside _kimi's functions: the optional [hf] extra, and the model registry imports this module. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_kimi.py | 13 ++++++------- tests/unit/torch/models/kimi_vl/__init__.py | 15 +++++++++++++++ .../{test_kimi_vl_ep.py => kimi_vl/test_ep.py} | 0 3 files changed, 21 insertions(+), 7 deletions(-) create mode 100644 tests/unit/torch/models/kimi_vl/__init__.py rename tests/unit/torch/models/{test_kimi_vl_ep.py => kimi_vl/test_ep.py} (100%) diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index 45de627e6f6..f48972306af 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -21,7 +21,8 @@ keep the vision tower and projector from the checkpoint's code. """ -# modelopt.torch.utils is imported inside the functions: the model registry imports this module. +# 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 @@ -60,10 +61,9 @@ def checkpoint_to_patch( def refuse_quantization_only_on_text_config(label: str, hf_config) -> None: """Refuse a checkpoint whose ``quantization_config`` only its ``text_config`` declares. - The adapter replaces ``text_config`` with a native config without it. Unless the experts - decompress as they load, the expert-parallel load then reads the checkpoint's quantization off - the top-level config, which Kimi-K2.5's config class copies it to; without that copy, the load - would read quantized experts as stored. + The adapter replaces ``text_config`` with a native config without it, and the expert-parallel + load would then read the quantized experts as stored. It loads unquantized checkpoints only, so + the way in is an unquantized copy of the model. """ if ( getattr(hf_config.text_config, "quantization_config", None) is not None @@ -72,8 +72,7 @@ def refuse_quantization_only_on_text_config(label: str, hf_config) -> None: raise ValueError( f"{label}: the checkpoint declares quantization_config on text_config only, which the " "native deepseek_v3 config replacing it drops, so the expert-parallel load would read " - "its experts as stored. Declare it on the top-level config too, or start from an " - "unquantized copy of the model." + "its experts as stored. Start from an unquantized copy of the model." ) diff --git a/tests/unit/torch/models/kimi_vl/__init__.py b/tests/unit/torch/models/kimi_vl/__init__.py new file mode 100644 index 00000000000..1275d78dff5 --- /dev/null +++ b/tests/unit/torch/models/kimi_vl/__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_vl_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py similarity index 100% rename from tests/unit/torch/models/test_kimi_vl_ep.py rename to tests/unit/torch/models/kimi_vl/test_ep.py From 5730f302b46c30ccb372fd2dc10bae366d15ffe1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:11:31 +0000 Subject: [PATCH 7/9] State the Kimi-VL adapter's contract, and mirror its module in its tests From meenchen-nv's review: kimi_vl/ep.py's module docstring is now the adapter's contract, and its walkthrough moved to the PR description. The adapter's tests move to tests/unit/torch/models/kimi_vl/test_ep.py, mirroring kimi_vl/ep.py, as CONTRIBUTING asks; the directory is a package, so it can hold a test_ep.py beside the other adapters'. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_vl/ep.py | 31 ++++------------------------- 1 file changed, 4 insertions(+), 27 deletions(-) diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index d1891f5c781..9d6899f0210 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -13,34 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""DP=EP a2a support for Kimi-VL-A3B (``kimi_vl``). +"""Expert-parallel preparation for remote-code Kimi-VL (``kimi_vl``). -a remote-code VLM whose text backbone is a Moonlight (DeepSeek-V3-lite: MLA, ``q_lora_rank=None``, 64 routed experts, 2 -shared) MoE and whose vision tower is MoonViT. The bf16 sibling of ``kimi_k25`` -- SAME shape, but NO compressed-tensors -decompressor (the checkpoint ships plain bf16 per-expert weights). - -Challenge (identical to kimi_k25): the bundled ``modeling_kimi_vl.py`` builds its language model from -an INLINE ``DeepseekV3ForCausalLM`` whose ``DeepseekV3MoE`` is an ``nn.ModuleList`` of per-expert -``gate_proj``/``up_proj``/``down_proj`` Linears -- incompatible with transformers' EP machinery, which -needs the FUSED 3-D ``gate_up_proj``/``down_proj`` Parameters the transformers-native ``deepseek_v3`` -modeling ships. - -Solution (mirrors kimi_k25): - 1. Shim the import symbols ``modeling_kimi_vl`` needs that newer transformers dropped (e.g. - ``is_torch_fx_available``) so the module IMPORTS -- they only guard torch.fx / legacy attn-mask - paths that our native-LM + MoonViT forward never execute. - 2. Swap ``hf_config.text_config`` to a transformers-native ``DeepseekV3Config`` (drops the remote - ``auto_map`` so the native class -- and its ``base_model_ep_plan`` -- is used for the LM). - 3. Apply the deepseek_v3 EP plan to that native ``text_config`` (post_init prefixes it to the VLM - root as ``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}``). - 4. Monkey-patch ``modeling_kimi_vl.DeepseekV3ForCausalLM`` -> the transformers-native class, so - ``KimiVLForConditionalGeneration.__init__`` (``self.language_model = DeepseekV3ForCausalLM( - config.text_config)``) builds the EP-capable native LM while MoonViT + the KimiVL wrapper stay - from the bundled code. Native deepseek_v3's per-expert WeightConverter fuses the checkpoint's - per-expert ``.weight`` tensors into the 3-D params on load (no decompressor -- weights are bf16). - -Version: no floor beyond the EP framework's 5.7; validated on transformers 5.9.0. The compat shims -stand in for the symbols the remote code imports that later releases dropped. +``prepare`` replaces the per-expert DeepSeek-V3 language model the checkpoint's own code builds with +transformers' native DeepSeek-V3, whose experts the expert-parallel load can shard, and keeps the +vision tower and projector from that code. It needs trusted remote code and transformers 5.7. """ from .. import _kimi, _onfly_decompress From fc735bebd9dd202c6d0d3928d7ce64d5a72ce373 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:14:57 +0000 Subject: [PATCH 8/9] Load a tiny Kimi-VL expert-parallel on GPUs, against the whole model From meenchen-nv's review: no test loaded checkpoint tensors through the Kimi-VL adapter. New 2-GPU case test_kimi_vl_loads_expert_parallel_and_ matches_the_whole_model writes a tiny Kimi-VL-shaped remote checkpoint (tests/_test_utils/torch/kimi_vl_tiny_modeling.py), its experts stored per expert, and loads it expert-parallel through the spec's prepare: each rank's native experts must equal the per-expert tensors fused, the vision tower and projector must load as stored, nothing may be left unplaced, and each rank's logits must match the whole model's on one GPU, with every rank's checks gathered under bounded collectives. Its non-distributed parts were checked on CPU: the swap, the exact fusion, the nested plan at the VLM root and the logits. Signed-off-by: Shengliang Xu --- .../torch/kimi_vl_tiny_modeling.py | 60 +++++++ .../torch/utils/test_expert_parallel_load.py | 161 +++++++++++++++++- 2 files changed, 218 insertions(+), 3 deletions(-) create mode 100644 tests/_test_utils/torch/kimi_vl_tiny_modeling.py diff --git a/tests/_test_utils/torch/kimi_vl_tiny_modeling.py b/tests/_test_utils/torch/kimi_vl_tiny_modeling.py new file mode 100644 index 00000000000..5a26aa8c050 --- /dev/null +++ b/tests/_test_utils/torch/kimi_vl_tiny_modeling.py @@ -0,0 +1,60 @@ +# 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. + +"""A tiny Kimi-VL-shaped model, as remote code: tests copy it into a checkpoint as ``modeling_kimi_vl.py``. + +It has what the Kimi-VL adapter patches, as the checkpoint's own code has it: a VLM class that builds +its language model, at construction, from the module-level name ``DeepseekV3ForCausalLM``, beside a +vision tower and a projector of its own, and whose ``tie_weights`` takes no arguments. The +checkpoint's own language model is a DeepSeek-V3 whose MoE holds per-expert Linears; here the name is +transformers' DeepSeek-V3 itself, since only the native class the adapter swaps in is built. +""" + +from torch import nn +from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig, PreTrainedModel + + +class KimiVLConfig(PretrainedConfig): + model_type = "kimi_vl" + sub_configs = {"text_config": DeepseekV3Config} + + def __init__(self, text_config=None, vision_hidden_size=16, **kwargs): + if text_config is None or isinstance(text_config, dict): + text_config = DeepseekV3Config(**(text_config or {})) + self.text_config = text_config + self.vision_hidden_size = vision_hidden_size + super().__init__(**kwargs) + + +class KimiVLForConditionalGeneration(PreTrainedModel): + config_class = KimiVLConfig + _no_split_modules = ["DeepseekV3DecoderLayer"] + + def __init__(self, config): + super().__init__(config) + width = config.vision_hidden_size + self.vision_tower = nn.Linear(width, width) + self.multi_modal_projector = nn.Linear(width, config.text_config.hidden_size) + self.language_model = DeepseekV3ForCausalLM(config.text_config) + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def forward(self, input_ids=None, **kwargs): + return self.language_model(input_ids=input_ids, **kwargs) + + def tie_weights(self): # as the checkpoint's own: no keyword arguments + return self.language_model.tie_weights() diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 533147ac1b0..3d50e5c3e5f 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -18,10 +18,14 @@ ``load_and_prepare_ep`` shards an MoE's routed experts over the ranks, and the ``torch_a2a`` experts forward sends each token to the rank that holds its expert. Every rank runs its own batch through the sharded model, and the result must match the whole model run on that batch alone. A checkpoint whose -experts are compressed-tensors compressed loads the same way, each rank decompressing its own. +experts are compressed-tensors compressed loads the same way, each rank decompressing its own, and +so does a remote-code VLM through its model spec's adapter. """ +import copy import json +import shutil +from pathlib import Path from types import SimpleNamespace import pytest @@ -37,8 +41,15 @@ import transformers.core_model_loading as cml import transformers.modeling_utils as mu +from _test_utils.torch import kimi_vl_tiny_modeling from torch.distributed.tensor import DTensor -from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + DeepseekV3Config, + DeepseekV3ForCausalLM, + Qwen3MoeConfig, +) from modelopt.torch.models import _onfly_decompress from modelopt.torch.utils.distributed import ( @@ -46,7 +57,7 @@ distribution_plan, set_distribution_plan, ) -from modelopt.torch.utils.moe import expert_parallel_mesh +from modelopt.torch.utils.moe import apply_distributed_support, expert_parallel_mesh from modelopt.torch.utils.plugins.model_load_utils import load_and_prepare_ep NUM_EXPERTS = 8 @@ -244,3 +255,147 @@ def test_expert_parallel_model_matches_the_whole_model( pytest.skip(f"{NUM_EXPERTS} experts do not shard over EP {ep_degree} on these GPUs") _save_tiny_qwen3_moe(tmp_path) dist_workers.run(_sharded_model_matches_the_whole_one, str(tmp_path), ep_degree) + + +# ------------------------------------------------------------------------------------- Kimi-VL + +# A tiny Kimi-VL: a vision tower and a projector beside a DeepSeek-V3 language model with a dense first +# layer, then one MoE layer of 8 routed experts, as the checkpoint's own code builds it (remote code, +# ``kimi_vl_tiny_modeling``), its experts stored per expert. +KIMI_VL_TEXT = { + "model_type": "deepseek_v3", + "vocab_size": 64, + "hidden_size": 64, + "intermediate_size": 64, + "moe_intermediate_size": 32, + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "n_routed_experts": NUM_EXPERTS, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "n_group": 1, + "topk_group": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "q_lora_rank": None, + "kv_lora_rank": 16, + "qk_nope_head_dim": 16, + "qk_rope_head_dim": 8, + "v_head_dim": 16, + "max_position_embeddings": 64, + "tie_word_embeddings": False, +} +KIMI_VL_MOE = "language_model.model.layers.1.mlp.experts" +KIMI_VL_OWN = ("vision_tower.weight", "multi_modal_projector.weight") + + +def _save_tiny_kimi_vl(ckpt_dir): + """A tiny Kimi-VL-shaped remote checkpoint whose experts are stored per expert, as Kimi-VL's are.""" + ckpt_dir = Path(ckpt_dir) + shutil.copy(kimi_vl_tiny_modeling.__file__, ckpt_dir / "modeling_kimi_vl.py") + remote = "modeling_kimi_vl.KimiVLForConditionalGeneration" + config = { + "model_type": "kimi_vl", + "architectures": ["KimiVLForConditionalGeneration"], + "auto_map": { + "AutoConfig": "modeling_kimi_vl.KimiVLConfig", + "AutoModel": remote, + "AutoModelForCausalLM": remote, + }, + "torch_dtype": "bfloat16", + "text_config": KIMI_VL_TEXT, + } + (ckpt_dir / "config.json").write_text(json.dumps(config)) + torch.manual_seed(0) + language_model = DeepseekV3ForCausalLM(DeepseekV3Config(**KIMI_VL_TEXT)).to(torch.bfloat16) + tensors = {} + inter = KIMI_VL_TEXT["moe_intermediate_size"] + for name, tensor in language_model.state_dict().items(): + name = f"language_model.{name}" + if name == f"{KIMI_VL_MOE}.gate_up_proj": # stored per expert, gate and up apart + for e in range(NUM_EXPERTS): + tensors[f"{KIMI_VL_MOE}.{e}.gate_proj.weight"] = tensor[e, :inter].contiguous() + tensors[f"{KIMI_VL_MOE}.{e}.up_proj.weight"] = tensor[e, inter:].contiguous() + elif name == f"{KIMI_VL_MOE}.down_proj": + for e in range(NUM_EXPERTS): + tensors[f"{KIMI_VL_MOE}.{e}.down_proj.weight"] = tensor[e].contiguous() + else: + tensors[name] = tensor.contiguous() + width, hidden = 16, KIMI_VL_TEXT["hidden_size"] + for name, shape in ( + ("vision_tower.weight", (width, width)), + ("vision_tower.bias", (width,)), + ("multi_modal_projector.weight", (hidden, width)), + ("multi_modal_projector.bias", (hidden,)), + ): + tensors[name] = torch.randn(shape).to(torch.bfloat16) + save_file(tensors, str(ckpt_dir / "model.safetensors"), metadata={"format": "pt"}) + + +def _kimi_vl_expert_parallel_job(rank, size, ckpt_dir): + device = torch.device("cuda", rank) + stored = load_file(str(Path(ckpt_dir) / "model.safetensors")) + per_rank = NUM_EXPERTS // size + owned = range(rank * per_rank, (rank + 1) * per_rank) + problems = [] + with bounded_collectives(seconds=120): + config = AutoConfig.from_pretrained(ckpt_dir, trust_remote_code=True) + plan = DistributionPlan(size, ep_size=size) + prepared = apply_distributed_support( + "kimi_vl", copy.deepcopy(config), ckpt_path=ckpt_dir, trust_remote_code=True + ) + model = load_and_prepare_ep( + ckpt_dir, plan, trust_remote_code=True, hf_config=copy.deepcopy(prepared) + ).eval() + whole = AutoModelForCausalLM.from_pretrained( + ckpt_dir, config=copy.deepcopy(prepared), trust_remote_code=True, dtype=torch.bfloat16 + ) + whole = whole.to(device).eval() + + # The native language model the adapter swaps in fuses this rank's experts, and only those, + # from the per-expert tensors as they load. + experts = model.language_model.model.layers[1].mlp.experts + gate_up = [ + torch.cat([stored[f"{KIMI_VL_MOE}.{e}.{part}_proj.weight"] for part in ("gate", "up")]) + for e in owned + ] + down = [stored[f"{KIMI_VL_MOE}.{e}.down_proj.weight"] for e in owned] + for name, expected in (("gate_up_proj", gate_up), ("down_proj", down)): + if not torch.equal(_local(getattr(experts, name)).cpu(), torch.stack(expected)): + problems.append(f"{name}: rank {rank}'s experts {list(owned)} differ") + # The vision tower and the projector, the checkpoint's own code's, load as stored. + problems.extend( + f"{name} differs from the checkpoint's" + for name in KIMI_VL_OWN + if not torch.equal(model.get_parameter(name).detach().cpu(), stored[name]) + ) + if model._modelopt_unplaced_source_keys: + problems.append(f"unplaced: {model._modelopt_unplaced_source_keys[:3]}") + + # Each rank its own batch through the experts, against the whole model on that batch. + generator = torch.Generator().manual_seed(rank) + input_ids = torch.randint(0, KIMI_VL_TEXT["vocab_size"], (2, 16), generator=generator) + with torch.no_grad(): + sharded = model(input_ids.to(device)).logits.float() + reference = whole(input_ids.to(device)).logits.float() + if not torch.allclose(sharded, reference, atol=2e-2, rtol=2e-2): + problems.append(f"logits differ by {(sharded - reference).abs().max().item():.3g}") + + gathered = [None] * size + dist.all_gather_object(gathered, problems) + found = [problem for rank_problems in gathered for problem in rank_problems] + assert not found, "\n".join(found[:20]) + + +def test_kimi_vl_loads_expert_parallel_and_matches_the_whole_model( + need_2_gpus, dist_workers, tmp_path +): + """A Kimi-VL-shaped remote checkpoint loaded expert-parallel through its model spec: the native + DeepSeek-V3 the adapter swaps in fuses each rank's own experts from the per-expert tensors, the + vision tower and projector load as stored, nothing is left unplaced, and each rank's logits match + the whole model's on one GPU. Each rank's checks are gathered before any rank fails, so a + regression fails instead of hanging.""" + if NUM_EXPERTS % dist_workers.world_size: + pytest.skip(f"{NUM_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs") + _save_tiny_kimi_vl(tmp_path) + dist_workers.run(_kimi_vl_expert_parallel_job, ckpt_dir=str(tmp_path)) From 4ff7f66e09826cf6556123934e0ec2307d715a64 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 01:09:17 +0000 Subject: [PATCH 9/9] Give the tiny Kimi-VL its own per-expert language model to swap From Claude's review: the stand-in's module-level DeepseekV3ForCausalLM was transformers' native class, so the GPU case would pass even if the adapter's swap did nothing. It is now the checkpoint's own: routed experts as per-expert Linears, which only a swap that misses builds, and the case's expert checks then fail. Signed-off-by: Shengliang Xu --- .../torch/kimi_vl_tiny_modeling.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/_test_utils/torch/kimi_vl_tiny_modeling.py b/tests/_test_utils/torch/kimi_vl_tiny_modeling.py index 5a26aa8c050..148598ca160 100644 --- a/tests/_test_utils/torch/kimi_vl_tiny_modeling.py +++ b/tests/_test_utils/torch/kimi_vl_tiny_modeling.py @@ -17,13 +17,27 @@ It has what the Kimi-VL adapter patches, as the checkpoint's own code has it: a VLM class that builds its language model, at construction, from the module-level name ``DeepseekV3ForCausalLM``, beside a -vision tower and a projector of its own, and whose ``tie_weights`` takes no arguments. The -checkpoint's own language model is a DeepSeek-V3 whose MoE holds per-expert Linears; here the name is -transformers' DeepSeek-V3 itself, since only the native class the adapter swaps in is built. +vision tower and a projector of its own, and whose ``tie_weights`` takes no arguments. The name is +the checkpoint's own DeepSeek-V3, whose MoE holds its routed experts as per-expert Linears that expert +parallelism cannot shard: only a swap that misses builds it, and a test then meets its layout. """ from torch import nn -from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig, PreTrainedModel +from transformers import DeepseekV3Config, PretrainedConfig, PreTrainedModel +from transformers import DeepseekV3ForCausalLM as _NativeDeepseekV3ForCausalLM + + +class DeepseekV3ForCausalLM(_NativeDeepseekV3ForCausalLM): + """The checkpoint's own: routed experts as per-expert Linears, which expert parallelism cannot shard.""" + + def __init__(self, config): + super().__init__(config) + for layer in self.model.layers: + if hasattr(layer.mlp, "experts"): + layer.mlp.experts = nn.ModuleList( + nn.Linear(config.hidden_size, config.moe_intermediate_size, bias=False) + for _ in range(config.n_routed_experts) + ) class KimiVLConfig(PretrainedConfig):