Training Gemma-4-31B on H100

Author

ATTUNE (cite)

GPUs

8 x H100 80 GB (AWS p5.48xlarge, 8 per node)

Layout

TP 2, PP 4 (13/16/16/15 layers), CP 1, DP 1, sequence parallel; one 8,192-token row per microbatch (BSHD); selective recompute of the core attention and the MLP; TE's auto attention backend (cuDNN fused attention on the sliding-window layers, unfused on the global layers); CPU optimizer offload of 80% with overlapped copies and the precision-aware optimizer; dense model, no experts

Tokens/s per GPU

1,178.10

$ / 1M tokens

$1.62

MFU

21.9% (N = 30,696,013,824 active parameters per token)

Step-time CV

2.19%

Convergence check

correlation 0.98649, final loss difference 0.0091, 0 pointwise failures over 20 steps

Component test

TE's auto attention backend against its unfused attention and an FP64 reference: H100, 2026-10-21. Shape: Gemma-4-31B attention at TP 2 (rank 0's heads), 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 58, global layers 5 and 59, 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.62e-2 against a bound of 4.07e-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

Component test

CPU optimizer offload fraction 0.8 against fraction 1.0 and an FP64 reference: H100, 2026-10-21. Shape: TP 2 rank 0's slices of Gemma-4-31B layer 0 (sliding), layer 5 (global) and embedding rows 0-8,191: 21 tensors, 556,051,968 parameters, 19.8% of them on the GPU at fraction 0.8 (the embedding slice, layer 0's attention projections and every norm gain); checkpoint weights, synthetic gradients scaled to the measured run's first three pre-clip norms. after three steps at the recipe's hyperparameters, the GPU partition's FP32 masters and updates are 0.62-1.00 x the default's relative L2 error against FP64 and its BF16 weights 1.00 x (worst tensor, the update of the norm gains, 7.24e-2 against a bound of 1.46e-1, 2 x the default's error + 1e-3, the same error as the default's); the CPU partition is bitwise identical to the default; a repeat of the default passes and the three negative controls (no bias correction, BF16 masters, no clipping) fail run on this recipe's training image. The Adam moments on the GPU path pass on the bound's 1e-3 term only: exp_avg 1.80e-7 (3.2 x the default's error) and exp_avg_sq 1.01e-6 (9.0 x). TE FusedAdam computes with beta1 and beta2 rounded to FP32; against an FP64 reference with the same rounded betas, its moments are 4.1e-8 and 4.7e-8, the FP32 rounding floor

Measured

2026-10-20

GPU-hours

5.69

Rows

gemma4-tulu3-fixed8192-v1 (rows_sha256 36956f40db9943482d88e9732d1a69ad9bc94ec00e822ab3df93c0a379e7effa)

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-31b-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 2 \
  --sequence-parallel \
  --pipeline-model-parallel-size 4 \
  --context-parallel-size 1 \
  --recompute-granularity selective \
  --recompute-modules core_attn mlp \
  --micro-batch-size 1 \
  --max-tokens-per-gpu 512 \
  --optimizer-cpu-offload \
  --overlap-cpu-optimizer-d2h-h2d \
  --use-precision-aware-optimizer \
  --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 \
  --decoder-first-pipeline-num-layers 13 \
  --decoder-last-pipeline-num-layers 15 \
  --optimizer-offload-fraction 0.8

# recipes/gemma-4-31b/h100/attune-tp2-pp4-split13.yaml
# An OpenRecipes recipe; the fields are documented in schema/recipe.schema.yaml.
schema_version: 1
model: Gemma-4-31B
gpu:
  type: H100
  memory_gb: 80
  count: 8
  per_node: 8
  instance: AWS p5.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, with the dense Gemma 4 support miles'
    launcher requires)
  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; a dense model 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), byte-identical to Gemma-4-31B-it's at
        revision 842da3794eaa0b77d5f08bae87a17459d91ff475
      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: TP2 PP4 DP1
  summary: TP 2, PP 4 (13/16/16/15 layers), CP 1, DP 1, sequence parallel; one 8,192-token row per microbatch (BSHD);
    selective recompute of the core attention and the MLP; TE's auto attention backend (cuDNN fused attention on the
    sliding-window layers, unfused on the global layers); CPU optimizer offload of 80% with overlapped copies and the
    precision-aware optimizer; dense model, no experts
  highlights: 13/16/16/15 layer split, CPU optimizer offload 0.8, cuDNN fused sliding-window attention
  tp: 2
  pp: 4
  vpp: 1
  cp: 1
  ep: 1
  etp: 1
  dp: 1
  sequence_parallel: true
  layer_split:
    first: 13
    last: 15
  max_tokens_per_gpu: null
  rows_per_microbatch: 1
  recompute:
    granularity: selective
    method: null
    num_layers: null
  offload:
    optimizer_fraction: 0.8
    overlap_d2h_h2d: true
  dispatcher: null
result:
  tokens_per_s_per_gpu: 1178.1
  usd_per_million_tokens: 1.62
  step_time_cv_percent: 2.19
  mfu: 0.2194
  mfu_note: null
  active_params:
    count: 30696013824
    source: https://huggingface.co/google/gemma-4-31B-it
    revision: 842da3794eaa0b77d5f08bae87a17459d91ff475
  measured: 2026-10-20
  steps_per_run: 20
  window:
  - 5
  - 19
  window_tokens: 15728640
  runs: 1
  profiler: false
  gpu_hours: 5.69
  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) and the optimizer offload fraction (1.0 to 0.8) passed
    their component tests at this model's shape
  reference: 'miles'' layout for Gemma-4-31B (`scripts/run_gemma_4_31b.py`): TP 4 with sequence parallel, TE''s unfused
    attention, `--qkv-format bshd` and Adam on the CPU (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d
    --use-precision-aware-optimizer`), here on 8 H100 with one row per microbatch in place of its dynamic batching,
    selective recompute of the core attention and the MLP in place of its full recompute, and PP 2 (30 layers per stage)
    in place of its PP 1: with the core attention recomputed alone, TP 4 ran out of memory at step 0 at PP 1 and at PP
    2. 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-31B attention at TP 2 (rank 0''s heads), 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 58, global
      layers 5 and 59, on the first row of the recipe''s rows)'
    gpu: H100
    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.62e-2 against a bound of 4.07e-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
  - test: CPU optimizer offload fraction 0.8 against fraction 1.0 and an FP64 reference
    shape: 'TP 2 rank 0''s slices of Gemma-4-31B layer 0 (sliding), layer 5 (global) and embedding rows 0-8,191: 21
      tensors, 556,051,968 parameters, 19.8% of them on the GPU at fraction 0.8 (the embedding slice, layer 0''s
      attention projections and every norm gain); checkpoint weights, synthetic gradients scaled to the measured run''s
      first three pre-clip norms'
    gpu: H100
    date: 2026-10-21
    result: pass
    detail: after three steps at the recipe's hyperparameters, the GPU partition's FP32 masters and updates are 0.62-1.00 x
      the default's relative L2 error against FP64 and its BF16 weights 1.00 x (worst tensor, the update of the norm
      gains, 7.24e-2 against a bound of 1.46e-1, 2 x the default's error + 1e-3, the same error as the default's); the
      CPU partition is bitwise identical to the default; a repeat of the default passes and the three negative controls
      (no bias correction, BF16 masters, no clipping) fail
    note: 'run on this recipe''s training image. The Adam moments on the GPU path pass on the bound''s 1e-3 term only:
      exp_avg 1.80e-7 (3.2 x the default''s error) and exp_avg_sq 1.01e-6 (9.0 x). TE FusedAdam computes with beta1 and
      beta2 rounded to FP32; against an FP64 reference with the same rounded betas, its moments are 4.1e-8 and 4.7e-8,
      the FP32 rounding floor'
  convergence:
    steps: 20
    result: pass
    correlation: 0.98649
    final_loss_difference: 0.0091
    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-31b-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 2 \
    --sequence-parallel \
    --pipeline-model-parallel-size 4 \
    --context-parallel-size 1 \
    --recompute-granularity selective \
    --recompute-modules core_attn mlp \
    --micro-batch-size 1 \
    --max-tokens-per-gpu 512 \
    --optimizer-cpu-offload \
    --overlap-cpu-optimizer-d2h-h2d \
    --use-precision-aware-optimizer \
    --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 \
    --decoder-first-pipeline-num-layers 13 \
    --decoder-last-pipeline-num-layers 15 \
    --optimizer-offload-fraction 0.8
notes: |
  - Starting point: miles' Gemma-4-31B launcher, `scripts/run_gemma_4_31b.py` at miles 0.1.0 (model page
    https://miles.radixark.com/docs/models/gemma/gemma-4). miles publishes it as a GRPO recipe for 8 H200 with no
    tokens/s figure. Ported to this workload with its TP 4 with sequence parallel, TE's unfused attention and Adam on
    the CPU on 8 H100, with PP 2, selective recompute of the core attention and the MLP and one row per microbatch, it
    measured 743.4 tokens/s/GPU; ATTUNE tuned it to 1,178.1 (+58%).
  - `--optimizer-offload-fraction 0.8` keeps 20% of the optimizer state on the GPU and shortens the CPU Adam step. At
    0.8 the first pipeline stage runs with less than 10% of its memory free.
  - At step 0 the loss on the first 128 rows is 2.27393, against 2.26972 for the reference recipe; Hugging Face
    transformers 5.12.1 gives 2.27153 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.