d0rj/t5gemma-2-1b-1b.dflare-large

T5Gemma 2 1B-1B DFlare Large

This repository contains a six-layer DFlare draft model with learned per-draft-layer fusion over eight verifier decoder states, trained for the frozen google/t5gemma-2-1b-1b verifier.

The checkpoint is a speculative draft model, not a standalone language model. It must be paired with the base verifier. This repository includes inference weights, complete optimizer and scheduler states, validation metrics, the exact training command, continuation state, TensorBoard events, and raw benchmark summaries.

"Large" distinguishes this 200,000-example mixed-data run from the earlier development checkpoint.

Model details

Item	Value
Draft architecture	DFlareDraftModel
Verifier	google/t5gemma-2-1b-1b
Draft transformer layers	6
Draft hidden / intermediate size	1,152 / 6,912
Draft vocabulary	32,000 tokens mapped from the verifier vocabulary
Speculative block	8 positions: one anchor and up to 7 proposed tokens
Default proposals per step	7
Runtime override	1-7 proposed tokens
Verifier decoder states	2, 5, 8, 11, 14, 17, 20, 22
Precision	BF16
Selected checkpoint	Epoch 5 (sixth/final epoch), global step 146,486

Usage with vLLM

Install a compatible vLLM build and the T5Gemma 2 vLLM plugin:

pip install "git+https://github.com/d0rj/t5gemma2-vllm-plugin.git"

The plugin is discovered through vLLM entry points. No VLLM_PLUGINS environment variable is required.

Serve with the checkpoint default (K=7)

vllm serve d0rj/t5gemma-2-1b-1b.dflare-large \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dflare \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 2048 \
  --max-num-seqs 1

The plugin reads the verifier reference from config.json and loads google/t5gemma-2-1b-1b automatically.

Limit the number of draft tokens

To override K, serve the verifier explicitly and reference this Hub repository. The adapter uses method: "dflash" for all DFlash-family schedulers, then dispatches to the concrete architecture recorded in the draft checkpoint.

vllm serve google/t5gemma-2-1b-1b \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dflare-k3 \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 2048 \
  --max-num-seqs 1 \
  --speculative-config '{"model":"d0rj/t5gemma-2-1b-1b.dflare-large","method":"dflash","num_speculative_tokens":3}'

Set num_speculative_tokens to an integer from 1 through 7. K=3 and K=7 were benchmarked below. CUDA Graph remains enabled when --enforce-eager is omitted.

curl http://127.0.0.1:8000/v1/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "t5gemma-2-1b-1b-dflare",
    "prompt": "Solve step by step.\n\nQuestion: A train travels 60 miles in 2 hours. What is its average speed?\nAnswer:",
    "max_tokens": 128,
    "temperature": 0
  }'

This draft is not intended to be loaded independently with Transformers AutoModelForSeq2SeqLM. The verifier supplies encoder execution, verification, embeddings, and target hidden states.

Training

Training used online hidden-state extraction. The frozen verifier and trainable draft shared one RTX 5070 Ti 16 GB GPU under WSL2; hidden states were generated in process rather than stored for the full dataset.

Data mixture

The mixture contained exactly 200,000 examples and was deterministically shuffled with seed 42.

Dataset	Examples	Share	Coverage
UltraChat 200k	60,000	30%	General instruction following and dialogue
ShareGPT Vicuna unfiltered	20,000	10%	Multi-turn conversation
Magicoder Evol Instruct	50,000	25%	Code and software tasks
xLAM function calling	40,000	20%	Tool and function calling
Glaive function calling v2	20,000	10%	Tool use and function calling
MathInstruct	10,000	5%	Mathematics and reasoning

The prepared data used a 90%/10% train/validation split. Encoder inputs were truncated to 2,048 tokens and decoder targets to 1,024 tokens. Exact retained counts and converters are in training/data_config.json.

Hyperparameters

Parameter	Value
Epochs	6
Optimizer	AdamW
Learning rate	6e-4
Weight decay	0.01
Scheduler	Linear; automatic 1% warmup
Packed token budget	2,048
Online extraction microbatch	32
Maximum anchors	256
Block size	8
Hidden-state dtype	BF16
Noise standard deviation	0.05
Loss	KL divergence
Verifier	Frozen
Dataloader workers	0
Seed	42
Training revision	a58573acefd975901bae8e83b6e8b5e6ace717ef
Software	Speculators 0.7.0.dev72; Transformers 5.12.1; PyTorch 2.12.0

The launcher is training/train.sh. train_command.txt records the resolved command and software versions.

TensorBoard

The requested cumulative comparison logs are included: dflash_t5gemma2_mixture_200k and dflare_t5gemma2_mixture_200k.

tensorboard --logdir tensorboard --port 6006

Validation metrics

These are teacher-forced metrics from the selected epoch-5 checkpoint. Position accuracy is not equivalent to runtime acceptance because later runtime positions depend on all earlier draft tokens being accepted.

Metric	Value
Validation loss	1.2540
Full-block token accuracy	33.84%
Position 1 accuracy	67.58%
Position 2 accuracy	48.25%
Position 3 accuracy	35.94%
Position 4 accuracy	27.88%
Position 5 accuracy	22.36%
Position 6 accuracy	18.83%
Position 7 accuracy	15.90%

Unrounded values are available in val_metrics.json.

vLLM performance

The raw verifier and all drafts used identical saved prompt snapshots on one RTX 5070 Ti 16 GB under WSL2. The vLLM development build was based on commit e4b3da3feb20c1854a4b23e431cfb787ee268f72.

Common settings: greedy decoding, seed 0, maximum model length 2,048, maximum output 128, --max-num-batched-tokens 4096, GPU memory utilization 0.85, chunked prefill disabled, CUDA Graph enabled, and 8 excluded warm-up requests. BS1 used concurrency/max-num-seqs 1; BS4 used 4. Dataset sizes: MATH-500 500, GSM8K 1,319, HumanEval 164, MBPP 500, MT-Bench 160 independent turns, and UltraChat test_sft 5,000.

Base tok/s is raw verifier throughput; Draft tok/s is end-to-end verified throughput. Speedup is their ratio. Latency, TTFT, and ITL are mean/p95 milliseconds. Acceptance rate is accepted tokens divided by all proposed draft tokens. Acceptance length is 1 + accepted_tokens / draft_steps. Values below 1.0x are genuine slowdowns for this setup.

Batch size 1

Dataset	K	Base tok/s	Draft tok/s	Speedup	Latency mean/p95 ms	TTFT mean/p95 ms	ITL mean/p95 ms	Accept rate	Accept p1/p2/p3	Accept len
math500	3	113.45	93.74	0.826x	1359.56/1919.75	150.99/254.30	9.56/13.64	39.18%	64.52%/35.08%/17.94%	2.175
math500	7	113.45	97.45	0.859x	1306.61/1831.94	149.97/241.51	9.16/13.16	18.55%	64.82%/35.30%/17.18%	2.298
gsm8k	3	115.56	80.95	0.701x	1577.24/2166.57	150.18/254.45	11.26/15.61	27.89%	53.39%/21.58%/8.69%	1.837
gsm8k	7	115.56	81.92	0.709x	1558.94/2191.44	152.78/255.01	11.09/15.86	12.80%	54.17%/21.95%/9.32%	1.896
humaneval	3	93.66	99.28	1.060x	700.23/1298.74	163.54/255.47	8.18/11.82	58.63%	77.59%/56.43%/41.86%	2.759
humaneval	7	93.66	114.43	1.222x	579.29/1124.12	153.89/249.90	7.29/11.29	32.60%	76.90%/55.26%/40.09%	3.282
mbpp	3	117.61	103.43	0.879x	1191.91/1663.97	151.14/261.67	8.62/12.05	46.96%	68.67%/44.38%/27.84%	2.409
mbpp	7	117.61	115.48	0.982x	1059.83/1511.53	149.30/254.47	7.82/11.13	24.50%	67.47%/42.50%/25.01%	2.715
mt_bench	3	116.89	85.55	0.732x	1468.51/2127.07	159.73/265.74	10.65/15.71	32.31%	57.74%/27.25%/11.94%	1.969
mt_bench	7	116.89	89.12	0.762x	1409.59/2124.10	146.91/259.67	10.23/15.34	14.63%	57.43%/27.01%/11.06%	2.024
ultrachat	3	62.83	62.40	0.993x	2037.94/2756.50	185.96/279.43	14.67/19.80	31.13%	57.22%/26.06%/10.11%	1.934
ultrachat	7	62.83	66.03	1.051x	1925.98/2635.48	180.71/272.82	13.82/18.90	14.39%	57.70%/26.66%/10.39%	2.008

Batch size 4

Dataset	K	Base tok/s	Draft tok/s	Speedup	Latency mean/p95 ms	TTFT mean/p95 ms	ITL mean/p95 ms	Accept rate	Accept p1/p2/p3	Accept len
math500	3	266.74	217.77	0.816x	2328.94/3331.75	279.51/506.23	16.27/23.21	39.35%	64.89%/35.28%/17.89%	2.181
math500	7	266.74	225.99	0.847x	2249.20/3338.34	263.03/475.01	15.76/24.25	18.69%	64.67%/35.48%/17.53%	2.309
gsm8k	3	280.58	201.03	0.716x	2536.44/3534.03	258.19/484.66	17.98/25.82	27.95%	53.64%/21.51%/8.69%	1.838
gsm8k	7	280.58	154.41	0.550x	3304.27/5885.10	1205.81/4089.58	16.56/27.05	12.83%	54.22%/21.89%/9.37%	1.898
humaneval	3	194.21	194.56	1.002x	1401.65/2758.01	307.73/627.48	17.78/33.04	58.74%	77.35%/56.42%/42.44%	2.762
humaneval	7	194.21	206.56	1.064x	1296.23/2684.54	395.19/956.47	16.23/32.95	33.05%	77.91%/56.44%/40.45%	3.313
mbpp	3	275.24	244.12	0.887x	2016.51/2791.48	280.07/557.88	14.26/19.44	47.05%	68.89%/44.50%/27.76%	2.411
mbpp	7	275.24	184.84	0.672x	2657.85/4488.95	1190.47/3259.46	12.45/22.86	24.66%	67.13%/42.74%/25.45%	2.726
mt_bench	3	309.23	225.31	0.729x	2231.25/3333.71	248.97/520.00	15.78/23.60	32.30%	57.50%/27.57%/11.81%	1.969
mt_bench	7	309.23	133.91	0.433x	3752.53/6717.97	1934.01/4820.78	14.32/29.83	14.58%	56.90%/26.94%/11.54%	2.021
ultrachat	3	155.46	161.25	1.037x	3155.82/4013.42	282.90/456.06	22.75/29.20	31.06%	57.10%/26.00%/10.07%	1.932
ultrachat	7	155.46	157.88	1.016x	3221.08/4123.29	352.64/612.31	22.73/29.63	14.37%	57.63%/26.59%/10.33%	2.006

Full results, all quantiles, request counts, prompt hashes, baseline rows, and complete suite comparisons are included in benchmarks/bs1 and benchmarks/bs4.

Repository contents

• model.safetensors: selected draft weights.
• config.json, config.py: architecture and verifier linkage.
• optimizer_state_dict.pt: complete AdamW state.
• scheduler_state_dict.pt: scheduler state.
• training_state.json: epoch/global-step continuation metadata.
• train_command.txt: exact training command and package versions.
• val_metrics.json: selected-checkpoint validation metrics.
• training/: data metadata and launcher.
• tensorboard/: cumulative TensorBoard events.
• benchmarks/: raw bs1 and bs4 reports.

For optimizer-level continuation, restore the root checkpoint files as an epoch directory in a compatible Speculators checkpoint tree and point checkpoint_best to it.

Limitations

• Performance depends on hardware, workload, batch size, and vLLM version; universal acceleration is not claimed.
• The plugin path was validated for single-GPU text serving. Tensor/pipeline parallelism, quantized serving, and the vision path require separate validation.
• Teacher-forced validation accuracy does not guarantee equal runtime acceptance.
• K=3 reduces draft work but changes block length relative to eight-position training; K=7 matches the trained block shape.
• Generated outputs retain the behavior and limitations of the verifier and training data.

License and attribution

This checkpoint is derived from T5Gemma 2 and is distributed under the Gemma license. Using the verifier requires accepting Google's Gemma terms. Each dataset retains its own license and conditions.

Training used the vLLM Speculators project with T5Gemma 2, DFlare, DSpark, online-extraction, and logging extensions from the accompanying development branch.