From 51d7aa92d15cba064091715b9617bbc019359c28 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:08:30 +0000 Subject: [PATCH 01/32] Add the Kimi-K3 expert-parallel adapter Kimi-K3's text backbone, kimi_linear, has no native transformers modeling to swap in, so the adapter patches the checkpoint's own code in place. It replaces the MoE block's per-expert ModuleList with fused experts that dispatch through the experts interface, keeping K3's latent expert space and its custom SwiGLU. It registers converters that fuse the checkpoint's per-expert w1/w3 into gate_up_proj and stack w2 into down_proj, and installs the on-the-fly decompressor for its MXFP4 experts. It registers torch_a2a before the plan decision builds the model on the meta device. Its spec states the plan prepare attaches and declares transformers 5.13, the release it was validated against. The default PTQ units keep K3's latent-expert projections and residual gates, and Kimi's mm_projector, unquantized, as vLLM and the published NVFP4 checkpoints expect. The fixes the checkpoint's own modeling needs on current transformers and fla-core now come from their own module, kimi_k3/_remote_fixes.py (#3132). The refusal of a Kimi-K3 left whole on every rank has its own test instead of a row in the Kimi-K2.5 adapter's. The adapter's earlier history is archived; this commit replaces it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/__init__.py | 1 + modelopt/torch/models/kimi_k3/__init__.py | 4 +- modelopt/torch/models/kimi_k3/ep.py | 381 +++++++++++ modelopt/torch/models/kimi_k3/specs.py | 38 ++ modelopt/torch/utils/moe.py | 5 +- .../units/default_disabled_quantizers.yaml | 17 +- tests/unit/recipe/test_presets.py | 44 ++ tests/unit/torch/models/test_kimi_k3_ep.py | 592 ++++++++++++++++++ .../unit/torch/utils/test_model_load_utils.py | 37 ++ 9 files changed, 1115 insertions(+), 4 deletions(-) create mode 100644 modelopt/torch/models/kimi_k3/ep.py create mode 100644 modelopt/torch/models/kimi_k3/specs.py create mode 100644 tests/unit/torch/models/test_kimi_k3_ep.py diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index bd748e773e0..e3794c17771 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -90,6 +90,7 @@ nemotron, nemotron_h, minimax_m3_vl, + kimi_k3, llama4_text, kimi_vl, kimi_k25, diff --git a/modelopt/torch/models/kimi_k3/__init__.py b/modelopt/torch/models/kimi_k3/__init__.py index b340a6b60b1..05b323e388c 100644 --- a/modelopt/torch/models/kimi_k3/__init__.py +++ b/modelopt/torch/models/kimi_k3/__init__.py @@ -13,4 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Kimi-K3 (HF model type ``kimi_k3``).""" +"""kimi_k3 specs (HF model type ``kimi_k3``).""" + +from . import specs diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py new file mode 100644 index 00000000000..f0adc73f1d5 --- /dev/null +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -0,0 +1,381 @@ +# 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 all-to-all support for Kimi-K3 (``kimi_k3``). + +a remote-code VLM whose text backbone is ``kimi_linear`` (KDA/linear attention + DeepSeek-V3-style routed MoE with a +LATENT expert space). + +Kimi-K3 has no transformers-native modeling (``CONFIG_MAPPING_NAMES`` has no ``kimi_linear``/``kimi_k3``), +so -- unlike ``kimi_k25``, which swaps to native ``deepseek_v3`` -- we cannot get EP-capable fused experts +by swapping in a native backbone. We must adapt the checkpoint's REMOTE code in place. Two structural +obstacles and how this plugin removes them: + +1. **Per-expert ``nn.ModuleList``.** The remote ``KimiSparseMoeBlock`` holds ``self.experts`` as an + ``nn.ModuleList`` of ``num_experts`` ``KimiBlockSparseMLP`` (each a separate ``w1``=gate / ``w3``=up / + ``w2``=down Linear), and its ``forward`` runs a hardcoded per-expert Python loop (``moe_infer``). The + transformers EP framework (and ModelOpt's ``torch_a2a`` transport) instead needs the FUSED 3-D grouped + params ``gate_up_proj [E, 2*I, H]`` / ``down_proj [E, H, I]`` behind ``ExpertsInterface``. This plugin + replaces ``self.experts`` with a fused ``KimiK3Experts`` module (decorated with + ``use_experts_implementation``) and rewrites ``KimiSparseMoeBlock.forward`` to dispatch through it -- so + the generic ``torch_a2a`` experts forward (NCCL all-to-all -> quantizable ``_grouped_linear``) takes over. + +2. **No expert fusion for the checkpoint's names.** The checkpoint stores per-expert ``w1``/``w3``/``w2`` + (gate/up/down), and transformers registers no expert fusion for custom-code models. This plugin + registers ``kimi_linear`` WeightConverters that fuse ``...experts..w1|w3`` into ``gate_up_proj`` and + stack ``w2`` into ``down_proj``, which the ``grouped_gemm`` plan entries then EP-shard. + +3. **Compressed (mxfp4) experts.** In this checkpoint the 896 routed experts are compressed-tensors + ``mxfp4-pack-quantized`` (4-bit, group 32; stored as ``w1/w3/w2.weight_packed`` + ``.weight_scale``) -- + everything else (attention, shared experts, latent projections, dense layer, embeddings, lm_head, + vision) is plain bf16. This plugin installs the format-agnostic on-the-fly decompressor + (``_onfly_decompress``) so from_pretrained dequantizes each per-expert packed weight -> bf16 INSIDE + the fused-expert conversion (EP-shard-aware -- only this rank's local experts), then drops the + ``quantization_config`` so transformers does not build a compressed-tensors HF quantizer. ModelOpt PTQ + then inserts fp8/nvfp4 quantizers on the resulting bf16 fused params. (Re-quantizing an already-4-bit + model to fp8/nvfp4 is lossy-on-lossy; that is the meaning of PTQ-ing this checkpoint.) + +Kimi-K3 specifics preserved as DENSE / replicated (NOT quantized through the EP experts path): + * the LATENT MoE wrapper -- ``routed_expert_down_proj`` (H->H_latent) / ``routed_expert_up_proj`` + (H_latent->H) and the optional ``routed_expert_norm`` -- the experts run in the H_latent + (``routed_expert_hidden_size``) space between them; the block forward keeps these around the fused call. + * the always-on ``shared_experts`` (a dense ``KimiMLP``) and the inline ``gate`` (``KimiMoEGate``, + sigmoid noaux_tc). The gate runs over the REPLICATED router weight and emits GLOBAL top-k ids, so it + stays a plain replicated module (no ``ep_router`` plan entry -- as for deepseek_v3's global-id + gate). The custom ``situ`` SwiGLU (``SituAndMul``) is preserved via the fused module's + ``_apply_gate``. + +VLM: the plan is set on ``text_config`` and bubbles up under ``language_model.model.*`` (the VLM builds +``self.language_model = KimiLinearForCausalLM(config.text_config)``); vision tower + mm projector stay bf16. + +Also shims ``OutputRecorder`` back into ``transformers.utils.generic`` (the remote code imports it there, +but it moved to ``transformers.modeling_utils``) so the modeling imports at all. + +Version floor: transformers >= 5.13, which it was developed and validated against (the spec declares it), +with flash-linear-attention (FLA) 0.5.0, which the remote code's KDA layers need at runtime. +""" + +import torch +import torch.nn as nn + +from .. import _onfly_decompress +from .._remote_code import stage_full_remote_closure +from ._remote_fixes import ( + _find_module_defining, + _install_output_recorder_shim, + _patch_kda_a_log, + _patch_kda_beta_sigmoid, + _patch_kimi_cache_get_mask_sizes, + _shim_create_causal_mask, + _shim_tie_weights_kwargs, +) + +# base-relative EP plan (bubbles to `language_model.model.` in the VLM). MoE lives at +# `layers.*.block_sparse_moe` (NOT `.mlp`; layer 0 is dense `.mlp`). Fused expert params shard dim-0 via +# grouped_gemm; the experts container gets the EP group via moe_tp_experts (-> torch_a2a_experts for a2a). +# No `ep_router` entry: KimiMoEGate runs inline over the replicated router weight and emits GLOBAL ids. +BASE_MODEL_EP_PLAN = { + "layers.*.block_sparse_moe.experts.gate_up_proj": "grouped_gemm", + "layers.*.block_sparse_moe.experts.down_proj": "grouped_gemm", + "layers.*.block_sparse_moe.experts": "moe_tp_experts", +} + + +def _find_remote_module(package: str): + """Return the loaded remote backbone module. + + the one in ``package`` that DEFINES ``KimiSparseMoeBlock`` (the Kimi MoE block this plugin patches) -- or + None. Name-independent; see ``_find_module_defining``. + """ + return _find_module_defining("KimiSparseMoeBlock", package) + + +def _build_fused_experts_class(kmod): + """Build the fused experts class the expert-parallel path needs. + + Build the fused ``KimiK3Experts`` class (bound to the remote module's ``SituAndMul``), decorated + with ``use_experts_implementation`` so its forward dispatches through transformers' ``ExpertsInterface`` + and honors ``config._experts_implementation`` (``torch_a2a`` at EP-PTQ time). Stores the standard fused + grouped params ``gate_up_proj [E, 2*I, H_latent]`` / ``down_proj [E, H_latent, I]`` (``is_transposed= + False``, ``[E, out, in]``). ``_apply_gate`` preserves Kimi's ``situ`` SwiGLU. + """ + from transformers.integrations.moe import use_experts_implementation + + SituAndMul = kmod.SituAndMul # noqa: N806 (a class, keeping its own name) + _get_situ = kmod._get_situ_activation_params + + class KimiK3Experts(nn.Module): + def __init__(self, config, hidden_dim): + super().__init__() + self.num_experts = config.num_experts + # The width the MoE block runs its experts at: the latent one under latent MoE, else hidden. + self.hidden_dim = hidden_dim + self.intermediate_dim = config.moe_intermediate_size + self.gate_up_proj = nn.Parameter( + torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim) + ) + self.down_proj = nn.Parameter( + torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim) + ) + # explicit attrs the torch_a2a experts forward reads (also set by use_experts_implementation). + self.has_gate = True + self.has_bias = False + self.is_transposed = False + if getattr(config, "hidden_act", None) == "situ": + beta, linear_beta = _get_situ(config) + self._situ = SituAndMul(beta=beta, linear_beta=linear_beta) + else: + self._situ = None + + def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: + # SituAndMul takes the FULL concat [gate; up] and splits internally -- matches the fused + # gate_up_proj row layout (rows [0:I]=gate from w1, [I:2I]=up from w3). + if self._situ is not None: + return self._situ(gate_up) + gate, up = gate_up.chunk(2, dim=-1) + return nn.functional.silu(gate) * up + + def forward(self, hidden_states, top_k_index, top_k_weights): + # Fallback (non-EP / eager) dense per-expert loop; the EP-PTQ path replaces this with the + # torch_a2a experts forward via config._experts_implementation. Mirrors the native + # DeepseekV3Experts/MiniMaxM3VLExperts reference loop, with situ gating. + final = torch.zeros_like(hidden_states) + with torch.no_grad(): + # One class past the experts: the id a router gives a token routed to no local expert. + mask = nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1).permute( + 2, 1, 0 + ) + hit = torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(mask[expert_idx]) + cur = self._apply_gate( + nn.functional.linear(hidden_states[token_idx], self.gate_up_proj[expert_idx]) + ) + cur = nn.functional.linear(cur, self.down_proj[expert_idx]) + cur = cur * top_k_weights[token_idx, top_k_pos, None] + final.index_add_(0, token_idx, cur.to(final.dtype)) + return final + + # is_transposed=False: [E, out, in] (like deepseek/qwen). has_gate=True: gate_up is FUSED (2*I) and + # split by _apply_gate. has_bias=False: bias-free experts. + use_experts_implementation(KimiK3Experts, is_transposed=False, has_bias=False, has_gate=True) + return KimiK3Experts + + +def _patch_kimi_moe_block(kmod) -> None: + """Replace the per-expert ModuleList with the fused experts module. + + Swap ``KimiSparseMoeBlock.experts`` (per-expert ModuleList) for the fused experts module and rewrite + its ``forward`` to dispatch through it (keeping the latent down/up projections, optional norm, and + shared experts around the fused call). Idempotent. + """ + from modelopt.torch.utils import print_rank_0 + + Block = kmod.KimiSparseMoeBlock # noqa: N806 (a class, keeping its own name) + if getattr(Block, "_kimi_k3_ep_patched", False): + return + FusedExperts = _build_fused_experts_class(kmod) # noqa: N806 (a class, keeping its own name) + + _orig_init = Block.__init__ + + def _init(self, config): + _orig_init(self, config) + # Replace the ModuleList of per-expert MLPs with the fused grouped-param module (EP-capable), as + # wide as the space the block runs its experts in. The block decides latent MoE itself (Kimi-K3's + # config has no use_latent_moe field) and the forward below follows it, so read it off the block. + del self.experts + width = config.routed_expert_hidden_size if self.use_latent_moe else config.hidden_size + self.experts = FusedExperts(config, width) + + def _forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + if self.use_latent_moe: + hidden_states = self.routed_expert_down_proj(hidden_states) + # Fused / a2a experts: returns the top-k-combined output in the latent space. + y = self.experts(hidden_states, topk_idx, topk_weight) + if self.use_latent_moe: + if self.latent_moe_use_norm: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + y = y.reshape(orig_shape) + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + Block.__init__ = _init + Block.forward = _forward + Block._kimi_k3_ep_patched = True + print_rank_0("[kimi_k3] patched KimiSparseMoeBlock: fused experts + ExpertsInterface forward") + + +def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bool) -> type: + """Import the remote modeling now, so its classes can be patched before the model is built. + + Force the remote modeling into ``sys.modules`` (so we can monkey-patch its classes BEFORE + ``from_pretrained`` instantiates the model) by resolving the model class the checkpoint's config + declares in its ``auto_map`` -- no hardcoded remote file/module name. Runs the OutputRecorder + closure + shims first. Returns that class, whose package holds this checkpoint's remote code. + """ + _install_output_recorder_shim() + stage_full_remote_closure() + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + auto_map = getattr(hf_config, "auto_map", None) or {} + class_ref = auto_map.get("AutoModelForCausalLM") or auto_map.get("AutoModel") + if not class_ref: + raise RuntimeError( + "kimi_k3: config.auto_map has no AutoModel/AutoModelForCausalLM entry to force-load the remote " + "modeling" + ) + return get_class_from_dynamic_module(class_ref, ckpt_path, trust_remote_code=trust_remote_code) + + +def _register_kimi_k3_conversions() -> None: + """Register the per-expert to fused-expert weight converters. + + Register the per-expert ``w1/w3/w2 -> gate_up_proj/down_proj`` fusion WeightConverters for Kimi-K3's + ``kimi_linear`` backbone. + + transformers 5.13's ``get_model_conversion_mapping`` SKIPS building any weight conversions for + custom-code (remote) models unless their class/model_type is in ``USER_REGISTERED_MAPPINGS`` + (``if submodule.is_custom_code() and class_name/model_type not in USER_REGISTERED_MAPPINGS: continue``). + Kimi-K3 is remote code, so the per-expert -> fused-expert fusion is NEVER built: the checkpoint's + ``experts..w1/w3/w2.weight_packed`` land as UNEXPECTED and the model's fused ``gate_up_proj`` / + ``down_proj`` as MISSING (newly-initialized garbage, and allocated at FULL 896-expert size on every + rank -> OOM at scale). Register the fusion converters ourselves (mirrors transformers' own native + ``mlp.experts.*.w1/w3/w2.weight`` fusion, adapted to K3's ``block_sparse_moe.experts``): w1=gate, + w3=up (Concatenate gate-then-up on dim 1), w2=down. The on-the-fly decompressor then rewrites these + ``.weight`` sources to the checkpoint's ``.weight_packed`` + scales (mxfp4 -> bf16). Idempotent. + """ + # transformers is optional. Its conversion API is in every version prepare() accepts, and a + # load without it would leave the experts MISSING, so an ImportError here must stop the load. + from transformers.conversion_mapping import register_checkpoint_conversion_mapping + from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter + + from modelopt.torch.utils import print_rank_0 + + def _mapping(): + return [ + WeightConverter( + source_patterns=[ + "block_sparse_moe.experts.*.w1.weight", + "block_sparse_moe.experts.*.w3.weight", + ], + target_patterns="block_sparse_moe.experts.gate_up_proj", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + source_patterns="block_sparse_moe.experts.*.w2.weight", + target_patterns="block_sparse_moe.experts.down_proj", + operations=[MergeModulelist(dim=0)], + ), + ] + + # Register for the backbone MODEL_TYPE "kimi_linear" ONLY (not the class names). The VLM nests + # KimiLinearForCausalLM (`language_model`) -> KimiLinearModel (`language_model.model`), both model_type + # kimi_linear; registering by model_type applies the fusion ONCE (at the outer `language_model.` scope) + # and the inner submodule is de-duplicated (ancestor already claimed the model_type). Registering the + # CLASS names too would apply it a second time at `language_model.model.` scope -> the same expert keys + # converted twice -> conflict at load. + register_checkpoint_conversion_mapping("kimi_linear", _mapping(), overwrite=True) + print_rank_0( + "[kimi_k3] registered per-expert w1/w3/w2 -> gate_up_proj/down_proj fusion for kimi_linear " + "(custom-code conversion; tf>=5.13 skips these otherwise)" + ) + + +def prepare( + hf_config, + *, + ckpt_path: str | None = None, + trust_remote_code: bool = False, + **_, +): + """Adapt Kimi-K3's remote MoE to the EP/a2a framework and set its EP plan. + + Steps: shim OutputRecorder -> force-load + monkey-patch the remote KimiSparseMoeBlock (fused experts + + ExpertsInterface forward) -> set base_model_ep_plan on text_config -> register the w1/w3/w2 fusion + converters and the decompressor. Returns the same config object, mutated in place. + + The 896 ``mxfp4-pack-quantized`` routed experts are DECOMPRESSED at load: ``_onfly_decompress`` + dequantizes each local expert's packed weight -> bf16 INSIDE the fused-expert conversion, stored as + bf16 fused params, and the compressed-tensors ``quantization_config`` is dropped. ModelOpt then + inserts/calibrates/folds the fp8/nvfp4 quantizers on those bf16 params. That is ~4x the packed HBM + (the memory bottleneck), but PTQ needs the bf16 weights. + """ + # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe + # imports the registry. + from modelopt.torch.utils.moe import register_torch_a2a + + if ckpt_path is None: + raise ValueError( + "kimi_k3: expert parallelism patches the checkpoint's own modeling code, so it needs " + "the checkpoint's path (apply_distributed_support's ckpt_path)." + ) + if not trust_remote_code: + # Patching below imports the checkpoint's own modeling code, and transformers' + # get_class_from_dynamic_module runs it without checking trust_remote_code itself. + raise ValueError( + "kimi_k3: expert parallelism imports the checkpoint's own modeling code; " + "load it with trust_remote_code=True." + ) + + text_cfg = getattr(hf_config, "text_config", None) or hf_config + + # This checkpoint's remote code, not whichever revision this process imported first. Every patch + # is idempotent, so a repeat prepare for the same checkpoint changes nothing. + package = _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code).__module__ + package = package.rpartition(".")[0] + kmod = _find_remote_module(package) + if kmod is None: + raise RuntimeError( + "kimi_k3: could not locate the loaded remote module defining KimiSparseMoeBlock to patch" + ) + _patch_kimi_moe_block(kmod) + _patch_kda_a_log(kmod) + _patch_kda_beta_sigmoid(kmod) + _shim_create_causal_mask(kmod) + _patch_kimi_cache_get_mask_sizes(kmod) + _shim_tie_weights_kwargs(package) + + if not getattr(text_cfg, "base_model_ep_plan", None): + # grouped_gemm shards the bf16 fused params; the experts container goes through moe_tp_experts + # (-> torch_a2a_experts). + text_cfg.base_model_ep_plan = dict(BASE_MODEL_EP_PLAN) + # transformers refuses to build a model whose experts implementation it has not registered, and + # the plan decision builds this one on the meta device before the expert-parallel load runs. + register_torch_a2a() + text_cfg._experts_implementation = "torch_a2a" + + # The 896 routed experts are compressed-tensors `mxfp4-pack-quantized` (4-bit, group 32); everything + # else (attention, shared experts, latent projections, dense layer, embeddings, lm_head, vision) is + # plain bf16. Register the per-expert w1/w3/w2 -> gate_up_proj/down_proj fusion converters (tf 5.13 + # skips these for custom-code models) BEFORE the decompressor patches get_model_conversion_mapping, so + # the decompressor has converters to rewrite into compressed-aware (weight_packed -> bf16) ones. Drop + # quantization_config afterwards so transformers builds no competing HF quantizer. + _register_kimi_k3_conversions() + # Plain bf16 experts come without a quantization_config: nothing to decompress. + if any(getattr(cfg, "quantization_config", None) is not None for cfg in (hf_config, text_cfg)): + _onfly_decompress.install_onfly_decompress(hf_config) + for cfg in (hf_config, text_cfg): + if getattr(cfg, "quantization_config", None) is not None: + cfg.quantization_config = None + return hf_config diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py new file mode 100644 index 00000000000..df584b6b1c3 --- /dev/null +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -0,0 +1,38 @@ +# 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_k3 specs (HF model type ``kimi_k3``). + +Kimi-K3: a remote-code VLM on the kimi_linear backbone, with packed mxfp4 experts that are +decompressed at load. +""" + +from ..specs import EPSpec, ModelSpec, register +from .ep import BASE_MODEL_EP_PLAN, prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="kimi_k3", + modeling_source="remote_code", + ep_spec=EPSpec( + prepare=prepare, + prepared_ep_plan=BASE_MODEL_EP_PLAN, + # prepare's patches to the checkpoint's remote code were validated on this release. + requires_transformers="5.13", + ), + ) +) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 08d40dc9ac4..72e8a41d36c 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -560,8 +560,9 @@ def expert_fn(sorted_tok, offsets): def register_torch_a2a(): """Register everything the torch_a2a path needs into stock transformers (idempotent). - Done lazily from ``load_and_prepare_ep``, not at import time, so importing this module has no - transformers side effect: + Done lazily, not at import time, so importing this module has no transformers side effect: by + ``load_and_prepare_ep``, and by a spec's ``prepare`` that names ``torch_a2a`` on the config a + model is built from before that load (Kimi-K3's). It registers: - the ``torch_a2a`` experts forward into ``ExpertsInterface`` (``ALL_EXPERTS_FUNCTIONS``); - two EP parallel styles into ``ALL_PARALLEL_STYLES``: diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 812b559ec78..93412495ae4 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -50,12 +50,25 @@ enable: false - quantizer_name: 'output.*' enable: false + # Kimi-K3: vLLM builds the latent-expert projections around the routed experts and the + # scalar residual gates unquantized (quant_config=None), so it cannot load a checkpoint + # that quantizes them. The [1, hidden] gates gain nothing from quantization either. + - quantizer_name: '*routed_expert_down_proj*' + enable: false + - quantizer_name: '*routed_expert_up_proj*' + enable: false + - quantizer_name: '*self_attention_res_proj*' + enable: false + - quantizer_name: '*mlp_res_proj*' + enable: false + - quantizer_name: '*output_attn_res_proj*' + enable: false # Multimodal vision branch: keep the vision encoder (SigLIP / ViT) and any # multimodal embedding projection in BF16 by default. Recipes that enable bare # `*weight_quantizer` / `*input_quantizer` or `*mlp*` wildcards otherwise also # match the vision tower (`model.vision_tower.*`, `model.visual.*`, # Llama-4 `vision_model.*`) and the embedding projection (`model.embed_vision.*`, - # Llama-4 `multi_modal_projector.*`); quantizing the vision branch crashes + # Llama-4 `multi_modal_projector.*`, Kimi `mm_projector.*`); quantizing the vision branch crashes # export / produces garbage image embeddings on VL models (gemma-4, # Qwen3.5-VL — NVBugs 6293731, 6293762, 6294017; Llama-4-Scout — NVBugs # 6359097, where `vision_model.patch_embedding.linear` has in_features=588, @@ -71,6 +84,8 @@ enable: false - quantizer_name: '*multi_modal_projector*' enable: false + - quantizer_name: '*mm_projector*' + enable: false - parent_class: 'nn.BatchNorm1d' quantizer_name: '*' enable: false diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 16f5d31a2d6..bab0c0a2477 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -22,9 +22,11 @@ """ import argparse +import copy import pytest import torch +import torch.nn as nn import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe, presets @@ -116,6 +118,48 @@ def test_default_exclusions_disable_mtp_at_any_depth(): assert not model.language_model.mtp.proj.input_quantizer.is_enabled +class _KimiK3Layer(nn.Module): + """Kimi-K3's decoder-layer projections that vLLM builds unquantized, beside one it quantizes.""" + + def __init__(self, hidden=32, latent=16): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.o_proj = nn.Linear(hidden, hidden, bias=False) + self.self_attention_res_proj = nn.Linear(hidden, 1, bias=False) + self.mlp_res_proj = nn.Linear(hidden, 1, bias=False) + self.block_sparse_moe = nn.Module() + self.block_sparse_moe.routed_expert_down_proj = nn.Linear(hidden, latent, bias=False) + self.block_sparse_moe.routed_expert_up_proj = nn.Linear(latent, hidden, bias=False) + + +@pytest.mark.parametrize("qformat", ["fp8", "nvfp4"]) +def test_default_presets_leave_the_kimi_k3_projections_vllm_builds_unquantized(qformat): + # vLLM's Kimi-K3 model passes quant_config=None for these, so it cannot load them quantized. + # Kimi-K2.5 and Kimi-K3 also name their multimodal projector ``mm_projector``; like the other + # projector spellings it stays BF16 unless a recipe re-enables it. + model = nn.Module() + model.layers = nn.ModuleList([_KimiK3Layer()]) + model.output_attn_res_proj = nn.Linear(32, 1, bias=False) + model.mm_projector = nn.Sequential(nn.Linear(32, 32), nn.GELU(), nn.Linear(32, 32)) + config = copy.deepcopy(presets.QUANT_CFG_CHOICES[qformat]) + config["algorithm"] = None # placement only, no calibration forward + mtq.quantize(model, config) + + layer = model.layers[0] + for module in ( + layer.block_sparse_moe.routed_expert_down_proj, + layer.block_sparse_moe.routed_expert_up_proj, + layer.self_attention_res_proj, + layer.mlp_res_proj, + model.output_attn_res_proj, + model.mm_projector[0], + model.mm_projector[2], + ): + assert not module.weight_quantizer.is_enabled + assert not module.input_quantizer.is_enabled + assert layer.self_attn.o_proj.weight_quantizer.is_enabled + + def test_local_hessian_layerwise_is_scoped_to_the_model_recipe(): # Layerwise calibration needs identifiable decoder layers, so it belongs to the model # recipe the published numbers come from -- not to the shared preset, which also backs diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py new file mode 100644 index 00000000000..7a025947ba3 --- /dev/null +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -0,0 +1,592 @@ +# 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-K3 (``kimi_k3``), whose kimi_linear backbone is remote code. + +The checkpoint's modeling is not available here, so the tests build a stand-in remote module -- the +classes and functions the adapter patches, with the per-expert ``nn.ModuleList`` MoE the remote code +ships -- and check what the adapter does to it on tiny CPU tensors. Every process-wide patch is +undone after each test. +""" + +import copy +import sys +import types +from concurrent.futures import Future +from types import SimpleNamespace + +import pytest +import torch +from _test_utils.torch.kimi_k3_remote import remote_module +from torch import nn + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.8"): + pytest.skip("Kimi-K3 expert parallelism needs transformers >= 5.8", allow_module_level=True) + +import transformers.conversion_mapping as cm +import transformers.core_model_loading as cml +import transformers.dynamic_module_utils as dmu +import transformers.modeling_utils as mu +import transformers.utils.generic as hf_generic +from transformers.core_model_loading import WeightConverter, build_glob_alternation +from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS + +from modelopt.torch.models import _onfly_decompress, get_spec +from modelopt.torch.models.kimi_k3 import ep as kimi_k3 +from modelopt.torch.utils import moe +from modelopt.torch.utils.moe import check_distributed_support + +ROUTES = torch.tensor([[0, 3], [1, 2], [3, 1], [2, 0], [1, 3]]) # top-2 experts of 5 tokens + + +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) + + +@pytest.fixture +def private_conversion_registry(monkeypatch): + """Let a test register conversion mappings without leaking them into transformers' registry.""" + monkeypatch.setattr(cm, "_checkpoint_conversion_mapping_cache", None) # rebuilt on first use + if hasattr(cm, "USER_REGISTERED_MAPPINGS"): # transformers >= 5.13 + monkeypatch.setattr(cm, "USER_REGISTERED_MAPPINGS", set(cm.USER_REGISTERED_MAPPINGS)) + + +@pytest.fixture +def restore_loader_hooks(monkeypatch): + """Undo the process-wide transformers hooks the on-the-fly decompressor installs.""" + for owner, name in ( + (WeightConverter, "add_tensor"), + (WeightConverter, "materialize_tensors"), + (cml, "_materialize_copy"), + (mu, "get_model_conversion_mapping"), + ): + monkeypatch.setattr(owner, name, getattr(owner, name)) + + +def _config(**overrides): + """A tiny kimi_linear MoE config: 4 routed experts, top-2, run in a 6-wide latent space.""" + fields = { + "num_experts": 4, + "num_experts_per_token": 2, + "hidden_size": 8, + "routed_expert_hidden_size": 6, + "moe_intermediate_size": 5, + "hidden_act": "silu", + "situ_beta": 1.5, + "situ_linear_beta": 0.25, + "latent_moe_use_norm": True, + "num_shared_experts": 1, + "_experts_implementation": "eager", + } + return SimpleNamespace(**{**fields, **overrides}) + + +def _randomize(module, seed=0): + gen = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for param in module.parameters(): + param.copy_(torch.randn(param.shape, generator=gen)) + + +def _load(converters, checkpoint): + """Route, collect and fuse ``checkpoint`` the way transformers' loader does.""" + per_target = {} + for key, tensor in checkpoint.items(): + for conv in converters: + target, pattern = conv.rename_source_key(key) + if pattern is not None: + break + else: + raise AssertionError(f"no converter claims {key}") + future = Future() + future.set_result(tensor) + per_target.setdefault(target, copy.deepcopy(conv)).add_tensor(target, key, pattern, future) + fused = {} + for target, conv in per_target.items(): + fused.update(conv.convert(target)) + return fused + + +# --------------------------------------------------------------------------- the fused experts + + +def test_fused_experts_dispatch_through_the_experts_interface(monkeypatch): + """``config._experts_implementation`` picks the forward (``torch_a2a`` under EP), so the module + must route through transformers' ``ExpertsInterface``; and a stock implementation reading the + fused params agrees with the eager loop, so the declared layout ([E, out, in], gated) is right.""" + routed = [] + batched_mm = ALL_EXPERTS_FUNCTIONS["batched_mm"] + + def probe(module, *args): + routed.append(module) + return batched_mm(module, *args) + + monkeypatch.setitem(ALL_EXPERTS_FUNCTIONS, "modelopt_test_probe", probe) + config = _config(hidden_act="situ") + experts = kimi_k3._build_fused_experts_class(remote_module())( + config, config.routed_expert_hidden_size + ) + _randomize(experts) + gen = torch.Generator().manual_seed(1) + tokens = torch.randn(len(ROUTES), config.routed_expert_hidden_size, generator=gen) + route_weights = torch.rand(ROUTES.shape, generator=gen) + + with torch.no_grad(): + eager = experts(tokens, ROUTES, route_weights) + config._experts_implementation = "modelopt_test_probe" + dispatched = experts(tokens, ROUTES, route_weights) + + assert routed == [experts] + torch.testing.assert_close(dispatched, eager) + + +@pytest.mark.parametrize( + "overrides", + [ + {}, + { + "latent_moe_use_norm": False, + "routed_expert_hidden_size": None, + "num_shared_experts": None, + }, + {"hidden_act": "situ"}, + ], + ids=["latent-norm-shared", "plain", "situ"], +) +def test_the_patched_block_reproduces_the_remote_per_expert_forward(overrides): + """Swapping the per-expert ModuleList for the fused module keeps the block's output: the latent + projections, the latent norm and the shared experts stay around the fused call, and w1 / w3 / + w2 are the gate rows, the up rows and the down of the fused params. A ``situ`` model hands the + full [gate; up] row to the remote ``SituAndMul``.""" + remote = remote_module() + config = _config(**overrides) + reference = remote.KimiSparseMoeBlock(config) + _randomize(reference) + gen = torch.Generator().manual_seed(1) + hidden_states = torch.randn(2, 3, config.hidden_size, generator=gen) + with torch.no_grad(): + expected = reference(hidden_states) + + kimi_k3._patch_kimi_moe_block(remote) + block = remote.KimiSparseMoeBlock(config) + assert not isinstance(block.experts, nn.ModuleList) + dense = {k: v for k, v in reference.state_dict().items() if not k.startswith("experts.")} + result = block.load_state_dict(dense, strict=False) + assert sorted(result.missing_keys) == ["experts.down_proj", "experts.gate_up_proj"] + assert not result.unexpected_keys + with torch.no_grad(): + gate_up = [torch.cat([m.w1.weight, m.w3.weight]) for m in reference.experts] + block.experts.gate_up_proj.copy_(torch.stack(gate_up)) + block.experts.down_proj.copy_(torch.stack([m.w2.weight for m in reference.experts])) + torch.testing.assert_close(block(hidden_states), expected) + + +# ------------------------------------------------------------------- weight conversion and plan + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_registered_conversions_fuse_w1_w3_into_gate_up_and_stack_w2(): + """transformers builds no weight conversions for remote code unless they are registered. The + registered pair fuses the checkpoint's per-expert w1 (gate) and w3 (up) into + ``gate_up_proj = [w1; w3]`` and stacks w2 into ``down_proj``, under the VLM's prefix.""" + kimi_k3._register_kimi_k3_conversions() + converters = cm.get_checkpoint_conversion_mapping("kimi_linear") + n_experts, inter, hidden = 3, 2, 4 + shapes = {"w1": (inter, hidden), "w3": (inter, hidden), "w2": (hidden, inter)} + gen = torch.Generator().manual_seed(0) + weights = { + (name, e): torch.randn(shape, generator=gen) + for e in range(n_experts) + for name, shape in shapes.items() + } + prefix = "language_model.model.layers.1.block_sparse_moe.experts." + + fused = _load(converters, {f"{prefix}{e}.{name}.weight": w for (name, e), w in weights.items()}) + + gate_up = [torch.cat([weights["w1", e], weights["w3", e]]) for e in range(n_experts)] + assert torch.equal(fused[prefix + "gate_up_proj"], torch.stack(gate_up)) + down = [weights["w2", e] for e in range(n_experts)] + assert torch.equal(fused[prefix + "down_proj"], torch.stack(down)) + + +def test_the_registration_raises_without_the_conversion_api(monkeypatch): + """Unregistered, the checkpoint's experts would load as MISSING: newly initialized, at full + size on every rank. So a missing conversion API stops the load instead of warning.""" + monkeypatch.setitem(sys.modules, "transformers.conversion_mapping", None) + + with pytest.raises(ImportError): + kimi_k3._register_kimi_k3_conversions() + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_the_plan_the_conversions_and_the_fused_params_name_the_same_tensors(): + """Three separately written names must agree, or EP silently shards nothing: the fused module's + parameters, the conversion targets that fill them, and the EP plan globs (matched the way + transformers matches them). The router, latent projections and shared experts stay replicated.""" + remote = remote_module() + kimi_k3._patch_kimi_moe_block(remote) + layer = nn.Module() + layer.block_sparse_moe = remote.KimiSparseMoeBlock(_config()) + model = nn.Module() + model.layers = nn.ModuleList([layer]) + kimi_k3._register_kimi_k3_conversions() + converters = cm.get_checkpoint_conversion_mapping("kimi_linear") + plan = kimi_k3.BASE_MODEL_EP_PLAN + matcher, glob_of_group, _ = build_glob_alternation(list(plan)) + + def style(name): + match = matcher.search(name) + return plan[glob_of_group[match.lastgroup]] if match else None + + sharded = {name for name, _ in model.named_parameters() if style(name) is not None} + + assert sharded == { + "layers.0.block_sparse_moe.experts.gate_up_proj", + "layers.0.block_sparse_moe.experts.down_proj", + } + assert {style(name) for name in sharded} == {"grouped_gemm"} + assert {name.removeprefix("layers.0.") for name in sharded} == { + target for conv in converters for target in conv.target_patterns + } + assert style("layers.0.block_sparse_moe.experts") == "moe_tp_experts" + + +# ------------------------------------------------------------------------------------- prepare + +_AUTO_MAP = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} + + +@pytest.fixture +def remote_checkpoint(monkeypatch): + """The checkpoint's remote code, a fresh stand-in that importing it registers; the shims + prepare installs on transformers are undone after the test.""" + remote = remote_module("transformers_modules.stub.modeling_kimi_k3") + + def import_remote_class(class_ref, ckpt_path, **kwargs): + monkeypatch.setitem(sys.modules, remote.__name__, remote) + return remote.KimiK3ForConditionalGeneration + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + _restore(monkeypatch, hf_generic, "OutputRecorder") + return remote + + +def _prepare(config): + return kimi_k3.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) + + +@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") +def test_prepare_sets_the_plan_and_a2a_experts_and_hands_the_quant_config_to_the_decompressor( + monkeypatch, +): + """The plan goes on the text config (its globs are written against the language model) and is + the one the spec states for a run that skips prepare, the experts forward is ``torch_a2a``, the + kimi_linear fusion is registered, and the decompressor sees the compressed-tensors config before + it is dropped (so transformers builds no competing quantizer).""" + seen = [] + monkeypatch.setattr( + _onfly_decompress, + "install_onfly_decompress", + lambda cfg: seen.append((cfg.quantization_config, cfg.text_config.quantization_config)), + ) + quant = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} + text = SimpleNamespace(quantization_config=quant.copy()) + config = SimpleNamespace(text_config=text, quantization_config=quant.copy(), auto_map=_AUTO_MAP) + + assert _prepare(config) is config + + assert text.base_model_ep_plan == kimi_k3.BASE_MODEL_EP_PLAN + assert text.base_model_ep_plan is not kimi_k3.BASE_MODEL_EP_PLAN + assert text.base_model_ep_plan == get_spec("kimi_k3").ep_spec.prepared_ep_plan + assert not hasattr(config, "base_model_ep_plan") + assert text._experts_implementation == "torch_a2a" + registered = cm.get_checkpoint_conversion_mapping("kimi_linear") + assert any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) + assert seen == [(quant, quant)] + assert config.quantization_config is None + assert text.quantization_config is None + + +@pytest.fixture +def torch_a2a_unregistered(monkeypatch): + """The process as the plan decision finds it, before anything registered torch_a2a; what the + test registers is dropped after it.""" + experts = type(ALL_EXPERTS_FUNCTIONS) + registered = {k: v for k, v in experts._global_mapping.items() if k != "torch_a2a"} + monkeypatch.setattr(experts, "_global_mapping", registered) + monkeypatch.setattr(moe, "_TORCH_A2A_REGISTERED", False) + + +@pytest.mark.usefixtures( + "private_conversion_registry", "remote_checkpoint", "torch_a2a_unregistered" +) +def test_a_model_builds_from_the_prepared_config_before_the_expert_parallel_load(): + """The plan decision builds the model from prepare's config on the meta device before the + expert-parallel load registers torch_a2a, and transformers refuses to build a model whose experts + implementation it has not registered: prepare registers the one it names.""" + text = SimpleNamespace() + _prepare(SimpleNamespace(text_config=text, auto_map=_AUTO_MAP)) + + # Kimi-K3's own model is the checkpoint's code; every transformers model runs the same check. + config = transformers.Qwen3MoeConfig( + hidden_size=8, + intermediate_size=8, + moe_intermediate_size=4, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + num_experts=4, + num_experts_per_tok=2, + vocab_size=16, + ) + config._experts_implementation = text._experts_implementation + with torch.device("meta"): + transformers.Qwen3MoeForCausalLM(config) + + +@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") +def test_prepare_keeps_a_plan_the_config_already_has(monkeypatch): + """A plan the checkpoint's config already carries describes its own layout and wins.""" + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + text = SimpleNamespace(base_model_ep_plan={"layers.*.custom": "grouped_gemm"}) + + _prepare(SimpleNamespace(text_config=text, auto_map=_AUTO_MAP)) + + assert text.base_model_ep_plan == {"layers.*.custom": "grouped_gemm"} + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_prepare_with_a_checkpoint_patches_the_loaded_remote_module_once(monkeypatch): + """Given a checkpoint, prepare imports the remote modeling through the class the config's + ``auto_map`` names, then patches the module that *defines* the MoE block (found by class, not + by file name, and not a VLM wrapper that merely imports it) -- a repeat prepare patches nothing + twice -- and leaves the remote-code staging shim installed.""" + remote = remote_module() + wrapper = types.ModuleType("kimi_k3_remote_stub_vlm") # imports the block, defines nothing + wrapper.KimiSparseMoeBlock = remote.KimiSparseMoeBlock + loads = [] + + def import_remote_class(class_ref, ckpt_path, **kwargs): + loads.append((class_ref, ckpt_path, kwargs)) + for module in (wrapper, remote): # what importing the remote code does + monkeypatch.setitem(sys.modules, module.__name__, module) + return remote.KimiK3ForConditionalGeneration + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + staging = dmu.get_cached_module_file + monkeypatch.setattr(dmu, "get_cached_module_file", staging) + _restore(monkeypatch, hf_generic, "OutputRecorder") + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + auto_map = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} + + for _ in range(2): + config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=auto_map) + kimi_k3.prepare(config, ckpt_path="/ckpt", trust_remote_code=True) + + assert loads == [(auto_map["AutoModelForCausalLM"], "/ckpt", {"trust_remote_code": True})] * 2 + assert dmu.get_cached_module_file is not staging + assert hasattr(hf_generic, "OutputRecorder") # the remote code imports it from there + # every incompatibility is patched on the module the checkpoint's code defined + assert not isinstance(remote.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + a_log = remote.KimiDeltaAttention(num_heads=3, head_dim=4).A_log + assert (a_log.shape, a_log.dtype) == ( + (4,), + torch.bfloat16, + ) # the checkpoint's shape, same dtype + remote.chunk_kda(1, 2, 3, 4, beta=torch.zeros(1), use_beta_sigmoid_in_kernel=True) + _, beta, flag = remote.calls.pop() + assert (beta.item(), flag) == (0.5, False) + remote.create_causal_mask( + config=None, input_embeds=None, attention_mask=None, past_key_values=None, cache_position=0 + ) + assert remote.KimiDynamicCache(seen_tokens=1).get_mask_sizes(2, 0) == (3, 0) + vlm = remote.KimiK3ForConditionalGeneration() + assert vlm.tie_weights(recompute_mapping=False) == "tied" + assert [call for call in remote.calls if call[0] == "tie_weights"] == [("tie_weights", vlm)] + + +def test_prepare_fails_loudly_when_no_loaded_module_defines_the_moe_block(monkeypatch): + """Skipping the patch silently would load the per-expert layout that EP cannot shard.""" + + class KimiK3ForConditionalGeneration: + __module__ = "transformers_modules.no_moe_block.modeling_kimi_k3" + + monkeypatch.setattr( + kimi_k3, "_force_load_remote_modeling", lambda *args: KimiK3ForConditionalGeneration + ) + + with pytest.raises(RuntimeError, match="KimiSparseMoeBlock"): + kimi_k3.prepare( + SimpleNamespace(text_config=SimpleNamespace()), + ckpt_path="/ckpt", + trust_remote_code=True, + ) + + +@pytest.mark.usefixtures("private_conversion_registry") +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + # Without the checkpoint there is no modeling code to patch: the plan and the fusion would + # name fused experts the unpatched per-expert modeling never builds. + ({"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. + ({"ckpt_path": "/ckpt"}, "trust_remote_code=True"), + ], + ids=["no-checkpoint", "untrusted-checkpoint"], +) +def test_prepare_refuses_before_changing_anything(monkeypatch, kwargs, match): + monkeypatch.setattr( + kimi_k3, + "_force_load_remote_modeling", + lambda *args: pytest.fail("the checkpoint's code was imported"), + ) + text = SimpleNamespace() + + with pytest.raises(ValueError, match=match): + kimi_k3.prepare(SimpleNamespace(text_config=text), **kwargs) + assert vars(text) == {} + registered = cm.get_checkpoint_conversion_mapping("kimi_linear") or [] + assert not any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) + + +def test_a_transformers_older_than_the_validated_one_is_refused_before_prepare(monkeypatch): + """The spec states the release the adapter's patches were validated against, so a run on an + older one is refused, with its way around, before prepare touches anything.""" + monkeypatch.setattr(transformers, "__version__", "5.12.0") + + with pytest.raises(RuntimeError, match=r"transformers >= 5\.13"): + check_distributed_support("kimi_k3") + + +@pytest.mark.usefixtures("private_conversion_registry") +def test_each_checkpoint_revision_patches_its_own_remote_module(monkeypatch): + """Two revisions of the remote code import as two packages, each defining its own MoE block. Each + checkpoint's prepare patches its own revision's module, whichever was imported first, and a + repeat for the first checkpoint replaces nothing that is already patched.""" + revisions = { + f"/ckpt/{rev}": remote_module(f"transformers_modules.{rev}.modeling_kimi") + for rev in ("rev_a", "rev_b") + } + + def import_remote_class(class_ref, ckpt_path, **kwargs): + remote = revisions[ckpt_path] + monkeypatch.setitem(sys.modules, remote.__name__, remote) + return remote.KimiK3ForConditionalGeneration + + monkeypatch.setattr(dmu, "get_class_from_dynamic_module", import_remote_class) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + _restore(monkeypatch, hf_generic, "OutputRecorder") + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", lambda cfg: None) + auto_map = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} + + def prepare(ckpt_path): + config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=auto_map) + kimi_k3.prepare(config, ckpt_path=ckpt_path, trust_remote_code=True) + + def patched_surface(remote): + return ( + remote.KimiSparseMoeBlock.__init__, + remote.KimiSparseMoeBlock.forward, + remote.KimiDeltaAttention.__init__, + remote.chunk_kda, + remote.fused_recurrent_kda, + remote.create_causal_mask, + remote.KimiDynamicCache.get_mask_sizes, + remote.KimiK3ForConditionalGeneration.tie_weights, + ) + + rev_a, rev_b = revisions["/ckpt/rev_a"], revisions["/ckpt/rev_b"] + shipped_b = patched_surface(rev_b) + prepare("/ckpt/rev_a") + assert patched_surface(rev_b) == shipped_b + prepare("/ckpt/rev_b") + patched = patched_surface(rev_a), patched_surface(rev_b) + prepare("/ckpt/rev_a") + + assert (patched_surface(rev_a), patched_surface(rev_b)) == patched + assert all(new is not old for new, old in zip(patched_surface(rev_b), shipped_b)) + for remote in (rev_a, rev_b): + assert not isinstance(remote.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + assert ( + remote.KimiK3ForConditionalGeneration().tie_weights(recompute_mapping=False) == "tied" + ) + + +def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares(monkeypatch): + """``AutoModelForCausalLM`` is preferred, then ``AutoModel``; a config declaring neither is an + error rather than a load that silently skips the patches.""" + loads = [] + monkeypatch.setattr( + dmu, "get_class_from_dynamic_module", lambda ref, path, **kwargs: loads.append(ref) + ) + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + _restore(monkeypatch, hf_generic, "OutputRecorder") + + for auto_map in ( + {"AutoModel": "remote.Base", "AutoModelForCausalLM": "remote.ForCausalLM"}, + {"AutoModel": "remote.Base"}, + ): + kimi_k3._force_load_remote_modeling(SimpleNamespace(auto_map=auto_map), "/ckpt", False) + + assert loads == ["remote.ForCausalLM", "remote.Base"] + config_only = SimpleNamespace(auto_map={"AutoConfig": "remote.Config"}) + with pytest.raises(RuntimeError, match="auto_map"): + kimi_k3._force_load_remote_modeling(config_only, "/ckpt", False) + + +@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") +def test_prepare_lets_an_error_reading_the_quant_config_through(monkeypatch): + """compressed_tensors raises pydantic's ValidationError, a ValueError, for a quantization_config + it cannot parse. Dropped after that, the config would load the experts newly initialized.""" + + def install(cfg): + raise ValueError("1 validation error for QuantizationConfig") + + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", install) + config = SimpleNamespace( + text_config=SimpleNamespace(), quantization_config={"format": "new"}, auto_map=_AUTO_MAP + ) + + with pytest.raises(ValueError, match="validation error"): + _prepare(config) + + +@pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks", "remote_checkpoint") +def test_prepare_leaves_the_loader_alone_for_bf16_experts(): + """Without a compressed-tensors config the experts are plain bf16: prepare completes and leaves + transformers' conversion mapping unwrapped, with or without ``compressed_tensors`` installed.""" + stock = mu.get_model_conversion_mapping + + _prepare(SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP)) + + assert mu.get_model_conversion_mapping is stock + + +def test_the_kimi_k3_spec_routes_expert_parallelism_through_this_adapter(): + """kimi_k3 is remote code, so its whole EP setup is this adapter's ``prepare``: another Kimi + adapter's would load it without its own patches.""" + assert get_spec("kimi_k3").ep_spec.prepare is kimi_k3.prepare diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 9e586769402..650203b9b99 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -911,6 +911,43 @@ def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead( assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp") +def _remote_kimi_k3(): + """A remote-code Kimi VLM as the checkpoint's own code builds it, unpatched by its spec's + ``prepare``: each MoE layer of its language model holds its routed experts as per-expert + modules in an ``nn.ModuleList``, at ``language_model.model.layers...experts``.""" + + def layer(): + block = nn.Module() + block.gate = nn.Linear(2, 4, bias=False) + block.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(4)) + return nn.ModuleDict({"block_sparse_moe": block}) + + model = nn.Module() + model.language_model = nn.Module() + model.language_model.model = nn.Module() + model.language_model.model.layers = nn.ModuleList(layer() for _ in range(2)) + model.tie_weights = lambda: None # a PreTrainedModel's; nothing here is tied + return model + + +def test_a_remote_kimi_k3_whose_prepare_attaches_its_plan_left_whole_is_an_error( + run_as, serve, monkeypatch, tmp_path +): + """Remote-code Kimi checkpoints get their plan from their spec's ``prepare``, which an EP degree + of 1 skips: it patches their code. The plan their spec states still finds the experts their own + code builds, so whole on every rank they are refused like Qwen3-MoE.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + monkeypatch.setattr( + get_spec("kimi_k3").ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") + ) + monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi_k3()) + run_as(4) + serve(SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace(n_routed_experts=8))) + + with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + resolve_distribution_plan(str(tmp_path)) + + @needs_qwen3_moe_ep_plan def test_an_ep_capable_moe_is_replicated_when_both_knobs_ask_for_it( run_as, serve, monkeypatch, tmp_path From 8ccca6f8592b7a77eea37b348a8ce9ffa469b241 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:29:10 +0000 Subject: [PATCH 02/32] Give the Kimi-K3 stand-in back the MoE block the adapter's tests build The remote-code fixes' review trimmed the shared stand-in to what those fixes patch. The adapter's tests also build the checkpoint's per-expert MoE block, its gate and situ activation, so they come back here. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/kimi_k3_remote.py | 101 +++++++++++++++++++++- 1 file changed, 98 insertions(+), 3 deletions(-) diff --git a/tests/_test_utils/torch/kimi_k3_remote.py b/tests/_test_utils/torch/kimi_k3_remote.py index 8c2128231f2..75261b5a778 100644 --- a/tests/_test_utils/torch/kimi_k3_remote.py +++ b/tests/_test_utils/torch/kimi_k3_remote.py @@ -13,21 +13,111 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""A stand-in for Kimi-K3's own remote modeling code: the classes and functions the fixes for it -patch, for tests on tiny CPU tensors. +"""A stand-in for Kimi-K3's own remote modeling code: the classes and functions its adapter patches, +with the per-expert ``nn.ModuleList`` MoE the checkpoint's code ships, for tests on tiny CPU tensors. """ import types import torch +import torch.nn.functional as F from torch import nn +def situ(gate_up, beta, linear_beta): + """The stand-in ``situ`` gated activation over a full [gate; up] row.""" + gate, up = gate_up.chunk(2, dim=-1) + return gate * torch.sigmoid(beta * gate) * (up + linear_beta) + + def remote_module(name="kimi_k3_remote_stub"): - """A stand-in for the remote modeling namespace: what the fixes patch, as shipped.""" + """A stand-in for the remote modeling namespace: everything the adapter patches, as shipped.""" remote = types.ModuleType(name) remote.calls = [] + def _get_situ_activation_params(config): + return config.situ_beta, config.situ_linear_beta + + class SituAndMul(nn.Module): + def __init__(self, beta, linear_beta): + super().__init__() + self.beta, self.linear_beta = beta, linear_beta + + def forward(self, gate_up): + return situ(gate_up, self.beta, self.linear_beta) + + class KimiBlockSparseMLP(nn.Module): + """One expert as the checkpoint stores it: w1 = gate, w3 = up, w2 = down.""" + + def __init__(self, config, hidden): + super().__init__() + inter = config.moe_intermediate_size + self.w1 = nn.Linear(hidden, inter, bias=False) + self.w3 = nn.Linear(hidden, inter, bias=False) + self.w2 = nn.Linear(inter, hidden, bias=False) + situ = config.hidden_act == "situ" + self.situ_params = _get_situ_activation_params(config) if situ else None + + def forward(self, x): + gate, up = self.w1(x), self.w3(x) + if self.situ_params is None: + return self.w2(F.silu(gate) * up) + return self.w2(situ(torch.cat([gate, up], dim=-1), *self.situ_params)) + + class KimiMoEGate(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_token + self.weight = nn.Parameter(torch.randn(config.num_experts, config.hidden_size)) + + def forward(self, hidden_states): + flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + weights, idx = torch.topk(torch.sigmoid(flat @ self.weight.t()), self.top_k, dim=-1) + return idx, weights + + class KimiSparseMoeBlock(nn.Module): + """Latent projections around a per-expert Python loop, plus always-on shared experts.""" + + def __init__(self, config): + super().__init__() + self.config = config + # Kimi-K3's config has no use_latent_moe field: the block derives it, as vLLM does. + self.use_latent_moe = config.routed_expert_hidden_size is not None + self.latent_moe_use_norm = config.latent_moe_use_norm + hidden = config.hidden_size + width = config.routed_expert_hidden_size if self.use_latent_moe else hidden + self.experts = nn.ModuleList( + KimiBlockSparseMLP(config, width) for _ in range(config.num_experts) + ) + self.gate = KimiMoEGate(config) + if self.use_latent_moe: + self.routed_expert_down_proj = nn.Linear(hidden, width, bias=False) + self.routed_expert_up_proj = nn.Linear(width, hidden, bias=False) + if config.latent_moe_use_norm: + self.routed_expert_norm = nn.LayerNorm(width) + if config.num_shared_experts is not None: + self.shared_experts = KimiBlockSparseMLP(config, hidden) + + def forward(self, hidden_states): + idx, weights = self.gate(hidden_states) + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + if self.use_latent_moe: + x = self.routed_expert_down_proj(x) + y = torch.stack( + [ + sum(w * self.experts[int(e)](token) for e, w in zip(top_e, top_w)) + for token, top_e, top_w in zip(x, idx, weights) + ] + ) + if self.use_latent_moe: + if self.latent_moe_use_norm: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + y = y.reshape(hidden_states.shape) + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(hidden_states) + return y + class KimiDeltaAttention(nn.Module): def __init__(self, num_heads, head_dim): super().__init__() @@ -70,6 +160,11 @@ def tie_weights(self): # the remote override, with the old no-kwarg signature return "tied" for obj in ( + _get_situ_activation_params, + SituAndMul, + KimiBlockSparseMLP, + KimiMoEGate, + KimiSparseMoeBlock, KimiDeltaAttention, chunk_kda, fused_recurrent_kda, From f06e8102dce265e425e840ec8d029334fd1cfbfd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:10:57 +0000 Subject: [PATCH 03/32] Spell the Kimi-K3 refusal fixture's expert count as Kimi-K3's config does Kimi-K3's text config names it num_experts; n_routed_experts is DeepSeek-V3's spelling, left from the Kimi-K2.5 copy of this test. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 650203b9b99..b88e8fe2b76 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -942,7 +942,7 @@ def test_a_remote_kimi_k3_whose_prepare_attaches_its_plan_left_whole_is_an_error ) monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi_k3()) run_as(4) - serve(SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace(n_routed_experts=8))) + serve(SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace(num_experts=8))) with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): resolve_distribution_plan(str(tmp_path)) From 83a8113ce0e6f360648b36ca4b466b1ff422d499 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:00:07 +0000 Subject: [PATCH 04/32] Keep the own forward for a Kimi-K3 MoE block built before the patch _patch_kimi_moe_block replaces KimiSparseMoeBlock's forward on the class, but swaps the per-expert ModuleList for the fused experts only in the patched __init__. A block built before the patch, as a model loaded earlier in the process holds, kept its ModuleList and then called it as the fused experts, which a ModuleList cannot run (NotImplementedError, no forward). The patched forward now runs the block's own forward for a block that still holds its per-expert experts. The patched-block test calls its reference block, built before the patch, again after patching. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 7 ++++++- tests/unit/torch/models/test_kimi_k3_ep.py | 5 ++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index f0adc73f1d5..a5a3815c6b4 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -182,7 +182,8 @@ def _patch_kimi_moe_block(kmod) -> None: Swap ``KimiSparseMoeBlock.experts`` (per-expert ModuleList) for the fused experts module and rewrite its ``forward`` to dispatch through it (keeping the latent down/up projections, optional norm, and - shared experts around the fused call). Idempotent. + shared experts around the fused call). A block built before the patch keeps its per-expert experts, + and so its own forward. Idempotent. """ from modelopt.torch.utils import print_rank_0 @@ -192,6 +193,7 @@ def _patch_kimi_moe_block(kmod) -> None: FusedExperts = _build_fused_experts_class(kmod) # noqa: N806 (a class, keeping its own name) _orig_init = Block.__init__ + _orig_forward = Block.forward def _init(self, config): _orig_init(self, config) @@ -203,6 +205,9 @@ def _init(self, config): self.experts = FusedExperts(config, width) def _forward(self, hidden_states): + if isinstance(self.experts, nn.ModuleList): + # Built before the patch: the per-expert experts only the block's own forward runs. + return _orig_forward(self, hidden_states) identity = hidden_states orig_shape = hidden_states.shape topk_idx, topk_weight = self.gate(hidden_states) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 7a025947ba3..f0d033ec996 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -178,7 +178,8 @@ def test_the_patched_block_reproduces_the_remote_per_expert_forward(overrides): """Swapping the per-expert ModuleList for the fused module keeps the block's output: the latent projections, the latent norm and the shared experts stay around the fused call, and w1 / w3 / w2 are the gate rows, the up rows and the down of the fused params. A ``situ`` model hands the - full [gate; up] row to the remote ``SituAndMul``.""" + full [gate; up] row to the remote ``SituAndMul``. A block built before the patch, as a model + loaded earlier in the process holds, keeps its per-expert experts and its own forward.""" remote = remote_module() config = _config(**overrides) reference = remote.KimiSparseMoeBlock(config) @@ -189,6 +190,8 @@ def test_the_patched_block_reproduces_the_remote_per_expert_forward(overrides): expected = reference(hidden_states) kimi_k3._patch_kimi_moe_block(remote) + with torch.no_grad(): + torch.testing.assert_close(reference(hidden_states), expected) block = remote.KimiSparseMoeBlock(config) assert not isinstance(block.experts, nn.ModuleList) dense = {k: v for k, v in reference.state_dict().items() if not k.startswith("experts.")} From 24cdc3c747e74e732442027939fcccbff802d526 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:01:54 +0000 Subject: [PATCH 05/32] Import print_rank_0 at module scope in the Kimi-K3 adapter Two functions imported it locally, with no import cycle or optional dependency to defer: the adapter already imports _remote_fixes at module scope, which imports print_rank_0 there. The transformers imports and register_torch_a2a stay local, for the reasons their comments give. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index a5a3815c6b4..4e68b9b74b9 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -70,6 +70,8 @@ import torch import torch.nn as nn +from modelopt.torch.utils import print_rank_0 + from .. import _onfly_decompress from .._remote_code import stage_full_remote_closure from ._remote_fixes import ( @@ -185,8 +187,6 @@ def _patch_kimi_moe_block(kmod) -> None: shared experts around the fused call). A block built before the patch keeps its per-expert experts, and so its own forward. Idempotent. """ - from modelopt.torch.utils import print_rank_0 - Block = kmod.KimiSparseMoeBlock # noqa: N806 (a class, keeping its own name) if getattr(Block, "_kimi_k3_ep_patched", False): return @@ -275,8 +275,6 @@ def _register_kimi_k3_conversions() -> None: from transformers.conversion_mapping import register_checkpoint_conversion_mapping from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter - from modelopt.torch.utils import print_rank_0 - def _mapping(): return [ WeightConverter( From d7f0fd9812048db2be3a27e3ca33d6748d6031de Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:10:30 +0000 Subject: [PATCH 06/32] Attribute the Kimi-K3 fused experts to transformers' MiniMaxM3VLExperts KimiK3Experts' fused parameter layout and its eager forward loop are adapted from transformers' MiniMaxM3VLExperts (v5.18.0, commit a906d3c), which is Apache-2.0. Following CONTRIBUTING.md's guidance on copied code: the file now carries the source link and the original copyright and license notice above the NVIDIA header, the insert-license hook leaves the file alone, and LICENSE lists the copyright holder among its third-party notices. Signed-off-by: Shengliang Xu --- .pre-commit-config.yaml | 1 + LICENSE | 1 + modelopt/torch/models/kimi_k3/ep.py | 21 +++++++++++++++++++-- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cbd1ddd9e38..b6f4ecc2174 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -114,6 +114,7 @@ repos: modelopt/torch/fastgen/plugins/qwen_image_pdd.py| modelopt/torch/kernels/quantization/linear_attention/fla_chunk_delta_h.py| modelopt/torch/kernels/quantization/linear_attention/fla_chunk_gated_delta_rule.py| + modelopt/torch/models/kimi_k3/ep.py| modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py| modelopt/torch/quantization/export_onnx.py| modelopt/torch/quantization/plugins/attention.py| diff --git a/LICENSE b/LICENSE index abc6a0cb761..634ffd50bf3 100644 --- a/LICENSE +++ b/LICENSE @@ -224,6 +224,7 @@ the following copyright holders, licensed under the Apache License, Version 2.0 Copyright (c) 2024 Heming Xia Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved. Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team + Copyright 2026 the MiniMax AI Team and HuggingFace Team. All rights reserved. Copyright (c) OpenMMLab. All rights reserved. Copyright (c) 2022 megvii-model. All Rights Reserved. diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 4e68b9b74b9..e908ce1067b 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -1,3 +1,20 @@ +# Adapted from https://github.com/huggingface/transformers/blob/a906d3c/src/transformers/models/minimax_m3_vl/modeling_minimax_m3_vl.py +# (``MiniMaxM3VLExperts``: the fused expert parameters and the eager ``forward`` loop of ``KimiK3Experts``), +# modified for Kimi-K3's ``situ`` gate and latent expert width. +# Copyright 2026 the MiniMax AI Team and HuggingFace Team. All rights reserved. +# +# 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. + # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -151,8 +168,8 @@ def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: def forward(self, hidden_states, top_k_index, top_k_weights): # Fallback (non-EP / eager) dense per-expert loop; the EP-PTQ path replaces this with the - # torch_a2a experts forward via config._experts_implementation. Mirrors the native - # DeepseekV3Experts/MiniMaxM3VLExperts reference loop, with situ gating. + # torch_a2a experts forward via config._experts_implementation. Adapted from transformers' + # MiniMaxM3VLExperts.forward (see the header), with situ gating. final = torch.zeros_like(hidden_states) with torch.no_grad(): # One class past the experts: the id a router gives a token routed to no local expert. From 0654eabfa3d711e73e9ac73dceadabc73c45dc48 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:14:13 +0000 Subject: [PATCH 07/32] Pin that Kimi-K3's packed experts export under the checkpoint's names The decompressor rewrites the registered w1/w3/w2 converters to read the packed experts, and the model keeps the rewritten converters as the conversions its load ran; the export names the experts back through them. The test rewrites the registered converters as the decompressor does for Kimi-K3's stored names and maps ModelOpt's per-expert names through the export's reverse name mapper: gate_proj, up_proj and down_proj go back to w1, w3 and w2. With the rewritten converters reversed as rewritten, the mapper refuses on an arity mismatch, and the export would keep the in-memory names. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k3_ep.py | 25 ++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index f0d033ec996..43ab6d80c56 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -233,6 +233,31 @@ def test_registered_conversions_fuse_w1_w3_into_gate_up_and_stack_w2(): assert torch.equal(fused[prefix + "down_proj"], torch.stack(down)) +@pytest.mark.usefixtures("private_conversion_registry") +def test_packed_experts_export_under_the_checkpoints_names(): + """The decompressor rewrites the registered converters to read the packed experts, and the model + keeps the rewritten ones as the conversions its load ran. The export names the experts back + through them, so they reverse as the fusion they rewrote: to the checkpoint's ``w1``/``w3``/``w2``, + not to ModelOpt's ``gate_proj``/``up_proj``/``down_proj``.""" + from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper + + kimi_k3._register_kimi_k3_conversions() + packed = [ + _onfly_decompress._compressedify(conv, None, None, ("weight_packed", "weight_scale")) + for conv in cm.get_checkpoint_conversion_mapping("kimi_linear") + ] + expert = "model.layers.1.block_sparse_moe.experts.0." + + to_checkpoint = build_reverse_name_mapper( + SimpleNamespace(_weight_conversions=packed), tensor_keys=True + ) + + exported = [ + to_checkpoint(f"{expert}{part}.weight") for part in ("gate_proj", "up_proj", "down_proj") + ] + assert exported == [f"{expert}{leaf}.weight" for leaf in ("w1", "w3", "w2")] + + def test_the_registration_raises_without_the_conversion_api(monkeypatch): """Unregistered, the checkpoint's experts would load as MISSING: newly initialized, at full size on every rank. So a missing conversion API stops the load instead of warning.""" From 9585987c70561af446043f3e74482c732fbdabba Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:20:32 +0000 Subject: [PATCH 08/32] Load, calibrate and name back a packed Kimi-K3 expert-parallel on 2 GPUs A bounded 2-rank GPU case for the Kimi-K3 adapter on a real distributed load. A tiny Kimi-K3-shaped remote checkpoint, written by the test from kimi_k3_tiny_modeling (the module tree, names and per-expert MoE the adapter patches, as the checkpoint's own code builds them), stores its 8 routed experts as packed MXFP4, as Kimi-K3 does. Loaded through the model spec and load_and_prepare_ep at an EP degree of the world size: - each rank decompresses only its own experts, exactly; - the sharded model's logits match the whole model loaded on one GPU; - FP8 calibration, each rank on its own batch, matches the whole model calibrated on every batch, expert weights and activations alike; - the export's reverse name mapper names the experts back to the checkpoint's w1/w3/w2. Each rank's findings are gathered before any rank fails, and the object collectives give up in two minutes, so a regression fails instead of hanging. One GPU skips it. Signed-off-by: Shengliang Xu --- .../torch/kimi_k3_tiny_modeling.py | 268 ++++++++++++++++++ .../torch/utils/test_expert_parallel_load.py | 242 ++++++++++++++++ 2 files changed, 510 insertions(+) create mode 100644 tests/_test_utils/torch/kimi_k3_tiny_modeling.py diff --git a/tests/_test_utils/torch/kimi_k3_tiny_modeling.py b/tests/_test_utils/torch/kimi_k3_tiny_modeling.py new file mode 100644 index 00000000000..5407da3e132 --- /dev/null +++ b/tests/_test_utils/torch/kimi_k3_tiny_modeling.py @@ -0,0 +1,268 @@ +# 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-K3-shaped model, as remote code: tests copy it into a checkpoint as ``modeling_kimi_k3.py``. + +It has the module tree, the names and the MoE the Kimi-K3 adapter patches, as the checkpoint's own +code builds them: a VLM-shaped ``KimiK3ForConditionalGeneration`` whose ``language_model`` is a +``KimiLinearForCausalLM`` (model type ``kimi_linear``), a dense first layer, and then +``block_sparse_moe`` blocks of per-expert ``w1``/``w3``/``w2`` in a latent space, with the latent norm, +shared experts and the ``situ`` gate. A per-token linear stands in for the attention, which the +expert-parallel path leaves alone. +""" + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig, PreTrainedModel +from transformers.modeling_outputs import CausalLMOutputWithPast + + +class KimiLinearConfig(PretrainedConfig): + model_type = "kimi_linear" + + def __init__( + self, + vocab_size=64, + hidden_size=64, + intermediate_size=64, + routed_expert_hidden_size=32, + moe_intermediate_size=32, + num_experts=8, + num_experts_per_token=2, + num_hidden_layers=3, + first_k_dense_replace=1, + num_shared_experts=1, + hidden_act="situ", + situ_beta=1.5, + situ_linear_beta=0.25, + latent_moe_use_norm=True, + rms_norm_eps=1e-6, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.routed_expert_hidden_size = routed_expert_hidden_size + self.moe_intermediate_size = moe_intermediate_size + self.num_experts = num_experts + self.num_experts_per_token = num_experts_per_token + self.num_hidden_layers = num_hidden_layers + self.first_k_dense_replace = first_k_dense_replace + self.num_shared_experts = num_shared_experts + self.hidden_act = hidden_act + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta + self.latent_moe_use_norm = latent_moe_use_norm + self.rms_norm_eps = rms_norm_eps + kwargs.setdefault("tie_word_embeddings", False) + super().__init__(**kwargs) + + +class KimiK3Config(PretrainedConfig): + model_type = "kimi_k3" + sub_configs = {"text_config": KimiLinearConfig} + + def __init__(self, text_config=None, **kwargs): + if isinstance(text_config, dict): + text_config = KimiLinearConfig(**text_config) + self.text_config = text_config if text_config is not None else KimiLinearConfig() + kwargs.setdefault("tie_word_embeddings", False) + super().__init__(**kwargs) + + +def _get_situ_activation_params(config): + return config.situ_beta, config.situ_linear_beta + + +class SituAndMul(nn.Module): + """``situ`` on a full ``[gate; up]`` row.""" + + def __init__(self, beta, linear_beta): + super().__init__() + self.beta, self.linear_beta = beta, linear_beta + + def forward(self, gate_up): + gate, up = gate_up.chunk(2, dim=-1) + return gate * torch.sigmoid(self.beta * gate) * (up + self.linear_beta) + + +class KimiRMSNorm(nn.Module): + def __init__(self, size, eps): + super().__init__() + self.weight = nn.Parameter(torch.ones(size)) + self.eps = eps + + def forward(self, x): + normed = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) + return (self.weight * normed).to(x.dtype) + + +class KimiMLP(nn.Module): + """The dense MLP: the first layer's, and the shared experts'.""" + + def __init__(self, hidden, inter): + super().__init__() + self.gate_proj = nn.Linear(hidden, inter, bias=False) + self.up_proj = nn.Linear(hidden, inter, bias=False) + self.down_proj = nn.Linear(inter, hidden, bias=False) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class KimiBlockSparseMLP(nn.Module): + """One routed expert, as the checkpoint stores it: ``w1`` = gate, ``w3`` = up, ``w2`` = down.""" + + def __init__(self, config, hidden): + super().__init__() + inter = config.moe_intermediate_size + self.w1 = nn.Linear(hidden, inter, bias=False) + self.w3 = nn.Linear(hidden, inter, bias=False) + self.w2 = nn.Linear(inter, hidden, bias=False) + situ = config.hidden_act == "situ" + self.act = SituAndMul(*_get_situ_activation_params(config)) if situ else None + + def forward(self, x): + gate_up = torch.cat([self.w1(x), self.w3(x)], dim=-1) + if self.act is not None: + return self.w2(self.act(gate_up)) + gate, up = gate_up.chunk(2, dim=-1) + return self.w2(F.silu(gate) * up) + + +class KimiMoEGate(nn.Module): + """Sigmoid scores over the replicated router weight, normalized over each token's top k.""" + + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_token + self.weight = nn.Parameter(torch.empty(config.num_experts, config.hidden_size)) + + def forward(self, hidden_states): + flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + scores = torch.sigmoid(F.linear(flat.float(), self.weight.float())) + weights, idx = torch.topk(scores, self.top_k, dim=-1) + return idx, (weights / weights.sum(dim=-1, keepdim=True)).to(hidden_states.dtype) + + +class KimiSparseMoeBlock(nn.Module): + """Latent projections around a per-expert loop, plus the always-on shared experts.""" + + def __init__(self, config): + super().__init__() + self.config = config + self.use_latent_moe = config.routed_expert_hidden_size is not None + self.latent_moe_use_norm = config.latent_moe_use_norm + hidden = config.hidden_size + width = config.routed_expert_hidden_size if self.use_latent_moe else hidden + self.experts = nn.ModuleList( + KimiBlockSparseMLP(config, width) for _ in range(config.num_experts) + ) + self.gate = KimiMoEGate(config) + if self.use_latent_moe: + self.routed_expert_down_proj = nn.Linear(hidden, width, bias=False) + self.routed_expert_up_proj = nn.Linear(width, hidden, bias=False) + if config.latent_moe_use_norm: + self.routed_expert_norm = KimiRMSNorm(width, config.rms_norm_eps) + if config.num_shared_experts is not None: + inter = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = KimiMLP(hidden, inter) + + def forward(self, hidden_states): + idx, weights = self.gate(hidden_states) + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + if self.use_latent_moe: + x = self.routed_expert_down_proj(x) + y = torch.zeros_like(x) + for e, expert in enumerate(self.experts): + token, slot = torch.where(idx == e) + if token.numel(): + y.index_add_(0, token, expert(x[token]) * weights[token, slot, None]) + if self.use_latent_moe: + if self.latent_moe_use_norm: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + y = y.reshape(hidden_states.shape) + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(hidden_states) + return y + + +class KimiDecoderLayer(nn.Module): + def __init__(self, config, layer_idx): + super().__init__() + hidden = config.hidden_size + self.input_layernorm = KimiRMSNorm(hidden, config.rms_norm_eps) + self.self_attn = nn.Linear(hidden, hidden, bias=False) # no token mixing: not under test + self.post_attention_layernorm = KimiRMSNorm(hidden, config.rms_norm_eps) + if layer_idx < config.first_k_dense_replace: + self.mlp = KimiMLP(hidden, config.intermediate_size) + else: + self.block_sparse_moe = KimiSparseMoeBlock(config) + + def forward(self, hidden_states): + hidden_states = hidden_states + self.self_attn(self.input_layernorm(hidden_states)) + ffn = self.mlp if hasattr(self, "mlp") else self.block_sparse_moe + return hidden_states + ffn(self.post_attention_layernorm(hidden_states)) + + +class KimiLinearPreTrainedModel(PreTrainedModel): + config_class = KimiLinearConfig + base_model_prefix = "model" + _no_split_modules = ["KimiDecoderLayer"] + + +class KimiLinearModel(KimiLinearPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + KimiDecoderLayer(config, i) for i in range(config.num_hidden_layers) + ) + self.norm = KimiRMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_init() + + def forward(self, input_ids): + hidden_states = self.embed_tokens(input_ids) + for layer in self.layers: + hidden_states = layer(hidden_states) + return self.norm(hidden_states) + + +class KimiLinearForCausalLM(KimiLinearPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.model = KimiLinearModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def forward(self, input_ids=None, **kwargs): + return CausalLMOutputWithPast(logits=self.lm_head(self.model(input_ids))) + + +class KimiK3PreTrainedModel(PreTrainedModel): + config_class = KimiK3Config + _no_split_modules = ["KimiDecoderLayer"] + + +class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.language_model = KimiLinearForCausalLM(config.text_config) + self.post_init() + + def forward(self, input_ids=None, **kwargs): + return self.language_model(input_ids=input_ids) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 3d50e5c3e5f..3a496021260 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -31,6 +31,7 @@ import pytest import torch import torch.distributed as dist +from _test_utils.torch import kimi_k3_tiny_modeling from _test_utils.torch.distributed.utils import bounded_collectives from packaging.version import Version from safetensors.torch import load_file, save_file @@ -51,9 +52,14 @@ Qwen3MoeConfig, ) +import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper from modelopt.torch.models import _onfly_decompress +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils.distributed import ( DistributionPlan, + ParallelState, + distributed_forward_loop, distribution_plan, set_distribution_plan, ) @@ -399,3 +405,239 @@ def test_kimi_vl_loads_expert_parallel_and_matches_the_whole_model( 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)) + + +# ------------------------------------------------------------------------------------- Kimi-K3 + +# A tiny Kimi-K3: a dense first layer, then one MoE layer of 8 routed experts, run in a 32-wide latent +# space, as Kimi-K3's own code builds it (remote code, ``kimi_k3_tiny_modeling``). +KIMI_K3_DIMS = { + "vocab_size": 64, + "hidden_size": 64, + "intermediate_size": 64, + "routed_expert_hidden_size": 32, + "moe_intermediate_size": 32, + "num_experts": 8, + "num_experts_per_token": 2, + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "num_shared_experts": 1, +} +KIMI_K3_EXPERTS = KIMI_K3_DIMS["num_experts"] +KIMI_K3_MOE = "language_model.model.layers.1.block_sparse_moe.experts." +# Kimi-K3's own compression: MXFP4 routed experts, everything else bf16. +KIMI_K3_MXFP4 = { + "quant_method": "compressed-tensors", + "format": "mxfp4-pack-quantized", + "quantization_status": "compressed", + "config_groups": { + "group_0": { + "format": "mxfp4-pack-quantized", + "targets": ["Linear"], + "weights": { + "num_bits": 4, + "type": "float", + "symmetric": True, + "strategy": "group", + "group_size": 32, + "scale_dtype": "torch.uint8", + "dynamic": False, + "observer": "minmax", + }, + } + }, + "ignore": [ + "re:.*self_attn.*", + "re:.*shared_experts.*", + "re:.*mlp\\.(gate|up|gate_up|down)_proj.*", + "re:.*lm_head.*", + ], +} + + +def _kimi_k3_expert(expert): + """Expert ``expert``'s ``w1``/``w3``/``w2``, the same on every rank.""" + gen = torch.Generator().manual_seed(expert) + + def draw(rows, cols): + # MXFP4 (e2m1) values at a scale of 1/16: packed at that power-of-two scale, they decompress + # exactly, and bf16 holds them exactly. + values = _E2M1[torch.randint(0, len(_E2M1), (rows, cols), generator=gen)] / 16 + return values * torch.where(torch.rand(rows, cols, generator=gen) < 0.5, -1.0, 1.0) + + inter, latent = KIMI_K3_DIMS["moe_intermediate_size"], KIMI_K3_DIMS["routed_expert_hidden_size"] + return {"w1": draw(inter, latent), "w3": draw(inter, latent), "w2": draw(latent, inter)} + + +def _save_tiny_kimi_k3(ckpt_dir): + """A tiny Kimi-K3-shaped remote checkpoint whose routed experts are packed MXFP4, as Kimi-K3's are.""" + ckpt_dir = Path(ckpt_dir) + shutil.copy(kimi_k3_tiny_modeling.__file__, ckpt_dir / "modeling_kimi_k3.py") + remote = "modeling_kimi_k3.KimiK3ForConditionalGeneration" + config = { + "model_type": "kimi_k3", + "architectures": ["KimiK3ForConditionalGeneration"], + "auto_map": { + "AutoConfig": "modeling_kimi_k3.KimiK3Config", + "AutoModel": remote, + "AutoModelForCausalLM": remote, + }, + "torch_dtype": "bfloat16", + "tie_word_embeddings": False, + # As Kimi-K3's: the scheme sits on the language model's config only. + "text_config": { + "model_type": "kimi_linear", + **KIMI_K3_DIMS, + "quantization_config": KIMI_K3_MXFP4, + }, + } + (ckpt_dir / "config.json").write_text(json.dumps(config)) + + compressor, scheme, stored_as, _ = _onfly_decompress._read_compression( + SimpleNamespace(quantization_config=KIMI_K3_MXFP4) + ) + with torch.device("meta"): + shapes = kimi_k3_tiny_modeling.KimiK3ForConditionalGeneration( + kimi_k3_tiny_modeling.KimiK3Config(text_config=dict(KIMI_K3_DIMS)) + ).state_dict() + gen = torch.Generator().manual_seed(0) + tensors = {} + for name, meta in shapes.items(): + if name.startswith(KIMI_K3_MOE): + continue + dense = torch.randn(meta.shape, generator=gen) * 0.1 + tensors[name] = (1 + dense if name.endswith("norm.weight") else dense).to(torch.bfloat16) + for e in range(KIMI_K3_EXPERTS): + for leaf, weight in _kimi_k3_expert(e).items(): + scale = torch.full((weight.shape[0], weight.shape[1] // 32), 1 / 16) + packed = compressor.compress({"weight": weight, "weight_scale": scale}, scheme) + for stored in stored_as: + tensors[f"{KIMI_K3_MOE}{e}.{leaf}.{stored}"] = packed[stored] + save_file(tensors, str(ckpt_dir / "model.safetensors"), metadata={"format": "pt"}) + + +def _whole_kimi_k3(ckpt_dir, config, device): + """The single-rank load: the whole model on one device, its experts decompressed as it loads.""" + prepared = apply_distributed_support( + "kimi_k3", copy.deepcopy(config), ckpt_path=ckpt_dir, trust_remote_code=True + ) + whole = AutoModelForCausalLM.from_pretrained( + ckpt_dir, config=prepared, trust_remote_code=True, dtype=torch.bfloat16 + ) + for module in whole.modules(): # the experts forward reads it at forward time + if hasattr(getattr(module, "config", None), "_experts_implementation"): + module.config._experts_implementation = "eager" + return whole.to(device).eval() + + +def _calibrated(model): + return { + name: module.amax.float().cpu() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.is_enabled and module.amax is not None + } + + +def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): + device = torch.device("cuda", rank) + batches = [ + torch.randint( + 0, KIMI_K3_DIMS["vocab_size"], (2, 16), generator=torch.Generator().manual_seed(r) + ) + for r in range(size) + ] + per_rank = KIMI_K3_EXPERTS // size + owned = range(rank * per_rank, (rank + 1) * per_rank) + problems = [] + # The decompressor's hooks are process-wide, and the workers outlive the job. + hooks = (mu.get_model_conversion_mapping, cml._materialize_copy) + try: + 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_k3", 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=prepared + ).eval() + set_distribution_plan(model, plan) + whole = _whole_kimi_k3(ckpt_dir, config, device) + for module in whole.modules(): + module._parallel_state = ParallelState(data_parallel_group=-1) + + # This rank's experts, decompressed as they loaded, and only those. + experts = model.language_model.model.layers[1].block_sparse_moe.experts + want = [_kimi_k3_expert(e) for e in owned] + gate_up = torch.stack([torch.cat([w["w1"], w["w3"]]) for w in want]) + down = torch.stack([w["w2"] for w in want]) + for name, expected in (("gate_up_proj", gate_up), ("down_proj", down)): + local = _local(getattr(experts, name)).float().cpu() + if not torch.equal(local, expected): + problems.append(f"{name}: rank {rank}'s experts {list(owned)} differ") + + # Each rank its own batch through the experts, against the whole model on that batch. + with torch.no_grad(): + sharded = model(batches[rank].to(device)).logits.float() + reference = whole(batches[rank].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}") + + # Calibrated over every rank's batch: each rank on its own, the whole model on all. + forward_loop = distributed_forward_loop( + model, + [{"input_ids": batches[rank]}], + device=device, + ep_mesh=expert_parallel_mesh(model), + ) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop) + mtq.quantize(whole, mtq.FP8_DEFAULT_CFG, lambda m: [m(b.to(device)) for b in batches]) + sharded_amax, reference_amax = _calibrated(model), _calibrated(whole) + if not any("_weight_quantizers." in name for name in sharded_amax): + problems.append("no expert weight quantizer was calibrated") + for name, amax in sharded_amax.items(): + if "_weight_quantizers." in name: # this rank's local expert -> the global one + prefix, index = name.rsplit(".", 1) + name = f"{prefix}.{rank * per_rank + int(index)}" + expected_amax = reference_amax.get(name) + if expected_amax is None: + problems.append(f"{name} is calibrated here but not in the whole model") + continue + # bf16 activations, summed in another order: compare by relative norm. + error = (amax - expected_amax).norm() / expected_amax.norm().clamp_min(1e-12) + if not error < 2e-2: + problems.append(f"{name}: amax off by {error.item():.3g} (relative)") + + # The export names each expert back to the checkpoint's w1/w3/w2. + to_checkpoint = build_reverse_name_mapper(model, tensor_keys=True) + for e in owned: + parts = [ + f"{KIMI_K3_MOE}{e}.{p}.weight" for p in ("gate_proj", "up_proj", "down_proj") + ] + leaves = [f"{KIMI_K3_MOE}{e}.{leaf}.weight" for leaf in ("w1", "w3", "w2")] + if [to_checkpoint(part) for part in parts] != leaves: + problems.append(f"expert {e} exports as {[to_checkpoint(p) for p in parts]}") + + gathered = [None] * size + dist.all_gather_object(gathered, problems) + finally: + mu.get_model_conversion_mapping, cml._materialize_copy = hooks + found = [problem for rank_problems in gathered for problem in rank_problems] + assert not found, "\n".join(found[:20]) + + +def test_kimi_k3_packed_experts_load_calibrate_and_name_back_expert_parallel( + need_2_gpus, dist_workers, tmp_path +): + """A Kimi-K3-shaped remote checkpoint with packed MXFP4 experts, loaded expert-parallel through + its model spec: each rank decompresses only its own experts, the sharded model matches the whole + one loaded on a single GPU, calibration over every rank's batch matches the whole model's, and + the export names the experts back to the checkpoint's ``w1``/``w3``/``w2``. Each rank's + checks are gathered before any rank fails, so a regression fails instead of hanging.""" + pytest.importorskip("compressed_tensors", reason="the checkpoint's experts are compressed") + if Version(transformers.__version__) < Version("5.16"): + pytest.skip("the torch_a2a forward wrappers are pre-hooks below transformers 5.16") + if KIMI_K3_EXPERTS % dist_workers.world_size: + pytest.skip(f"{KIMI_K3_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs") + _save_tiny_kimi_k3(tmp_path) + dist_workers.run(_kimi_k3_expert_parallel_job, ckpt_dir=str(tmp_path)) From 90d79ae82a79db4ad6061a99cfbcc078dc0044ae Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:52:31 +0000 Subject: [PATCH 09/32] Restore the decompressor's spawn hook after the Kimi-K3 tests too The decompressor keeps integer sources uncast at the loader's spawn of each source tensor, which covers an expert-parallel load's shard path, where it hooked only the copy before. The unit tests' hook-restoring fixture and the 2-GPU job's save and restore now cover spawn_materialize as well as _materialize_copy, so neither patch leaks into later tests or the pooled GPU workers, whichever of the two the decompressor patches. The GPU file imports the tiny Kimi-K3 modeling below its transformers guard, and the Kimi-K3 case's skip reason names the package it needs. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_expert_parallel_load.py | 12 +++++++----- tests/unit/torch/models/test_kimi_k3_ep.py | 4 +++- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 3a496021260..6f4ca1bfbcf 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -31,7 +31,6 @@ import pytest import torch import torch.distributed as dist -from _test_utils.torch import kimi_k3_tiny_modeling from _test_utils.torch.distributed.utils import bounded_collectives from packaging.version import Version from safetensors.torch import load_file, save_file @@ -42,7 +41,7 @@ import transformers.core_model_loading as cml import transformers.modeling_utils as mu -from _test_utils.torch import kimi_vl_tiny_modeling +from _test_utils.torch import kimi_k3_tiny_modeling, kimi_vl_tiny_modeling from torch.distributed.tensor import DTensor from transformers import ( AutoConfig, @@ -550,7 +549,7 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): owned = range(rank * per_rank, (rank + 1) * per_rank) problems = [] # The decompressor's hooks are process-wide, and the workers outlive the job. - hooks = (mu.get_model_conversion_mapping, cml._materialize_copy) + hooks = (mu.get_model_conversion_mapping, cml._materialize_copy, cml.spawn_materialize) try: with bounded_collectives(seconds=120): config = AutoConfig.from_pretrained(ckpt_dir, trust_remote_code=True) @@ -621,7 +620,7 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): gathered = [None] * size dist.all_gather_object(gathered, problems) finally: - mu.get_model_conversion_mapping, cml._materialize_copy = hooks + mu.get_model_conversion_mapping, cml._materialize_copy, cml.spawn_materialize = hooks found = [problem for rank_problems in gathered for problem in rank_problems] assert not found, "\n".join(found[:20]) @@ -634,7 +633,10 @@ def test_kimi_k3_packed_experts_load_calibrate_and_name_back_expert_parallel( one loaded on a single GPU, calibration over every rank's batch matches the whole model's, and the export names the experts back to the checkpoint's ``w1``/``w3``/``w2``. Each rank's checks are gathered before any rank fails, so a regression fails instead of hanging.""" - pytest.importorskip("compressed_tensors", reason="the checkpoint's experts are compressed") + pytest.importorskip( + "compressed_tensors", + reason="compressed-tensors is needed: the checkpoint's experts are compressed", + ) if Version(transformers.__version__) < Version("5.16"): pytest.skip("the torch_a2a forward wrappers are pre-hooks below transformers 5.16") if KIMI_K3_EXPERTS % dist_workers.world_size: diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 43ab6d80c56..c20a82862bc 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -74,11 +74,13 @@ def private_conversion_registry(monkeypatch): @pytest.fixture def restore_loader_hooks(monkeypatch): - """Undo the process-wide transformers hooks the on-the-fly decompressor installs.""" + """Undo the process-wide transformers hooks the on-the-fly decompressor installs, whichever of + them its release patches: the copy, or the spawn that also covers the shard path.""" for owner, name in ( (WeightConverter, "add_tensor"), (WeightConverter, "materialize_tensors"), (cml, "_materialize_copy"), + (cml, "spawn_materialize"), (mu, "get_model_conversion_mapping"), ): monkeypatch.setattr(owner, name, getattr(owner, name)) From 2ace2206c4d6760b30fa5fa74afaa94ea89dd81c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:39:28 +0000 Subject: [PATCH 10/32] Expect the left-whole refusal's arm this transformers takes for remote Kimi-K3 The plan decision now raises the generic left-whole refusal as a ValueError, and keeps the RuntimeError for a transformers below what the model's spec needs, which names the release. The remote Kimi-K3 case expects whichever arm its transformers takes, as the Kimi-VL case does. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b88e8fe2b76..b58bca618ce 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -935,7 +935,8 @@ def test_a_remote_kimi_k3_whose_prepare_attaches_its_plan_left_whole_is_an_error ): """Remote-code Kimi checkpoints get their plan from their spec's ``prepare``, which an EP degree of 1 skips: it patches their code. The plan their spec states still finds the experts their own - code builds, so whole on every rank they are refused like Qwen3-MoE.""" + code builds, so whole on every rank they are refused like Qwen3-MoE: a ValueError, or the + RuntimeError naming the release when this transformers is below what their spec needs.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") monkeypatch.setattr( get_spec("kimi_k3").ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") @@ -943,8 +944,12 @@ def test_a_remote_kimi_k3_whose_prepare_attaches_its_plan_left_whole_is_an_error monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi_k3()) run_as(4) serve(SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace(num_experts=8))) + below_floor = Version(transformers.__version__) < Version(_ep_floor("kimi_k3")) - with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + with pytest.raises( + RuntimeError if below_floor else ValueError, + match="can be expert-parallel, but the plan resolved", + ): resolve_distribution_plan(str(tmp_path)) From 1e87e88f3c6362bf677e882e2ddecab3f1871915 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:43:33 +0000 Subject: [PATCH 11/32] Pin that a packed Kimi-K3 loaded through another model class gets no prepare The plan decision applies a model's spec only for the classes the sharded loads build; a model class they refuse is planned whole on every rank, never expert-parallel. For a packed Kimi-K3, the spec's prepare patches the checkpoint's own code for the process and marks only its copy of the config for the experts' decompression, so running it for such a load would leave the whole load with fused blocks its unmarked config cannot fill. The test plans a packed Kimi-K3 on four ranks through AutoModel: prepare never runs, the plan is the whole model on every rank, and the config returned is the one read, unmarked and still holding its compressed-tensors scheme. It fails, on prepare, when the spec is applied regardless of the class. Signed-off-by: Shengliang Xu --- .../unit/torch/utils/test_model_load_utils.py | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b58bca618ce..24513a89755 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -75,7 +75,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist import modelopt.torch.utils.moe as moe -from modelopt.torch.models import get_spec +from modelopt.torch.models import _onfly_decompress, get_spec from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -953,6 +953,38 @@ def test_a_remote_kimi_k3_whose_prepare_attaches_its_plan_left_whole_is_an_error resolve_distribution_plan(str(tmp_path)) +def test_a_packed_kimi_k3_loaded_through_another_class_is_planned_without_its_spec( + run_as, serve, monkeypatch, tmp_path +): + """A model class the sharded loads refuse is planned whole on every rank, never expert-parallel, + so a packed Kimi-K3 loaded through one (``AutoModel``) gets no ``prepare``, which patches the + checkpoint's own code for the process and marks only its copy of the config for the experts' + decompression. The whole load then reads the config as the checkpoint holds it: unmarked, and + with the compressed-tensors scheme transformers' own quantizer loads the experts by.""" + monkeypatch.setattr( + get_spec("kimi_k3").ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") + ) + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") + ) + # As on a transformers new enough for Kimi-K3's expert parallelism, which would rule it out here. + monkeypatch.setattr(model_load_utils, "_spec_refuses_ep", lambda model_type: False) + run_as(4) + scheme = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} + text = SimpleNamespace(num_experts=8, quantization_config=scheme) + # Kimi-K3's own config class copies its language model's scheme to the top level. + config = serve( + SimpleNamespace(model_type="kimi_k3", text_config=text, quantization_config=scheme) + ) + + plan, planned = resolve_distribution_plan(str(tmp_path), auto_class=AutoModel) + + assert plan == DistributionPlan(4) + assert planned is config and not _onfly_decompress.is_marked(planned) + assert planned.quantization_config is scheme + assert planned.text_config.quantization_config is scheme + + @needs_qwen3_moe_ep_plan def test_an_ep_capable_moe_is_replicated_when_both_knobs_ask_for_it( run_as, serve, monkeypatch, tmp_path From 47f8884a670903bc1cb371cd3ff686826d1d68ff Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:43:38 +0000 Subject: [PATCH 12/32] Copy the tiny Kimi-K3 config's scheme to the top level, as Kimi-K3's does Kimi-K3's own config class copies its language model's quantization_config to the top-level config, where the sharded loads' refusal of a quantized checkpoint reads it. The tiny Kimi-K3 the GPU case writes now does the same. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/kimi_k3_tiny_modeling.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/_test_utils/torch/kimi_k3_tiny_modeling.py b/tests/_test_utils/torch/kimi_k3_tiny_modeling.py index 5407da3e132..c81990eb82e 100644 --- a/tests/_test_utils/torch/kimi_k3_tiny_modeling.py +++ b/tests/_test_utils/torch/kimi_k3_tiny_modeling.py @@ -79,6 +79,9 @@ def __init__(self, text_config=None, **kwargs): if isinstance(text_config, dict): text_config = KimiLinearConfig(**text_config) self.text_config = text_config if text_config is not None else KimiLinearConfig() + # As Kimi-K3's own config class: its language model's scheme is the checkpoint's too. + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config kwargs.setdefault("tie_word_embeddings", False) super().__init__(**kwargs) From 1577b6d637c17522e3ad4f886243424eabec2c1e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:43:40 +0000 Subject: [PATCH 13/32] Import the export's reverse name mapper at module scope in the Kimi-K3 tests It has no import cycle or optional dependency to defer; it now sits with the other ModelOpt imports, below the transformers guard. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_kimi_k3_ep.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index c20a82862bc..a03a42dbbc4 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -47,6 +47,7 @@ from transformers.core_model_loading import WeightConverter, build_glob_alternation from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS +from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper from modelopt.torch.models import _onfly_decompress, get_spec from modelopt.torch.models.kimi_k3 import ep as kimi_k3 from modelopt.torch.utils import moe @@ -241,8 +242,6 @@ def test_packed_experts_export_under_the_checkpoints_names(): keeps the rewritten ones as the conversions its load ran. The export names the experts back through them, so they reverse as the fusion they rewrote: to the checkpoint's ``w1``/``w3``/``w2``, not to ModelOpt's ``gate_proj``/``up_proj``/``down_proj``.""" - from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper - kimi_k3._register_kimi_k3_conversions() packed = [ _onfly_decompress._compressedify(conv, None, None, ("weight_packed", "weight_scale")) From c052f0c82c0b49d47b33b55ac207e390f2a1b0d9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:58:54 +0000 Subject: [PATCH 14/32] Restore only the loader hooks the decompressor patches in the Kimi-K3 tests The decompressor patches the loader's spawn_materialize and the conversion mapping, on every transformers release it runs on. The Kimi-K3 unit fixture and the 2-GPU job also saved and restored the converter's add_tensor and materialize_tensors and _materialize_copy, which nothing patches any more. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_expert_parallel_load.py | 4 ++-- tests/unit/torch/models/test_kimi_k3_ep.py | 13 +++---------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 6f4ca1bfbcf..ee240d2e6c8 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -549,7 +549,7 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): owned = range(rank * per_rank, (rank + 1) * per_rank) problems = [] # The decompressor's hooks are process-wide, and the workers outlive the job. - hooks = (mu.get_model_conversion_mapping, cml._materialize_copy, cml.spawn_materialize) + hooks = (mu.get_model_conversion_mapping, cml.spawn_materialize) try: with bounded_collectives(seconds=120): config = AutoConfig.from_pretrained(ckpt_dir, trust_remote_code=True) @@ -620,7 +620,7 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): gathered = [None] * size dist.all_gather_object(gathered, problems) finally: - mu.get_model_conversion_mapping, cml._materialize_copy, cml.spawn_materialize = hooks + mu.get_model_conversion_mapping, cml.spawn_materialize = hooks found = [problem for rank_problems in gathered for problem in rank_problems] assert not found, "\n".join(found[:20]) diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index a03a42dbbc4..990f0910b1d 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -44,7 +44,7 @@ import transformers.dynamic_module_utils as dmu import transformers.modeling_utils as mu import transformers.utils.generic as hf_generic -from transformers.core_model_loading import WeightConverter, build_glob_alternation +from transformers.core_model_loading import build_glob_alternation from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper @@ -75,15 +75,8 @@ def private_conversion_registry(monkeypatch): @pytest.fixture def restore_loader_hooks(monkeypatch): - """Undo the process-wide transformers hooks the on-the-fly decompressor installs, whichever of - them its release patches: the copy, or the spawn that also covers the shard path.""" - for owner, name in ( - (WeightConverter, "add_tensor"), - (WeightConverter, "materialize_tensors"), - (cml, "_materialize_copy"), - (cml, "spawn_materialize"), - (mu, "get_model_conversion_mapping"), - ): + """Undo the process-wide transformers hooks the on-the-fly decompressor installs.""" + for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) From 6bf6d48f8cf9d957655e0e81489c2cd89e575548 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:06:29 +0000 Subject: [PATCH 15/32] Say in Kimi-K3's prepare that its patches outlast the load they were made for A process that already ran a Kimi-K3 expert-parallel load keeps the patched remote classes, so a later load of the same checkpoint through a class the sharded loads refuse, such as AutoModel, runs on them. The plan decision itself never makes such a load. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index e908ce1067b..adaf504cde8 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -340,6 +340,10 @@ def prepare( bf16 fused params, and the compressed-tensors ``quantization_config`` is dropped. ModelOpt then inserts/calibrates/folds the fp8/nvfp4 quantizers on those bf16 params. That is ~4x the packed HBM (the memory bottleneck), but PTQ needs the bf16 weights. + + The patches are process-wide and stay after the load: a later load of the same checkpoint in this + process, through a class the sharded loads refuse (such as ``AutoModel``), runs on the patched + code. The plan decision itself never makes such a load. """ # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe # imports the registry. From 5d84282b44db774e2c269a8319972d8fb54ff0cc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:31:05 +0000 Subject: [PATCH 16/32] Install Kimi-K3's decompressor before prepare patches anything From Claude's review: prepare patched the remote MoE block and registered the fusion converters before the decompressor's install, which can still refuse the checkpoint (compressed-tensors missing, several config groups, an unknown format, a dtype other than bf16). A process that went on after the refusal kept the patches. The install now runs first: it validates before it installs anything, and its hooks rewrite only a config it marked. The converters need only be registered before the load, which the comment now says. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 28 ++++++++++++---------- tests/unit/torch/models/test_kimi_k3_ep.py | 11 ++++++--- 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index adaf504cde8..d21d6d3059e 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -331,9 +331,10 @@ def prepare( ): """Adapt Kimi-K3's remote MoE to the EP/a2a framework and set its EP plan. - Steps: shim OutputRecorder -> force-load + monkey-patch the remote KimiSparseMoeBlock (fused experts + - ExpertsInterface forward) -> set base_model_ep_plan on text_config -> register the w1/w3/w2 fusion - converters and the decompressor. Returns the same config object, mutated in place. + Steps: install the decompressor -> shim OutputRecorder -> force-load + monkey-patch the remote + KimiSparseMoeBlock (fused experts + ExpertsInterface forward) -> set base_model_ep_plan on + text_config -> register the w1/w3/w2 fusion converters. Returns the same config object, mutated in + place. The 896 ``mxfp4-pack-quantized`` routed experts are DECOMPRESSED at load: ``_onfly_decompress`` dequantizes each local expert's packed weight -> bf16 INSIDE the fused-expert conversion, stored as @@ -364,6 +365,14 @@ def prepare( text_cfg = getattr(hf_config, "text_config", None) or hf_config + # The 896 routed experts are compressed-tensors `mxfp4-pack-quantized` (4-bit, group 32); everything + # else (attention, shared experts, latent projections, dense layer, embeddings, lm_head, vision) is + # plain bf16. The decompressor is the step that can still refuse the checkpoint, so it comes before + # anything below patches the process: a refusal leaves it unchanged. Its hooks rewrite only the + # converters of a config it marked. Plain bf16 experts come without a quantization_config. + if any(getattr(cfg, "quantization_config", None) is not None for cfg in (hf_config, text_cfg)): + _onfly_decompress.install_onfly_decompress(hf_config) + # This checkpoint's remote code, not whichever revision this process imported first. Every patch # is idempotent, so a repeat prepare for the same checkpoint changes nothing. package = _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code).__module__ @@ -389,16 +398,11 @@ def prepare( register_torch_a2a() text_cfg._experts_implementation = "torch_a2a" - # The 896 routed experts are compressed-tensors `mxfp4-pack-quantized` (4-bit, group 32); everything - # else (attention, shared experts, latent projections, dense layer, embeddings, lm_head, vision) is - # plain bf16. Register the per-expert w1/w3/w2 -> gate_up_proj/down_proj fusion converters (tf 5.13 - # skips these for custom-code models) BEFORE the decompressor patches get_model_conversion_mapping, so - # the decompressor has converters to rewrite into compressed-aware (weight_packed -> bf16) ones. Drop - # quantization_config afterwards so transformers builds no competing HF quantizer. + # The per-expert w1/w3/w2 -> gate_up_proj/down_proj fusion converters (tf 5.13 skips these for + # custom-code models). The decompressor rewrites them into compressed-aware (weight_packed -> bf16) + # ones at the load, so both need only come before it. Drop quantization_config afterwards so + # transformers builds no competing HF quantizer. _register_kimi_k3_conversions() - # Plain bf16 experts come without a quantization_config: nothing to decompress. - if any(getattr(cfg, "quantization_config", None) is not None for cfg in (hf_config, text_cfg)): - _onfly_decompress.install_onfly_decompress(hf_config) for cfg in (hf_config, text_cfg): if getattr(cfg, "quantization_config", None) is not None: cfg.quantization_config = None diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 990f0910b1d..6c7940b47e5 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -580,10 +580,12 @@ def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares kimi_k3._force_load_remote_modeling(config_only, "/ckpt", False) -@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") -def test_prepare_lets_an_error_reading_the_quant_config_through(monkeypatch): +@pytest.mark.usefixtures("private_conversion_registry") +def test_prepare_lets_an_error_reading_the_quant_config_through(monkeypatch, remote_checkpoint): """compressed_tensors raises pydantic's ValidationError, a ValueError, for a quantization_config - it cannot parse. Dropped after that, the config would load the experts newly initialized.""" + it cannot parse. Dropped after that, the config would load the experts newly initialized. The + refusal comes before prepare patches the remote block or registers its converters, so a process + that goes on after it is unchanged.""" def install(cfg): raise ValueError("1 validation error for QuantizationConfig") @@ -595,6 +597,9 @@ def install(cfg): with pytest.raises(ValueError, match="validation error"): _prepare(config) + assert isinstance(remote_checkpoint.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + registered = cm.get_checkpoint_conversion_mapping("kimi_linear") or [] + assert not any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks", "remote_checkpoint") From 30847b20395bcf7b44963138bc0c7a73e1cf8904 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:32:50 +0000 Subject: [PATCH 17/32] Tell a compressed Kimi-K3 checkpoint by its own config too From Claude's review: prepare decided whether to install the decompressor from the config's quantization_config, which prepare itself drops, while the decompressor clears its mark at the load. A config prepared and loaded before, passed to a later load, so skipped the decompressor, and the fused experts loaded newly initialized from a checkpoint storing packed weights. The checkpoint's own config.json now counts too, at its top or text level: such a config reaches the install, which refuses it, since nothing would decompress its experts. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 25 ++++++++++- tests/unit/torch/models/test_kimi_k3_ep.py | 48 ++++++++++++++++++++-- 2 files changed, 67 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index d21d6d3059e..bc7b52d81bf 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -270,6 +270,22 @@ def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bo return get_class_from_dynamic_module(class_ref, ckpt_path, trust_remote_code=trust_remote_code) +def _checkpoint_quantized(ckpt_path: str) -> bool: + """Whether the checkpoint's own ``config.json`` has a ``quantization_config``, top or text level. + + The config ``prepare`` is handed cannot always say: ``prepare`` drops its quantization_config, + and the decompressor clears its mark at the load, so a config prepared and loaded before has + neither. + """ + from transformers import PretrainedConfig + + stored, _ = PretrainedConfig.get_config_dict(ckpt_path) + return any( + isinstance(level, dict) and level.get("quantization_config") is not None + for level in (stored, stored.get("text_config")) + ) + + def _register_kimi_k3_conversions() -> None: """Register the per-expert to fused-expert weight converters. @@ -369,8 +385,13 @@ def prepare( # else (attention, shared experts, latent projections, dense layer, embeddings, lm_head, vision) is # plain bf16. The decompressor is the step that can still refuse the checkpoint, so it comes before # anything below patches the process: a refusal leaves it unchanged. Its hooks rewrite only the - # converters of a config it marked. Plain bf16 experts come without a quantization_config. - if any(getattr(cfg, "quantization_config", None) is not None for cfg in (hf_config, text_cfg)): + # converters of a config it marked. Plain bf16 experts come without a quantization_config. The + # checkpoint's own config counts too: a config prepared and loaded before has lost both its + # quantization_config and the mark, and the install refuses it rather than let the experts load + # newly initialized. + if any( + getattr(cfg, "quantization_config", None) is not None for cfg in (hf_config, text_cfg) + ) or _checkpoint_quantized(ckpt_path): _onfly_decompress.install_onfly_decompress(hf_config) # This checkpoint's remote code, not whichever revision this process imported first. Every patch diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 6c7940b47e5..a037b535604 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -22,6 +22,7 @@ """ import copy +import json import sys import types from concurrent.futures import Future @@ -297,6 +298,14 @@ def style(name): # ------------------------------------------------------------------------------------- prepare _AUTO_MAP = {"AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration"} +_CHECKPOINT_QUANTIZED = kimi_k3._checkpoint_quantized + + +@pytest.fixture(autouse=True) +def unquantized_checkpoints_on_disk(monkeypatch): + """The checkpoints these tests name are stand-ins with no ``config.json``: prepare reads each as + an unquantized one. The test of that read puts the real one back.""" + monkeypatch.setattr(kimi_k3, "_checkpoint_quantized", lambda ckpt_path: False) @pytest.fixture @@ -602,13 +611,44 @@ def install(cfg): assert not any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) +@pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks") +@pytest.mark.parametrize("level", ["top", "text"]) +def test_a_config_that_lost_its_quantization_config_is_refused_by_the_checkpoints( + monkeypatch, remote_checkpoint, tmp_path, level +): + """prepare drops the config's quantization_config, and the decompressor clears its mark at the + load, so a config prepared and loaded before says nothing of the compression; loaded again, its + experts would load newly initialized. The checkpoint's own config.json says it, at either level, + and the decompressor refuses the config before prepare patches anything.""" + monkeypatch.setattr(kimi_k3, "_checkpoint_quantized", _CHECKPOINT_QUANTIZED) + scheme = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} + stored = ( + {"quantization_config": scheme} + if level == "top" + else {"text_config": {"quantization_config": scheme}} + ) + (tmp_path / "config.json").write_text(json.dumps(stored)) + config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP) + + with pytest.raises(ValueError, match="no compressed-tensors quantization_config"): + kimi_k3.prepare(config, ckpt_path=str(tmp_path), trust_remote_code=True) + assert isinstance(remote_checkpoint.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + + @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks", "remote_checkpoint") -def test_prepare_leaves_the_loader_alone_for_bf16_experts(): - """Without a compressed-tensors config the experts are plain bf16: prepare completes and leaves - transformers' conversion mapping unwrapped, with or without ``compressed_tensors`` installed.""" +def test_prepare_leaves_the_loader_alone_for_bf16_experts(monkeypatch, tmp_path): + """Without a compressed-tensors config, in the config or the checkpoint's, the experts are plain + bf16: prepare completes and leaves transformers' conversion mapping unwrapped, with or without + ``compressed_tensors`` installed.""" + monkeypatch.setattr(kimi_k3, "_checkpoint_quantized", _CHECKPOINT_QUANTIZED) + (tmp_path / "config.json").write_text(json.dumps({"text_config": {}})) stock = mu.get_model_conversion_mapping - _prepare(SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP)) + kimi_k3.prepare( + SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP), + ckpt_path=str(tmp_path), + trust_remote_code=True, + ) assert mu.get_model_conversion_mapping is stock From 01afeba6f59f7127c9228f3f35b18932e644d501 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:03:42 +0000 Subject: [PATCH 18/32] Pin that a Kimi-K3 refusal changes nothing, and that a returned config is refused From cjluo-nv's review: the refusal test now has the real decompressor refuse, for a scheme compressed-tensors cannot parse and for a dtype other than bf16, and checks that the remote block, the config, the converter registry, the registered experts implementations and the loader hooks are all as before. The lost-scheme test builds the config as a load returns it, and the 2-GPU Kimi-K3 job passes the returned model.config to another load of the checkpoint, which must refuse it. Signed-off-by: Shengliang Xu --- .../torch/utils/test_expert_parallel_load.py | 16 +++++- tests/unit/torch/models/test_kimi_k3_ep.py | 50 ++++++++++++++----- 2 files changed, 51 insertions(+), 15 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index ee240d2e6c8..a7a35147f2e 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -63,7 +63,7 @@ set_distribution_plan, ) 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 +from modelopt.torch.utils.plugins.model_load_utils import load_and_prepare_ep, load_pretrained NUM_EXPERTS = 8 @@ -617,6 +617,17 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): if [to_checkpoint(part) for part in parts] != leaves: problems.append(f"expert {e} exports as {[to_checkpoint(p) for p in parts]}") + # The returned model's config has lost its scheme and the decompressor's mark. Passed + # to another load of the checkpoint, it is refused, not read as bf16 experts. + try: + load_pretrained( + ckpt_dir, config=model.config, distribution=plan, trust_remote_code=True + ) + problems.append("a load with the returned model.config was not refused") + except ValueError as e: + if "no compressed-tensors quantization_config" not in str(e): + problems.append(f"a load with the returned model.config failed otherwise: {e}") + gathered = [None] * size dist.all_gather_object(gathered, problems) finally: @@ -632,7 +643,8 @@ def test_kimi_k3_packed_experts_load_calibrate_and_name_back_expert_parallel( its model spec: each rank decompresses only its own experts, the sharded model matches the whole one loaded on a single GPU, calibration over every rank's batch matches the whole model's, and the export names the experts back to the checkpoint's ``w1``/``w3``/``w2``. Each rank's - checks are gathered before any rank fails, so a regression fails instead of hanging.""" + checks are gathered before any rank fails, so a regression fails instead of hanging. Passed to + another load, the returned model's config is refused rather than read as bf16 experts.""" pytest.importorskip( "compressed_tensors", reason="compressed-tensors is needed: the checkpoint's experts are compressed", diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index a037b535604..df3c7d31dc0 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -589,26 +589,44 @@ def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares kimi_k3._force_load_remote_modeling(config_only, "/ckpt", False) -@pytest.mark.usefixtures("private_conversion_registry") -def test_prepare_lets_an_error_reading_the_quant_config_through(monkeypatch, remote_checkpoint): +@pytest.mark.usefixtures( + "private_conversion_registry", "restore_loader_hooks", "torch_a2a_unregistered" +) +@pytest.mark.parametrize( + ("dtype", "error", "match"), + [ + (None, ValueError, "validation error"), + (torch.float16, NotImplementedError, "decompress to bfloat16"), + ], + ids=["unparsable-scheme", "not-bf16"], +) +def test_prepare_lets_an_error_reading_the_quant_config_through( + remote_checkpoint, dtype, error, match +): """compressed_tensors raises pydantic's ValidationError, a ValueError, for a quantization_config - it cannot parse. Dropped after that, the config would load the experts newly initialized. The - refusal comes before prepare patches the remote block or registers its converters, so a process - that goes on after it is unchanged.""" - - def install(cfg): - raise ValueError("1 validation error for QuantizationConfig") - - monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", install) + it cannot parse, here one with no config groups, and the experts decompress to bf16 only. + Dropped after that, the config would load the experts newly initialized. The refusal comes + before prepare changes anything, so a process that goes on after it is unchanged: the remote + block, the config, the converter registry, the experts implementations and the loader hooks.""" + pytest.importorskip("compressed_tensors") + scheme = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} config = SimpleNamespace( - text_config=SimpleNamespace(), quantization_config={"format": "new"}, auto_map=_AUTO_MAP + text_config=SimpleNamespace(quantization_config=scheme), + quantization_config=scheme, + auto_map=_AUTO_MAP, + dtype=dtype, ) + before = copy.deepcopy(config) + hooks = (mu.get_model_conversion_mapping, cml.spawn_materialize) - with pytest.raises(ValueError, match="validation error"): + with pytest.raises(error, match=match): _prepare(config) assert isinstance(remote_checkpoint.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + assert config == before registered = cm.get_checkpoint_conversion_mapping("kimi_linear") or [] assert not any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) + assert "torch_a2a" not in type(ALL_EXPERTS_FUNCTIONS)._global_mapping + assert (mu.get_model_conversion_mapping, cml.spawn_materialize) == hooks @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks") @@ -628,7 +646,13 @@ def test_a_config_that_lost_its_quantization_config_is_refused_by_the_checkpoint else {"text_config": {"quantization_config": scheme}} ) (tmp_path / "config.json").write_text(json.dumps(stored)) - config = SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP) + # What the load returns as model.config: prepared, with its scheme dropped, and the mark cleared. + text = SimpleNamespace( + quantization_config=None, + base_model_ep_plan=dict(kimi_k3.BASE_MODEL_EP_PLAN), + _experts_implementation="torch_a2a", + ) + config = SimpleNamespace(text_config=text, quantization_config=None, auto_map=_AUTO_MAP) with pytest.raises(ValueError, match="no compressed-tensors quantization_config"): kimi_k3.prepare(config, ckpt_path=str(tmp_path), trust_remote_code=True) From b7e0fc5abaab1381069f9ff80d75d932c48e6e8a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:40:14 +0000 Subject: [PATCH 19/32] Make Kimi-K3's code importable for the plan decision, and prepare it for an EP plan alone #2748 has the plan decision build the model the checkpoint's own code builds, and run a spec's prepare only once the plan is expert-parallel. Kimi-K3's spec now names its import step (the OutputRecorder its code imports from where transformers no longer keeps it, and its imports staged), which patches no class, config, converter or loader hook, and puts its stated plan on text_config, where the expert count is. Its prepare, and so its torch_a2a registration, runs only on the expert-parallel route. The note on patches outliving their load narrows to a later load in the same process after an expert-parallel one. The end-to-end test runs the real Kimi-K3 prepare through the decision: FSDP2-sharded, or left whole and refused, the remote MoE block, the converter registry, the experts implementations, the loader and the config as read are untouched; an expert-parallel plan fuses the experts and registers torch_a2a. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 28 +++++--- modelopt/torch/models/kimi_k3/specs.py | 4 +- tests/unit/torch/models/test_kimi_k3_ep.py | 81 +++++++++++++++++++++- 3 files changed, 100 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index bc7b52d81bf..3931f999d2a 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -248,16 +248,26 @@ def _forward(self, hidden_states): print_rank_0("[kimi_k3] patched KimiSparseMoeBlock: fused experts + ExpertsInterface forward") +def make_importable() -> None: + """What Kimi-K3's own code needs to import on this transformers, and nothing more. + + The ``OutputRecorder`` it imports from where transformers no longer keeps it, and its relative + imports staged. It patches no class, config, converter or loader hook, so the plan decision runs + it before ``prepare``, whatever plan comes out. Idempotent. + """ + _install_output_recorder_shim() + stage_full_remote_closure() + + def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bool) -> type: """Import the remote modeling now, so its classes can be patched before the model is built. Force the remote modeling into ``sys.modules`` (so we can monkey-patch its classes BEFORE ``from_pretrained`` instantiates the model) by resolving the model class the checkpoint's config - declares in its ``auto_map`` -- no hardcoded remote file/module name. Runs the OutputRecorder + closure - shims first. Returns that class, whose package holds this checkpoint's remote code. + declares in its ``auto_map`` -- no hardcoded remote file/module name. Makes the code importable + first. Returns that class, whose package holds this checkpoint's remote code. """ - _install_output_recorder_shim() - stage_full_remote_closure() + make_importable() from transformers.dynamic_module_utils import get_class_from_dynamic_module auto_map = getattr(hf_config, "auto_map", None) or {} @@ -358,9 +368,9 @@ def prepare( inserts/calibrates/folds the fp8/nvfp4 quantizers on those bf16 params. That is ~4x the packed HBM (the memory bottleneck), but PTQ needs the bf16 weights. - The patches are process-wide and stay after the load: a later load of the same checkpoint in this - process, through a class the sharded loads refuse (such as ``AutoModel``), runs on the patched - code. The plan decision itself never makes such a load. + Only an expert-parallel plan runs it, but its patches are process-wide and stay after that load: + a later load of the same checkpoint in the same process, after an expert-parallel load, runs on + the patched code. """ # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe # imports the registry. @@ -414,8 +424,8 @@ def prepare( # grouped_gemm shards the bf16 fused params; the experts container goes through moe_tp_experts # (-> torch_a2a_experts). text_cfg.base_model_ep_plan = dict(BASE_MODEL_EP_PLAN) - # transformers refuses to build a model whose experts implementation it has not registered, and - # the plan decision builds this one on the meta device before the expert-parallel load runs. + # transformers refuses to build a model whose experts implementation it has not registered: on the + # expert-parallel route, the only one that runs prepare, before anything builds this config. register_torch_a2a() text_cfg._experts_implementation = "torch_a2a" diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index df584b6b1c3..7c5fcfd5e78 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -20,7 +20,7 @@ """ from ..specs import EPSpec, ModelSpec, register -from .ep import BASE_MODEL_EP_PLAN, prepare +from .ep import BASE_MODEL_EP_PLAN, make_importable, prepare __all__: list[str] = [] @@ -31,6 +31,8 @@ ep_spec=EPSpec( prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="text_config", + make_importable=make_importable, # prepare's patches to the checkpoint's remote code were validated on this release. requires_transformers="5.13", ), diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index df3c7d31dc0..23e04c8e3f8 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -51,8 +51,10 @@ from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper from modelopt.torch.models import _onfly_decompress, get_spec from modelopt.torch.models.kimi_k3 import ep as kimi_k3 +from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import moe from modelopt.torch.utils.moe import check_distributed_support +from modelopt.torch.utils.plugins import model_load_utils ROUTES = torch.tensor([[0, 3], [1, 2], [3, 1], [2, 0], [1, 3]]) # top-2 experts of 5 tokens @@ -374,9 +376,9 @@ def torch_a2a_unregistered(monkeypatch): "private_conversion_registry", "remote_checkpoint", "torch_a2a_unregistered" ) def test_a_model_builds_from_the_prepared_config_before_the_expert_parallel_load(): - """The plan decision builds the model from prepare's config on the meta device before the - expert-parallel load registers torch_a2a, and transformers refuses to build a model whose experts - implementation it has not registered: prepare registers the one it names.""" + """transformers refuses to build a model whose experts implementation it has not registered, so + prepare, which only the expert-parallel route runs, registers the one it names: a model built + from its config before the expert-parallel load registers torch_a2a builds.""" text = SimpleNamespace() _prepare(SimpleNamespace(text_config=text, auto_map=_AUTO_MAP)) @@ -659,6 +661,79 @@ def test_a_config_that_lost_its_quantization_config_is_refused_by_the_checkpoint assert isinstance(remote_checkpoint.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) +def _built_by_its_own_code(): + """Kimi-K3 as its own code builds it, on the meta device: each MoE layer of its language model + holds its 8 routed experts as per-expert modules, at + ``language_model.model.layers..block_sparse_moe.experts``.""" + + def layer(): + moe_block = nn.Module() + moe_block.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(8)) + return nn.ModuleDict({"block_sparse_moe": moe_block}) + + model = nn.Module() + model.language_model = nn.Module() + model.language_model.model = nn.Module() + model.language_model.model.layers = nn.ModuleList(layer() for _ in range(2)) + model.tie_weights = lambda: None + return model + + +@pytest.mark.usefixtures( + "private_conversion_registry", "restore_loader_hooks", "torch_a2a_unregistered" +) +@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_kimi_k3s_code( + monkeypatch, remote_checkpoint, tmp_path, ranks, dense, outcome +): + """Deciding the plan builds Kimi-K3 as its own code builds it, importable but unpatched, and + finds its 8 experts through the plan the spec states. Only an expert-parallel plan (4 ranks) + then runs ``prepare``. FSDP2-sharded, or left whole on 3 ranks, which no EP degree of the 8 + experts divides (refused), the run keeps the remote MoE block, the converter registry, the + experts implementations, the loader and the config as read.""" + as_read = SimpleNamespace( + model_type="kimi_k3", text_config=SimpleNamespace(num_experts=8), auto_map=_AUTO_MAP + ) + 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") + expert_parallel = outcome == "expert-parallel" + assert hasattr(hf_generic, "OutputRecorder") # importable on every plan + fused = not isinstance(remote_checkpoint.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) + assert fused == expert_parallel + registered = cm.get_checkpoint_conversion_mapping("kimi_linear") or [] + fusion = any("block_sparse_moe.experts.*.w1.weight" in c.source_patterns for c in registered) + assert fusion == expert_parallel + assert ("torch_a2a" in type(ALL_EXPERTS_FUNCTIONS)._global_mapping) == expert_parallel + if expert_parallel: + assert config.text_config.base_model_ep_plan == kimi_k3.BASE_MODEL_EP_PLAN + assert config.text_config._experts_implementation == "torch_a2a" + else: + assert (mu.get_model_conversion_mapping, cml.spawn_materialize) == hooks + if outcome == "fsdp2": + assert vars(config.text_config) == {"num_experts": 8} + + @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks", "remote_checkpoint") def test_prepare_leaves_the_loader_alone_for_bf16_experts(monkeypatch, tmp_path): """Without a compressed-tensors config, in the config or the checkpoint's, the experts are plain From 7c7c208c57ae956ece95eca108dcabee66d8f294 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:57:40 +0000 Subject: [PATCH 20/32] Count a packed Kimi-K3 checkpoint as one the expert-parallel load takes The plan decision now runs before prepare installs the decompressor, and reads a packed Kimi-K3 checkpoint's compressed-tensors format on the config as read. Kimi-K3's spec names that format among those its prepare decompresses, as #2748 lets a spec state, so a decided run of a packed Kimi-K3 is still planned expert-parallel. The end-to-end test gains a packed row, which fails without it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/specs.py | 1 + tests/unit/torch/models/test_kimi_k3_ep.py | 22 +++++++++++++++++----- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index 7c5fcfd5e78..50a7e0801f1 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -33,6 +33,7 @@ prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", make_importable=make_importable, + decompressed_formats=("compressed-tensors",), # prepare's patches to the checkpoint's remote code were validated on this release. requires_transformers="5.13", ), diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/test_kimi_k3_ep.py index 23e04c8e3f8..8c83c251d2c 100644 --- a/tests/unit/torch/models/test_kimi_k3_ep.py +++ b/tests/unit/torch/models/test_kimi_k3_ep.py @@ -683,21 +683,32 @@ def layer(): "private_conversion_registry", "restore_loader_hooks", "torch_a2a_unregistered" ) @pytest.mark.parametrize( - ("ranks", "dense", "outcome"), - [(3, "fsdp", "fsdp2"), (3, "replicate", "refused"), (4, "replicate", "expert-parallel")], - ids=["fsdp2", "whole", "expert-parallel"], + ("ranks", "dense", "outcome", "packed"), + [ + (3, "fsdp", "fsdp2", False), + (3, "replicate", "refused", False), + (4, "replicate", "expert-parallel", False), + (4, "replicate", "expert-parallel", True), + ], + ids=["fsdp2", "whole", "expert-parallel", "expert-parallel-packed"], ) def test_only_an_expert_parallel_plan_patches_kimi_k3s_code( - monkeypatch, remote_checkpoint, tmp_path, ranks, dense, outcome + monkeypatch, remote_checkpoint, tmp_path, ranks, dense, outcome, packed ): """Deciding the plan builds Kimi-K3 as its own code builds it, importable but unpatched, and finds its 8 experts through the plan the spec states. Only an expert-parallel plan (4 ranks) then runs ``prepare``. FSDP2-sharded, or left whole on 3 ranks, which no EP degree of the 8 experts divides (refused), the run keeps the remote MoE block, the converter registry, the - experts implementations, the loader and the config as read.""" + experts implementations, the loader and the config as read. A packed checkpoint is planned + expert-parallel too, before ``prepare`` installs the decompressor its spec names.""" as_read = SimpleNamespace( model_type="kimi_k3", text_config=SimpleNamespace(num_experts=8), auto_map=_AUTO_MAP ) + installs = [] + if packed: + scheme = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} + as_read.quantization_config = as_read.text_config.quantization_config = scheme + monkeypatch.setattr(_onfly_decompress, "install_onfly_decompress", installs.append) monkeypatch.setattr(dist, "size", lambda group=None: ranks) monkeypatch.setattr( transformers.AutoConfig, "from_pretrained", lambda *a, **k: copy.deepcopy(as_read) @@ -732,6 +743,7 @@ def test_only_an_expert_parallel_plan_patches_kimi_k3s_code( assert (mu.get_model_conversion_mapping, cml.spawn_materialize) == hooks if outcome == "fsdp2": assert vars(config.text_config) == {"num_experts": 8} + assert len(installs) == packed @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks", "remote_checkpoint") From 721064806cd576d7599e9a8cf5051fb7ad4a05ef Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:17:15 +0000 Subject: [PATCH 21/32] Mirror the Kimi-K3 adapter's module in its tests From cjluo-nv's review: as CONTRIBUTING asks, the Kimi-K3 adapter's tests move to tests/unit/torch/models/kimi_k3/test_ep.py, mirroring kimi_k3/ep.py, as the other Kimi adapters' do. The GPU integration case keeps its place in the expert-parallel load's file. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/kimi_k3/__init__.py | 15 +++++++++++++++ .../{test_kimi_k3_ep.py => kimi_k3/test_ep.py} | 0 2 files changed, 15 insertions(+) create mode 100644 tests/unit/torch/models/kimi_k3/__init__.py rename tests/unit/torch/models/{test_kimi_k3_ep.py => kimi_k3/test_ep.py} (100%) diff --git a/tests/unit/torch/models/kimi_k3/__init__.py b/tests/unit/torch/models/kimi_k3/__init__.py new file mode 100644 index 00000000000..1275d78dff5 --- /dev/null +++ b/tests/unit/torch/models/kimi_k3/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + diff --git a/tests/unit/torch/models/test_kimi_k3_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py similarity index 100% rename from tests/unit/torch/models/test_kimi_k3_ep.py rename to tests/unit/torch/models/kimi_k3/test_ep.py From 36e572515c317ad23021ac8761c7834e798eb829 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:37:53 +0000 Subject: [PATCH 22/32] Build Kimi-K3's code before the plan decision, refuse a load its prepare spoiled - The plan decision builds the checkpoint's own code before prepare, and Kimi-K3's VLM class takes none of the keyword arguments post_init passes tie_weights. Its spec's make_buildable now also loads the code and adapts tie_weights, as #2748's Kimi steps do. - Kimi-K3's prepare swaps its remote MoE block's class for good, so its spec says prepare_patches_every_load: once an expert-parallel load of it ran, a whole-model or FSDP2 load in the process is refused before anything is built (cjluo-nv's review). The note on patches outliving their load says so. The unit test fails without the refusal, and the 2-GPU job now runs the review's exact continuation, a whole-model load through AutoModel with the returned model.config, which every rank must refuse. - From Claude's review, _checkpoint_quantized reads every nested config spelling the load reads, llm_config too, and the lost-scheme test gains that level. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 41 ++++++++++++++----- modelopt/torch/models/kimi_k3/specs.py | 5 ++- modelopt/torch/models/specs.py | 6 +++ modelopt/torch/utils/moe.py | 22 ++++++++++ .../torch/utils/plugins/model_load_utils.py | 3 ++ .../torch/utils/test_expert_parallel_load.py | 18 ++++++++ tests/unit/torch/models/kimi_k3/test_ep.py | 21 +++++++--- .../unit/torch/utils/test_model_load_utils.py | 33 +++++++++++++++ 8 files changed, 130 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 3931f999d2a..a177d2e0b0b 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -248,17 +248,33 @@ def _forward(self, hidden_states): print_rank_0("[kimi_k3] patched KimiSparseMoeBlock: fused experts + ExpertsInterface forward") -def make_importable() -> None: - """What Kimi-K3's own code needs to import on this transformers, and nothing more. +def _make_importable() -> None: + """Make Kimi-K3's code import: the ``OutputRecorder`` it imports, and its imports staged. - The ``OutputRecorder`` it imports from where transformers no longer keeps it, and its relative - imports staged. It patches no class, config, converter or loader hook, so the plan decision runs - it before ``prepare``, whatever plan comes out. Idempotent. + transformers no longer keeps ``OutputRecorder`` where the code imports it from, and stages only + the direct relative imports. """ _install_output_recorder_shim() stage_full_remote_closure() +def make_buildable( + hf_config, *, ckpt_path: str | None = None, trust_remote_code: bool = False, **_ +) -> None: + """What Kimi-K3's own code needs to import and build on this transformers, and nothing more. + + Its imports (:func:`_make_importable`), and its 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. Code this load + may not run has nothing to build. Idempotent. + """ + _make_importable() + if ckpt_path is None or not trust_remote_code: + return + package = _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code).__module__ + _shim_tie_weights_kwargs(package.rpartition(".")[0]) + + def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bool) -> type: """Import the remote modeling now, so its classes can be patched before the model is built. @@ -267,7 +283,7 @@ def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bo declares in its ``auto_map`` -- no hardcoded remote file/module name. Makes the code importable first. Returns that class, whose package holds this checkpoint's remote code. """ - make_importable() + _make_importable() from transformers.dynamic_module_utils import get_class_from_dynamic_module auto_map = getattr(hf_config, "auto_map", None) or {} @@ -281,7 +297,7 @@ def _force_load_remote_modeling(hf_config, ckpt_path: str, trust_remote_code: bo def _checkpoint_quantized(ckpt_path: str) -> bool: - """Whether the checkpoint's own ``config.json`` has a ``quantization_config``, top or text level. + """Whether the checkpoint's own ``config.json`` has a ``quantization_config``, at any level. The config ``prepare`` is handed cannot always say: ``prepare`` drops its quantization_config, and the decompressor clears its mark at the load, so a config prepared and loaded before has @@ -289,10 +305,13 @@ def _checkpoint_quantized(ckpt_path: str) -> bool: """ from transformers import PretrainedConfig + # Not at module scope: the model registry imports this module, and moe imports the registry. + from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS + stored, _ = PretrainedConfig.get_config_dict(ckpt_path) return any( isinstance(level, dict) and level.get("quantization_config") is not None - for level in (stored, stored.get("text_config")) + for level in (stored, *(stored.get(attr) for attr in _SUB_CONFIG_ATTRS)) ) @@ -368,9 +387,9 @@ def prepare( inserts/calibrates/folds the fp8/nvfp4 quantizers on those bf16 params. That is ~4x the packed HBM (the memory bottleneck), but PTQ needs the bf16 weights. - Only an expert-parallel plan runs it, but its patches are process-wide and stay after that load: - a later load of the same checkpoint in the same process, after an expert-parallel load, runs on - the patched code. + Only an expert-parallel plan runs it, but its patches are process-wide and stay after that load, + which the spec states (``prepare_patches_every_load``): a later whole-model or FSDP2-sharded load + of Kimi-K3 in the same process is refused before anything is built. """ # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe # imports the registry. diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index 50a7e0801f1..7d651e0bb05 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -20,7 +20,7 @@ """ from ..specs import EPSpec, ModelSpec, register -from .ep import BASE_MODEL_EP_PLAN, make_importable, prepare +from .ep import BASE_MODEL_EP_PLAN, make_buildable, prepare __all__: list[str] = [] @@ -32,7 +32,8 @@ prepare=prepare, prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", - make_importable=make_importable, + make_buildable=make_buildable, + prepare_patches_every_load=True, decompressed_formats=("compressed-tensors",), # prepare's patches to the checkpoint's remote code were validated on this release. requires_transformers="5.13", diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 253dae6fa4c..a307ef778f9 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -230,6 +230,12 @@ class EPSpec(SpecSection): 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.""" + prepare_patches_every_load: bool = False + """Whether ``prepare`` patches the model's code for every later load of it in the process, not + only for the expert-parallel load it ran for (Kimi-K3's swaps its remote MoE block's class). + Once such a ``prepare`` has run, a whole-model or FSDP2-sharded load of the model in the process + is refused before anything is built.""" + 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 diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 72e8a41d36c..6f8aa5048b1 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -220,6 +220,26 @@ def _spec_refuses_ep(model_type: str) -> bool: return False +# The model types whose spec's prepare ran in this process and patched every later load of them. +_PATCHED_EVERY_LOAD: set[str] = set() + + +def check_not_patched_by_prepare(model_type: str) -> None: + """Refuse a load other than an expert-parallel one of a model whose ``prepare`` patched it for good. + + A spec whose ``prepare`` patches every later load of its model + (:attr:`~modelopt.torch.models.EPSpec.prepare_patches_every_load`) leaves code only an + expert-parallel load builds from: once it ran in this process, a whole-model or FSDP2-sharded + load of the model is refused before anything is built. + """ + if model_type in _PATCHED_EVERY_LOAD: + raise NotImplementedError( + f"An expert-parallel load of {model_type} in this process patched its modeling code " + "for every later load of it, which a whole-model or FSDP2-sharded load cannot build " + "from. Load it in a new process." + ) + + def apply_distributed_support( model_type: str, hf_config, @@ -289,6 +309,8 @@ def apply_distributed_support( if prepare and ep.prepare is not None: hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) + if ep.prepare_patches_every_load: + _PATCHED_EVERY_LOAD.add(model_type) 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 4eb82feac63..2a5f5a6f14b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -82,6 +82,7 @@ _spec_refuses_ep, apply_distributed_support, check_distributed_support, + check_not_patched_by_prepare, config_levels, register_torch_a2a, ) @@ -995,6 +996,8 @@ def handed() -> DistributionPlan | None: model_type = getattr(chosen, "model_type", None) or ( _read_config(ckpt_path, auto_class, trust_remote_code, **hub_options).model_type ) + if route != _EXPERT_PARALLEL: + check_not_patched_by_prepare(model_type) load_kwargs = ( {} if route is None diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index a7a35147f2e..4975de811dd 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -45,6 +45,7 @@ from torch.distributed.tensor import DTensor from transformers import ( AutoConfig, + AutoModel, AutoModelForCausalLM, DeepseekV3Config, DeepseekV3ForCausalLM, @@ -55,6 +56,7 @@ from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper from modelopt.torch.models import _onfly_decompress from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.utils import moe from modelopt.torch.utils.distributed import ( DistributionPlan, ParallelState, @@ -627,11 +629,27 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): except ValueError as e: if "no compressed-tensors quantization_config" not in str(e): problems.append(f"a load with the returned model.config failed otherwise: {e}") + # Nor may a whole-model load of it in this process, through a class the sharded loads + # refuse: prepare patched Kimi-K3's MoE block for every later load, so it is refused + # before anything is built. + try: + load_pretrained( + ckpt_dir, + auto_class=AutoModel, + config=model.config, + distribution=DistributionPlan(size, 1, "replicate"), + trust_remote_code=True, + ) + problems.append("a whole-model load after the expert-parallel one was not refused") + except NotImplementedError as e: + if "patched its modeling code" not in str(e): + problems.append(f"a whole-model load after it failed otherwise: {e}") gathered = [None] * size dist.all_gather_object(gathered, problems) finally: mu.get_model_conversion_mapping, cml.spawn_materialize = hooks + moe._PATCHED_EVERY_LOAD.discard("kimi_k3") # the workers outlive the job found = [problem for rank_problems in gathered for problem in rank_problems] assert not found, "\n".join(found[:20]) diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index 8c83c251d2c..ac2265162a9 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -303,6 +303,12 @@ def style(name): _CHECKPOINT_QUANTIZED = kimi_k3._checkpoint_quantized +@pytest.fixture(autouse=True) +def no_load_patched_for_good(monkeypatch): + """Each test starts in a process no expert-parallel Kimi-K3 load has patched, and leaves none.""" + monkeypatch.setattr(moe, "_PATCHED_EVERY_LOAD", set()) + + @pytest.fixture(autouse=True) def unquantized_checkpoints_on_disk(monkeypatch): """The checkpoints these tests name are stand-ins with no ``config.json``: prepare reads each as @@ -632,7 +638,7 @@ def test_prepare_lets_an_error_reading_the_quant_config_through( @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks") -@pytest.mark.parametrize("level", ["top", "text"]) +@pytest.mark.parametrize("level", ["top", "text", "llm"]) def test_a_config_that_lost_its_quantization_config_is_refused_by_the_checkpoints( monkeypatch, remote_checkpoint, tmp_path, level ): @@ -642,11 +648,11 @@ def test_a_config_that_lost_its_quantization_config_is_refused_by_the_checkpoint and the decompressor refuses the config before prepare patches anything.""" monkeypatch.setattr(kimi_k3, "_checkpoint_quantized", _CHECKPOINT_QUANTIZED) scheme = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} - stored = ( - {"quantization_config": scheme} - if level == "top" - else {"text_config": {"quantization_config": scheme}} - ) + stored = { + "top": {"quantization_config": scheme}, + "text": {"text_config": {"quantization_config": scheme}}, + "llm": {"llm_config": {"quantization_config": scheme}}, + }[level] (tmp_path / "config.json").write_text(json.dumps(stored)) # What the load returns as model.config: prepared, with its scheme dropped, and the mark cleared. text = SimpleNamespace( @@ -730,6 +736,9 @@ def test_only_an_expert_parallel_plan_patches_kimi_k3s_code( assert plan.is_expert_parallel == (outcome == "expert-parallel") expert_parallel = outcome == "expert-parallel" assert hasattr(hf_generic, "OutputRecorder") # importable on every plan + tie = vars(remote_checkpoint.KimiK3ForConditionalGeneration)["tie_weights"] + assert getattr(tie, "_kimi_k3_kwargs_shim", False) # and buildable on every plan + assert ("kimi_k3" in moe._PATCHED_EVERY_LOAD) == expert_parallel fused = not isinstance(remote_checkpoint.KimiSparseMoeBlock(_config()).experts, nn.ModuleList) assert fused == expert_parallel registered = cm.get_checkpoint_conversion_mapping("kimi_linear") or [] diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 24513a89755..f5d2b8cdd12 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -895,6 +895,39 @@ def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( resolve_distribution_plan(str(tmp_path)) +@pytest.mark.skipif( + Version(transformers.__version__) < Version(_ep_floor("kimi_k3")), + reason="Kimi-K3's expert parallelism needs a later transformers", +) +@pytest.mark.parametrize( + ("distribution", "auto_class"), + [(DistributionPlan(2, 1, "replicate"), AutoModel), (DistributionPlan(2, dense="fsdp"), None)], + ids=["whole", "fsdp2"], +) +def test_a_load_after_a_prepare_that_patched_every_load_is_refused( + run_as, serve, loaders, monkeypatch, tmp_path, distribution, auto_class +): + """Kimi-K3's ``prepare`` patches its remote MoE block for every later load in the process, and + its spec says so. Once it ran, as it does for an expert-parallel load, a whole-model or + FSDP2-sharded load of Kimi-K3 in the process is refused before anything is built: it would build + the patched block and load the checkpoint's experts into it wrong.""" + monkeypatch.setattr(moe, "_PATCHED_EVERY_LOAD", set()) + monkeypatch.setattr( + get_spec("kimi_k3").ep_spec, "prepare", lambda hf_config, **kwargs: hf_config + ) + run_as(2) + config = serve(SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace())) + moe.apply_distributed_support( + "kimi_k3", copy.deepcopy(config), ckpt_path=str(tmp_path), trust_remote_code=True + ) + + with pytest.raises(NotImplementedError, match="patched its modeling code"): + load_pretrained( + str(tmp_path), auto_class, distribution=distribution, trust_remote_code=True + ) + assert loaders == {} + + @pytest.mark.parametrize("ep_degree", [None, "1"], ids=["auto", "set-to-1"]) def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead( run_as, serve, monkeypatch, tmp_path, ep_degree From 03f43e2237e3f2a2909d14ae9a795681ca6ed5a5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 01:12:05 +0000 Subject: [PATCH 23/32] Record Kimi-K3's lasting patch where it is made, and say what revision it reads From Claude's review of the refusal: - The record that a prepare patched every later load landed only once prepare returned, so a step after the MoE block's patch that raised left the process patched but unrecorded. Kimi-K3's prepare now calls moe.mark_patched_every_load right after that patch; apply_distributed_support still records a prepare that completes. test_a_prepare_that_raises_after_patching_still_records_it fails without the early mark. - _checkpoint_quantized is the one place prepare reads the checkpoint itself, and prepare is handed no hub options: its docstring now says a hub id is read at its default revision, a local checkpoint as it is. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 7 +++++-- modelopt/torch/utils/moe.py | 11 ++++++++++- tests/unit/torch/models/kimi_k3/test_ep.py | 15 +++++++++++++++ 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index a177d2e0b0b..789b2a64edf 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -301,7 +301,8 @@ def _checkpoint_quantized(ckpt_path: str) -> bool: The config ``prepare`` is handed cannot always say: ``prepare`` drops its quantization_config, and the decompressor clears its mark at the load, so a config prepared and loaded before has - neither. + neither. ``prepare`` is handed no hub options, so a hub id is read at its default revision; a + local checkpoint, the usual case, is read as it is. """ from transformers import PretrainedConfig @@ -393,7 +394,7 @@ def prepare( """ # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe # imports the registry. - from modelopt.torch.utils.moe import register_torch_a2a + from modelopt.torch.utils.moe import mark_patched_every_load, register_torch_a2a if ckpt_path is None: raise ValueError( @@ -433,6 +434,8 @@ def prepare( "kimi_k3: could not locate the loaded remote module defining KimiSparseMoeBlock to patch" ) _patch_kimi_moe_block(kmod) + # Recorded at the first patch that lasts, so a step below that raises still leaves the record. + mark_patched_every_load("kimi_k3") _patch_kda_a_log(kmod) _patch_kda_beta_sigmoid(kmod) _shim_create_causal_mask(kmod) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 6f8aa5048b1..ae5718c0578 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -224,6 +224,15 @@ def _spec_refuses_ep(model_type: str) -> bool: _PATCHED_EVERY_LOAD: set[str] = set() +def mark_patched_every_load(model_type: str) -> None: + """Record that a ``prepare`` patched ``model_type`` for every later load in this process. + + :func:`apply_distributed_support` records it once such a ``prepare`` returns; a ``prepare`` calls + this itself at its first patch that lasts, so a raise after it still leaves the record. + """ + _PATCHED_EVERY_LOAD.add(model_type) + + def check_not_patched_by_prepare(model_type: str) -> None: """Refuse a load other than an expert-parallel one of a model whose ``prepare`` patched it for good. @@ -310,7 +319,7 @@ def apply_distributed_support( if prepare and ep.prepare is not None: hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) if ep.prepare_patches_every_load: - _PATCHED_EVERY_LOAD.add(model_type) + mark_patched_every_load(model_type) elif not prepare and ep.make_buildable is not None: ep.make_buildable(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) return hf_config diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index ac2265162a9..bf3ce78d035 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -597,6 +597,21 @@ def test_the_remote_modeling_is_imported_through_the_class_the_auto_map_declares kimi_k3._force_load_remote_modeling(config_only, "/ckpt", False) +@pytest.mark.usefixtures("private_conversion_registry", "remote_checkpoint") +def test_a_prepare_that_raises_after_patching_still_records_it(monkeypatch): + """Kimi-K3's prepare records its lasting patch where it makes it, so a step after it that raises + still leaves the process refusing a load other than an expert-parallel one.""" + + def fails(kmod): + raise RuntimeError("no kernels to wrap") + + monkeypatch.setattr(kimi_k3, "_patch_kda_beta_sigmoid", fails) + + with pytest.raises(RuntimeError, match="no kernels to wrap"): + _prepare(SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP)) + assert "kimi_k3" in moe._PATCHED_EVERY_LOAD + + @pytest.mark.usefixtures( "private_conversion_registry", "restore_loader_hooks", "torch_a2a_unregistered" ) From 75a6781cb134b064c9c2c51ed718037081bb8246 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 03:48:13 +0000 Subject: [PATCH 24/32] Key Kimi-K3's early record on the spec's model type, and say what the refusal precedes From Claude's review: - The early record of prepare's lasting patch spelled kimi_k3 as a literal, which had to equal the spec's model_type. kimi_k3.ep.MODEL_TYPE is now the one definition the spec registers and the record keys on, and the test checks the record against the spec. - The refusal precedes the load's build, not the plan decision's throwaway meta build, which runs earlier: the docs now say a refused load is refused before it builds or loads the model, and that under torchrun an EP degree of 1 meets the decision's own refusal first. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 7 +++++-- modelopt/torch/models/kimi_k3/specs.py | 4 ++-- modelopt/torch/models/specs.py | 2 +- modelopt/torch/utils/moe.py | 3 ++- tests/unit/torch/models/kimi_k3/test_ep.py | 2 +- 5 files changed, 11 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 789b2a64edf..c1e6d03fef1 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -105,6 +105,9 @@ # `layers.*.block_sparse_moe` (NOT `.mlp`; layer 0 is dense `.mlp`). Fused expert params shard dim-0 via # grouped_gemm; the experts container gets the EP group via moe_tp_experts (-> torch_a2a_experts for a2a). # No `ep_router` entry: KimiMoEGate runs inline over the replicated router weight and emits GLOBAL ids. +# The model type the spec registers, and the early record of prepare's lasting patch keys on. +MODEL_TYPE = "kimi_k3" + BASE_MODEL_EP_PLAN = { "layers.*.block_sparse_moe.experts.gate_up_proj": "grouped_gemm", "layers.*.block_sparse_moe.experts.down_proj": "grouped_gemm", @@ -390,7 +393,7 @@ def prepare( Only an expert-parallel plan runs it, but its patches are process-wide and stay after that load, which the spec states (``prepare_patches_every_load``): a later whole-model or FSDP2-sharded load - of Kimi-K3 in the same process is refused before anything is built. + of Kimi-K3 in the same process is refused before it builds or loads the model. """ # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe # imports the registry. @@ -435,7 +438,7 @@ def prepare( ) _patch_kimi_moe_block(kmod) # Recorded at the first patch that lasts, so a step below that raises still leaves the record. - mark_patched_every_load("kimi_k3") + mark_patched_every_load(MODEL_TYPE) _patch_kda_a_log(kmod) _patch_kda_beta_sigmoid(kmod) _shim_create_causal_mask(kmod) diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index 7d651e0bb05..2800c166034 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -20,13 +20,13 @@ """ from ..specs import EPSpec, ModelSpec, register -from .ep import BASE_MODEL_EP_PLAN, make_buildable, prepare +from .ep import BASE_MODEL_EP_PLAN, MODEL_TYPE, make_buildable, prepare __all__: list[str] = [] register( ModelSpec( - model_type="kimi_k3", + model_type=MODEL_TYPE, modeling_source="remote_code", ep_spec=EPSpec( prepare=prepare, diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index a307ef778f9..52eae22ab49 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -234,7 +234,7 @@ class EPSpec(SpecSection): """Whether ``prepare`` patches the model's code for every later load of it in the process, not only for the expert-parallel load it ran for (Kimi-K3's swaps its remote MoE block's class). Once such a ``prepare`` has run, a whole-model or FSDP2-sharded load of the model in the process - is refused before anything is built.""" + is refused before it builds or loads one.""" decompressed_formats: tuple[str, ...] = () """The checkpoint formats (a ``quantization_config``'s ``quant_method``) whose experts diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index ae5718c0578..863bed2fa3c 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -239,7 +239,8 @@ def check_not_patched_by_prepare(model_type: str) -> None: A spec whose ``prepare`` patches every later load of its model (:attr:`~modelopt.torch.models.EPSpec.prepare_patches_every_load`) leaves code only an expert-parallel load builds from: once it ran in this process, a whole-model or FSDP2-sharded - load of the model is refused before anything is built. + load of the model is refused before it builds or loads one. The plan decision's throwaway meta + build runs earlier, so under torchrun an EP degree of 1 meets the decision's own refusal first. """ if model_type in _PATCHED_EVERY_LOAD: raise NotImplementedError( diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index bf3ce78d035..58bc48cc657 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -609,7 +609,7 @@ def fails(kmod): with pytest.raises(RuntimeError, match="no kernels to wrap"): _prepare(SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP)) - assert "kimi_k3" in moe._PATCHED_EVERY_LOAD + assert get_spec("kimi_k3").model_type in moe._PATCHED_EVERY_LOAD @pytest.mark.usefixtures( From 528e65c75dfd13d316279cd93f4831ddaca62d85 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 19:52:46 +0000 Subject: [PATCH 25/32] Skip the Kimi-K3 tests that read a compressed scheme below compressed-tensors 0.16 Signed-off-by: Shengliang Xu #3137 now refuses an older compressed-tensors, whose compressors take no scheme: the Kimi-K3 tests that parse or load a compressed checkpoint through the real library skip there, as the decompressor's own do. --- tests/gpu/torch/utils/test_expert_parallel_load.py | 1 + tests/unit/torch/models/kimi_k3/test_ep.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 4975de811dd..5892e2cc442 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -665,6 +665,7 @@ def test_kimi_k3_packed_experts_load_calibrate_and_name_back_expert_parallel( another load, the returned model's config is refused rather than read as bf16 experts.""" pytest.importorskip( "compressed_tensors", + minversion="0.16", reason="compressed-tensors is needed: the checkpoint's experts are compressed", ) if Version(transformers.__version__) < Version("5.16"): diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index 58bc48cc657..0914471780d 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -631,7 +631,7 @@ def test_prepare_lets_an_error_reading_the_quant_config_through( Dropped after that, the config would load the experts newly initialized. The refusal comes before prepare changes anything, so a process that goes on after it is unchanged: the remote block, the config, the converter registry, the experts implementations and the loader hooks.""" - pytest.importorskip("compressed_tensors") + pytest.importorskip("compressed_tensors", minversion="0.16") scheme = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} config = SimpleNamespace( text_config=SimpleNamespace(quantization_config=scheme), From 0b5eba3360a47e7a23375d931ae34cbd815ff8f7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 19:54:53 +0000 Subject: [PATCH 26/32] Compare the Kimi-K3 test's calibrated amaxes to one bf16 step An amax is a max, and the expert-parallel and whole models take it over the same bf16 activations: all that precedes Kimi-K3's one MoE layer is dense and replicated. The relative-norm tolerance of 2e-2 came from a test whose amaxes are summed over other orders; here one bf16 step (2**-7 of the value) covers the experts' down_proj input, which their grouped and eager GEMMs may round apart. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_expert_parallel_load.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 5892e2cc442..9607097ce1b 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -604,10 +604,14 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): if expected_amax is None: problems.append(f"{name} is calibrated here but not in the whole model") continue - # bf16 activations, summed in another order: compare by relative norm. - error = (amax - expected_amax).norm() / expected_amax.norm().clamp_min(1e-12) - if not error < 2e-2: - problems.append(f"{name}: amax off by {error.item():.3g} (relative)") + # An amax is a max, and both models take it over the same bf16 activations: all that + # precedes the experts is dense and replicated here. Allow one bf16 step (2**-7 of + # the value) for the experts' down_proj input, which their grouped and eager GEMMs + # may round apart. + if not torch.allclose(amax, expected_amax, rtol=2**-7, atol=0): + problems.append( + f"{name}: amax {amax.tolist()}, the whole model's {expected_amax.tolist()}" + ) # The export names each expert back to the checkpoint's w1/w3/w2. to_checkpoint = build_reverse_name_mapper(model, tensor_keys=True) From 6791db7ba2300d23eb3ab14f1d30ecd539977e0b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 22:30:22 +0000 Subject: [PATCH 27/32] Plan a Kimi-K3 checkpoint at an EP degree of 1 through its own code's real build The tiny Kimi-K3 remote modeling now has the no-keyword tie_weights the checkpoint ships, and a test plans a tiny Kimi-K3 checkpoint at ep_size=1 through the real meta build of that code: refused whole on 4 ranks, or planned FSDP2. Without the build step at that degree (#2748) the build fails on tie_weights(recompute_mapping=False). Signed-off-by: Shengliang Xu --- .../torch/kimi_k3_tiny_modeling.py | 3 + tests/unit/torch/models/kimi_k3/test_ep.py | 58 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/tests/_test_utils/torch/kimi_k3_tiny_modeling.py b/tests/_test_utils/torch/kimi_k3_tiny_modeling.py index c81990eb82e..6bdf1cca406 100644 --- a/tests/_test_utils/torch/kimi_k3_tiny_modeling.py +++ b/tests/_test_utils/torch/kimi_k3_tiny_modeling.py @@ -269,3 +269,6 @@ def __init__(self, config): def forward(self, input_ids=None, **kwargs): return self.language_model(input_ids=input_ids) + + def tie_weights(self): # as Kimi-K3's own: no keyword arguments + return self.language_model.tie_weights() diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index 0914471780d..2be3d11f619 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -23,6 +23,7 @@ import copy import json +import shutil import sys import types from concurrent.futures import Future @@ -30,6 +31,7 @@ import pytest import torch +from _test_utils.torch import kimi_k3_tiny_modeling from _test_utils.torch.kimi_k3_remote import remote_module from torch import nn @@ -770,6 +772,62 @@ def test_only_an_expert_parallel_plan_patches_kimi_k3s_code( assert len(installs) == packed +@pytest.mark.parametrize( + ("dense", "fsdp2"), [("auto", False), ("fsdp", True)], ids=["whole", "fsdp2"] +) +def test_an_ep_degree_of_1_builds_kimi_k3s_own_code_to_decide(monkeypatch, tmp_path, dense, fsdp2): + """With expert parallelism ruled out, deciding the plan still builds Kimi-K3's own code on the + meta device, its ``tie_weights`` without the keyword arguments transformers' ``post_init`` + passes: the build step runs all the same, and no ``prepare``. The decision reaches its outcome: + the 8 experts left whole on 4 ranks are refused, or the model is planned FSDP2-sharded.""" + _restore(monkeypatch, hf_generic, "OutputRecorder") + monkeypatch.setattr(dmu, "get_cached_module_file", dmu.get_cached_module_file) + monkeypatch.setattr( + get_spec("kimi_k3").ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") + ) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + shutil.copy(kimi_k3_tiny_modeling.__file__, tmp_path / "modeling_kimi_k3.py") + remote = "modeling_kimi_k3.KimiK3ForConditionalGeneration" + auto_map = {"AutoConfig": "modeling_kimi_k3.KimiK3Config", "AutoModelForCausalLM": remote} + text = { + "model_type": "kimi_linear", + "hidden_size": 16, + "intermediate_size": 16, + "routed_expert_hidden_size": 8, + "moe_intermediate_size": 8, + "num_experts": 8, + "num_experts_per_token": 2, + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "num_shared_experts": 1, + } + stored = {"model_type": "kimi_k3", "auto_map": auto_map, "text_config": text} + (tmp_path / "config.json").write_text(json.dumps(stored)) + below_floor = Version(transformers.__version__) < Version(moe._ep_floor("kimi_k3")) + + def decide(): + return model_load_utils.resolve_distribution_plan( + str(tmp_path), ep_size=1, dense=dense, trust_remote_code=True + ) + + loaded = set(sys.modules) + try: + if fsdp2: + plan, _ = decide() + assert (plan.ep_size, plan.dense) == (1, "fsdp") + else: + with pytest.raises( + RuntimeError if below_floor else ValueError, + match="can be expert-parallel, but the plan resolved", + ): + decide() + finally: + # The checkpoint's code stays imported otherwise, for later tests to find as Kimi-K3's. + for name in set(sys.modules) - loaded: + if name.startswith("transformers_modules."): + del sys.modules[name] + + @pytest.mark.usefixtures("private_conversion_registry", "restore_loader_hooks", "remote_checkpoint") def test_prepare_leaves_the_loader_alone_for_bf16_experts(monkeypatch, tmp_path): """Without a compressed-tensors config, in the config or the checkpoint's, the experts are plain From 1b4d445134777d7e8c846066e9fe614b959c5106 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 22:47:36 +0000 Subject: [PATCH 28/32] Build the GPU test's tiny Kimi-K3 with the build-step tie_weights shim The tiny Kimi-K3 remote modeling now has the no-keyword tie_weights Kimi-K3's own code ships, and the GPU test's checkpoint writer built that class directly to read its state-dict shapes: on transformers 5.18, post_init's tie_weights(recompute_mapping=False) raised. The writer now builds it with the adapter's tie_weights shim, patched onto the class for that build only, so the loads in the test still meet the unshimmed checkpoint code. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_expert_parallel_load.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 9607097ce1b..588b2bb48d4 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -27,6 +27,7 @@ import shutil from pathlib import Path from types import SimpleNamespace +from unittest import mock import pytest import torch @@ -55,6 +56,7 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper from modelopt.torch.models import _onfly_decompress +from modelopt.torch.models._kimi import _drop_tie_weights_kwargs from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils import moe from modelopt.torch.utils.distributed import ( @@ -497,8 +499,12 @@ def _save_tiny_kimi_k3(ckpt_dir): compressor, scheme, stored_as, _ = _onfly_decompress._read_compression( SimpleNamespace(quantization_config=KIMI_K3_MXFP4) ) - with torch.device("meta"): - shapes = kimi_k3_tiny_modeling.KimiK3ForConditionalGeneration( + # Its tie_weights takes no keyword arguments, as Kimi-K3's own does, and transformers passes some + # at construction: built with the adapter's build-step shim, kept to this build. + vlm = kimi_k3_tiny_modeling.KimiK3ForConditionalGeneration + shim = _drop_tie_weights_kwargs(vlm.tie_weights) + with torch.device("meta"), mock.patch.object(vlm, "tie_weights", shim): + shapes = vlm( kimi_k3_tiny_modeling.KimiK3Config(text_config=dict(KIMI_K3_DIMS)) ).state_dict() gen = torch.Generator().manual_seed(0) From 2f3ee4796a33f7be05b05c40ed9fb820fb14c162 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 23:06:57 +0000 Subject: [PATCH 29/32] Install nothing from Kimi-K3's build step when there is nothing to build Kimi-K3's make_buildable made the code importable, installing the OutputRecorder shim and the remote-code staging for the process, before it checked for a checkpoint path and trusted code, so a plan decision with neither still patched the process. Since the decision runs the build step at an EP degree of 1 too, the Kimi-K3 left-whole refusal test left the staging installed, and five Kimi prepare tests that compare it failed when they ran later in the same process. The build step now returns first, as the Kimi-VL and Kimi-K2.5 one does; the remote-module load it goes on to makes the code importable itself. A test pins that it installs nothing then, and the left-whole refusal test records the build step instead of running it, checking that the decision runs it once. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 7 ++++--- tests/unit/torch/models/kimi_k3/test_ep.py | 20 +++++++++++++++++++ .../unit/torch/utils/test_model_load_utils.py | 9 ++++++--- 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index c1e6d03fef1..257c411092d 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -268,12 +268,13 @@ def make_buildable( Its imports (:func:`_make_importable`), and its 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. Code this load - may not run has nothing to build. Idempotent. + loads, so the plan decision runs it before ``prepare``, whatever plan comes out. With no + checkpoint, or code this load may not run, there is nothing to build, and it installs nothing. + Idempotent. """ - _make_importable() if ckpt_path is None or not trust_remote_code: return + # Makes the code importable first. package = _force_load_remote_modeling(hf_config, ckpt_path, trust_remote_code).__module__ _shim_tie_weights_kwargs(package.rpartition(".")[0]) diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index 2be3d11f619..50e94d77e19 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -772,6 +772,26 @@ def test_only_an_expert_parallel_plan_patches_kimi_k3s_code( assert len(installs) == packed +@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 Kimi-K3's is imported or built, + so the build step installs no import shim and no staging either: a plan decision that runs it + leaves the process as it found it.""" + _restore(monkeypatch, hf_generic, "OutputRecorder") + monkeypatch.delattr(hf_generic, "OutputRecorder", raising=False) + staging = dmu.get_cached_module_file + monkeypatch.setattr(dmu, "get_cached_module_file", staging) + + kimi_k3.make_buildable(SimpleNamespace(auto_map=_AUTO_MAP), **kwargs) + + assert dmu.get_cached_module_file is staging + assert not hasattr(hf_generic, "OutputRecorder") + + @pytest.mark.parametrize( ("dense", "fsdp2"), [("auto", False), ("fsdp", True)], ids=["whole", "fsdp2"] ) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index f5d2b8cdd12..550dd983b53 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -971,9 +971,11 @@ def test_a_remote_kimi_k3_whose_prepare_attaches_its_plan_left_whole_is_an_error 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("kimi_k3").ep_spec, "prepare", lambda *a, **k: pytest.fail("ran prepare") - ) + spec = get_spec("kimi_k3").ep_spec + monkeypatch.setattr(spec, "prepare", lambda *a, **k: pytest.fail("ran prepare")) + # The build step's own work is tested with Kimi-K3's code; here, that the decision runs it. + built = [] + monkeypatch.setattr(spec, "make_buildable", lambda hf_config, **kwargs: built.append(kwargs)) monkeypatch.setattr(model_load_utils, "build_meta_causal_lm", lambda *a, **k: _remote_kimi_k3()) run_as(4) serve(SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace(num_experts=8))) @@ -984,6 +986,7 @@ def test_a_remote_kimi_k3_whose_prepare_attaches_its_plan_left_whole_is_an_error match="can be expert-parallel, but the plan resolved", ): resolve_distribution_plan(str(tmp_path)) + assert len(built) == 1 def test_a_packed_kimi_k3_loaded_through_another_class_is_planned_without_its_spec( From dc8ac99b665222842e5d69d4d721a24de9045a49 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 23:46:16 +0000 Subject: [PATCH 30/32] Refuse a later kimi_linear load too, and pass the Kimi-K3 block what else it takes Kimi-K3's prepare registers its fused-expert converters under its backbone's model type, kimi_linear, for the whole process, but recorded only kimi_k3 as spoiled. A later whole-model load of a text-only kimi_linear checkpoint then built its own unpatched block, whose per-expert experts the converters fuse into parameters it does not have, so they loaded newly initialized with only a warning. EPSpec.prepare_patches_every_load now names the model types a prepare spoils; Kimi-K3's are kimi_k3 and kimi_linear, recorded where prepare first patches and again once it returns. The refusal says what loaded expert-parallel without naming the refused type as that load's. The patched block's __init__ passes on any argument beyond the config, as a revision adding a layer index would need. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 17 +++++++++---- modelopt/torch/models/kimi_k3/specs.py | 4 ++-- modelopt/torch/models/specs.py | 11 +++++---- modelopt/torch/utils/moe.py | 10 ++++---- .../torch/utils/test_expert_parallel_load.py | 3 ++- tests/unit/torch/models/kimi_k3/test_ep.py | 22 ++++++++++++++++- .../unit/torch/utils/test_model_load_utils.py | 24 ++++++++++++------- 7 files changed, 63 insertions(+), 28 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index 257c411092d..f76e0f9a652 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -107,6 +107,11 @@ # No `ep_router` entry: KimiMoEGate runs inline over the replicated router weight and emits GLOBAL ids. # The model type the spec registers, and the early record of prepare's lasting patch keys on. MODEL_TYPE = "kimi_k3" +# Its kimi_linear backbone's model type, which the fused-expert converters register under. +BACKBONE_MODEL_TYPE = "kimi_linear" +# The model types every later load of which prepare spoils: the converters apply to any +# kimi_linear checkpoint, a text-only one too, whose own MoE block was never patched. +PATCHED_EVERY_LOAD = (MODEL_TYPE, BACKBONE_MODEL_TYPE) BASE_MODEL_EP_PLAN = { "layers.*.block_sparse_moe.experts.gate_up_proj": "grouped_gemm", @@ -215,8 +220,8 @@ def _patch_kimi_moe_block(kmod) -> None: _orig_init = Block.__init__ _orig_forward = Block.forward - def _init(self, config): - _orig_init(self, config) + def _init(self, config, *args, **kwargs): + _orig_init(self, config, *args, **kwargs) # Replace the ModuleList of per-expert MLPs with the fused grouped-param module (EP-capable), as # wide as the space the block runs its experts in. The block decides latent MoE itself (Kimi-K3's # config has no use_latent_moe field) and the forward below follows it, so read it off the block. @@ -365,7 +370,7 @@ def _mapping(): # and the inner submodule is de-duplicated (ancestor already claimed the model_type). Registering the # CLASS names too would apply it a second time at `language_model.model.` scope -> the same expert keys # converted twice -> conflict at load. - register_checkpoint_conversion_mapping("kimi_linear", _mapping(), overwrite=True) + register_checkpoint_conversion_mapping(BACKBONE_MODEL_TYPE, _mapping(), overwrite=True) print_rank_0( "[kimi_k3] registered per-expert w1/w3/w2 -> gate_up_proj/down_proj fusion for kimi_linear " "(custom-code conversion; tf>=5.13 skips these otherwise)" @@ -394,7 +399,8 @@ def prepare( Only an expert-parallel plan runs it, but its patches are process-wide and stay after that load, which the spec states (``prepare_patches_every_load``): a later whole-model or FSDP2-sharded load - of Kimi-K3 in the same process is refused before it builds or loads the model. + of Kimi-K3, or of any kimi_linear checkpoint, in the same process is refused before it builds or + loads the model. """ # Not at module scope: the model registry imports this module, and modelopt.torch.utils.moe # imports the registry. @@ -439,7 +445,8 @@ def prepare( ) _patch_kimi_moe_block(kmod) # Recorded at the first patch that lasts, so a step below that raises still leaves the record. - mark_patched_every_load(MODEL_TYPE) + for patched in PATCHED_EVERY_LOAD: + mark_patched_every_load(patched) _patch_kda_a_log(kmod) _patch_kda_beta_sigmoid(kmod) _shim_create_causal_mask(kmod) diff --git a/modelopt/torch/models/kimi_k3/specs.py b/modelopt/torch/models/kimi_k3/specs.py index 2800c166034..ea8e9ef676e 100644 --- a/modelopt/torch/models/kimi_k3/specs.py +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -20,7 +20,7 @@ """ from ..specs import EPSpec, ModelSpec, register -from .ep import BASE_MODEL_EP_PLAN, MODEL_TYPE, make_buildable, prepare +from .ep import BASE_MODEL_EP_PLAN, MODEL_TYPE, PATCHED_EVERY_LOAD, make_buildable, prepare __all__: list[str] = [] @@ -33,7 +33,7 @@ prepared_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="text_config", make_buildable=make_buildable, - prepare_patches_every_load=True, + prepare_patches_every_load=PATCHED_EVERY_LOAD, decompressed_formats=("compressed-tensors",), # prepare's patches to the checkpoint's remote code were validated on this release. requires_transformers="5.13", diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 52eae22ab49..2d42b3c816c 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -230,11 +230,12 @@ class EPSpec(SpecSection): 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.""" - prepare_patches_every_load: bool = False - """Whether ``prepare`` patches the model's code for every later load of it in the process, not - only for the expert-parallel load it ran for (Kimi-K3's swaps its remote MoE block's class). - Once such a ``prepare`` has run, a whole-model or FSDP2-sharded load of the model in the process - is refused before it builds or loads one.""" + prepare_patches_every_load: tuple[str, ...] = () + """The model types ``prepare`` patches every later load of in the process, not only the + expert-parallel load it ran for: the model's own, and any other its patches reach (Kimi-K3's + swaps its remote MoE block's class, and registers weight converters under its ``kimi_linear`` + backbone's type). Once such a ``prepare`` has run, a whole-model or FSDP2-sharded load of any of + them in the process is refused before it builds or loads one.""" decompressed_formats: tuple[str, ...] = () """The checkpoint formats (a ``quantization_config``'s ``quant_method``) whose experts diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 863bed2fa3c..db631f7f080 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -244,9 +244,9 @@ def check_not_patched_by_prepare(model_type: str) -> None: """ if model_type in _PATCHED_EVERY_LOAD: raise NotImplementedError( - f"An expert-parallel load of {model_type} in this process patched its modeling code " - "for every later load of it, which a whole-model or FSDP2-sharded load cannot build " - "from. Load it in a new process." + f"An expert-parallel load in this process, of {model_type} or of a model built on it, " + "patched its modeling code for every later load of it, which a whole-model or " + "FSDP2-sharded load cannot build from. Load it in a new process." ) @@ -319,8 +319,8 @@ def apply_distributed_support( if prepare and ep.prepare is not None: hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) - if ep.prepare_patches_every_load: - mark_patched_every_load(model_type) + for patched in ep.prepare_patches_every_load: + mark_patched_every_load(patched) elif not prepare and ep.make_buildable is not None: ep.make_buildable(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) return hf_config diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 588b2bb48d4..33f271e091c 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -659,7 +659,8 @@ def _kimi_k3_expert_parallel_job(rank, size, ckpt_dir): dist.all_gather_object(gathered, problems) finally: mu.get_model_conversion_mapping, cml.spawn_materialize = hooks - moe._PATCHED_EVERY_LOAD.discard("kimi_k3") # the workers outlive the job + # The workers outlive the job: prepare recorded Kimi-K3 and its kimi_linear backbone. + moe._PATCHED_EVERY_LOAD.difference_update({"kimi_k3", "kimi_linear"}) found = [problem for rank_problems in gathered for problem in rank_problems] assert not found, "\n".join(found[:20]) diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index 50e94d77e19..cf09b3df70e 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -611,7 +611,27 @@ def fails(kmod): with pytest.raises(RuntimeError, match="no kernels to wrap"): _prepare(SimpleNamespace(text_config=SimpleNamespace(), auto_map=_AUTO_MAP)) - assert get_spec("kimi_k3").model_type in moe._PATCHED_EVERY_LOAD + # Its backbone's too: the fused-expert converters apply to any kimi_linear load. + assert {get_spec("kimi_k3").model_type, "kimi_linear"} <= moe._PATCHED_EVERY_LOAD + + +def test_the_patched_block_takes_what_else_the_remote_one_takes(): + """A revision of the code whose MoE block takes more than its config, a layer index as many + DeepSeek-V3-style blocks do, builds through the patch as it does without it.""" + remote = remote_module("transformers_modules.stub_layer_idx.modeling_kimi_k3") + own_init = remote.KimiSparseMoeBlock.__init__ + + def init(self, config, layer_idx=None): + own_init(self, config) + self.layer_idx = layer_idx + + remote.KimiSparseMoeBlock.__init__ = init + kimi_k3._patch_kimi_moe_block(remote) + + block = remote.KimiSparseMoeBlock(_config(), layer_idx=3) + + assert block.layer_idx == 3 + assert not isinstance(block.experts, nn.ModuleList) @pytest.mark.usefixtures( diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 550dd983b53..fd2fea9e760 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -900,28 +900,34 @@ def test_a_remote_moe_whose_prepare_attaches_its_plan_left_whole_is_an_error( reason="Kimi-K3's expert parallelism needs a later transformers", ) @pytest.mark.parametrize( - ("distribution", "auto_class"), - [(DistributionPlan(2, 1, "replicate"), AutoModel), (DistributionPlan(2, dense="fsdp"), None)], - ids=["whole", "fsdp2"], + ("model_type", "distribution", "auto_class"), + [ + ("kimi_k3", DistributionPlan(2, 1, "replicate"), AutoModel), + ("kimi_k3", DistributionPlan(2, dense="fsdp"), None), + ("kimi_linear", DistributionPlan(2, 1, "replicate"), AutoModel), + ], + ids=["whole", "fsdp2", "its-backbone-whole"], ) def test_a_load_after_a_prepare_that_patched_every_load_is_refused( - run_as, serve, loaders, monkeypatch, tmp_path, distribution, auto_class + run_as, serve, loaders, monkeypatch, tmp_path, model_type, distribution, auto_class ): """Kimi-K3's ``prepare`` patches its remote MoE block for every later load in the process, and - its spec says so. Once it ran, as it does for an expert-parallel load, a whole-model or - FSDP2-sharded load of Kimi-K3 in the process is refused before anything is built: it would build - the patched block and load the checkpoint's experts into it wrong.""" + registers its fused-expert converters under its backbone's type, kimi_linear, and its spec says + so. Once it ran, as it does for an expert-parallel load, a whole-model or FSDP2-sharded load of + Kimi-K3 in the process is refused before anything is built: it would build the patched block and + load the checkpoint's experts into it wrong. So is one of a text-only kimi_linear checkpoint, + whose own unpatched block the converters would load newly initialized experts into.""" monkeypatch.setattr(moe, "_PATCHED_EVERY_LOAD", set()) monkeypatch.setattr( get_spec("kimi_k3").ep_spec, "prepare", lambda hf_config, **kwargs: hf_config ) run_as(2) - config = serve(SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace())) + config = serve(SimpleNamespace(model_type=model_type, text_config=SimpleNamespace())) moe.apply_distributed_support( "kimi_k3", copy.deepcopy(config), ckpt_path=str(tmp_path), trust_remote_code=True ) - with pytest.raises(NotImplementedError, match="patched its modeling code"): + with pytest.raises(NotImplementedError, match=r"patched its modeling code.*a new process"): load_pretrained( str(tmp_path), auto_class, distribution=distribution, trust_remote_code=True ) From 188533b7390ca74eaeb37457592b7aa92b9a4df4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 1 Nov 2026 00:11:37 +0000 Subject: [PATCH 31/32] Put the Kimi-K3 plan's comment back above the plan The model-type constants added above the plan left its comment, on the plan's base-relative globs, its block_sparse_moe experts and its missing ep_router entry, standing over MODEL_TYPE instead. Comment only. Signed-off-by: Shengliang Xu --- modelopt/torch/models/kimi_k3/ep.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/models/kimi_k3/ep.py b/modelopt/torch/models/kimi_k3/ep.py index f76e0f9a652..bad4c460f30 100644 --- a/modelopt/torch/models/kimi_k3/ep.py +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -101,10 +101,6 @@ _shim_tie_weights_kwargs, ) -# base-relative EP plan (bubbles to `language_model.model.` in the VLM). MoE lives at -# `layers.*.block_sparse_moe` (NOT `.mlp`; layer 0 is dense `.mlp`). Fused expert params shard dim-0 via -# grouped_gemm; the experts container gets the EP group via moe_tp_experts (-> torch_a2a_experts for a2a). -# No `ep_router` entry: KimiMoEGate runs inline over the replicated router weight and emits GLOBAL ids. # The model type the spec registers, and the early record of prepare's lasting patch keys on. MODEL_TYPE = "kimi_k3" # Its kimi_linear backbone's model type, which the fused-expert converters register under. @@ -113,6 +109,10 @@ # kimi_linear checkpoint, a text-only one too, whose own MoE block was never patched. PATCHED_EVERY_LOAD = (MODEL_TYPE, BACKBONE_MODEL_TYPE) +# base-relative EP plan (bubbles to `language_model.model.` in the VLM). MoE lives at +# `layers.*.block_sparse_moe` (NOT `.mlp`; layer 0 is dense `.mlp`). Fused expert params shard dim-0 via +# grouped_gemm; the experts container gets the EP group via moe_tp_experts (-> torch_a2a_experts for a2a). +# No `ep_router` entry: KimiMoEGate runs inline over the replicated router weight and emits GLOBAL ids. BASE_MODEL_EP_PLAN = { "layers.*.block_sparse_moe.experts.gate_up_proj": "grouped_gemm", "layers.*.block_sparse_moe.experts.down_proj": "grouped_gemm", From 140cd265399280004c53843504422fa31e9d79b0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 1 Nov 2026 08:06:26 +0000 Subject: [PATCH 32/32] Import the tiny Kimi-K3 modeling after the Kimi-K3 tests' transformers guard The tiny modeling imports transformers at module scope, so importing it above the guard made collecting the file without the optional [hf] extra fail on ModuleNotFoundError instead of skipping. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/kimi_k3/test_ep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/models/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py index cf09b3df70e..077f142e3ea 100644 --- a/tests/unit/torch/models/kimi_k3/test_ep.py +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -31,7 +31,6 @@ import pytest import torch -from _test_utils.torch import kimi_k3_tiny_modeling from _test_utils.torch.kimi_k3_remote import remote_module from torch import nn @@ -47,6 +46,7 @@ import transformers.dynamic_module_utils as dmu import transformers.modeling_utils as mu import transformers.utils.generic as hf_generic +from _test_utils.torch import kimi_k3_tiny_modeling from transformers.core_model_loading import build_glob_alternation from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS