d0rj/t5gemma-2-1b-1b.dspark-base

T5Gemma 2 1B-1B DSpark Base

This repository contains a three-layer DSpark speculative draft model trained for the frozen google/t5gemma-2-1b-1b verifier.

"Base" denotes the reduced three-layer draft. DSpark Large uses five draft transformer layers with the same data, vocabulary, target-state selection, Markov head, confidence head, and training budget. The smaller model trades a modest amount of acceptance accuracy for substantially lower draft compute.

This is a speculative draft checkpoint, not a standalone language model. The repository includes inference weights, complete optimizer/scheduler state, validation metrics, continuation metadata, exact training commands, both DSpark TensorBoard runs, and raw benchmark artifacts for the Base and Large comparisons.

Model details

Item	Value
Draft architecture	DSparkDraftModel
Verifier	google/t5gemma-2-1b-1b
Draft transformer layers	3
Large comparison	5 draft transformer layers
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
Markov head	Vanilla first-order bias, rank 256
Confidence head	Enabled; conditioned on the Markov embedding
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 automatically through vLLM entry points; VLLM_PLUGINS is not required.

Serve the Hub checkpoint with its default K=7

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

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

Override the number of proposed tokens

The adapter uses vLLM's DFlash scheduler for every DFlash-family checkpoint and selects DSpark from the checkpoint architecture.

vllm serve google/t5gemma-2-1b-1b \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dspark-base-k3 \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 2048 \
  --max-num-seqs 1 \
  --speculative-config '{"model":"d0rj/t5gemma-2-1b-1b.dspark-base","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-dspark-base",
    "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 checkpoint is not intended to be loaded independently through Transformers AutoModelForSeq2SeqLM.

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 and were not stored for the complete dataset.

Data mixture

The deterministically shuffled seed-42 mixture contained exactly 200,000 examples.

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 source counts and converters are in training/data_config.json.

Hyperparameters

Parameter	Value
Draft layers	3
Epochs	6
Optimizer	AdamW
Learning rate / weight decay	6e-4 / 0.01
Scheduler	Linear; automatic 1% warmup
Packed token budget	2,048
Online extraction microbatch	32
Maximum anchors / block size	256 / 8
Hidden-state dtype	BF16
Noise standard deviation	0.05
Main loss	0.1 cross-entropy + 0.9 total variation
Confidence loss weight	1.0
Verifier	Frozen
Dataloader workers	0
Seed	42
Training revision	4bd884b204230bbbec44dbb6538eb82593c4016f
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

Both requested DSpark runs are included:

• dspark_base_t5gemma2_mixture_200k: this three-layer model;
• dspark_t5gemma2_mixture_200k: the five-layer Large model.

tensorboard --logdir tensorboard --port 6006

Validation comparison

These are teacher-forced validation metrics from each selected epoch-5 checkpoint. Runtime acceptance is lower because later speculative positions depend on every preceding draft token being accepted.

Metric	DSpark Large (5L)	DSpark Base (3L)
Validation loss	0.6738	0.6961
Full-block token accuracy	45.73%	43.67%
Teacher-forced acceptance rate	42.77%	40.55%
Teacher-forced acceptance length	2.557	2.418
Confidence loss	0.2815	0.2832
Confidence absolute error	0.1886	0.1896
Position 1 accuracy	66.60%	64.26%
Position 2 accuracy	55.52%	53.21%
Position 3 accuracy	47.79%	45.51%
Position 4 accuracy	42.48%	40.30%
Position 5 accuracy	38.69%	36.74%
Position 6 accuracy	35.74%	33.96%
Position 7 accuracy	33.21%	31.62%

Unrounded Base values are in val_metrics.json; both runs can be inspected in TensorBoard.

vLLM performance

Both campaigns used identical prompt snapshots: their per-dataset SHA-256 hashes match. Settings were also identical within each batch size: 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 were MATH-500 500, GSM8K 1,319, HumanEval 164, MBPP 500, MT-Bench 160 independent turns, and UltraChat test_sft 5,000.

The tables below compare DSpark Base only with the verifier baseline measured in the same benchmark campaign. Latency, TTFT, and ITL show mean/p95 milliseconds. Acceptance rate is accepted draft tokens divided by all proposals; p1 is first-position acceptance; acceptance length is 1 + accepted_tokens / draft_steps.

Batch size 1

Dataset	K	Verifier tok/s	DSpark Base tok/s	Speedup	Latency mean/p95 ms	TTFT mean/p95 ms	ITL mean/p95 ms	Accept rate	p1	Accept len
math500	3	123.60	152.19	1.231x	837.38/1146.01	119.19/213.45	5.68/8.04	41.56%	63.26%	2.247
math500	7	123.60	149.95	1.213x	849.09/1242.33	129.85/256.74	5.69/8.78	20.32%	62.06%	2.423
gsm8k	3	127.30	128.83	1.012x	990.98/1346.59	111.31/207.46	6.94/9.58	27.13%	48.78%	1.814
gsm8k	7	127.30	121.81	0.957x	1048.34/1480.12	127.43/265.25	7.26/10.53	12.43%	47.85%	1.870
humaneval	3	112.12	162.20	1.447x	428.57/818.54	100.60/181.40	5.12/7.48	57.68%	75.54%	2.730
humaneval	7	112.12	149.24	1.331x	444.15/816.37	160.06/265.69	4.94/8.73	33.09%	73.07%	3.317
mbpp	3	130.38	175.87	1.349x	700.90/919.01	108.79/204.45	5.01/6.63	50.96%	70.10%	2.529
mbpp	7	130.38	175.67	1.347x	696.66/978.64	144.75/269.57	4.79/6.83	27.43%	67.37%	2.920
mt_bench	3	127.47	144.01	1.130x	872.29/1313.48	97.45/185.92	6.34/10.21	32.28%	54.41%	1.968
mt_bench	7	127.47	132.15	1.037x	950.60/1560.65	132.69/258.32	6.66/11.56	14.79%	52.76%	2.035
ultrachat	3	66.23	83.93	1.267x	1515.15/2205.47	154.38/250.98	10.78/15.62	30.25%	53.35%	1.907
ultrachat	7	66.23	83.77	1.265x	1518.37/2203.91	165.57/262.97	10.72/15.63	14.19%	53.12%	1.993

Batch size 4

Dataset	K	Verifier tok/s	DSpark Base tok/s	Speedup	Latency mean/p95 ms	TTFT mean/p95 ms	ITL mean/p95 ms	Accept rate	p1	Accept len
math500	3	327.25	424.44	1.297x	1195.68/1773.15	143.30/297.68	8.35/12.35	41.55%	63.32%	2.246
math500	7	327.25	394.37	1.205x	1286.87/1958.87	169.23/310.46	8.87/13.72	20.28%	61.78%	2.419
gsm8k	3	344.96	376.84	1.092x	1352.79/1901.42	146.36/300.51	9.52/13.59	27.04%	48.59%	1.811
gsm8k	7	344.96	361.79	1.049x	1409.89/2024.30	149.38/303.52	9.94/14.45	12.44%	47.59%	1.871
humaneval	3	265.07	371.59	1.402x	731.68/1438.71	153.84/371.52	9.22/16.54	57.23%	75.49%	2.717
humaneval	7	265.07	388.96	1.467x	695.93/1258.20	170.66/440.41	9.07/18.30	33.28%	73.49%	3.329
mbpp	3	345.00	478.08	1.386x	1031.81/1450.64	149.58/290.25	7.49/10.49	51.15%	70.32%	2.535
mbpp	7	345.00	482.65	1.399x	1016.76/1519.82	157.02/310.51	7.46/11.04	27.60%	67.48%	2.932
mt_bench	3	359.97	395.71	1.099x	1267.36/1889.96	149.93/296.09	8.95/13.28	32.00%	54.34%	1.960
mt_bench	7	359.97	356.55	0.991x	1404.46/2271.47	182.32/363.33	9.92/17.31	14.60%	51.91%	2.022
ultrachat	3	169.47	231.21	1.364x	2200.76/2860.88	203.06/327.66	15.82/20.66	30.23%	53.29%	1.907
ultrachat	7	169.47	230.66	1.361x	2205.89/2923.48	208.65/335.64	15.81/21.20	14.16%	53.16%	1.991

Full Base reports are in benchmarks/base. The raw Large reports remain available in benchmarks/large for reproducibility, but are intentionally omitted from the speed tables above. The reports include all quantiles, prompt hashes, request counts, baseline rows, CSV tables, and raw JSON.

Repository contents

• model.safetensors: selected three-layer DSpark 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/: Base and Large DSpark TensorBoard events.
• benchmarks/: Base and Large bs1/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.
• Base and Large were measured in separate benchmark campaigns. Prompt hashes and settings match, but the two co-run verifier measurements demonstrate system-level drift.
• 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, DSpark, online-extraction, CUDA Graph, and logging extensions from the accompanying development branch.