From 646693bb8df0da8326caf5d66165819c90d3c9e4 Mon Sep 17 00:00:00 2001 From: fares Date: Fri, 17 Jul 2026 15:01:28 +0000 Subject: [PATCH] feat(models): add Kimi K2 + K2.7 (text-only) support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kimi K2 reuses DeepSeek-V3's architecture verbatim (MLA + MoE hyperparameters, weight-key naming, even auto_map points at transformers' own DeepseekV3ForCausalLM), so this subclasses DeepseekV3Config directly and adds a plain (non-sparse) MLA attention layer (layers/mla.py) rather than redefining ~20 config fields. Reuses GLM-4 MoE's existing HF<->prime conversion functions as-is since the MoE layout (per-expert gate/up/down, shared-experts, e_score_correction_bias) matches exactly. K2.7 (and K2.5/K2.6) ship as a multimodal wrapper with the same text backbone, language_model.-prefixed. Loading one strips that prefix and drops the vision tower entirely — text-only, no VLM support. Verified via scripts/mini_moe.py's HF<->prime roundtrip check (max logits diff 1e-6, using transformers' own DeepseekV3ForCausalLM as the reference — Kimi's own checkpoints use this exact class). No real Kimi checkpoint was available to validate against directly. vLLM already serves real Kimi checkpoints natively (architectures=DeepseekV3ForCausalLM routes into its existing DeepSeek-V3 support), so no inference-side changes are needed. Co-Authored-By: Claude Sonnet 5 --- docs/advanced.md | 3 + scripts/mini_moe.py | 38 +++ src/prime_rl/trainer/models/__init__.py | 3 + .../trainer/models/kimi_k2/__init__.py | 9 + .../models/kimi_k2/configuration_kimi_k2.py | 38 +++ .../models/kimi_k2/converting_kimi_k2.py | 43 +++ .../models/kimi_k2/modeling_kimi_k2.py | 300 ++++++++++++++++++ src/prime_rl/trainer/models/layers/mla.py | 226 +++++++++++++ 8 files changed, 660 insertions(+) create mode 100644 src/prime_rl/trainer/models/kimi_k2/__init__.py create mode 100644 src/prime_rl/trainer/models/kimi_k2/configuration_kimi_k2.py create mode 100644 src/prime_rl/trainer/models/kimi_k2/converting_kimi_k2.py create mode 100644 src/prime_rl/trainer/models/kimi_k2/modeling_kimi_k2.py create mode 100644 src/prime_rl/trainer/models/layers/mla.py diff --git a/docs/advanced.md b/docs/advanced.md index e07b401891..0d8796cc2c 100644 --- a/docs/advanced.md +++ b/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: diff --git a/scripts/mini_moe.py b/scripts/mini_moe.py index 0aca7cf6cf..a3206b8bd7 100644 --- a/scripts/mini_moe.py +++ b/scripts/mini_moe.py @@ -16,6 +16,7 @@ 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, @@ -23,6 +24,8 @@ 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, diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py index 14b5233182..49eee7e01a 100644 --- a/src/prime_rl/trainer/models/__init__.py +++ b/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) diff --git a/src/prime_rl/trainer/models/kimi_k2/__init__.py b/src/prime_rl/trainer/models/kimi_k2/__init__.py new file mode 100644 index 0000000000..b2a531e5e5 --- /dev/null +++ b/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", +] diff --git a/src/prime_rl/trainer/models/kimi_k2/configuration_kimi_k2.py b/src/prime_rl/trainer/models/kimi_k2/configuration_kimi_k2.py new file mode 100644 index 0000000000..4b1357182d --- /dev/null +++ b/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") diff --git a/src/prime_rl/trainer/models/kimi_k2/converting_kimi_k2.py b/src/prime_rl/trainer/models/kimi_k2/converting_kimi_k2.py new file mode 100644 index 0000000000..4aaec6f41a --- /dev/null +++ b/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", +] diff --git a/src/prime_rl/trainer/models/kimi_k2/modeling_kimi_k2.py b/src/prime_rl/trainer/models/kimi_k2/modeling_kimi_k2.py new file mode 100644 index 0000000000..eff864d781 --- /dev/null +++ b/src/prime_rl/trainer/models/kimi_k2/modeling_kimi_k2.py @@ -0,0 +1,300 @@ +import math +from typing import Optional, Union + +import torch +from torch import Tensor, nn +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, auto_docstring +from transformers.utils.deprecation import deprecate_kwarg + +from prime_rl.trainer.models.base import PreTrainedModelPrimeRL +from prime_rl.trainer.models.kimi_k2.configuration_kimi_k2 import KimiK2Config +from prime_rl.trainer.models.kimi_k2.converting_kimi_k2 import ( + convert_hf_layer_to_tt, + convert_hf_to_tt_moe, + convert_tt_layer_to_hf, + convert_tt_to_hf_moe, + strip_multimodal_wrapper, +) +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.mla import MLA_IMPL2CLASS, MLAConfig +from prime_rl.trainer.models.layers.mlp import MLP, MLPConfig +from prime_rl.trainer.models.layers.moe import MoE, MoEArgs +from prime_rl.trainer.models.layers.rms_norm import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig +from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class KimiK2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: KimiK2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + softmax_scale = qk_head_dim**-0.5 + if config.rope_scaling is not None: + mscale_all_dim = config.rope_scaling.get("mscale_all_dim", 0) + if mscale_all_dim: + mscale = yarn_get_mscale(config.rope_scaling["factor"], mscale_all_dim) + softmax_scale = softmax_scale * mscale * mscale + + mla_config = MLAConfig( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + attention_bias=config.attention_bias, + rms_norm_eps=config.rms_norm_eps, + softmax_scale=softmax_scale, + rope_interleave=config.rope_interleave, + ) + self.self_attn = MLA_IMPL2CLASS[config._attn_implementation](mla_config) + + moe_args = MoEArgs( + num_experts=config.n_routed_experts, + num_shared_experts=config.n_shared_experts, + score_func=config.scoring_func, + route_norm=config.norm_topk_prob, + route_scale=config.routed_scaling_factor, + score_before_experts=False, + top_k=config.num_experts_per_tok, + load_balance_coeff=1e-3, + use_grouped_mm=config.use_grouped_mm, + fp8=getattr(config, "fp8", False), + ) + mlp_config = MLPConfig( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + gate_act=config.hidden_act, + bias=False, + ) + + if layer_idx >= config.first_k_dense_replace: + self.mlp = MoE(moe_args, dim=config.hidden_size, hidden_dim=config.moe_intermediate_size) + else: + self.mlp = MLP(mlp_config) + + self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + self.post_attention_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + cu_seqlens: Optional[torch.LongTensor] = None, + max_seqlen: Optional[int] = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class KimiK2PreTrainedModel(PreTrainedModelPrimeRL): + config: KimiK2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["KimiK2DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": KimiK2DecoderLayer, + } + + def _init_weights(self, module): + super()._init_weights(module) + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.1.up_proj" in module_name for module_name in state_dict.keys()) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.w1" in module_name for module_name in state_dict.keys()) + + @classmethod + def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + convert_tt_to_hf_moe(state_dict) + return state_dict + + @classmethod + def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + strip_multimodal_wrapper(state_dict) # no-op for a plain (non-K2.7) checkpoint + convert_hf_to_tt_moe(state_dict) + return state_dict + + @classmethod + def convert_layer_to_hf(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_tt_layer_to_hf(state_dict, layer_idx) + return state_dict + + @classmethod + def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: + convert_hf_layer_to_tt(state_dict, layer_idx) + return state_dict + + +@auto_docstring +class KimiK2Model(KimiK2PreTrainedModel): + def __init__(self, config: KimiK2Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [KimiK2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + + if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): + rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + rope_type = "default" + rotary_config = RotaryEmbeddingConfig( + max_position_embeddings=config.max_position_embeddings, + rope_type=rope_type, + model_config=config, + ) + self.rotary_emb = RotaryEmbedding(rotary_config) + self.gradient_checkpointing = False + + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if self.config._attn_implementation in ("flash_attention_2", "flash_attention_3", "flash_attention_4"): + cu_seqlens, max_seqlen = get_cu_seqlens_from_position_ids(position_ids) + torch._dynamo.mark_dynamic(cu_seqlens, 0) + else: + max_seqlen = None + cu_seqlens = None + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states) + + +@auto_docstring +class KimiK2ForCausalLM(KimiK2PreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = KimiK2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.post_init() + + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + temperature: Optional[torch.Tensor] = None, + **kwargs: Unpack[TransformersKwargs], + ) -> PrimeLmOutput: + r""" + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices of input tokens in the KV cache. Accepted only for HuggingFace API + compatibility — prime-rl asserts `use_cache is None` since training does not + perform autoregressive decoding, so this argument is unused. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels used by PrimeRL's wrapped LM head to optionally compute per-token logprobs/entropy. + temperature (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Per-token temperatures for logprobs/entropy computation when `labels` are provided. + """ + assert use_cache is None, "use_cache is not supported for custom kimi_k2 for now" + assert past_key_values is None, "past_key_values is not supported for custom kimi_k2 for now" + + if position_ids is None: + if inputs_embeds is not None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + else: + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) + + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + return self.lm_head( + hidden_states[:, slice_indices, :], + labels[:, slice_indices] if labels is not None else None, + temperature=temperature, + ) + + def init_buffers_post_meta(self): + buffer_names = [name for name, _ in self.named_buffers()] + if "model.rotary_emb.inv_freq" in buffer_names: + rotary_emb = self.model.rotary_emb + inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn( + rotary_emb.config, rotary_emb.inv_freq.device + ) + rotary_emb.inv_freq.copy_(inv_freq) + + +__all__ = ["KimiK2Config", "KimiK2PreTrainedModel", "KimiK2Model", "KimiK2ForCausalLM"] diff --git a/src/prime_rl/trainer/models/layers/mla.py b/src/prime_rl/trainer/models/layers/mla.py new file mode 100644 index 0000000000..9791c953b5 --- /dev/null +++ b/src/prime_rl/trainer/models/layers/mla.py @@ -0,0 +1,226 @@ +"""Plain (dense, non-absorbed) Multi-head Latent Attention (MLA) — DeepSeek-V3's attention +mechanism, shared by any family built on it (Kimi K2 reuses it verbatim; DeepSeek-V3 itself +would too). No sparsity/indexer awareness — see `layers/dsa.py` for the DSA-capable sibling +used to convert a checkpoint built on this module into DeepSeek Sparse Attention. +""" + +import functools + +import torch +import torch.nn.functional as F +from torch import nn + +from .rms_norm import RMSNorm, RMSNormConfig +from .rotary_emb import apply_rotary_pos_emb, rotate_half + +try: + from flash_attn import flash_attn_varlen_func +except ImportError: + flash_attn_varlen_func = None # type: ignore + +try: + from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func +except ImportError: + flash_attn_3_varlen_func = None # type: ignore + +try: + from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func +except ImportError: + flash_attn_4_varlen_func = None # type: ignore + + +class MLAConfig: + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + q_lora_rank: int | None, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + attention_bias: bool, + rms_norm_eps: float, + softmax_scale: float, + rope_interleave: bool, + ): + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.attention_bias = attention_bias + self.rms_norm_eps = rms_norm_eps + self.softmax_scale = softmax_scale + self.rope_interleave = rope_interleave + + +def apply_rotary_pos_emb_interleaved(q, k, cos, sin, unsqueeze_dim=1): + """DeepSeek-V3-style interleaved RoPE: view as (d//2, 2) -> transpose -> standard rotate_half.""" + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class _MLAProjections(nn.Module): + """Shared parameter set for the MLA attention variants below (q/kv LoRA bottleneck).""" + + def __init__(self, config: MLAConfig): + super().__init__() + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.qk_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.softmax_scale = config.softmax_scale + self.rope_interleave = config.rope_interleave + + if self.q_lora_rank is None: + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias) + self.q_a_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.q_lora_rank, eps=config.rms_norm_eps)) + self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) + + self.kv_a_proj_with_mqa = nn.Linear( + config.hidden_size, config.kv_lora_rank + config.qk_rope_head_dim, bias=config.attention_bias + ) + self.kv_a_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.kv_lora_rank, eps=config.rms_norm_eps)) + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, self.num_heads * (config.qk_nope_head_dim + config.v_head_dim), bias=False + ) + + self.o_proj = nn.Linear(self.num_heads * config.v_head_dim, config.hidden_size, bias=config.attention_bias) + + def _project_qkv(self, hidden_states: torch.Tensor): + batch_size, seq_length = hidden_states.shape[:2] + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(batch_size, seq_length, self.num_heads, self.qk_head_dim).transpose(1, 2) + q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + kv_compressed, k_rope = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + + kv = self.kv_b_proj(self.kv_a_layernorm(kv_compressed)) + kv = kv.view(batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2) + k_nope, value_states = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + k_rope = k_rope.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + return q_nope, q_rope, k_nope, k_rope, value_states + + def _apply_rope(self, q_rope, k_rope, position_embeddings): + cos, sin = position_embeddings + if self.rope_interleave: + return apply_rotary_pos_emb_interleaved(q_rope, k_rope, cos, sin) + return apply_rotary_pos_emb(q_rope, k_rope, cos, sin) + + +class MLAFlashAttention(_MLAProjections): + """MLA via flash-attn's varlen kernels.""" + + _funcs = { + 2: flash_attn_varlen_func, + 3: flash_attn_3_varlen_func, + 4: flash_attn_4_varlen_func, + } + + def __init__(self, config: MLAConfig, flash_attn_version: int = 2): + super().__init__(config) + self._flash_attn_version = flash_attn_version + self.func = self._funcs[flash_attn_version] + self._flash_attn_call = self.func + if self._flash_attn_version == 4: + self._flash_attn_call = torch._dynamo.disable(self.func) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + q_nope, q_rope, k_nope, k_rope, value_states = self._project_qkv(hidden_states) + q_rope, k_rope = self._apply_rope(q_rope, k_rope, position_embeddings) + + k_rope = k_rope.expand(*k_nope.shape[:-1], -1) + query_states = torch.cat([q_nope, q_rope], dim=-1) + key_states = torch.cat([k_nope, k_rope], dim=-1) + + # Flash attention expects (batch, seq, heads, dim); we take [0] below for varlen. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + # Pad V to qk_head_dim (flash attention requires the same head dim for Q/K/V). + if self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + + args = [query_states[0], key_states[0], value_states[0], cu_seqlens, cu_seqlens] + if self._flash_attn_version != 4: + args.extend([max_seqlen, max_seqlen]) + + out = self._flash_attn_call(*args, causal=True, softmax_scale=self.softmax_scale) + if isinstance(out, tuple): + out = out[0] + + if self.qk_head_dim != self.v_head_dim: + out = out[..., : self.v_head_dim] # truncate padding back to v_head_dim + + out = out.contiguous() + attn_output = out.view(1, out.shape[0], -1) + attn_output = self.o_proj(attn_output) + return attn_output, None + + +class MLASDPAAttention(_MLAProjections): + """MLA via `F.scaled_dot_product_attention`.""" + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + q_nope, q_rope, k_nope, k_rope, value_states = self._project_qkv(hidden_states) + q_rope, k_rope = self._apply_rope(q_rope, k_rope, position_embeddings) + + k_rope = k_rope.expand(*k_nope.shape[:-1], -1) + query_states = torch.cat([q_nope, q_rope], dim=-1) + key_states = torch.cat([k_nope, k_rope], dim=-1) + + # SDPA supports a different V head dim natively — no padding needed. + out = F.scaled_dot_product_attention( + query_states, key_states, value_states, is_causal=True, scale=self.softmax_scale + ) + out = out.transpose(1, 2).contiguous() + attn_output = out.view(out.shape[0], out.shape[1], -1) + attn_output = self.o_proj(attn_output) + return attn_output, None + + +MLA_IMPL2CLASS = { + "flash_attention_2": functools.partial(MLAFlashAttention, flash_attn_version=2), + "sdpa": MLASDPAAttention, + "flash_attention_3": functools.partial(MLAFlashAttention, flash_attn_version=3), + "flash_attention_4": functools.partial(MLAFlashAttention, flash_attn_version=4), +} + +__all__ = ["MLAConfig", "MLA_IMPL2CLASS", "MLAFlashAttention", "MLASDPAAttention", "apply_rotary_pos_emb_interleaved"]