You can deploy this model and start earning money today!

inference-optimization/
gemma-4-31B-it-speculator.dflash.swa.causal.fullvocab.muon

Name: gemma-4-31B-it-speculator.dflash.swa.causal.fullvocab.muon
Availability: InStock

safetensors

custom_code

Views

No views yet

Training

Prepare the data:

python scripts/prepare_data.py \
  --model google/gemma-4-31B-it \
  --data magpie_gemma-4-31B-it.jsonl \
  --data ultrachat_gemma-4-31B-it.jsonl \
  --output ./data \
  --num-preprocessing-workers 32

Launch the hidden-state server (GPUs 0-3):

CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/launch_vllm.py \
  google/gemma-4-31B-it \
  --hidden-states-path ./hidden_states \
  --target-layer-ids 1 17 29 47 58 \
  -- \
  --port 8887 \
  --max-model-len 16384 \
  -tp 2 \
  --data-parallel-size 2 \
  --no-enable-chunked-prefill

Train (GPUs 4-7):

CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --standalone --nproc_per_node=4 scripts/train.py \
  --speculator-type dflash \
  --verifier-name-or-path google/gemma-4-31B-it \
  --data-path ./data \
  --hidden-states-path ./hidden_states \
  --on-missing generate --on-generate delete \
  --vllm-endpoint http://localhost:8887/v1 \
  --target-layer-ids 1 17 29 47 58 \
  --num-layers 5 \
  --max-anchors 512 \
  --block-size 8 \
  --sliding-window 2048 \
  --scheduler-type cosine \
  --epochs 3 \
  --noise-std 0.0 \
  --optimizer muon \
  --loss-fn '{"ce": 0.1, "tv": 0.9}' \
  --fsdp-shard \
  --logger trackio \
  --run-name gemma-4-31B-it-dflash-swa-causal-fullvocab-muon \
  --save-path ./checkpoints

Deployment

Serve the model:

CUDA_VISIBLE_DEVICES=0,1 python -m vllm.entrypoints.openai.api_server \
  --model google/gemma-4-31B-it \
  --speculative-config '{"model": "inference-optimization/gemma-4-31B-it-speculator.dflash.swa.causal.fullvocab.muon", "method": "dflash", "num_speculative_tokens": 7}' \
  -tp 2 \
  --max-model-len 16384 \
  --port 8001

Run benchmarks:

cd speculators/scripts/evaluate && python evaluate.py \
  --target http://localhost:8001/v1 \
  --output-dir ./epoch3_throughput \
  throughput

Evaluation results

Dataset	Pos 0	Pos 1	Pos 2	Pos 3	Pos 4	Pos 5	Pos 6	Avg. Length
HumanEval	85.8%	71.1%	59.1%	49.5%	41.6%	34.5%	27.9%	4.69
math_reasoning	91.4%	78.9%	67.2%	55.8%	45.9%	38.0%	29.6%	5.07
qa	68.6%	42.5%	24.7%	14.4%	9.0%	5.5%	3.4%	2.68
question	74.5%	50.1%	33.3%	23.8%	17.3%	13.1%	9.8%	3.22
rag	81.2%	62.0%	45.9%	34.0%	25.7%	17.6%	10.6%	3.77
summarization	72.4%	46.7%	28.4%	16.8%	9.3%	4.6%	2.2%	2.80
tool_call	71.8%	49.5%	34.8%	24.2%	17.7%	12.2%	8.3%	3.19
translation	78.0%	55.3%	38.3%	26.1%	17.1%	10.1%	5.7%	3.31
writing	73.3%	49.0%	32.9%	23.0%	16.8%	12.5%	9.2%	3.17
Mean	77.4%	56.1%	40.5%	29.7%	22.3%	16.5%	11.9%	3.54