---
language:
- en
license: gemma
library_name: vllm
base_model: google/t5gemma-2-1b-1b
datasets:
- HuggingFaceH4/ultrachat_200k
- Aeala/ShareGPT_Vicuna_unfiltered
- ise-uiuc/Magicoder-Evol-Instruct-110K
- Salesforce/xlam-function-calling-60k
- glaiveai/glaive-function-calling-v2
- TIGER-Lab/MathInstruct
tags:
- speculative-decoding
- speculators
- vllm
- t5gemma2
- dspark
---

# 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](https://huggingface.co/google/t5gemma-2-1b-1b) verifier.

"Base" denotes the reduced three-layer draft. [DSpark Large](https://huggingface.co/d0rj/t5gemma-2-1b-1b.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](https://github.com/d0rj/t5gemma2-vllm-plugin):

```bash
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

```bash
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.

```bash
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.

```bash
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](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) | 60,000 | 30% | General instruction following and dialogue |
| [ShareGPT Vicuna unfiltered](https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered) | 20,000 | 10% | Multi-turn conversation |
| [Magicoder Evol Instruct](https://huggingface.co/datasets/ise-uiuc/Magicoder-Evol-Instruct-110K) | 50,000 | 25% | Code and software tasks |
| [xLAM function calling](https://huggingface.co/datasets/Salesforce/xlam-function-calling-60k) | 40,000 | 20% | Tool and function calling |
| [Glaive function calling v2](https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2) | 20,000 | 10% | Tool use and function calling |
| [MathInstruct](https://huggingface.co/datasets/TIGER-Lab/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](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](training/train.sh). [train_command.txt](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.

```bash
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](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](benchmarks/base). The raw Large reports remain available in [benchmarks/large](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](https://github.com/vllm-project/speculators) with T5Gemma 2, DSpark, online-extraction, CUDA Graph, and logging extensions from the accompanying development branch.