Initial K3 snapshot: 0.5B KDA/MLA/MoE train path

Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
2026-08-25 14:43:17 +08:00
commit 584f7e9e73
140 changed files with 21592 additions and 0 deletions
@@ -0,0 +1,128 @@
+"""K3Config — Kimi K3 架构的小规模复现配置 (KDA + Gated MLA + Stable LatentMoE).
+
+对照 learning/kimi-k3-notes §尺寸速查 (真实 K3 → 本 toy 缩比):
+    hidden 7168 → 256;  L 93 → 4;  H=HV 96 → 8;  K=V 128 → 16;
+    MLA kv_lora 512 → 32, q_lora 1536 → 64, nope/v 128 → 16;
+    MoE ℓ=d/2=3584 → 128, 896/16 → 16/2, shared 2, d_ff 3072 → 96.
+
+Hybrid Attention (K3): 每 4 层 1 次 Gated MLA, 末层强制 MLA.
+
+Presets:
+    toy  — ~8M, 自训 8k SP, 本地过拟合
+    0.5b — ~482M, Qwen3 词表, 32–40GB bf16；默认 step 是冒烟，翻译前置用 --max-tokens
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Qwen3 config.json; train_k3 overrides with len(tokenizer).
+QWEN3_VOCAB_SIZE = 151936
+
+
+@dataclass
+class K3Config:
+    # 主干
+    hidden_size: int = 256
+    num_hidden_layers: int = 4
+    vocab_size: int = 8192               # toy: data/spm_4k; 0.5b: Qwen3
+    initializer_range: float = 0.02
+    norm_eps: float = 1e-6
+    tie_word_embeddings: bool = False
+    max_position_embeddings: int = 2048   # NoPE, 仅语义保留
+
+    # KDA (K3: H = HV = 96, 无 GVA)
+    num_heads: int = 8
+    head_dim: int = 16
+    chunk_size: int = 16
+    lower_bound: float | None = -5.0
+    use_gate_in_kernel: bool = True
+    use_qk_l2norm_in_kernel: bool = True
+    use_beta_sigmoid_in_kernel: bool = True
+
+    # Gated MLA (NoPE)
+    kv_lora_rank: int = 32
+    q_lora_rank: int = 64
+    qk_nope_head_dim: int = 16
+    v_head_dim: int = 16
+
+    # Stable LatentMoE
+    moe_latent_size: int = 128           # ℓ = d/2
+    n_routed: int = 16
+    top_k: int = 2
+    n_shared: int = 2
+    moe_d_ff: int = 96
+    situ_beta1: float = 4.0
+    situ_beta2: float = 25.0
+
+    kda_backend: str = "reference"
+
+    # Depth mixer. off = DecoderBlock residual; block matches K3.
+    attnres: str = "off"  # off | full | block
+    attnres_block_size: int | None = None  # DecoderBlocks / AttnRes block; None ≈ L/8
+    attnres_zero_init_queries: bool = True
+    attnres_final_aggregate: bool = True
+    gradient_checkpointing: bool = False
+
+    def __post_init__(self):
+        from ..layers.attn_res import validate_attnres
+
+        validate_attnres(self.attnres, self.attnres_block_size)
+
+    @classmethod
+    def preset(cls, name: str) -> K3Config:
+        if name == "toy":
+            return cls()
+        if name in {"0.5b", "500m"}:
+            # H * head_dim == hidden. Routed 16: LatentMoE still runs every expert.
+            # ~482M with tied Qwen3 embeddings. 6×(3 KDA + 1 MLA).
+            return cls(
+                hidden_size=768,
+                num_hidden_layers=24,
+                vocab_size=QWEN3_VOCAB_SIZE,
+                tie_word_embeddings=True,
+                max_position_embeddings=2048,
+                num_heads=12,
+                head_dim=64,
+                chunk_size=64,
+                kv_lora_rank=192,
+                q_lora_rank=512,
+                qk_nope_head_dim=64,
+                v_head_dim=64,
+                moe_latent_size=384,
+                n_routed=16,
+                top_k=2,
+                n_shared=2,
+                moe_d_ff=512,
+                # The pure-PyTorch reference is far too slow at this size.
+                kda_backend="triton",
+                gradient_checkpointing=True,
+            )
+        raise ValueError(f"unknown preset: {name}")
+
+    @property
+    def H(self) -> int:
+        return self.num_heads
+
+    @property
+    def HV(self) -> int:
+        return self.num_heads
+
+    @property
+    def K(self) -> int:
+        return self.head_dim
+
+    @property
+    def V(self) -> int:
+        return self.head_dim
+
+    def layer_types(self) -> list[str]:
+        """Hybrid pattern: 每 4 层 1 次 MLA (0-based 层 3,7,...), 末层强制 MLA."""
+        types = ["kda"] * self.num_hidden_layers
+        for i in range(self.num_hidden_layers):
+            if i % 4 == 3:
+                types[i] = "mla"
+        types[-1] = "mla"
+        return types
+
+    def layer_specs(self) -> list[tuple[str, str]]:
+        return [(kind, "moe") for kind in self.layer_types()]