RedHatAI/gemma-4-31B-it-speculator.dspark

This is a DSpark speculator model for google/gemma-4-31B-it.

Training Details

This model was warm-started from a DFlash checkpoint and trained using the Speculators library. Training compute for this model was generously provided by Lambda, a leading cloud platform for AI training and inference.

Commands

Using the Speculators library and the helper scripts provided in the repo.

Prepare data

# In virtual environment with speculators installed
python scripts/prepare_data.py \
  --model google/gemma-4-31B-it \
  --data ./regenerated_data.jsonl \
  --output ./output \
  --seq-length 16384

Launch vLLM

# In (separate) virtual environment with vllm installed
CUDA_VISIBLE_DEVICES=0,1 vllm_venv/bin/python scripts/launch_vllm.py \
  google/gemma-4-31B-it \
  --target-layer-ids 1 17 29 47 58 \
  --max-model-len 32768 \
  --max-num-batched-tokens 32768 \
  --tensor-parallel-size 2 \
  --no-enable-chunked-prefill

Launch training

Must be run once vLLM has finished launching and is running in the background.

# In virtual environment with speculators installed
CUDA_VISIBLE_DEVICES=2,3 torchrun \
  --standalone \
  --nproc_per_node 2 \
  scripts/train.py \
  --verifier-name-or-path google/gemma-4-31B-it \
  --speculator-type dspark \
  --from-pretrained <dflash-checkpoint> \
  --data-path ./output \
  --on-missing generate \
  --on-generate delete \
  --scheduler-type cosine \
  --draft-vocab-size 32000 \
  --max-anchors 1024 \
  --target-layer-ids 1 17 29 47 58 \
  --num-layers 5 \
  --lr 1e-4 \
  --backbone-lr-scale 0.1 \
  --epochs 3 \
  --sliding-window 2048 \
  --markov-rank 256 \
  --markov-head-type vanilla \
  --enable-confidence-head \
  --confidence-head-with-markov \
  --draft-hidden-act silu

Model Specifications


Base Model	google/gemma-4-31B-it
Chat Template	google/gemma-4-31B-it (use `/chat/completions` endpoint)
Format	Safetensors
License	Apache 2.0
Validation Hardware	Nvidia h100

Deployment

# Deploy with speculative decoding
vllm serve google/gemma-4-31B-it -tp 4 --spec-model RedHatAI/gemma-4-31B-it-speculator.dspark --spec-tokens 7 --spec-method dspark

Acceptance Rates

Per-position token acceptance rates across datasets:

Dataset	Pos 0	Pos 1	Pos 2	Pos 3	Pos 4	Pos 5	Pos 6	Avg. Length
HumanEval	85.6%	72.6%	61.3%	52.2%	44.7%	37.6%	30.8%	4.85
math_reasoning	89.2%	78.6%	67.7%	59.2%	50.7%	42.7%	35.3%	5.24
qa	66.4%	41.8%	26.0%	15.8%	10.1%	6.5%	4.4%	2.71
question	73.4%	51.2%	35.6%	26.1%	19.7%	15.1%	11.8%	3.33
rag	75.1%	55.1%	40.5%	29.4%	20.9%	14.0%	8.5%	3.44
summarization	66.9%	40.3%	24.7%	15.0%	8.7%	4.4%	2.6%	2.63
tool_call	69.1%	49.4%	35.1%	25.2%	19.0%	13.8%	9.3%	3.21
translation	72.3%	52.0%	36.6%	26.5%	18.8%	12.6%	8.1%	3.27
writing	73.4%	51.3%	35.6%	25.9%	19.6%	15.0%	11.7%	3.32