# 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.