Gemma

Training Gemma-4-26B-A4B on B200

Ranked by tokens/s per GPU · 8,192-token SFT rows · global batch 128 · BF16

#	Tokens/s per GPU	$ / 1M tokens	GPUs	Layout	Measured	Author
#1	3,763.4	$1.05	8	TP1 PP1 EP8 DP8 one row per microbatch, selective recompute, no offload	2026-10-20	ATTUNE	Recipe ›

Notes for Gemma-4-26B-A4B

Weights: the BF16 checkpoint google/gemma-4-26B-A4B-it (revision 4d7ae498), with no conversion: miles loads it through Megatron-Bridge (--megatron-to-hf-mode bridge), so --hf-checkpoint and --ref-load are both the checkpoint directory.
MODEL_ARGS come from miles' scripts/models/gemma-4-26b-a4b-it.py at commit 78527b91, printed by miles/utils/external_utils/model_args_utils.py gemma-4-26b-a4b-it; run the command from a miles checkout at that commit, in the radixark/miles v0.1.0 image, whose Megatron-Bridge (commit 7f0fb345) holds the Gemma 4 provider.
It trains on the Gemma 4 rows, built with --loss-mask-type gemma4 (data/README.md).
Selective recompute and --micro-batch-size 1 replace the launcher's full recompute and dynamic batching: full recompute stops at step 0, where save_for_backward receives Gemma 4's rotary embedding as a tuple, and miles refuses dynamic batching with --qkv-format bshd. --max-tokens-per-gpu stays from the launcher and has no effect without dynamic batching. Of the selective recompute modules, moe_act and layernorm stop in the first backward.
On B200, keep --attention-backend unfused. With auto, TE 2.17 selects FlashAttention 4 for the sliding-window layers, whose sm100 kernel does not support a local window at head dim 256, and the run stops at step 0; with fused, no cuDNN backend covers the layers. On H200, auto runs cuDNN fused attention on the 25 sliding-window layers and TE's unfused attention on the 5 global layers (head dim 512).

The steps every recipe shares: Running a recipe.

Recipes

#1attune-tp1-pp1-ep83,763.4 tokens/s/GPUTP1 PP1 EP8 DP8one row per microbatch, selective recompute, no offload

Starting point: miles' Gemma-4-26B-A4B launcher, scripts/run_gemma_4_26b_a4b.py at miles 0.1.0 (model page https://miles.radixark.com/docs/models/gemma/gemma-4). miles publishes it as a GRPO recipe with no tokens/s figure. Ported to this workload with its layout (TP 4 with sequence parallel, EP 8, TE's unfused attention) on 8 B200, with selective recompute of the core attention and one row per microbatch, it measured 1,989.9 tokens/s/GPU; ATTUNE tuned it to 3,763.4 (+89%).
Keep PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True: without it, a tuning run of this layout ran out of memory at step 1, asking for the 8 GiB of FP32 logits with 149 GB allocated.
At step 0 the loss on the first 128 rows is 4.28136, against 4.28294 for the reference recipe; Hugging Face transformers 5.12.1 gives 4.28956 on the same rows.

Author: ATTUNE (cite)
GPUs: 8 x B200 180 GB (AWS p6-b200.48xlarge, 8 per node)
Layout: TP 1, PP 1, CP 1, EP 8, ETP 1, DP 8; one 8,192-token row per microbatch (BSHD); selective recompute of the core attention; TE's unfused attention; distributed optimizer; alltoall dispatcher; no offload
Tokens/s per GPU: 3,763.40
$ / 1M tokens: $1.05
MFU: 3.8% (N = 3,821,830,144 active parameters per token)
Step-time CV: 1.39%
Convergence check: correlation 0.99891, final loss difference 0.0015, 0 pointwise failures over 20 steps
Measured: 2026-10-20
GPU-hours: 2.16
Rows: gemma4-tulu3-fixed8192-v1 (rows_sha256 36956f40db9943482d88e9732d1a69ad9bc94ec00e822ab3df93c0a379e7effa)

Stack

framework: miles 0.1.0 (radixark/miles commit 78527b91, a slime fork), no local changes
megatron lm: radixark/Megatron-LM commit 37f91617 (megatron-core 0.16.0rc0), no local changes
megatron bridge: 0.5.0 (commit 7f0fb345; the Gemma 4 provider and bridge)
transformer engine: 2.17.0
pytorch: 2.11.0 (CUDA 13.0 build)
cuda: 13.0
nccl: 2.28.9 (CUDA 13.0 build; printed by the measured run)
deepep: 1.2.1 (installed; the alltoall dispatcher does not use it)
python: 3.12.3
network: one node of 8 GPUs (NVLink); no inter-node traffic

Environment

CUDA_DEVICE_MAX_CONNECTIONS=1
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

CommandHow to run

export CUDA_DEVICE_MAX_CONNECTIONS=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py gemma-4-26b-a4b-it)"
read -ra MODEL_ARGS <<< "$MODEL_ARGS_LINE"
python3 train.py \
  --actor-num-nodes 1 \
  --actor-num-gpus-per-node 8 \
  --num-gpus-per-node 8 \
  "${MODEL_ARGS[@]}" \
  --hf-checkpoint "$HF_CKPT" \
  --ref-load "$HF_CKPT" \
  --megatron-to-hf-mode bridge \
  --loss-type sft_loss \
  --calculate-per-token-loss \
  --disable-compute-advantages-and-returns \
  --debug-train-only \
  --rollout-function-path pretokenized_sft_rollout.generate_rollout \
  --prompt-data "$SFT_DATA" \
  --input-key messages \
  --metadata-key metadata \
  --num-rollout 20 \
  --rollout-batch-size 128 \
  --global-batch-size 128 \
  --tensor-model-parallel-size 1 \
  --pipeline-model-parallel-size 1 \
  --context-parallel-size 1 \
  --expert-model-parallel-size 8 \
  --expert-tensor-parallel-size 1 \
  --recompute-granularity selective \
  --micro-batch-size 1 \
  --max-tokens-per-gpu 1024 \
  --optimizer adam --lr 1e-5 --lr-decay-style cosine --min-lr 1e-6 \
  --lr-warmup-fraction 0.1 --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \
  --attention-dropout 0.0 --hidden-dropout 0.0 \
  --accumulate-allreduce-grads-in-fp32 \
  --no-gradient-accumulation-fusion \
  --no-check-for-nan-in-loss-and-grad \
  --attention-softmax-in-fp32 \
  --attention-backend unfused \
  --qkv-format bshd \
  --distributed-timeout-minutes 30

recipes/gemma-4-26b-a4b/b200/attune-tp1-pp1-ep8.yamlDownload GitHub

# recipes/gemma-4-26b-a4b/b200/attune-tp1-pp1-ep8.yaml
# An OpenRecipes recipe; the fields are documented in schema/recipe.schema.yaml.
schema_version: 1
model: Gemma-4-26B-A4B
gpu:
  type: B200
  memory_gb: 180
  count: 8
  per_node: 8
  instance: AWS p6-b200.48xlarge
author:
  name: ATTUNE
  url: https://arxiv.org/abs/2610.17528
status: verified
stack:
  framework: miles 0.1.0 (radixark/miles commit 78527b91, a slime fork), no local changes
  megatron_lm: radixark/Megatron-LM commit 37f91617 (megatron-core 0.16.0rc0), no local changes
  megatron_bridge: 0.5.0 (commit 7f0fb345; the Gemma 4 provider and bridge)
  transformer_engine: 2.17.0
  pytorch: 2.11.0 (CUDA 13.0 build)
  cuda: '13.0'
  nccl: 2.28.9 (CUDA 13.0 build; printed by the measured run)
  deepep: 1.2.1 (installed; the alltoall dispatcher does not use it)
  python: 3.12.3
  network: one node of 8 GPUs (NVLink); no inter-node traffic
  source: 'version check of the measured run''s training image (same digest), 2026-10-20: miles and Megatron-LM git
    commits (clean miles tree), installed TE, PyTorch, megatron-core and Megatron-Bridge, and miles'' Gemma 4 model and
    launcher scripts byte-identical to miles v0.1.0; version dump of the same image, 2026-10-13: DeepEP and Python; NCCL
    from the measured run''s log. The image is radixark/miles v0.1.0 (amd64 digest
    sha256:45ae7833f97a01369dbd8f336784830a87d94587bc010c7b8065291cc15cd7ca) with EFA 1.49 added'
workload:
  row_tokens: 8192
  rows: SFT conversations
  global_batch: 128
  precision: bf16
  dataset:
    name: Tulu 3 SFT mixture
    source: allenai/tulu-3-sft-mixture
    revision: b14afda60f1bbebe55d5d2fa1e4df5042f97f8be
    license: ODC-By 1.0
    rows_sha256: 36956f40db9943482d88e9732d1a69ad9bc94ec00e822ab3df93c0a379e7effa
    build:
      name: gemma4-tulu3-fixed8192-v1
      tokenizer: Gemma-4-26B-A4B-it (revision 4d7ae4984b7db7de8f8457170b3f1a419ee76d52; tokenizer.json sha256
        cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f; chat_template.jinja sha256
        ae53464bf3be25802b3a5b37def7fd89667067d7577049b3b2d74c4d8de4c6d4)
      loss_mask: gemma4
      rows: 6400
      row_tokens: 8192
      padding_tokens: 0
      seed: 42
      packing: source conversations in a seeded global permutation, concatenated with an EOS token between them and cut
        into consecutive 8,192-token rows; rows without supervised targets are skipped
      supervised_token_ratio: 0.67
      parquet_sha256: 4357fac051c7968a0c53f5bb885686aa66fc0197dea89fbfc9fb9244df64506c
      parquet_bytes: 71707115
      builder: data/generate_fixed_sft.py with --loss-mask-type gemma4 (gemma4_mask.py; data/README.md)
layout:
  label: TP1 PP1 EP8 DP8
  summary: TP 1, PP 1, CP 1, EP 8, ETP 1, DP 8; one 8,192-token row per microbatch (BSHD); selective recompute of the
    core attention; TE's unfused attention; distributed optimizer; alltoall dispatcher; no offload
  highlights: one row per microbatch, selective recompute, no offload
  tp: 1
  pp: 1
  vpp: 1
  cp: 1
  ep: 8
  etp: 1
  dp: 8
  sequence_parallel: false
  layer_split: null
  max_tokens_per_gpu: null
  rows_per_microbatch: 1
  recompute:
    granularity: selective
    method: null
    num_layers: null
  offload: null
  dispatcher: alltoall
result:
  tokens_per_s_per_gpu: 3763.4
  usd_per_million_tokens: 1.05
  step_time_cv_percent: 1.39
  mfu: 0.0384
  mfu_note: null
  active_params:
    count: 3821830144
    source: https://huggingface.co/google/gemma-4-26B-A4B-it
    revision: 4d7ae4984b7db7de8f8457170b3f1a419ee76d52
  measured: 2026-10-20
  steps_per_run: 20
  window:
  - 5
  - 19
  window_tokens: 15728640
  runs: 1
  profiler: false
  gpu_hours: 2.16
  gpu_hours_note: GPUs x (scheduler admission to finish), including startup; the same run is the candidate arm of the
    convergence check
correctness:
  result: pass
  summary: passed; the recipe changes only the tensor-parallel size and sequence parallelism, placement arguments, from
    its reference recipe
  reference: 'miles'' layout for Gemma-4-26B-A4B (`scripts/run_gemma_4_26b_a4b.py`): TP 4 with sequence parallel, EP 8,
    TE''s unfused attention and `--qkv-format bshd`, with selective recompute of the core attention in place of its full
    recompute and one row per microbatch in place of its dynamic batching, the two changes this stack needs to run it,
    here on 8 B200. Model arguments, rows and schedule are the same in both recipes.'
  configuration: pass
  component_tests: []
  convergence:
    steps: 20
    result: pass
    correlation: 0.99891
    final_loss_difference: 0.0015
    pointwise_failures: 0
    note: 'candidate arm: the measured run; reference arm: a 20-step run of the reference recipe on the same image and
      rows; the parity probe off in both'
  rules: argument classes of ATTUNE's skill/scripts/correctness/rules/numerics_list.yaml (amazon-science/attune aff5b81);
    20-step check by Megatron-Bridge validate_convergence with its pipeline thresholds
env:
  CUDA_DEVICE_MAX_CONNECTIONS: '1'
  PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
patches: []
run: |
  export CUDA_DEVICE_MAX_CONNECTIONS=1
  export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

  MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py gemma-4-26b-a4b-it)"
  read -ra MODEL_ARGS <<< "$MODEL_ARGS_LINE"
  python3 train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 8 \
    --num-gpus-per-node 8 \
    "${MODEL_ARGS[@]}" \
    --hf-checkpoint "$HF_CKPT" \
    --ref-load "$HF_CKPT" \
    --megatron-to-hf-mode bridge \
    --loss-type sft_loss \
    --calculate-per-token-loss \
    --disable-compute-advantages-and-returns \
    --debug-train-only \
    --rollout-function-path pretokenized_sft_rollout.generate_rollout \
    --prompt-data "$SFT_DATA" \
    --input-key messages \
    --metadata-key metadata \
    --num-rollout 20 \
    --rollout-batch-size 128 \
    --global-batch-size 128 \
    --tensor-model-parallel-size 1 \
    --pipeline-model-parallel-size 1 \
    --context-parallel-size 1 \
    --expert-model-parallel-size 8 \
    --expert-tensor-parallel-size 1 \
    --recompute-granularity selective \
    --micro-batch-size 1 \
    --max-tokens-per-gpu 1024 \
    --optimizer adam --lr 1e-5 --lr-decay-style cosine --min-lr 1e-6 \
    --lr-warmup-fraction 0.1 --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \
    --attention-dropout 0.0 --hidden-dropout 0.0 \
    --accumulate-allreduce-grads-in-fp32 \
    --no-gradient-accumulation-fusion \
    --no-check-for-nan-in-loss-and-grad \
    --attention-softmax-in-fp32 \
    --attention-backend unfused \
    --qkv-format bshd \
    --distributed-timeout-minutes 30
notes: |
  - Starting point: miles' Gemma-4-26B-A4B launcher, `scripts/run_gemma_4_26b_a4b.py` at miles 0.1.0 (model page
    https://miles.radixark.com/docs/models/gemma/gemma-4). miles publishes it as a GRPO recipe with no tokens/s figure.
    Ported to this workload with its layout (TP 4 with sequence parallel, EP 8, TE's unfused attention) on 8 B200, with
    selective recompute of the core attention and one row per microbatch, it measured 1,989.9 tokens/s/GPU; ATTUNE
    tuned it to 3,763.4 (+89%).
  - Keep `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`: without it, a tuning run of this layout ran out of memory at
    step 1, asking for the 8 GiB of FP32 logits with 149 GB allocated.
  - At step 0 the loss on the first 128 rows is 4.28136, against 4.28294 for the reference recipe; Hugging Face
    transformers 5.12.1 gives 4.28956 on the same rows.

Citation

If you use a recipe by ATTUNE, please cite the paper.

BibTeX

@article{fan2026attune,
  title   = {{ATTUNE: Agentic Throughput Tuning for LLM Training}},
  author  = {Leo Fan and Hongyang Wang and Wenbo Chen and Jingguo Liu and Xian Li and Xuan Wang and
             Ruochen Wen and Yiming Yan and Hao Wu and Zheshen Wang and Benoit Dumoulin},
  journal = {arXiv preprint arXiv:2610.17528},
  year    = {2026},
  url     = {https://arxiv.org/abs/2610.17528}
}

Gemma-4-26B-A4B on other GPUs

H100

–

No recipe yet

H200 · 8 GPUs

4,023.2

tokens/s/GPU · $0.55 / 1M tokens

B200 · 8 GPUs · this page

3,763.4

tokens/s/GPU · $1.05 / 1M tokens