ProtonPrat/FlowDraft-checkpoints

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:

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.