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

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.

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

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.

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

Batch size 4

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.