K3_draft / README.md
slippedJim's picture
Add Kimi-K3 DSpark draft trained with on-policy speculative distillation (step 5413)
3930efe verified
|
Raw
History Blame Contribute Delete
3.39 kB
---
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.