diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index b87f3872f73..bd748e773e0 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -92,6 +92,7 @@ minimax_m3_vl, llama4_text, kimi_vl, + kimi_k25, nemotron_h_omni, qwen2_moe, qwen3, diff --git a/modelopt/torch/models/_kimi.py b/modelopt/torch/models/_kimi.py index f48972306af..3914104dfdd 100644 --- a/modelopt/torch/models/_kimi.py +++ b/modelopt/torch/models/_kimi.py @@ -23,7 +23,7 @@ # modelopt.torch.utils and transformers are imported inside the functions: the model registry imports # this module, and it must import without transformers, the optional [hf] extra. -from ._remote_code import keep_patched +from ._remote_code import _apply, keep_patched, stage_full_remote_closure def checkpoint_to_patch( @@ -114,6 +114,55 @@ def _prepare_4d_causal_attention_mask(*args, **kwargs): # pragma: no cover print_rank_0(f"[kimi] compat-shim warning (attn_mask_utils): {type(e).__name__}: {e}") +def make_importable() -> None: + """The names newer transformers dropped, and the code's two-deep relative imports staged.""" + install_compat_import_shims() + # transformers stages only 1-hop relative imports; Kimi's remote code is 2 hops deep. + stage_full_remote_closure() + + +def make_buildable( + modeling: str, vlm_class: str, hf_config, ckpt_path: str | None, trust_remote_code: bool +) -> None: + """What both checkpoints' own code needs to import and build on this transformers, and no more. + + The names newer transformers dropped, the code's two-deep relative imports staged, and the VLM + class's ``tie_weights`` taking the keyword arguments transformers' ``post_init`` passes it. None + of it changes what the model computes or loads, so the plan decision runs it before ``prepare``, + whatever plan comes out. A config of transformers' own port, or code this load may not run, has + nothing of the checkpoint's to build. Idempotent. + """ + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + if type(hf_config).__module__.startswith("transformers.") or ckpt_path is None: + return + if not trust_remote_code: + return + make_importable() + get_class_from_dynamic_module( + f"{modeling}.{vlm_class}", ckpt_path, trust_remote_code=trust_remote_code + ) + # Remote code only: transformers' own port of a model can share the module's name. + _apply(f".{modeling}", lambda module: _shim_tie_weights(module, vlm_class)) + + +def _shim_tie_weights(module, vlm_class: str) -> bool: + """Have ``module``'s VLM class take the keyword arguments transformers passes ``tie_weights``. + + The remote ``tie_weights()`` takes none, but transformers' ``post_init`` and load call + ``tie_weights(recompute_mapping=False)``: a TypeError at construction. Returns whether it changed + the class. + """ + kimi_cls = getattr(module, vlm_class, None) + if kimi_cls is None: + return False + own_tie = vars(kimi_cls).get("tie_weights") + if own_tie is None or getattr(own_tie, "_kimi_tie_weights_shim", False): + return False + kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie) + return True + + def _drop_tie_weights_kwargs(orig_tie): """Wrap a remote ``tie_weights()`` that predates the kwargs newer transformers pass to it. @@ -134,10 +183,10 @@ def _native_for_the_prepared_config(own_class, native_class): """The ``DeepseekV3ForCausalLM`` the patched module builds its language model with. transformers' native class for the native ``DeepseekV3Config`` ``prepare`` puts on - ``text_config``, and the checkpoint's own class for its own config. Deciding the plan prepares a - copy of the config before the plan is known, and the patch outlives that copy: a load the plan - keeps off expert parallelism builds from the config as read, and so gets the checkpoint's own - language model, as it would with no patch. + ``text_config``, and the checkpoint's own class for its own config. The patch outlives the load + it was made for: a later load in the process that its plan keeps off expert parallelism builds + from the config as read, and so gets the checkpoint's own language model, as it would with no + patch. """ from transformers import DeepseekV3Config @@ -183,13 +232,7 @@ def to_native(module): module.DeepseekV3ForCausalLM = _native_for_the_prepared_config( current, NativeDeepseekV3ForCausalLM ) - # The remote tie_weights() takes no kwargs, but transformers' post_init and load call - # tie_weights(recompute_mapping=False) -> TypeError. - kimi_cls = getattr(module, vlm_class, None) - if kimi_cls is not None: - own_tie = vars(kimi_cls).get("tie_weights") - if own_tie is not None and not getattr(own_tie, "_kimi_tie_weights_shim", False): - kimi_cls.tie_weights = _drop_tie_weights_kwargs(own_tie) + _shim_tie_weights(module, vlm_class) return True # The dynamic-module loader names the module like diff --git a/modelopt/torch/models/kimi_k25/__init__.py b/modelopt/torch/models/kimi_k25/__init__.py new file mode 100644 index 00000000000..2cf9aee48e6 --- /dev/null +++ b/modelopt/torch/models/kimi_k25/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""kimi_k25 specs (HF model type ``kimi_k25``).""" + +from . import specs diff --git a/modelopt/torch/models/kimi_k25/ep.py b/modelopt/torch/models/kimi_k25/ep.py new file mode 100644 index 00000000000..b9b6451360b --- /dev/null +++ b/modelopt/torch/models/kimi_k25/ep.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DP=EP a2a support for Kimi-K2.5 (``kimi_k25``). + +a remote-code VLM whose text backbone is a DeepSeek-V3 variant with compressed-tensors-compressed experts. +Kimi-K2.6 and Kimi-K2.7-Code ship the same model type and modeling classes. + +The challenge: Kimi's bundled ``modeling_deepseek.py`` defines ``DeepseekV3MoE`` as an +``nn.ModuleList`` of per-expert ``DeepseekV3MLP`` (separate ``gate_proj`` / ``up_proj`` / +``down_proj`` Linears per expert). That layout is incompatible with transformers' EP machinery, +which expects the **fused** 3-D ``gate_up_proj`` / ``down_proj`` Parameters that the +transformers-native ``deepseek_v3`` modeling ships. We need the native modeling to get EP-capable +experts, but we also want the rest of the VLM (vision tower + mm projector) to come from Kimi's +custom code. + +Solution -- monkey-patch ``modeling_kimi_k25``'s ``DeepseekV3ForCausalLM`` import to point at the +transformers-native class, AND swap ``hf_config.text_config`` to a transformers-native +``DeepseekV3Config`` (the native class won't accept Kimi's custom config subclass). Then +``KimiK25ForConditionalGeneration.__init__`` (which does ``self.language_model = +DeepseekV3ForCausalLM(config.text_config)``) builds the native EP-capable model for the language +sub-model, while vision_tower + mm_projector are still built from Kimi's custom code. + +On-the-fly compressed-tensors expert decompress happens during from_pretrained's +``WeightConverter`` machinery (installed by ``_onfly_decompress.install_onfly_decompress``); the +fused per-expert source ``.weight`` tensors are decompressed + merged into the 3-D +``gate_up_proj`` / ``down_proj`` parameters. Format-agnostic (compressed_tensors decides the +dequant from the checkpoint's quantization_config). + +EP plan: attached to the (now native) ``text_config``. transformers' ``post_init`` aggregates +child-module plans up to the root with name prefix, so the final root ``_ep_plan`` keys end up as +``language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj,...}`` -- matching the actual +parameter paths in the loaded VLM. + +Version floor: transformers >= 5.8, which the spec declares. The on-the-fly decompressor, installed +for a compressed-tensors checkpoint, needs it: earlier releases' ``Concatenate`` trips over the +sibling keys it has already consumed. The decompressor couples tightly to +transformers' still-evolving ``WeightConverter`` conversion-ops API. +""" + +from .. import _kimi, _onfly_decompress + +# The checkpoint formats prepare has the experts decompress as they load; the spec states them for +# the plan decision, which runs before prepare. +DECOMPRESSED_FORMATS = ("compressed-tensors",) + + +def make_buildable( + hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_ +) -> None: + """Have the checkpoint's own code import and build on this transformers. + + The plan decision needs it before ``prepare``; it changes nothing the model computes or loads. + """ + _kimi.make_buildable( + "modeling_kimi_k25", + "KimiK25ForConditionalGeneration", + hf_config, + ckpt_path, + trust_remote_code, + ) + + +def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): + """Load Kimi-K2.5's language model as the transformers-native deepseek_v3. + + Swap text_config to native deepseek_v3 + monkey-patch Kimi's modeling to use the native + DeepseekV3ForCausalLM + apply the EP plan + install the on-the-fly decompressor for + compressed-tensors experts. + """ + # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe + # imports the registry. + from modelopt.torch.utils.moe import apply_distributed_support + + ckpt_path = _kimi.checkpoint_to_patch("kimi_k25", hf_config, ckpt_path, trust_remote_code) + + # 1. Install on-the-fly compressed-tensors expert decompressor. Operates on transformers-native + # deepseek_v3's per-expert WeightConverter (which fuses per-expert .weight tensors into the + # 3-D gate_up_proj/down_proj Parameters). Format-agnostic; driven by checkpoint's + # quantization_config, which Kimi-K2.5 carries only on the text_config step 2 replaces. It + # goes first, so a checkpoint it refuses leaves everything as it was. Experts stored + # otherwise are left to the expert-parallel load's format handling, which reads their scheme + # from the copy of quantization_config Kimi's config class keeps on the top-level config; + # step 2 leaves that copy alone. A checkpoint without that copy is refused. + quantization_config = _onfly_decompress._find_quantization_config(hf_config) or {} + if quantization_config.get("quant_method") in DECOMPRESSED_FORMATS: + _onfly_decompress.install_onfly_decompress(hf_config) + else: + _kimi.refuse_quantization_only_on_text_config("kimi_k25", hf_config) + + # Before any remote-code load: this plugin's get_class_from_dynamic_module below, and + # from_pretrained later in load_and_prepare_ep. + _kimi.make_importable() # the swap below loads the code and adapts its tie_weights + + # 2. Replace text_config with a transformers-native DeepseekV3Config built from the same fields. + # Kimi's __init__ passes config.text_config to DeepseekV3ForCausalLM; the native class rejects + # Kimi's custom config subclass. to_native_backbone_config also drops auto_map + + # quantization_config so transformers builds the native modeling instead of re-routing back + # to Kimi's remote code via auto_map. + hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config) + + # 3. Apply the EP plan to the now-native text_config. transformers' post_init walks children + # and prefixes their _ep_plan onto the parent's: text_config's plan ends up at the VLM root + # as `language_model.model.layers.*.mlp.experts.{gate_up_proj,down_proj}`. + hf_config.text_config = apply_distributed_support("deepseek_v3", hf_config.text_config) + + # 4. Force-load Kimi's modeling_kimi_k25 module (so we can patch its DeepseekV3ForCausalLM + # symbol) and swap to the native class. + _kimi.patch_language_model_to_native( + "kimi_k25", + "modeling_kimi_k25", + "KimiK25ForConditionalGeneration", + ckpt_path, + trust_remote_code, + ) + return hf_config diff --git a/modelopt/torch/models/kimi_k25/specs.py b/modelopt/torch/models/kimi_k25/specs.py new file mode 100644 index 00000000000..93035a9d3b6 --- /dev/null +++ b/modelopt/torch/models/kimi_k25/specs.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""kimi_k25 specs (HF model type ``kimi_k25``). + +Kimi-K2.5 (Kimi-K2.6 and Kimi-K2.7-Code share its model type): the compressed-expert sibling of +``kimi_vl`` -- same swap to a native deepseek_v3 language model, plus on-the-fly decompression of +its packed experts at load. + +Only the expert-parallel section is registered under ``kimi_k25``: the MoE and export sections stay +``deepseek_v3``'s, so lookups by ``kimi_k25`` find neither yet. 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 DECOMPRESSED_FORMATS, make_buildable, prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_k25", + # transformers ships Kimi-K2.5 itself from 5.13, so modeling_source stays "transformers": the + # registry pairs it with that release, and a "remote_code" claim would be stale from 5.13 on. + # Expert parallelism still runs on the checkpoint's own code (``trust_remote_code=True``), + # which ``prepare`` patches, refusing the port. So the version describes transformers' port, + # not the model this spec's load builds: a MoE or export section added here must name the + # native deepseek_v3 language model's modules. + min_transformers_version="5.13", + # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan, + # and decompresses the experts as they load, which needs transformers 5.8. + ep_spec=EPSpec( + prepare=prepare, + prepared_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="text_config", + make_buildable=make_buildable, + decompressed_formats=DECOMPRESSED_FORMATS, + requires_transformers="5.8", + ), + ) +) diff --git a/modelopt/torch/models/kimi_vl/ep.py b/modelopt/torch/models/kimi_vl/ep.py index 9d6899f0210..1b80f6ce066 100644 --- a/modelopt/torch/models/kimi_vl/ep.py +++ b/modelopt/torch/models/kimi_vl/ep.py @@ -21,7 +21,22 @@ """ from .. import _kimi, _onfly_decompress -from .._remote_code import stage_full_remote_closure + + +def make_buildable( + hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_ +) -> None: + """Have the checkpoint's own code import and build on this transformers. + + The plan decision needs it before ``prepare``; it changes nothing the model computes or loads. + """ + _kimi.make_buildable( + "modeling_kimi_vl", + "KimiVLForConditionalGeneration", + hf_config, + ckpt_path, + trust_remote_code, + ) def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_): @@ -38,8 +53,7 @@ def prepare(hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool ckpt_path = _kimi.checkpoint_to_patch("kimi_vl", hf_config, ckpt_path, trust_remote_code) _kimi.refuse_quantization_only_on_text_config("kimi_vl", hf_config) - _kimi.install_compat_import_shims() - stage_full_remote_closure() + _kimi.make_importable() # the swap below loads the code and adapts its tie_weights # Native deepseek_v3 config for the LM (drops the remote auto_map so the native class is used). hf_config.text_config = _onfly_decompress.to_native_backbone_config(hf_config) diff --git a/modelopt/torch/models/kimi_vl/specs.py b/modelopt/torch/models/kimi_vl/specs.py index 73ba5e78b2c..bc511f806c9 100644 --- a/modelopt/torch/models/kimi_vl/specs.py +++ b/modelopt/torch/models/kimi_vl/specs.py @@ -27,7 +27,7 @@ from ..deepseek_v3.specs import BASE_MODEL_EP_PLAN from ..specs import EPSpec, ModelSpec, register -from .ep import prepare +from .ep import make_buildable, prepare __all__: list[str] = [] @@ -36,6 +36,11 @@ model_type="kimi_vl", modeling_source="remote_code", # prepare swaps a native deepseek_v3 in as the language model, with deepseek_v3's plan. - ep_spec=EPSpec(prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN), + ep_spec=EPSpec( + prepare=prepare, + prepared_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="text_config", + make_buildable=make_buildable, + ), ) ) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index fba161de33d..566bebd4922 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -213,19 +213,37 @@ class EPSpec(SpecSection): fields above have been applied, so a ``prepare`` that returns another config than the one it was handed must carry over the ``base_model_ep_plan`` they attached. - It runs before the run's plan is decided, wherever the plan could still be expert-parallel, so - a run whose EP degree comes out at 1 has run it too: it must leave a load of the whole model, - on every rank or FSDP2-sharded, as it would be without it, and may run more than once.""" + It runs only once a plan is decided expert-parallel: on every rank, in the plan decision's + collective "model spec prepared" step, on a copy of the config as read, so a refusal there + raises on every rank; or, for an expert-parallel plan handed in, in the load, after the load's + own checks, the GPU refusal included. It may patch the model's code, which is why nothing else + runs it. It may run more than once in a process.""" + + make_buildable: "Callable[..., Any] | None" = None + """For a checkpoint carrying its own modeling code: what that code needs to import and build on + the installed transformers, so the plan decision can build the model before ``prepare`` runs. + + Called as ``make_buildable(hf_config, ckpt_path=..., trust_remote_code=...)``, on every rank, + before the plan is decided, for every plan that could still be expert-parallel, whatever plan + comes out. So it must be idempotent and change nothing the model computes or loads: it may shim + names transformers dropped, wrap the remote-code staging to reach the modules the code imports, + and adapt a method to the arguments transformers now passes it, but patches no MoE or expert + class, config, weight converter or weight-loading hook.""" + + decompressed_formats: tuple[str, ...] = () + """The checkpoint formats (a ``quantization_config``'s ``quant_method``) whose experts + ``prepare`` has decompress as they load, which the expert-parallel load then takes beside + unquantized checkpoints. The plan decision runs before ``prepare``, and reads this to count such + a checkpoint one the expert-parallel load takes.""" prepared_ep_plan: dict[str, str] | None = None """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare`` supplies. - Stated, not applied: ``prepare`` still attaches it. A run that rules expert parallelism out - skips ``prepare``, which patches the model's code, yet must still tell an expert-parallel MoE - it would hold whole on every rank from a dense model. It finds the routed experts through this - plan, in the model as the checkpoint's own code builds it: ``prepare`` replaces the experts in - place, so the plan names their containers there too.""" + Deciding the plan attaches it with no ``prepare`` run, so the decision counts the routed experts + in the model the checkpoint's own code builds. A run that rules expert parallelism out finds + them the same way, to tell an expert-parallel MoE it would hold whole on every rank from a dense + model. ``prepare`` replaces the experts in place, so the plan names their containers there too.""" unsupported: str | None = None """Why this model cannot be quantized expert-parallel yet, when its config would route it there. diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index b9ed3863c63..08d40dc9ac4 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -226,6 +226,7 @@ def apply_distributed_support( *, ckpt_path: "str | None" = None, trust_remote_code: bool = False, + prepare: bool = True, ): """Prepare ``hf_config`` so ``model_type`` can be loaded expert-parallel, from its model spec. @@ -241,6 +242,11 @@ def apply_distributed_support( one (:func:`_ep_config`), not only on the config the spec would take: on the root, or on another nested config. + With ``prepare=False``, the spec's ``prepare`` does not run: the plan it would attach + (``prepared_ep_plan``) goes on the config instead, and only what makes the checkpoint's own code + import and build runs (``make_buildable``). That is what a plan decision reads, and it changes + nothing the model computes or loads, as a load that is not expert-parallel must find the process. + Returns the config to load, unchanged for a model type with no spec or no EP section. The plan is written into the config it is given, so a caller that keeps its own config passes a copy, as every caller here does. The generic all-to-all transport is not wired here -- it self-registers @@ -277,11 +283,14 @@ def apply_distributed_support( f"config.{' or config.'.join(attrs)}, which this config does not have, nor a plan " "of its own." ) - if ep.base_model_ep_plan and not existing: - target.base_model_ep_plan = dict(ep.base_model_ep_plan) + plan = ep.base_model_ep_plan or (None if prepare else ep.prepared_ep_plan) + if plan and not existing: + target.base_model_ep_plan = dict(plan) - if ep.prepare is not None: + if prepare and ep.prepare is not None: hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) + elif not prepare and ep.make_buildable is not None: + ep.make_buildable(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) return hf_config diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 00e2b4e9379..866e76b1d20 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1849,11 +1849,12 @@ def resolve_distribution_plan( Unless the call or rank 0's knobs rule expert parallelism out (a single process, an ``auto_class`` the sharded loads refuse, or an EP degree of 1, which auto is for a model whose spec refuses expert parallelism on the installed transformers), every rank applies the model's - spec to a copy of the config first - (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what - supplies the expert-parallel plan, and an expert-parallel plan returns that copy. Any other - plan returns the config without it, which the export would otherwise write out; the spec's - ``prepare`` has run for it all the same, which its contract allows + spec to a copy of the config first, in the collective "model spec" step: its plan, and what + makes the checkpoint's own code importable, with no ``prepare`` run + (:func:`~modelopt.torch.utils.moe.apply_distributed_support` with ``prepare=False``). For + several models it is what supplies the expert-parallel plan. An expert-parallel plan then has + every rank prepare a copy of the config as read, in the collective "model spec prepared" step, + and returns that copy. Any other plan returns the config as read, with no ``prepare`` run (:attr:`~modelopt.torch.models.EPSpec.prepare`). A plan that leaves an expert-parallel MoE whole on every rank raises, unless both knobs ask for exactly that: ``ep_size=1`` with ``dense="replicate"`` (``MODELOPT_EP_DEGREE=1`` and @@ -1896,7 +1897,7 @@ def checkpoint_read() -> bool: # class the sharded loads refuse rules it out too: the decision then plans the whole model on # every rank, or refuses, without the spec, whose prepare could fail, or patch the process, # for a load that never shards. - prepared = [hf_config] + stated = [hf_config] if dist.broadcast_from( "spec", lambda: ( @@ -1905,28 +1906,45 @@ def checkpoint_read() -> bool: and _ep_degree(ep_size, hf_config.model_type) != 1 ), ): - # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. A - # rank that cannot apply it fails every rank, as a checkpoint read does above. - def spec_applied() -> bool: - prepared[0] = apply_distributed_support( + # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. Its + # plan only: its prepare, which may patch the process, runs for an expert-parallel plan + # alone. A rank that cannot apply it fails every rank, as a checkpoint read does above. + def spec_stated() -> bool: + stated[0] = apply_distributed_support( hf_config.model_type, copy.deepcopy(hf_config), ckpt_path=ckpt_path, trust_remote_code=trust_remote_code, + prepare=False, ) return True - dist.all_reduce("model spec", spec_applied, all) + dist.all_reduce("model spec", spec_stated, all) plan = dist.broadcast_from( "plan", lambda: _decide_distribution( - ckpt_path, prepared[0], world_size, ep_size, dense, trust_remote_code, auto_class + ckpt_path, stated[0], world_size, ep_size, dense, trust_remote_code, auto_class ), ) if plan != DistributionPlan(): print_rank_0(f"Distribution: {plan!r}") - # Only the expert-parallel load shards by the plan the spec attached. - return plan, prepared[0] if plan.is_expert_parallel else hf_config + if not plan.is_expert_parallel: + return plan, hf_config + # Only the expert-parallel load shards by the plan the spec attached, and only it gets the spec + # prepared, on every rank. + prepared = [hf_config] + + def spec_applied() -> bool: + prepared[0] = apply_distributed_support( + hf_config.model_type, + copy.deepcopy(hf_config), + ckpt_path=ckpt_path, + trust_remote_code=trust_remote_code, + ) + return True + + dist.all_reduce("model spec prepared", spec_applied, all) + return plan, prepared[0] def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None: @@ -2008,7 +2026,11 @@ def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> Exception | None: parallelism out, or the transformers release it needs, is not asked here: the plan decision and the load ask the spec first (:func:`~modelopt.torch.utils.moe.check_distributed_support`). """ - if (method := _ckpt_quant_method(hf_config)) not in _EP_SOURCE_FORMATS: + # The formats it takes include those the spec's prepare, which runs only once the plan is + # decided, has the experts decompress as they load. + ep = _ep_spec(hf_config.model_type) + takes = _EP_SOURCE_FORMATS | frozenset(ep.decompressed_formats if ep is not None else ()) + if (method := _ckpt_quant_method(hf_config)) not in takes: return NotImplementedError( f"The expert-parallel load cannot dequantize a {method!r} checkpoint." ) diff --git a/tests/unit/torch/models/kimi_k25/__init__.py b/tests/unit/torch/models/kimi_k25/__init__.py new file mode 100644 index 00000000000..1275d78dff5 --- /dev/null +++ b/tests/unit/torch/models/kimi_k25/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + diff --git a/tests/unit/torch/models/kimi_k25/test_ep.py b/tests/unit/torch/models/kimi_k25/test_ep.py new file mode 100644 index 00000000000..6e51d9aba44 --- /dev/null +++ b/tests/unit/torch/models/kimi_k25/test_ep.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Expert-parallel adapter for Kimi-K2.5 (``kimi_k25``), a remote-code VLM on DeepSeek-V3. + +The adapter rebuilds the language model as transformers' native deepseek_v3 and installs the +on-the-fly expert decompressor. The checkpoint's remote code is not available here, so the tests +stub the dynamic-module loader with a stand-in remote module. Every process-wide patch is undone +after each test. +""" + +import sys +import types +from types import SimpleNamespace + +import pytest + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.8"): + pytest.skip( + "Kimi-K2.5 expert parallelism needs transformers >= 5.8, for the decompression", + allow_module_level=True, + ) + +import transformers.dynamic_module_utils as dmu +from transformers import DeepseekV3Config, PretrainedConfig + +from modelopt.torch.models import _kimi, _onfly_decompress, _remote_code, get_spec +from modelopt.torch.models.kimi_k25 import ep as kimi_k25 +from modelopt.torch.utils.moe import _routed_expert_keys, check_distributed_support + +QUANT = {"quant_method": "compressed-tensors", "format": "pack-quantized"} +REMOTE = "transformers_modules.Kimi_hyphen_K2_dot_5.0123abcd" # how the loader names remote code + + +class _RemoteTextConfig(PretrainedConfig): + """Stands in for the remote DeepSeek-V3 config Kimi-K2.5 bundles as its ``text_config``.""" + + model_type = "kimi_k2_remote_text" + + +def _vlm_config(): + """Kimi-K2.5's config: the compressed-tensors scheme sits only on the text config.""" + text = _RemoteTextConfig( + hidden_size=64, + num_hidden_layers=2, + n_routed_experts=8, + moe_intermediate_size=16, + auto_map={"AutoConfig": "configuration_deepseek.DeepseekV3Config"}, + quantization_config=QUANT.copy(), + ) + return SimpleNamespace(text_config=text) + + +@pytest.fixture +def remote(monkeypatch): + """Stand in for the checkpoint's remote code: importing it registers ``modeling_kimi_k25``, + whose ``DeepseekV3ForCausalLM`` is the per-expert remote language model. Also restores the + loader hooks and kept patches prepare installs, and leaves out its compat import shims, which + the Kimi-VL tests cover for both adapters.""" + module = types.ModuleType(f"{REMOTE}.modeling_kimi_k25") + module.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + module.loads = [] + + def import_remote_class(class_ref, ckpt_path, **kwargs): + module.loads.append((class_ref, ckpt_path, kwargs)) + monkeypatch.setitem(sys.modules, module.__name__, module) + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + monkeypatch.setattr(dmu, "get_class_in_module", dmu.get_class_in_module) + monkeypatch.setattr(_remote_code, "_KEPT_PATCHES", {}) + monkeypatch.setattr(_kimi, "install_compat_import_shims", lambda: None) + return module + + +@pytest.mark.parametrize( + "quantization_config", + [ + QUANT, + # Left to the expert-parallel load's format handling: bf16 experts load as they are, and + # any other scheme is read from the copy of quantization_config Kimi's config class keeps + # on the top level. + None, + {"quant_method": "fp8", "weight_block_size": [128, 128]}, + ], + ids=["compressed-tensors", "bf16", "block-fp8"], +) +def test_prepare_rebuilds_the_language_model_config_as_native_deepseek_v3( + monkeypatch, remote, quantization_config +): + """The text config becomes transformers' own ``DeepseekV3Config`` (no remote ``auto_map``, no + stale ``quantization_config``) with an EP plan for the fused experts, naming the experts the + spec states for a run that skips prepare; the decompressor is installed for compressed-tensors + experts only, while the text config it replaces still carries their scheme, and exactly for the + formats the spec states it decompresses, which the plan decision, made before prepare runs, does + not hold against the expert-parallel load; the remote-code staging shim is left installed.""" + installs = [] + monkeypatch.setattr( + _onfly_decompress, + "install_onfly_decompress", + lambda cfg: installs.append(_onfly_decompress._find_quantization_config(cfg)), + ) + staging = dmu.get_cached_module_file + config = _vlm_config() + config.text_config.quantization_config = quantization_config + if quantization_config is not QUANT: + config.quantization_config = quantization_config # the top-level copy the load reads + + assert kimi_k25.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) is config + + text = config.text_config + assert type(text) is DeepseekV3Config + assert getattr(text, "auto_map", None) is None + assert getattr(text, "quantization_config", None) is None + assert (text.hidden_size, text.n_routed_experts, text.moe_intermediate_size) == (64, 8, 16) + assert text.base_model_ep_plan["layers.*.mlp.experts.gate_up_proj"] == "grouped_gemm" + assert text.base_model_ep_plan["layers.*.mlp.experts.down_proj"] == "grouped_gemm" + assert _routed_expert_keys(text.base_model_ep_plan) == _routed_expert_keys( + get_spec("kimi_k25").ep_spec.prepared_ep_plan + ) + decompressed = get_spec("kimi_k25").ep_spec.decompressed_formats + expected = (quantization_config or {}).get("quant_method") in decompressed + assert installs == ([QUANT] if expected else []) + assert dmu.get_cached_module_file is not staging + + +def test_the_swap_leaves_a_sibling_remote_module_alone(monkeypatch, remote): + """Kimi-K2.5's ``modeling_kimi_k25`` imports ``modeling_deepseek``, which defines a + ``DeepseekV3ForCausalLM`` of its own; the swap, scoped to ``modeling_kimi_k25``, must not reach + it. The swap itself is the shared adapter tests'.""" + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + bystander = types.ModuleType(f"{REMOTE}.modeling_deepseek") + remote_lm = bystander.DeepseekV3ForCausalLM = type("DeepseekV3ForCausalLM", (), {}) + monkeypatch.setitem(sys.modules, bystander.__name__, bystander) + + kimi_k25.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + + assert bystander.DeepseekV3ForCausalLM is remote_lm + + +def test_the_build_step_leaves_transformers_own_port_alone(monkeypatch, remote): + """From 5.13 transformers ships its own Kimi-K2.5, whose module shares the remote one's name + (``transformers.models.kimi_k25.modeling_kimi_k25``). The build step adapts the remote VLM + class's ``tie_weights`` and must not reach the port's, which a later load of it would use.""" + + def own_tie(self): + return "tied" + + port = types.ModuleType("transformers.models.kimi_k25.modeling_kimi_k25") + port.KimiK25ForConditionalGeneration = type("Port", (), {"tie_weights": own_tie}) + monkeypatch.setitem(sys.modules, port.__name__, port) + remote.KimiK25ForConditionalGeneration = type("Remote", (), {"tie_weights": own_tie}) + + kimi_k25.make_buildable(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) + + assert vars(port.KimiK25ForConditionalGeneration)["tie_weights"] is own_tie + assert vars(remote.KimiK25ForConditionalGeneration)["tie_weights"]._kimi_tie_weights_shim + + +def test_a_checkpoint_the_decompressor_refuses_is_refused_before_anything_changes( + monkeypatch, remote +): + """The decompressor refuses a checkpoint it cannot decompress as it loads before prepare + replaces the language model's config or wraps the loader's remote-code staging.""" + + def refuse(cfg): + raise NotImplementedError("the checkpoint has 2 compressed-tensors config groups") + + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", refuse) + staging = dmu.get_cached_module_file + config = _vlm_config() + text = config.text_config + + with pytest.raises(NotImplementedError, match="config groups"): + kimi_k25.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) + + assert config.text_config is text + assert dmu.get_cached_module_file is staging + + +def test_a_transformers_too_old_to_decompress_is_refused_before_prepare(monkeypatch): + """The experts decompress as they load from transformers 5.8: the spec refuses an older + release before prepare swaps the language model's config.""" + monkeypatch.setattr(transformers, "__version__", "5.7.0") + + with pytest.raises(RuntimeError, match=r"transformers >= 5\.8"): + check_distributed_support("kimi_k25") + + +def test_prepare_refuses_transformers_own_port(): + """From 5.13 a load without ``trust_remote_code`` builds transformers' port, which the patch + never reaches; ``prepare`` must refuse it rather than run unpatched.""" + native = pytest.importorskip("transformers.models.kimi_k25.configuration_kimi_k25") + with pytest.raises(ValueError, match="runs on the checkpoint's own modeling code"): + kimi_k25.prepare(native.Kimi_K25Config(), ckpt_path="/ckpt", trust_remote_code=True) diff --git a/tests/unit/torch/models/kimi_vl/test_ep.py b/tests/unit/torch/models/kimi_vl/test_ep.py index bd5158077a9..c2d7c820991 100644 --- a/tests/unit/torch/models/kimi_vl/test_ep.py +++ b/tests/unit/torch/models/kimi_vl/test_ep.py @@ -21,8 +21,11 @@ after each test. """ +import contextlib import copy import importlib +import json +import shutil import sys import types from types import SimpleNamespace @@ -37,13 +40,26 @@ if Version(transformers.__version__) < Version("5.7"): pytest.skip("Kimi-VL expert parallelism needs transformers >= 5.7", allow_module_level=True) +import transformers.core_model_loading as cml import transformers.dynamic_module_utils as dmu +import transformers.modeling_utils as mu import transformers.utils.import_utils as import_utils -from transformers import DeepseekV3Config, DeepseekV3ForCausalLM, PretrainedConfig +from _test_utils.torch import kimi_vl_tiny_modeling +from torch import nn +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + DeepseekV3Config, + DeepseekV3ForCausalLM, + PretrainedConfig, +) from modelopt.torch.models import _onfly_decompress, _remote_code, get_spec +from modelopt.torch.models.kimi_k25 import ep as kimi_k25 from modelopt.torch.models.kimi_vl import ep as kimi_vl +from modelopt.torch.utils import distributed as dist from modelopt.torch.utils.moe import _routed_expert_keys +from modelopt.torch.utils.plugins import model_load_utils # Symbols the remote modeling imports that newer transformers dropped. COMPAT_SYMBOLS = ( @@ -159,11 +175,14 @@ def test_prepare_rebuilds_the_language_model_config_without_a_decompressor(monke assert dmu.get_cached_module_file is not staging -# What the Kimi adapters share, through each adapter's prepare. +# Kimi-K2.5's adapter is Kimi-VL's with on-the-fly decompression added; these hold for both. _BOTH_ADAPTERS = pytest.mark.parametrize( ("adapter", "remote"), - [(kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration"))], - ids=["vl"], + [ + (kimi_vl, ("modeling_kimi_vl", "KimiVLForConditionalGeneration")), + (kimi_k25, ("modeling_kimi_k25", "KimiK25ForConditionalGeneration")), + ], + ids=["vl", "k25"], indirect=["remote"], ) @@ -196,10 +215,10 @@ def test_prepare_with_a_checkpoint_swaps_in_native_deepseek_v3_and_shims_tie_wei def test_a_load_off_expert_parallelism_builds_the_checkpoint_s_own_language_model( monkeypatch, adapter, remote ): - """Deciding the plan prepares a copy of the config on every rank before the plan is known, and - the patch outlives that copy. A load the plan keeps off expert parallelism (replicated, or - FSDP2) builds from the config as read, so its VLM gets the checkpoint's own language model, as - with no patch; only the native config ``prepare`` makes builds transformers' native one.""" + """The patch outlives the load it was made for. A later load in the process that its plan keeps + off expert parallelism (replicated, or FSDP2) builds from the config as read, so its VLM gets + the checkpoint's own language model, as with no patch; only the native config ``prepare`` makes + builds transformers' native one.""" monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) as_read = _vlm_config() prepared = adapter.prepare(copy.deepcopy(as_read), ckpt_path="/ckpt", trust_remote_code=True) @@ -209,6 +228,131 @@ def test_a_load_off_expert_parallelism_builds_the_checkpoint_s_own_language_mode assert type(remote.vlm(as_read).language_model) is remote.remote_lm +def _built_by_its_own_code(): + """The VLM the checkpoint's own code builds, on the meta device: each MoE layer of its language + model holds its routed experts as per-expert modules, at ``language_model.model.layers..mlp.experts``.""" + + def layer(): + mlp = nn.Module() + mlp.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(8)) + return nn.ModuleDict({"mlp": mlp}) + + model = nn.Module() + model.language_model = nn.Module() + model.language_model.model = nn.Module() + model.language_model.model.layers = nn.ModuleList(layer() for _ in range(2)) + model.tie_weights = lambda: None + return model + + +@_BOTH_ADAPTERS +@pytest.mark.parametrize( + ("ranks", "dense", "outcome"), + [(3, "fsdp", "fsdp2"), (3, "replicate", "refused"), (4, "replicate", "expert-parallel")], + ids=["fsdp2", "whole", "expert-parallel"], +) +def test_only_an_expert_parallel_plan_patches_the_checkpoint_s_code( + monkeypatch, tmp_path, adapter, remote, ranks, dense, outcome +): + """Deciding the plan builds the model the checkpoint's own code builds, importable but unpatched, + and finds its 8 experts through the plan the spec states. Only an expert-parallel plan (4 ranks) + then runs ``prepare``, which swaps in the native language model. FSDP2-sharded, or left whole on + 3 ranks, which no EP degree of the 8 experts divides (refused), the run keeps the checkpoint's + own class, the loader and the config as read.""" + model_type = "kimi_vl" if adapter is kimi_vl else "kimi_k25" + as_read = _vlm_config() + as_read.model_type = model_type + monkeypatch.setattr(dist, "size", lambda group=None: ranks) + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: copy.deepcopy(as_read) + ) + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: _built_by_its_own_code() + ) + hooks = (mu.get_model_conversion_mapping, cml.spawn_materialize) + + if outcome == "refused": + with pytest.raises(ValueError, match="can be expert-parallel, but the plan resolved"): + model_load_utils.resolve_distribution_plan( + str(tmp_path), dense=dense, trust_remote_code=True + ) + else: + plan, config = model_load_utils.resolve_distribution_plan( + str(tmp_path), dense=dense, trust_remote_code=True + ) + assert plan.is_expert_parallel == (outcome == "expert-parallel") + assert import_utils.is_torch_fx_available() is False # importable on every plan + assert getattr(vars(remote.vlm)["tie_weights"], "_kimi_tie_weights_shim", False) # buildable + assert (remote.DeepseekV3ForCausalLM is remote.remote_lm) == (outcome != "expert-parallel") + if outcome == "expert-parallel": + assert type(config.text_config) is DeepseekV3Config + assert config.text_config.base_model_ep_plan + else: + assert (mu.get_model_conversion_mapping, cml.spawn_materialize) == hooks + if outcome == "fsdp2": + assert type(config.text_config) is _RemoteTextConfig + assert config.text_config.base_model_ep_plan is None + + +@pytest.mark.parametrize( + "kwargs", + [{"ckpt_path": None, "trust_remote_code": True}, {"ckpt_path": "/ckpt"}], + ids=["no-checkpoint", "untrusted"], +) +def test_the_build_step_installs_nothing_when_there_is_nothing_to_build(monkeypatch, kwargs): + """With no checkpoint, or code this load may not run, nothing of the checkpoint's is imported or + built, so the build step installs no import shim and no staging either.""" + for module_name, name in COMPAT_SYMBOLS: + with contextlib.suppress(ImportError): + _restore(monkeypatch, importlib.import_module(module_name), name) + monkeypatch.delattr(import_utils, "is_torch_fx_available", raising=False) + staging = dmu.get_cached_module_file + monkeypatch.setattr(dmu, "get_cached_module_file", staging) + + kimi_vl.make_buildable(_vlm_config(), **kwargs) + + assert dmu.get_cached_module_file is staging + assert not hasattr(import_utils, "is_torch_fx_available") + + +def test_the_build_step_lets_the_checkpoint_s_own_vlm_build_on_this_transformers( + monkeypatch, tmp_path +): + """The plan decision builds the checkpoint's own code on the meta device before ``prepare`` + runs, and Kimi-VL's VLM class has a ``tie_weights`` that takes none of the keyword arguments + transformers' ``post_init`` passes it. ``make_buildable`` adapts it, and swaps nothing: the + module still builds the language model from its own ``DeepseekV3ForCausalLM``.""" + for module_name, name in COMPAT_SYMBOLS: + with contextlib.suppress(ImportError): + _restore(monkeypatch, importlib.import_module(module_name), name) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + shutil.copy(kimi_vl_tiny_modeling.__file__, tmp_path / "modeling_kimi_vl.py") + remote = "modeling_kimi_vl.KimiVLForConditionalGeneration" + auto_map = {"AutoConfig": "modeling_kimi_vl.KimiVLConfig", "AutoModelForCausalLM": remote} + text = {"model_type": "deepseek_v3", "hidden_size": 16, "num_hidden_layers": 1} + stored = {"model_type": "kimi_vl", "auto_map": auto_map, "text_config": text} + (tmp_path / "config.json").write_text(json.dumps(stored)) + loaded = set(sys.modules) + try: + config = AutoConfig.from_pretrained(tmp_path, trust_remote_code=True) + + kimi_vl.make_buildable(config, ckpt_path=str(tmp_path), trust_remote_code=True) + model = model_load_utils.build_meta_causal_lm( + str(tmp_path), True, None, copy.deepcopy(config), auto_class=AutoModelForCausalLM + ) + + assert type(model).tie_weights._kimi_tie_weights_shim + # Nothing swapped: the language model is the checkpoint's own, its experts per expert. + own = sys.modules[type(model).__module__].DeepseekV3ForCausalLM + assert not getattr(own, "_kimi_native_swap", False) + assert type(model.language_model) is own + finally: + # The checkpoint's code stays imported otherwise, for later tests to find as a Kimi module. + for name in set(sys.modules) - loaded: + if name.startswith("transformers_modules."): + del sys.modules[name] + + @_BOTH_ADAPTERS @pytest.mark.parametrize( ("missing", "match"), @@ -233,7 +377,7 @@ def test_the_swap_fails_loudly_when_the_remote_modeling_has_nothing_to_swap( adapter.prepare(_vlm_config(), ckpt_path="/ckpt", trust_remote_code=True) -@pytest.mark.parametrize("adapter", [kimi_vl], ids=["vl"]) +@pytest.mark.parametrize("adapter", [kimi_vl, kimi_k25], ids=["vl", "k25"]) @pytest.mark.parametrize( ("config", "kwargs", "match"), [ @@ -282,7 +426,9 @@ def unstaged_loader(*args, **kwargs): assert dmu.get_cached_module_file is unstaged_loader -@pytest.mark.parametrize(("model_type", "adapter"), [("kimi_vl", kimi_vl)], ids=["vl"]) +@pytest.mark.parametrize( + ("model_type", "adapter"), [("kimi_vl", kimi_vl), ("kimi_k25", kimi_k25)], ids=["vl", "k25"] +) def test_the_spec_routes_expert_parallelism_through_its_own_adapter(model_type, adapter): """The checkpoint's code builds no EP-capable experts, so its whole EP setup is this adapter's ``prepare``: another Kimi adapter's would load it without its own swap or decompression.""" diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 7a4bd9ef182..76b23b2a26b 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -87,6 +87,17 @@ def test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal(model_type): ) +@pytest.mark.parametrize("model_type", EP_MODEL_TYPES) +def test_a_spec_with_a_prepare_states_its_plan(model_type): + """Deciding the plan runs no prepare, so a plan only prepare attaches goes unseen there: the + decision reads the MoE as dense and plans it without expert parallelism, with no refusal.""" + ep = get_spec(model_type).ep_spec + assert not ep.prepare or ep.base_model_ep_plan or ep.prepared_ep_plan, ( + f"{model_type} has a prepare and states no plan, which the plan decision needs, as it runs " + "no prepare: state the plan prepare attaches as prepared_ep_plan" + ) + + @pytest.mark.parametrize("model_type", EP_MODEL_TYPES) def test_a_plan_goes_where_the_load_reads_one(model_type): """A spelling the load does not read is read as no plan at all, and the MoE planned as dense.""" diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 817da74a8c6..95a7007db71 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -825,6 +825,33 @@ def test_an_moe_the_expert_parallel_load_refuses_is_decided_without_it( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2) +@pytest.mark.skipif( + Version(transformers.__version__) < Version(_ep_floor("kimi_k25")), + reason="Kimi-K2.5's expert parallelism needs a later transformers", +) +@pytest.mark.parametrize( + ("model_type", "expert_parallel"), [("kimi_k25", True), ("kimi_vl", False)] +) +def test_a_checkpoint_its_spec_decompresses_is_decided_expert_parallel( + run_as, serve, monkeypatch, tmp_path, model_type, expert_parallel +): + """The decision runs before ``prepare``, which installs the decompressor. A compressed-tensors + checkpoint whose spec decompresses its experts as they load (Kimi-K2.5's) is planned + expert-parallel all the same; one whose spec does not (Kimi-VL's) is planned without it, as the + expert-parallel load would refuse it.""" + spec = get_spec(model_type).ep_spec + monkeypatch.setattr(spec, "prepare", lambda hf_config, **kwargs: hf_config) + monkeypatch.setattr(spec, "make_buildable", lambda *args, **kwargs: None) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi()) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: None) + run_as(4) + scheme = {"quant_method": "compressed-tensors"} + text = SimpleNamespace(n_routed_experts=4, quantization_config=scheme) + serve(SimpleNamespace(model_type=model_type, quantization_config=scheme, text_config=text)) + + assert resolve_distribution_plan(str(tmp_path))[0].is_expert_parallel == expert_parallel + + def _remote_kimi(): """A remote-code Kimi VLM as the checkpoint's own code builds it, unpatched by its spec's ``prepare``: each MoE layer of its language model holds its routed experts as per-expert @@ -844,7 +871,7 @@ def layer(): return model -@pytest.mark.parametrize("model_type", ["kimi_vl"]) +@pytest.mark.parametrize("model_type", ["kimi_vl", "kimi_k25"]) def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( run_as, serve, monkeypatch, tmp_path, model_type ): @@ -1735,8 +1762,9 @@ def test_a_load_prepares_its_config_once( ) prepared = [] - def apply_distributed_support(model_type, hf_config, **kwargs): - prepared.append(model_type) + def apply_distributed_support(model_type, hf_config, prepare=True, **kwargs): + if prepare: # the plan alone, which deciding reads, prepares nothing + prepared.append(model_type) hf_config.base_model_ep_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan return hf_config @@ -1961,9 +1989,12 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, def _one_rank_cannot_apply_the_spec_job(rank, size, ckpt, step): apply = model_load_utils.apply_distributed_support + prepares = [] def apply_on_rank0_only(*args, **kwargs): - if rank == 1: + prepares.append(kwargs.get("prepare", True)) + # "prepared": rank 1 states the spec as rank 0 does, and fails only where prepare runs. + if rank == 1 and (step != "prepared" or prepares[-1]): raise OSError("this rank cannot import the checkpoint's code") return apply(*args, **kwargs) @@ -1983,7 +2014,7 @@ def apply_on_rank0_only(*args, **kwargs): mock.patch.object(model_load_utils, "load_and_prepare_ep", _loaded_in_step) ) try: - if step == "decision": + if step in ("decision", "prepared"): resolve_distribution_plan(ckpt) else: load_pretrained(ckpt, distribution=DistributionPlan(2, ep_size=2)) @@ -1992,14 +2023,17 @@ def apply_on_rank0_only(*args, **kwargs): else: raised = "no error" assert "cannot import the checkpoint's code" in raised, raised + # Deciding an expert-parallel plan, every rank reached the step that prepares it. + assert (True in prepares) == (step != "decision"), prepares @needs_ep_transformers -@pytest.mark.parametrize("step", ["decision", "load"]) +@pytest.mark.parametrize("step", ["decision", "prepared", "load"]) def test_a_rank_that_cannot_apply_the_spec_fails_every_rank(cpu_dist_workers, tmp_path, step): """A spec's preparation reads the checkpoint and imports its code on each rank, which can fail - on one alone; the others must not wait for it, in the plan's broadcast when deciding, nor in the - load a handed expert-parallel plan starts.""" + on one alone; the others must not wait for it, in the plan's broadcast when deciding, in the + step that prepares the expert-parallel plan decided, nor in the load a handed expert-parallel + plan starts.""" _deepseek_v3(n_routed_experts=8).save_pretrained(tmp_path) cpu_dist_workers.run(_one_rank_cannot_apply_the_spec_job, ckpt=str(tmp_path), step=step)