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

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.

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

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.

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

Batch size 4

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.