Add Kimi-K3 DSpark draft trained with on-policy speculative distillation (step 5413)
3930efe verified | license: other | |
| base_model: moonshotai/Kimi-K3 | |
| base_model_relation: finetune | |
| tags: | |
| - speculative-decoding | |
| - draft-model | |
| - dspark | |
| - k3_dspark | |
| - mla | |
| - torchspec | |
| - kimi-k3 | |
| - safetensors | |
| pipeline_tag: text-generation | |
| # Kimi-K3-DSpark-SDDD | |
| A DSpark draft model for speculative decoding with [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3), | |
| trained with on-policy speculative distillation. | |
| The file layout, tensor names and `config.json` are identical to | |
| [Inferact/Kimi-K3-DSpark](https://huggingface.co/Inferact/Kimi-K3-DSpark), so this | |
| is a drop-in replacement for it — only the weights differ. | |
| ## Architecture | |
| 3,562,312,961 parameters in bfloat16 across 68 tensors: | |
| | Component | Shape | Note | | |
| | --- | --- | --- | | |
| | `embed_tokens.weight` | 163840 x 7168 | frozen, copied from Kimi-K3 | | |
| | `context_proj.weight` | 7168 x 35840 | fuses the 5 captured teacher layers | | |
| | 5 x decoder layer | — | MLA, `q_lora_rank=1536`, `kv_lora_rank=512` | | |
| | `markov_head.markov_w{1,2}` | 163840 x 256 | low-rank output head | | |
| | `confidence_head.proj` | 1 x 7424 | per-position confidence | | |
| The draft consumes auxiliary hidden states from teacher layers `[2, 23, 47, 71, 89]` | |
| (93 layers total) and predicts `block_size=7` positions ahead. There is no `lm_head`; | |
| `embed_tokens` is frozen and shared with the teacher, so only 2,387,907,841 | |
| parameters were actually trained. | |
| ## Training | |
| On-policy speculative distillation: the teacher generates its own continuations, | |
| and the draft is trained to predict the teacher's next tokens from the teacher's | |
| hidden states. This keeps the training distribution equal to the deployment | |
| distribution, which offline distillation on a fixed corpus does not. | |
| | | | | |
| | --- | --- | | |
| | Framework | LumenRL, batch-alternating scheduler | | |
| | Teacher backend | ATOM `AsyncLLMEngine`, tensor parallel 8 | | |
| | Hardware | 1 node, 8 x AMD Instinct MI350 | | |
| | Steps | 5413 | | |
| | Global batch / micro batch | 64 / 1 | | |
| | Learning rate | 7.5e-5 cosine to 7.5e-6, 4% warmup | | |
| | Gradient clipping | 0.5 | | |
| | Weight decay | 0 | | |
| | Precision | bfloat16 | | |
| | Rollouts | greedy, up to 992 new tokens, reasoning traces enabled | | |
| | Data | 476,904 conversations | | |
| Teacher rollouts and draft updates alternate in rounds of 50 batches: the teacher | |
| engine generates and captures hidden states, then releases the GPUs so the draft | |
| can train on the cached states. | |
| ## Results | |
| Simulated acceptance length over 7 speculative positions, measured on a held-out | |
| set of 32 prompts: | |
| | Step | Acceptance length | Eval loss | | |
| | --- | --- | --- | | |
| | 1700 | 1.1491 | — | | |
| | 5100 | 1.86586 | 2.41002 | | |
| | 5200 | 1.87052 | 2.39803 | | |
| | 5300 | 1.88224 | 2.39511 | | |
| | 5400 | **1.88704** | 2.37878 | | |
| Gains had flattened to roughly +0.005 per 100 steps by the end of the run, in step | |
| with the cosine schedule reaching its floor. | |
| Note on the metric: 32 prompts is a small evaluation set. Early in training the | |
| round-to-round spread was around 0.10, comparable to the signal; the late-training | |
| numbers above are far more stable but should still be read as approximate. | |
| ## Usage | |
| Load it exactly as you would `Inferact/Kimi-K3-DSpark` — same `model_type` | |
| (`k3_dspark`), same tensor names, same config keys. The draft requires a runtime | |
| that can expose the teacher's auxiliary hidden states at layers | |
| `[2, 23, 47, 71, 89]`. | |
| ## License | |
| Inherits the license of the Kimi-K3 base model. | |