Gemma

Training Gemma-4-26B-A4B on H200

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	4,023.2	$0.55	8	TP1 PP1 EP8 DP8 selective recompute of attention and MoE, cuDNN fused sliding-window attention, 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-ep84,023.2 tokens/s/GPUTP1 PP1 EP8 DP8selective recompute of attention and MoE, cuDNN fused sliding-window attention, 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 H200, with selective recompute of the core attention and one row per microbatch, it measured 1,878.2 tokens/s/GPU; ATTUNE tuned it to 4,023.2 (+114%).
TP 1 fits with both --recompute-modules core_attn moe and --log-probs-chunk-size 1024. With the core attention recomputed alone, a tuning run ran out of memory at step 0 in the FP32 logits of the 262,144-token vocabulary (8 GiB per row); with chunked log-probs but no MoE recompute, another ran out at step 1 in the expert activations.
At step 0 the loss on the first 128 rows is 4.28773, against 4.28841 for the reference recipe; Hugging Face transformers 5.12.1 gives 4.28956 on the same rows.
The convergence check covers 20 steps. This recipe runs cuDNN fused attention on the sliding-window layers, where the reference recipe runs TE's unfused attention, and on another MoE model cuDNN fused attention passed the 20-step check while its loss drifted up later (ATTUNE's lever catalog). Compare the loss over a longer run before a full training run.

Author: ATTUNE (cite)
GPUs: 8 x H200 141 GB (AWS p5en.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 and the MoE layers; log-probs in chunks of 1,024 tokens; TE's auto attention backend (cuDNN fused attention on the sliding-window layers, unfused on the global layers); distributed optimizer; alltoall dispatcher; no offload
Tokens/s per GPU: 4,023.25
$ / 1M tokens: $0.55
MFU: 9.3% (N = 3,821,830,144 active parameters per token)
Step-time CV: 1.21%
Convergence check: correlation 0.99848, final loss difference 0.008, 0 pointwise failures over 20 steps
Component test: TE's auto attention backend against its unfused attention and an FP64 reference: H200, 2026-10-21. Shape: Gemma-4-26B-A4B attention at TP 1, forward and backward: 16 query heads, sequence 8,192, micro-batch 1, BF16; sliding-window layers with 8 key-value groups, head dim 256 and window 1,024, global layers with 2 groups and head dim 512; four synthetic cases and four from the checkpoint (sliding layers 0 and 28, global layers 5 and 29, on the first row of the recipe's rows). on the sliding-window layers TE runs cuDNN fused attention, whose relative L2 error against FP64 is 0.15-0.82 x the unfused path's (worst sliding-layer tensor, layer 0 dq, 1.24e-2 against a bound of 3.11e-2, 2 x the unfused error + 1e-3); the global layers stay on the unfused path, bitwise identical; a repeat of the default passes and the two negative controls (a wrong window, a wrong softmax scale) fail run on this recipe's training image; one attention call per case, not a training run
Measured: 2026-10-20
GPU-hours: 2.02
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 auto \
  --qkv-format bshd \
  --distributed-timeout-minutes 30 \
  --log-probs-chunk-size 1024 \
  --recompute-modules core_attn moe

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

# recipes/gemma-4-26b-a4b/h200/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: H200
  memory_gb: 141
  count: 8
  per_node: 8
  instance: AWS p5en.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 and the MoE layers; log-probs in chunks of 1,024 tokens; TE's auto attention backend (cuDNN fused
    attention on the sliding-window layers, unfused on the global layers); distributed optimizer; alltoall dispatcher;
    no offload
  highlights: selective recompute of attention and MoE, cuDNN fused sliding-window attention, 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: 4023.25
  usd_per_million_tokens: 0.55
  step_time_cv_percent: 1.21
  mfu: 0.0933
  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.02
  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 attention backend change (unfused to auto) passed its component test at this model's shape
  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.
    Model arguments, rows and schedule are the same in both recipes.'
  configuration: pass
  component_tests:
  - test: TE's auto attention backend against its unfused attention and an FP64 reference
    shape: 'Gemma-4-26B-A4B attention at TP 1, forward and backward: 16 query heads, sequence 8,192, micro-batch 1, BF16;
      sliding-window layers with 8 key-value groups, head dim 256 and window 1,024, global layers with 2 groups and head
      dim 512; four synthetic cases and four from the checkpoint (sliding layers 0 and 28, global layers 5 and 29, on
      the first row of the recipe''s rows)'
    gpu: H200
    date: 2026-10-21
    result: pass
    detail: on the sliding-window layers TE runs cuDNN fused attention, whose relative L2 error against FP64 is 0.15-0.82 x
      the unfused path's (worst sliding-layer tensor, layer 0 dq, 1.24e-2 against a bound of 3.11e-2, 2 x the unfused error + 1e-3);
      the global layers stay on the unfused path, bitwise identical; a repeat of the default passes and the two negative
      controls (a wrong window, a wrong softmax scale) fail
    note: run on this recipe's training image; one attention call per case, not a training run
  convergence:
    steps: 20
    result: pass
    correlation: 0.99848
    final_loss_difference: 0.008
    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 auto \
    --qkv-format bshd \
    --distributed-timeout-minutes 30 \
    --log-probs-chunk-size 1024 \
    --recompute-modules core_attn moe
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 H200, with
    selective recompute of the core attention and one row per microbatch, it measured 1,878.2 tokens/s/GPU; ATTUNE
    tuned it to 4,023.2 (+114%).
  - TP 1 fits with both `--recompute-modules core_attn moe` and `--log-probs-chunk-size 1024`. With the core attention
    recomputed alone, a tuning run ran out of memory at step 0 in the FP32 logits of the 262,144-token vocabulary
    (8 GiB per row); with chunked log-probs but no MoE recompute, another ran out at step 1 in the expert activations.
  - At step 0 the loss on the first 128 rows is 4.28773, against 4.28841 for the reference recipe; Hugging Face
    transformers 5.12.1 gives 4.28956 on the same rows.
  - The convergence check covers 20 steps. This recipe runs cuDNN fused attention on the sliding-window layers, where
    the reference recipe runs TE's unfused attention, and on another MoE model cuDNN fused attention passed the 20-step
    check while its loss drifted up later (ATTUNE's lever catalog). Compare the loss over a longer run before a full
    training run.

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 · this page

4,023.2

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

B200 · 8 GPUs

3,763.4

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