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/__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..9f4e3797ccc --- /dev/null +++ b/modelopt/torch/models/kimi_k3/ep.py @@ -0,0 +1,473 @@ +# Adapted from https://github.com/huggingface/transformers/blob/a906d3c4b65095f2308b6a6a193e934d03b8eb5d/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 +# +# 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 modelopt.torch.utils import print_rank_0 + +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, +) + +# 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-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. 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. + 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). A block built before the patch keeps its per-expert experts, + and so its own forward. Idempotent. + """ + 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__ + _orig_forward = Block.forward + + 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. + 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): + 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) + 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 _make_importable() -> None: + """Make Kimi-K3's code import: the ``OutputRecorder`` it imports, and its imports staged. + + 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. With no + checkpoint, or code this load may not run, there is nothing to build, and it installs nothing. + Idempotent. + """ + 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]) + + +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. Makes the code importable + first. Returns that class, whose package holds this checkpoint's remote code. + """ + _make_importable() + 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 _checkpoint_quantized(ckpt_path: str) -> bool: + """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 + 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 + + # 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(attr) for attr in _SUB_CONFIG_ATTRS)) + ) + + +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 + + 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(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)" + ) + + +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: 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 + 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. + + 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, 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. + from modelopt.torch.utils.moe import mark_patched_every_load, 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 + + # 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. 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 + # 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) + # Recorded at the first patch that lasts, so a step below that raises still leaves the record. + 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) + _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: 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" + + # 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() + 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..ea8e9ef676e --- /dev/null +++ b/modelopt/torch/models/kimi_k3/specs.py @@ -0,0 +1,42 @@ +# 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, MODEL_TYPE, PATCHED_EVERY_LOAD, make_buildable, prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type=MODEL_TYPE, + modeling_source="remote_code", + ep_spec=EPSpec( + prepare=prepare, + prepared_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="text_config", + make_buildable=make_buildable, + 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 253dae6fa4c..06caa0eaf8d 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -230,6 +230,15 @@ 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: 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. They are recorded once ``prepare`` + returns, so a ``prepare`` that can still raise after a lasting patch calls + :func:`~modelopt.torch.utils.moe.mark_patched_every_load` at that patch, as Kimi-K3's does.""" + 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 c7c87adf8f3..ba7b70ffefe 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -220,6 +220,36 @@ 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 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. + + 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 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( + 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." + ) + + def apply_distributed_support( model_type: str, hf_config, @@ -289,6 +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) + for patched in ep.prepare_patches_every_load: + mark_patched_every_load(patched) # The load builds from the config prepare returns: a bad one would fail only there, far # from the spec at fault. if hf_config is None: @@ -574,8 +606,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/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/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/_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, 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..6bdf1cca406 --- /dev/null +++ b/tests/_test_utils/torch/kimi_k3_tiny_modeling.py @@ -0,0 +1,274 @@ +# 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() + # 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) + + +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) + + def tie_weights(self): # as Kimi-K3's own: no keyword arguments + return self.language_model.tie_weights() diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 3d50e5c3e5f..33f271e091c 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 @@ -41,24 +42,32 @@ 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, + AutoModel, AutoModelForCausalLM, DeepseekV3Config, DeepseekV3ForCausalLM, 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.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 ( DistributionPlan, + ParallelState, + distributed_forward_loop, distribution_plan, 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 @@ -399,3 +408,280 @@ 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) + ) + # 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) + 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.spawn_materialize) + 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 + # 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) + 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]}") + + # 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}") + # 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 + # 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]) + + +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. Passed to + 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"): + 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)) 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/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/kimi_k3/test_ep.py b/tests/unit/torch/models/kimi_k3/test_ep.py new file mode 100644 index 00000000000..077f142e3ea --- /dev/null +++ b/tests/unit/torch/models/kimi_k3/test_ep.py @@ -0,0 +1,892 @@ +# 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 json +import shutil +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 _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 + +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 + + +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 ((cml, "spawn_materialize"), (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``. 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) + _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) + 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.")} + 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)) + + +@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``.""" + 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.""" + 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"} +_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 + 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 +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(): + """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)) + + # 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_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)) + # 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( + "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, 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", minversion="0.16") + scheme = {"quant_method": "compressed-tensors", "format": "mxfp4-pack-quantized"} + config = SimpleNamespace( + 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(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") +@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 +): + """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 = { + "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( + 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) + 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", "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, 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. 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) + ) + 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 + 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 [] + 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} + 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"] +) +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 + 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 + + 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 + + +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 a28c4a620b1..c9fcdb942cc 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 @@ -908,6 +908,67 @@ 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( + ("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, model_type, distribution, auto_class +): + """Kimi-K3's ``prepare`` patches its remote MoE block for every later load in the process, and + 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", _prepares_as_its_spec_states("kimi_k3") + ) + run_as(2) + 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=r"patched its modeling code.*a new process"): + load_pretrained( + str(tmp_path), auto_class, distribution=distribution, trust_remote_code=True + ) + assert loaders == {} + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version(_ep_floor("kimi_k3")), + reason="Kimi-K3's expert parallelism needs a later transformers", +) +def test_a_prepare_that_patched_every_load_is_recorded_though_it_returns_no_config( + monkeypatch, tmp_path +): + """A ``prepare`` that patches every later load stays recorded when it returns no config to load: + the refusal of what it returned comes after the record, as its patches stay in the process, so a + whole-model or FSDP2-sharded load of the model is still refused after it.""" + monkeypatch.setattr(moe, "_PATCHED_EVERY_LOAD", set()) + spec = get_spec("kimi_k3").ep_spec + monkeypatch.setattr(spec, "prepare", lambda hf_config, **kwargs: None) + config = SimpleNamespace(model_type="kimi_k3", text_config=SimpleNamespace()) + + with pytest.raises(ValueError, match="prepare returned None"): + moe.apply_distributed_support( + "kimi_k3", config, ckpt_path=str(tmp_path), trust_remote_code=True + ) + assert set(spec.prepare_patches_every_load) <= moe._PATCHED_EVERY_LOAD + + @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 @@ -924,6 +985,83 @@ 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: a ValueError, or the + RuntimeError naming the release when this transformers is below what their spec needs.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + 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))) + below_floor = Version(transformers.__version__) < Version(_ep_floor("kimi_k3")) + + with pytest.raises( + RuntimeError if below_floor else ValueError, + 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( + 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