gemma-4-31B-it-speculator.dflash

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

A DFlash speculative decoding draft model for google/gemma-4-31B-it, trained with Speculators.

It was trained using the Speculators library on the mlabonne/open-perfectblend dataset. Training data used responses from this dataset for the gemma-4-31B-it model (with reasoning).

Model configuration

The published config.json declares draft_vocab_size as 262144. The recorded training command used --draft-vocab-size 128000; this difference is retained here for transparency and should be resolved against the training/export configuration before changing the published config.

Use with vLLM

Use a vLLM release/build with DFlash support and the Gemma-4 model support enabled. The draft model is used with the matching Gemma-4 verifier:

vllm serve google/gemma-4-31B-it \
  --tensor-parallel-size 2 \
  --speculative-config '{
    "model": "RedHatAI/gemma-4-31B-it-speculator.dflash",
    "num_speculative_tokens": 7,
    "method": "dflash"
  }'

Training details

The checkpoint's recorded training invocation used the following DFlash recipe (paths are replaced with placeholders):

torchrun --standalone --nproc_per_node=4 \
  /path/to/speculators/train/__main__.py \
  --draft-vocab-size 128000 \
  --verifier-name-or-path google/gemma-4-31b-it \
  --data-path /path/to/prepared-open-perfectblend-data \
  --save-path /path/to/checkpoints \
  --hidden-states-path /path/to/hidden-states \
  --log-dir /path/to/training-logs \
  --run-name dflash_5layers_5epochs \
  --speculator-type dflash \
  --block-size 8 \
  --max-anchors 1536 \
  --num-layers 5 \
  --target-layer-ids 1 17 29 47 59 \
  --sliding-window 2048 \
  --epochs 3 \
  --lr 6e-4 \
  --weight-decay 0.01 \
  --scheduler-type cosine \
  --optimizer muon \
  --muon-lr 0.003 \
  --dpace-alpha 0.5 \
  --per-position-loss-weight dpace \
  --loss-fn ce \
  --checkpoint-freq 0.1 \
  --log-freq 100 \
  --num-workers 12 \
  --prefetch-factor 2 \
  --fsdp-shard \
  --vllm-endpoint http://localhost:8906/v1 \
  --on-missing generate \
  --on-generate delete \
  --request-timeout 180 \
  --total-seq-len 8192 \
  --max-retries 20

The training run used 4 processes (world size 4), Speculators 0.8.0, vLLM 0.29.0, Transformers 5.14.1, and PyTorch 2.13.0. The data-path is the prepared tokenized data derived from mlabonne/open-perfectblend. The recorded run name contains 5epochs, but the actual --epochs setting was 3.

Evaluation

Average acceptance length and per-position acceptance rates from the reported evaluation:

License

This model is distributed under the Apache-2.0 license. The Gemma-4 verifier is subject to its own license and terms.