feat(speculative): add Kimi K3 EAGLE-3 training (#3286)
* feat: add Kimi K3 model support Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * ci: add Kimi K3 recipe metadata Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor: keep Kimi K3 changes model-scoped Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(kimi-k3): isolate fp32 kernel parameters Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * chore(kimi-k3): remove bundled license file Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(kimi-k3): address review feedback Avoid import-time transformer mutation, remove the dead tokenizer fallback, and cover K3 MoE router-weight placement against reference inference for both loop and grouped-mm expert backends. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * test(kimi-k3): cover chat effort and config resolution Accept the documented medium thinking effort with durable input validation, and prove both Kimi K3 model types resolve through local configs when the installed Transformers mapping has no builtin entry. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(kimi-k3): guard optional attention residual init Keep checkpoint-free initialization valid when attention residual mixing is disabled while preserving pruned pipeline-stage handling, with a CPU forward regression test for the no-residual configuration. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * chore(kimi-k3): add NVIDIA copyright notices Keep the upstream Moonshot attribution and license while marking NVIDIA modifications so the repository copyright check recognizes both files. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * refactor(kimi-k3): simplify decoder initialization Remove the identical MoE and dense branches around MLP weight initialization so the shared behavior is explicit. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * feat(speculative): add Kimi K3 EAGLE-3 training Kimi K3 interleaves KDA layers with NoPE MLA full-attention layers, which neither the dense EAGLE-3 draft nor the DeepSeek MLA draft (which always rotates the rope slice) can represent. Add a dedicated draft that mirrors one K3 MLA layer: Q-LoRA, the compressed KV latent, the unrotated rope slice, the sigmoid output gate, K3 RMSNorm and the SiTU MLP, with KDA, the routed experts and the attention-residual mixer left in the target. Everything EAGLE-3 specific is inherited from the DeepSeek MLA draft through a new _build_inner_model hook, so the TTT recurrence, the vocab-remap buffers and the trainer are reused unchanged. freeze_embeddings, which the recipe calls for every draft, is added to that shared MLA base. The frozen target is loaded through AutoModel's custom expert-parallel path (K3 has no HuggingFace implementation and its 896 routed experts do not fit on one GPU), and the combinations K3 cannot serve (non-colocated backends, context or tensor parallelism, sequence packing, parallel drafting) are rejected before the target is loaded. Signed-off-by: khazic <khazzz1c@gmail.com> * fix(speculative): gate Kimi K3 pipeline parallelism and default its dispatcher to torch Online EAGLE-3 supervision hooks one complete target forward, so a pipelined K3 target (whose layer ModuleDict holds only its own stage, keyed by global index) either raises a KeyError in the aux-layer capture or silently captures nothing. The README already said pp_size must stay 1; enforce it with the other gates, before the 93-layer target is loaded. The expert dispatcher defaulted to hybridep, which forces the DeepEP path even where that build is absent (the world_size==1 fallback hides it until the real multi-node run). Default to torch, the dispatcher K3's own SFT reference config runs expert parallelism with, and keep deepep/hybridep opt-in. Signed-off-by: khazic <khazzz1c@gmail.com> * fix(speculative): guard the Kimi K3 gate against a missing draft config The non-colocated EAGLE-3 target backends call _setup_online_target with draft_base_config=None, so reading .model_type unconditionally raised AttributeError before the backend dispatch. Signed-off-by: khazic <khazzz1c@gmail.com> * docs(speculative): document the fused tensor layouts in the K3 EAGLE-3 draft Signed-off-by: khazic <khazzz1c@gmail.com> * fix(eagle3): reshard Kimi K3 target after forward Signed-off-by: Huiying Li <willwin.lee@gmail.com> --------- Signed-off-by: HuiyingLi <willwin.lee@gmail.com> Signed-off-by: khazic <khazzz1c@gmail.com> Signed-off-by: Huiying Li <willwin.lee@gmail.com> Co-authored-by: HuiyingLi <willwin.lee@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
K
khazzz1c committed
c852b16ff97221bd97a408e224f011a817cbc8e9
Parent: 03c7e70
Committed by GitHub <noreply@github.com>
on 8/11/2026, 4:03:16 AM