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..f48972306af --- /dev/null +++ b/modelopt/torch/models/_kimi.py @@ -0,0 +1,206 @@ +# 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 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 + + +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, 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 + 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. 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. 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. + return orig_tie(self) + + setattr(_tie_weights, "_kimi_tie_weights_shim", True) + 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``, 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 + + 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 getattr(current, "_kimi_native_swap", False): + return False + 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) + 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..9d6899f0210 --- /dev/null +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -0,0 +1,56 @@ +# 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 preparation for remote-code Kimi-VL (``kimi_vl``). + +``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 .._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..73ba5e78b2c --- /dev/null +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -0,0 +1,41 @@ +# 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. Nor do ``is_moe``'s name and structure +fallbacks recognize the native ``DeepseekV3MoE`` blocks ``prepare`` installs, so the HF export's +MoE-block handling skips them. Sharing the sections is left to a follow-up. +""" + +from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN +from ..specs import EPSpec, ModelSpec, register +from .ep import 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/_test_utils/torch/kimi_vl_tiny_modeling.py b/tests/_test_utils/torch/kimi_vl_tiny_modeling.py new file mode 100644 index 00000000000..148598ca160 --- /dev/null +++ b/tests/_test_utils/torch/kimi_vl_tiny_modeling.py @@ -0,0 +1,74 @@ +# 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 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, 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): + 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)) 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/kimi_vl/test_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py new file mode 100644 index 00000000000..bd5158077a9 --- /dev/null +++ b/tests/unit/torch/models/kimi_vl/test_ep.py @@ -0,0 +1,289 @@ +# 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 copy +import importlib +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + +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.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, (), {"__init__": build, "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 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 +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"), + [ + ("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..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, @@ -825,6 +825,49 @@ 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(): + """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..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({"mlp": 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", ["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 +): + """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: 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") + ) + 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 if below_floor else ValueError, + 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