feat(models): add Kimi K2 + K2.7 (text-only) support

docs/advanced.md

@@ -35,9 +35,12 @@ impl = "custom"        # or "hf" to force the HF path
 | Trinity (AFMoE) | `arcee-ai/Trinity-Mini`, … | ✅ | ✅ |
 | GLM-4 / GLM-4.5 / INTELLECT-3 | `THUDM/GLM-4-9B-0414`, `zai-org/GLM-4.5`, `PrimeIntellect/INTELLECT-3`, … | ✅ | ✅ |
 | GPT-OSS (HF MoE) | `openai/gpt-oss-20b`, `openai/gpt-oss-120b` | ❌ | ✅ |
+| Kimi K2 (`kimi_k2`) | `moonshotai/Kimi-K2-Instruct`, … | ✅ | ❌ |
 
 The custom path enables you to set EP, CP, selective activation checkpointing, low-precision training (`[trainer.model.quantization]`), and faster MoE kernels (`moe_use_grouped_mm = true`, default). Forcing `impl = "hf"` is mostly useful when debugging — it's slower and disables most MoE-specific knobs.
 
+`kimi_k2` reuses DeepSeek-V3's architecture directly (`layers/mla.py`, dense — not sparse — Multi-head Latent Attention), since Kimi K2 is architecturally identical, down to the HF weight-key naming. It also covers Kimi K2.7's text backbone: point `model.name` at a K2.7 checkpoint and the `language_model.`-prefixed weights load with the vision tower dropped — there is no multimodal/VLM support for it.
+
 ### Low-precision training
 
 Set `[trainer.model.quantization]` to train dense linears and MoE expert GEMMs in low precision. Two backends are available via the `type` discriminator:

scripts/mini_moe.py

@@ -16,13 +16,16 @@
 
 import torch
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers import DeepseekV3ForCausalLM as HFDeepseekV3ForCausalLM
 from transformers import Glm4MoeForCausalLM as HFGlm4MoeForCausalLM
 from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import (
     Qwen3_5MoeForConditionalGeneration as HFQwen3_5MoeVLM,
 )
 
 from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig
 from prime_rl.trainer.models.glm4_moe import Glm4MoeForCausalLM as PrimeRLGlm4MoeForCausalLM
+from prime_rl.trainer.models.kimi_k2 import KimiK2Config
+from prime_rl.trainer.models.kimi_k2 import KimiK2ForCausalLM as PrimeRLKimiK2ForCausalLM
 from prime_rl.trainer.models.laguna import LagunaConfig
 from prime_rl.trainer.models.laguna import LagunaForCausalLM as PrimeRLLagunaForCausalLM
 from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head
@@ -185,6 +188,41 @@ def _qwen3_5_moe_vlm_config():
         "prime_model_class": PrimeRLLagunaForCausalLM,
         "tokenizer_source": "poolside/Laguna-XS.2",
     },
+    "kimi_k2": {
+        "config_class": KimiK2Config,
+        "config_kwargs": dict(
+            vocab_size=4096,
+            hidden_size=256,
+            intermediate_size=512,
+            num_hidden_layers=4,
+            num_attention_heads=8,
+            num_key_value_heads=8,
+            q_lora_rank=128,
+            kv_lora_rank=64,
+            qk_nope_head_dim=32,
+            qk_rope_head_dim=16,
+            v_head_dim=32,
+            hidden_act="silu",
+            max_position_embeddings=4096,
+            rms_norm_eps=1e-6,
+            rope_theta=50000,
+            rope_interleave=True,
+            attention_bias=False,
+            moe_intermediate_size=128,
+            n_routed_experts=8,
+            num_experts_per_tok=4,
+            n_shared_experts=1,
+            first_k_dense_replace=1,
+            norm_topk_prob=True,
+            routed_scaling_factor=2.827,
+            n_group=1,
+            topk_group=1,
+            use_grouped_mm=False,
+        ),
+        "hf_model_class": HFDeepseekV3ForCausalLM,
+        "prime_model_class": PrimeRLKimiK2ForCausalLM,
+        "tokenizer_source": "deepseek-ai/DeepSeek-V3",
+    },
     "qwen3_5_moe_vlm": {
         "config_fn": _qwen3_5_moe_vlm_config,
         "hf_model_class": HFQwen3_5MoeVLM,

src/prime_rl/trainer/models/__init__.py

@@ -15,6 +15,7 @@
 from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM
 from prime_rl.trainer.models.glm_moe_dsa import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
 from prime_rl.trainer.models.gpt_oss import GptOssConfig, GptOssForCausalLM
+from prime_rl.trainer.models.kimi_k2 import KimiK2Config, KimiK2ForCausalLM
 from prime_rl.trainer.models.laguna import LagunaConfig, LagunaForCausalLM
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, cast_float_and_contiguous
 from prime_rl.trainer.models.llama import LlamaForCausalLM
@@ -29,6 +30,7 @@
 AutoConfig.register("afmoe", AfmoeConfig, exist_ok=True)
 AutoConfig.register("glm4_moe", Glm4MoeConfig, exist_ok=True)
 AutoConfig.register("glm_moe_dsa", GlmMoeDsaConfig, exist_ok=True)
+AutoConfig.register("kimi_k2", KimiK2Config, exist_ok=True)
 AutoConfig.register("laguna", LagunaConfig, exist_ok=True)
 AutoConfig.register("minimax_m2", MiniMaxM2Config, exist_ok=True)
 AutoConfig.register("nemotron_h", NemotronHConfig, exist_ok=True)
@@ -43,6 +45,7 @@
 _CUSTOM_CAUSAL_LM_MAPPING.register(AfmoeConfig, AfmoeForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(Glm4MoeConfig, Glm4MoeForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(GlmMoeDsaConfig, GlmMoeDsaForCausalLM, exist_ok=True)
+_CUSTOM_CAUSAL_LM_MAPPING.register(KimiK2Config, KimiK2ForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(LagunaConfig, LagunaForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(MiniMaxM2Config, MiniMaxM2ForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(NemotronHConfig, NemotronHForCausalLM, exist_ok=True)

src/prime_rl/trainer/models/kimi_k2/__init__.py

@@ -0,0 +1,9 @@
+from prime_rl.trainer.models.kimi_k2.configuration_kimi_k2 import KimiK2Config
+from prime_rl.trainer.models.kimi_k2.modeling_kimi_k2 import KimiK2ForCausalLM, KimiK2Model, KimiK2PreTrainedModel
+
+__all__ = [
+    "KimiK2Config",
+    "KimiK2ForCausalLM",
+    "KimiK2Model",
+    "KimiK2PreTrainedModel",
+]

src/prime_rl/trainer/models/kimi_k2/configuration_kimi_k2.py

@@ -0,0 +1,38 @@
+import warnings
+
+from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
+
+
+class KimiK2Config(DeepseekV3Config):
+    r"""
+    Configuration class for Kimi K2 (and K2.7's text backbone).
+
+    Kimi K2 reuses DeepSeek-V3's architecture verbatim (same MLA/MoE hyperparameters, same
+    HF weight-key naming — Moonshot ships `transformers`' own `DeepseekV3ForCausalLM` as
+    `auto_map` target). This subclasses `DeepseekV3Config` directly rather than redefining
+    its ~20 MLA/MoE fields; only the three fields below are prime-rl/Kimi-specific.
+
+    K2.7 (and K2.5, K2.6) are released as a multimodal wrapper (`model_type="kimi_k25"`)
+    with a nested `text_config` of this same shape plus a vision tower; this config
+    describes only the text backbone — see `converting_kimi_k2.py` for how a K2.7
+    checkpoint's `language_model.`-prefixed weights get loaded (vision ignored entirely).
+
+    Args:
+        use_grouped_mm (`bool`, defaults to `True`):
+            Whether to use grouped matrix multiplication for MoE.
+        topk_method (`str`, defaults to `"noaux_tc"`):
+            MoE routing top-k method (bias-based load balancing, no auxiliary loss).
+        scoring_func (`str`, defaults to `"sigmoid"`):
+            Scoring function for the MoE router.
+    """
+
+    model_type = "kimi_k2"
+
+    def __init__(self, use_grouped_mm=True, topk_method="noaux_tc", scoring_func="sigmoid", **kwargs):
+        super().__init__(**kwargs)
+        self.use_grouped_mm = use_grouped_mm
+        self.topk_method = topk_method
+        self.scoring_func = scoring_func
+
+        if not self.use_grouped_mm:
+            warnings.warn("not using grouped mm for moe is very slow, should only be used for debugging")

src/prime_rl/trainer/models/kimi_k2/converting_kimi_k2.py

@@ -0,0 +1,43 @@
+"""Kimi K2 weight conversion. The MoE layout (per-expert gate/up/down + shared-experts +
+`e_score_correction_bias` load-balancing term) is byte-identical to what GLM-4 MoE already
+converts, since both follow the standard DeepSeek-style naming — reused as-is, no adaptation.
+"""
+
+from torch import Tensor
+
+from prime_rl.trainer.models.glm4_moe.converting_glm4_moe import (
+    convert_hf_layer_to_tt,
+    convert_hf_to_tt_moe,
+    convert_tt_layer_to_hf,
+    convert_tt_to_hf_moe,
+)
+
+_LANGUAGE_MODEL_PREFIX = "language_model."
+
+
+def strip_multimodal_wrapper(state_dict: dict[str, Tensor]) -> None:
+    """Drop everything but the text backbone from a K2.5/K2.6/K2.7-style multimodal
+    checkpoint, in-place.
+
+    Those checkpoints wrap this exact text backbone (`language_model.model.layers.N...`)
+    alongside a vision tower. We only support the text backbone (see `KimiK2Config`'s
+    docstring) — un-prefix the language-model keys and drop everything else (vision tower,
+    projector, etc.) whatever they're named, since we never need them. A no-op for a plain
+    (non-multimodal) Kimi-K2 checkpoint, which carries no `language_model.`-prefixed keys.
+    """
+    prefixed = {k: v for k, v in state_dict.items() if k.startswith(_LANGUAGE_MODEL_PREFIX)}
+    if not prefixed:
+        return  # plain Kimi-K2 checkpoint, nothing to strip
+
+    state_dict.clear()
+    for key, value in prefixed.items():
+        state_dict[key[len(_LANGUAGE_MODEL_PREFIX) :]] = value
+
+
+__all__ = [
+    "convert_hf_layer_to_tt",
+    "convert_hf_to_tt_moe",
+    "convert_tt_layer_to_hf",
+    "convert_tt_to_hf_moe",
+    "strip_multimodal_wrapper",
+]