| --- |
| license: mit |
| base_model: Qwen/Qwen3-4B |
| tags: |
| - speculative-decoding |
| - dflash |
| - flow-matching |
| - drafter |
| --- |
| |
| # FlowDraft checkpoints |
|
|
| Drafter weights for **FlowDraft**: flow-map drafters for speculative decoding with **Qwen3-4B**. |
| Code, reports and full instructions: https://github.com/ProtonPratt/FlowDraft-Main |
|
|
| A drafter guesses the next 15 tokens of Qwen3-4B in one forward pass. Qwen3-4B checks all 15 guesses in one pass |
| and keeps the longest correct prefix. The output is exactly what Qwen3-4B writes alone. The score is **τ**: tokens |
| written per guess-and-check round. |
|
|
| ## Layout |
|
|
| This repo has the same layout as the GitHub repo. Download into the clone and every file is at the path the code |
| reads: |
|
|
| ```bash |
| git clone https://github.com/ProtonPratt/FlowDraft-Main.git && cd FlowDraft-Main |
| python scripts/hf_checkpoints.py fetch # DFlash-100K + the best drafter (4.1 GB) |
| python scripts/hf_checkpoints.py fetch --set all # everything here (~62 GB) |
| sha256sum -c WEIGHTS_SHA256SUMS --ignore-missing |
| ``` |
|
|
| Run the code from the repo root. The export configs (`dspark_config.json`) hold repo-relative paths. |
|
|
| ## Exports (drafters) |
|
|
| Results: 480 benchmark answers (80 per set: GSM8K, MATH-500, HumanEval, MBPP, MT-Bench), one H100, batch size 1, |
| overhead-free benchmark loop. |
|
|
| | Path | What it is | greedy τ | t1 τ | greedy speed | t1 speed | |
| |---|---|---:|---:|---:|---:| |
| | `init/dflash100k_epoch6` | DFlash-100K (baseline; also the start of every run) | 4.78 | 4.08 | 3.59x | 3.05x | |
| | `runs/fmlm2_v3_6ep_bv/export/epoch_1` | **best**: fmlm2 V3, 6 passes + BV pass | **5.31** | **4.70** | 3.50x | **3.11x** | |
| | `runs/fmlm2_v3_6ep_bv_lt/export/epoch_1` | BV pass + Lafa learned time per slot | 5.30 | 4.71 | | | |
| | `runs/fmlm2_v3_6ep_bv_emb/export/epoch_1` | BV pass + Lafa embedding loss | 5.27 | 4.65 | | | |
| | `runs/fmlm2_v3_6ep_bv_lt_sc/export/epoch_1` | Lafa A + self-conditioning, two-step drafting | | | | | |
| | `runs/fmlm2_v3_6ep_tv/export/epoch_1` | 6 passes + TV pass | 5.26 | 4.70 | | | |
| | `runs/fmlm2_v3_6ep_ce/export/epoch_1` | 6 passes + CE pass (control) | 5.17 | 4.37 | | | |
| | `runs/fmlm2_v3_6ep/export/epoch_6` | fmlm2 V3 after 6 passes (start of every extra pass) | | | | | |
| | `runs/fmlm2_v3/export/epoch_1` | fmlm2 V3 at 1 pass | | | | | |
| | `runs/dspark_joint_6p/export/epoch_6` | DSpark (no flow) after 6 passes | | | | | |
| | `runs/dspark_joint_6p_tv/export/epoch_1` | DSpark 6 passes + TV pass | 5.20 | 4.64 | 3.48x | 3.09x | |
| | `runs/dspark_a/export/epoch_1` | Markov head on frozen DFlash-100K (head only) | 5.04 | 4.36 | 3.37x | 2.90x | |
| | `runs/full_scratch/export/epoch_6` | plain FMLM+ from scratch, 6 passes | | | | | |
|
|
| Empty cells: see the reports in the GitHub repo (`reports/README.md`). An fmlm2 export has `model.safetensors` |
| (DFlash layout), `flow.safetensors` (flow parts) and `head.safetensors` (joint head). A DSpark export has |
| `model.safetensors` + `head.safetensors`. |
|
|
| ## Resume checkpoints |
|
|
| Full training state (weights, optimizer, schedule, data cursor). Use them only to continue training. Resume needs |
| 8 GPUs and the same training settings; `src/train.py` checks both. |
|
|
| | Path | Size | What it is | |
| |---|---:|---| |
| | `runs/fmlm2_v3_6ep/ckpt/epoch_6_step_17156.pt` | 12 GB | fmlm2 V3 at the end of pass 6 | |
| | `runs/dspark_joint_6p/ckpt/epoch_6_step_17156.pt` | 6.8 GB | DSpark at the end of pass 6 | |
| | `runs/fmlm2_v3_tv1p/ckpt/step_1250.pt` | 12 GB | 1-pass matched TV run, paused at step 1,250 | |
|
|
| ## Limits |
|
|
| - Trained on 99,999 **greedy** Qwen3-4B answers, one seed per run. |
| - The target model, Qwen3-4B, is not in this repo. Download it from `Qwen/Qwen3-4B` (Apache-2.0). |
| - The model code is upstream DFlash (z-lab/dflash, MIT). |
|
|
| Team TokenYahu: Vipen Loka, Dheeraj Hegde, Pratyush Jena. |
|
|
