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

T5Gemma 2 1B-1B DFlash Large

This repository contains a five-layer DFlash draft model that predicts a seven-token speculative block from selected 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	DFlashDraftModel
Verifier	google/t5gemma-2-1b-1b
Draft transformer layers	5
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, 7, 13, 18, 21
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.dflash-large \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dflash \
  --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-dflash-k3 \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 2048 \
  --max-num-seqs 1 \
  --speculative-config '{"model":"d0rj/t5gemma-2-1b-1b.dflash-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-dflash",
    "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.

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.2283
Full-block token accuracy	34.51%
Position 1 accuracy	67.90%
Position 2 accuracy	49.17%
Position 3 accuracy	36.79%
Position 4 accuracy	28.76%
Position 5 accuracy	23.30%
Position 6 accuracy	19.33%
Position 7 accuracy	16.21%

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	64.68	0.570x	1970.50/2394.33	176.98/348.85	14.18/17.33	6.32%	16.82%/1.89%/0.26%	1.190
math500	7	113.45	74.16	0.654x	1716.92/2122.72	138.95/232.84	12.49/15.42	2.73%	16.89%/1.92%/0.26%	1.191
gsm8k	3	115.56	65.68	0.568x	1943.86/2365.31	178.53/355.62	13.93/16.93	6.84%	18.35%/1.95%/0.21%	1.205
gsm8k	7	115.56	73.26	0.634x	1743.27/2099.80	152.56/246.42	12.55/15.04	2.75%	17.19%/1.83%/0.20%	1.192
humaneval	3	93.66	58.98	0.630x	1178.81/2207.20	162.99/323.05	15.05/20.25	5.99%	16.43%/1.45%/0.10%	1.180
humaneval	7	93.66	64.81	0.692x	1022.86/2083.64	157.66/238.51	13.24/16.24	2.71%	17.46%/1.41%/0.08%	1.189
mbpp	3	117.61	67.35	0.573x	1830.21/2291.14	181.88/374.19	13.37/16.29	6.52%	18.39%/1.10%/0.06%	1.196
mbpp	7	117.61	74.82	0.636x	1635.87/2016.28	152.67/244.47	12.10/14.41	2.84%	18.66%/1.14%/0.06%	1.199
mt_bench	3	116.89	65.85	0.563x	1907.75/2463.49	175.25/376.67	13.83/17.45	6.04%	16.21%/1.76%/0.15%	1.181
mt_bench	7	116.89	73.27	0.627x	1714.64/2135.27	139.66/237.62	12.69/15.55	2.35%	14.88%/1.46%/0.09%	1.164
ultrachat	3	62.83	52.06	0.829x	2442.49/3377.26	168.36/252.91	18.01/24.71	7.08%	18.36%/2.59%/0.30%	1.212
ultrachat	7	62.83	51.13	0.814x	2487.86/3443.12	172.78/258.77	18.33/25.19	3.06%	18.49%/2.58%/0.32%	1.215

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	185.18	0.694x	2740.52/3396.06	232.73/422.78	19.89/24.71	6.29%	16.64%/1.97%/0.24%	1.189
math500	7	266.74	188.44	0.706x	2690.23/3411.27	223.15/410.68	19.61/24.83	2.74%	16.86%/2.02%/0.26%	1.192
gsm8k	3	280.58	188.79	0.673x	2702.38/3317.10	235.61/424.62	19.48/24.18	6.81%	18.37%/1.86%/0.20%	1.204
gsm8k	7	280.58	188.46	0.672x	2705.30/3477.26	301.18/834.45	18.98/23.59	2.74%	17.14%/1.81%/0.20%	1.192
humaneval	3	194.21	147.26	0.758x	1856.84/3610.00	250.89/497.85	24.55/42.01	5.95%	16.41%/1.35%/0.08%	1.179
humaneval	7	194.21	141.91	0.731x	1877.28/3978.09	343.90/982.06	22.96/33.21	2.68%	17.40%/1.31%/0.09%	1.188
mbpp	3	275.24	189.94	0.690x	2595.66/3320.97	234.11/439.53	19.14/23.96	6.46%	18.25%/1.11%/0.04%	1.194
mbpp	7	275.24	170.29	0.619x	2875.58/4074.32	634.04/1885.92	18.20/24.69	2.85%	18.78%/1.13%/0.05%	1.200
mt_bench	3	309.23	195.62	0.633x	2573.23/3166.33	221.35/448.16	18.68/23.10	6.07%	16.45%/1.71%/0.06%	1.182
mt_bench	7	309.23	180.37	0.583x	2753.90/3791.87	381.83/1026.59	21.71/25.58	2.48%	15.47%/1.60%/0.12%	1.173
ultrachat	3	155.46	135.15	0.869x	3763.95/4515.90	257.60/405.56	27.76/33.28	7.09%	18.40%/2.57%/0.30%	1.213
ultrachat	7	155.46	134.03	0.862x	3794.92/4543.27	255.94/399.63	28.03/33.56	3.07%	18.54%/2.58%/0.31%	1.215

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.