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

T5Gemma 2 1B-1B DSpark Large

This repository contains a five-layer DSpark draft model with a rank-256 first-order Markov logit-bias head and a confidence head, 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	DSparkDraftModel
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.dspark-large \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dspark \
  --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-dspark-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-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-dspark",
    "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	0.1 cross-entropy + 0.9 total variation, plus confidence loss (weight 1.0)
Verifier	Frozen
Dataloader workers	0
Seed	42
Training revision	a58573acefd975901bae8e83b6e8b5e6ace717ef
Software	Speculators 0.7.0.dev72; Transformers 5.12.1; PyTorch 2.12.0
Markov head	Vanilla, rank 256
Confidence head	Enabled; conditioned on Markov embedding

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, dflare_t5gemma2_mixture_200k, and dspark_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	0.6738
Full-block token accuracy	45.73%
Position 1 accuracy	66.60%
Position 2 accuracy	55.52%
Position 3 accuracy	47.79%
Position 4 accuracy	42.48%
Position 5 accuracy	38.69%
Position 6 accuracy	35.74%
Position 7 accuracy	33.21%
Teacher-forced acceptance rate	42.77%
Teacher-forced mean acceptance length	2.557
Confidence loss	0.2815
Confidence absolute error	0.1886

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	131.01	1.155x	972.82/1383.58	139.16/243.83	6.59/9.67	44.80%	65.41%/42.19%/26.82%	2.344
math500	7	113.45	136.48	1.203x	932.87/1415.98	142.46/239.27	6.26/10.08	22.84%	64.00%/39.61%/24.04%	2.599
gsm8k	3	115.56	111.95	0.969x	1140.44/1627.45	137.90/232.76	7.91/11.64	31.19%	53.50%/25.78%/14.30%	1.936
gsm8k	7	115.56	111.09	0.961x	1149.49/1690.77	142.42/238.71	7.94/11.94	14.63%	53.36%/25.01%/13.27%	2.024
humaneval	3	93.66	130.99	1.399x	530.74/905.70	146.04/236.92	6.10/8.98	60.43%	77.12%/58.94%/45.24%	2.813
humaneval	7	93.66	145.41	1.553x	455.86/852.21	143.63/241.35	5.51/9.66	34.97%	73.45%/55.20%/39.37%	3.448
mbpp	3	117.61	146.73	1.248x	840.11/1144.96	133.81/238.02	5.95/8.07	52.37%	71.03%/51.06%/35.03%	2.571
mbpp	7	117.61	158.24	1.346x	773.38/1069.82	144.88/240.50	5.46/7.70	28.85%	68.24%/46.95%/30.63%	3.020
mt_bench	3	116.89	121.09	1.036x	1037.41/1565.51	141.77/240.41	7.31/11.71	37.06%	59.00%/32.88%/19.30%	2.112
mt_bench	7	116.89	122.24	1.046x	1027.68/1596.33	141.25/255.56	7.28/11.79	17.42%	56.41%/30.24%/17.08%	2.220
ultrachat	3	62.83	78.10	1.243x	1628.21/2287.42	171.98/258.65	11.54/16.24	32.41%	55.57%/28.18%/13.47%	1.972
ultrachat	7	62.83	78.70	1.253x	1616.26/2290.96	176.72/263.78	11.40/16.28	15.45%	55.33%/28.03%/13.31%	2.081

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	286.34	1.073x	1774.14/2512.93	237.01/434.45	12.18/17.88	44.99%	65.67%/42.25%/27.05%	2.350
math500	7	266.74	294.75	1.105x	1725.25/2646.27	247.76/461.03	11.72/18.88	22.89%	63.98%/39.64%/24.22%	2.602
gsm8k	3	280.58	267.82	0.954x	1904.86/2717.59	232.41/436.56	13.20/19.22	31.15%	53.33%/25.75%/14.35%	1.934
gsm8k	7	280.58	256.27	0.913x	1991.49/2946.94	240.23/457.41	13.82/21.07	14.63%	53.19%/24.90%/13.31%	2.024
humaneval	3	194.21	234.79	1.209x	1172.97/2107.93	265.30/502.88	14.78/29.63	59.91%	76.53%/58.74%/44.47%	2.797
humaneval	7	194.21	210.60	1.084x	1276.63/2616.22	564.26/1681.50	12.49/25.28	35.09%	74.01%/55.03%/39.85%	3.457
mbpp	3	275.24	310.80	1.129x	1592.68/2178.77	240.73/455.66	11.35/15.68	52.29%	71.25%/50.82%/34.79%	2.569
mbpp	7	275.24	274.50	0.997x	1794.33/3171.44	643.14/2119.01	10.10/14.88	28.83%	68.32%/47.02%/30.55%	3.018
mt_bench	3	309.23	282.93	0.915x	1771.67/2604.84	252.12/503.80	12.36/19.15	36.54%	58.03%/32.54%/19.06%	2.096
mt_bench	7	309.23	224.25	0.725x	2225.17/3657.94	838.84/2441.89	11.43/22.87	17.22%	55.30%/29.97%/16.74%	2.205
ultrachat	3	155.46	190.92	1.228x	2665.16/3441.95	264.37/429.11	19.01/24.80	32.40%	55.58%/28.18%/13.44%	1.972
ultrachat	7	155.46	189.60	1.220x	2683.34/3469.80	287.66/433.69	18.97/25.07	15.40%	55.21%/27.94%/13.27%	2.078

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.