samf.sh/talks/2026/07/14

Pre-Training AuroraGPT at Scale on Aurora

Less a story of peak FLOP/s than of what works, what breaks, and how cheaply you recover.

Sam Foreman¹, Nathan Nichols, Varuni Sastry, Samuel Wheeler, Khalid Hossain, Huihuo Zheng, Murali Emani, Filippo Simini, Marieme Ngom, Ethan Wong, Venkat Vishwanath

2026-07-14

Outline

Motivation

• How to do production training on a rapidly evolving software stack?

  • across {Intel, NVIDIA, AMD, …} hardware?
  • while also mitigating failures ?
    • {hardware, system, network, lustre, …}

• Tension between:

  • long running pre-training jobs
  • typical facility scheduling policies
    • INCITE (>= 20% of machine, 2k nodes on Aurora)
      -> very large batches
      -> bad model, unstable training³⁴⁵

Since TPC’26: where production stands

• 20B/512N beats the 2B/256N baseline on every eval, per token (HellaSwag-norm 0.635 vs 0.555)
• The wins and the wall are both new since June. This talk: what changed.

The stack

AuroraGPT-2B: the reference run on Aurora

This is the pre-torchtitan reference. Everything that follows is the migration story: same scale, same data, what changed and what broke when we cut over.

Why MDS (Megatron-DeepSpeed) first: the only option at the time

When AuroraGPT kicked off, MDS was the only LLM pre-training framework that ran at scale and supported:

• Intel XPU
• Model, pipeline parallelism
• DeepSpeed ZeRO Offloading

Supporting context:

• PyTorch FSDP1 had Intel XPU gaps (collectives, AC patterns, optimizer-state sharding)
• torchtitan existed as a research project (not tested)
• MDS was the pragmatic choice

By early 2026, the calculus changed: torchtitan + DTensor + FSDP2 closed the gap and the MDS fork’s maintenance cost crossed over.

Why SophiaG: large-batch stability at 50M tok/batch (256N)

W&B Report: AuroraGPT-2B Pre-Training⁸

LR-finder — exponential sweep, blow-up / 10

Classic exponential LR sweep⁹¹⁰; pick the steepest descent (blow-up / 10 as a safe default). See also our cross-optimizer LR scaling work¹¹.

Why we moved to torchtitan

All production runs: training loss

Evals: 2B + 20B vs MDS, per token

MoE on Intel XPU: where we are (honestly)

DeepSeek-style MLA + MoE, 500M → 10B. FSDP-only today (TP=1, no-compile, AC=none, seq_len=4096).

Open knobs we’re sweeping: EP 12 · TP 4 · context-parallel · per-block compile · DeepEP / HybridEP backends · Float8. This is exactly where cross-vendor MoE experience would help (see the ask).

mano: Muon quality at AdamW speed

mano¹² normalizes updates on a rotating Oblique manifold with $O(\text{dim})$ vector-norm ops (no Newton-Schulz iterations). So it matches Muon’s loss without Muon’s throughput tax.

Optimizer	Loss	TPS
Muon	3.557	4,556
mano	3.631	7,048
AdamW	3.801	7,245
SophiaG	4.719	7,208

• Muon and mano tie on loss (~3.6); mano runs at AdamW speed (~7,000 vs ~4,600 TPS) → wins on wall-clock.
• Caveat: at large batch (GBS=384) AdamW still wins (2.71 vs 2.88); mano’s LR was tuned at GBS=48 and needs re-tuning.

80B: one bf16 cliff, three dead optimizers

At 80B (dim=9216), every constant-finder LR NaN-ed in the first ~dozen steps: grad_norm → inf one step before the loss, while loss was still flat (~12.9):

• mano (our 2B speed win) died first (step 5), despite its 80B LR-finder band looking the safest of the three. Early-step ranking does not predict sustained stability.
• Shared failure across 3 optimizers ⇒ a corner-level instability (bf16 at dim=9216), not tuning. Two grad-path triggers: LBS>1 and large dp_degree.
• Fix in flight: long warmup (≥200 steps) + grad clipping, possibly an fp32 grad path. Stable corner today: TP=4 · LBS=1 · GBS=372, validated 100/100 steps NaN-free (loss 12.93 → 7.72) at ~9.8% MFU.

80B cliff: the GBS ladder + why LR-scaling backfires

The stable corner (TP=4 · LBS=1 · bf16 · GBS=372) held NaN-free under batch ramp, until it didn’t:

Batch	GBS	Result
1×	372	✓ 100/100 steps (12.93→7.72)
4×	1,488	✓ 46/46 steps
8×	2,976	✓ 34/34 steps
16×	5,952	❌ NaN at step 29
16× + LR-scaled	5,952	❌ NaN at step 7 (worse!)

• Do not linearly scale LR with batch here: it moves the cliff earlier.
• dp_degree bisect: dp=192 and dp=264 both ran 30 steps clean, so the old “dp>186 unsafe” warning was over-conservative.
• 2048N (dp=6138) SIGSEGVed in set_determinism (seed broadcast faulted at 24,864 ranks); 512N proven, 1024N is the missing measurement.

Beyond pretraining: CPT · SFT · RL

The base run is done, so the pipeline now extends past pretraining:

• CPT: stages 2-3 of the SophiaG reference (dolmino, then math+code) → 7.77T tokens
• SFT: checkpoint-729-hf (on that CPT’d base) is now a first-class production asset
• GRPO: on-policy RL (TRL + vLLM-XPU) on the SFT checkpoint, verified multi-node

CPT: lower loss ≠ better model

Two ~300B-token pilots (256N, GBS=6,144) vs the olmo-100 base plateau:

⚠ but downstream evals reversed it

ARC-Easy: base 0.651 → dolmino-100 0.547 (−10.4pp) vs olmo50-dolmino50 0.610 (holds). HellaSwag fell for both. No mixing ratio declared a winner. Hypothesis: re-warm to peak LR caused forgetting; gentle-LR retry pending.

SFT: tulu_math_uc_mix on the CPT’d 2B

• checkpoint-729-hf is a first-class production asset: it feeds every downstream GRPO experiment.
• Survived 3 SIGABRTs across 4 mpiexec auto-retry attempts.
• Run 3 in progress on the full ~54B-token mix (~93.1M rows); the v2-base variant is blocked by a scale-only oneCCL fault at 32N.

GRPO: arithmetic RL on Intel XPU

On-policy RL via TRL GRPOTrainer + vLLM-XPU, on the SFT’d 2B (sum_digits task, 1,000 steps, 8N × 12 = 96 ranks, LR 1e-6):

• Multi-node GRPO now works (2026-07-06): the apparent “desync” was two ordinary bugs (a oneCCL MPI-transport SIGSEGV + an AVG→SUM reduction), not an XPU pathology.
• Still experimental (verified on XPU, not production-ready).

The Restarts: At Scale, Failure is the Default

• Llama 3 405B — 16K H100s · 54 days · 419 failures (≈ 1 every 3h); 99% recovered via automation¹³
• OPT-175B — 35 manual restarts + 100+ cycled hosts in 2 mo on ~1K A100s¹⁴
• BLOOM-176B — frequent loss spikes; embedding-norm + checkpoint cadence on 384 A100s × 3.5 mo¹⁵
• GLM-130B — loss spikes “increasingly frequent”; some recover, others go to NaN¹⁶

“The Restarts”: three layers of recovery

Bad-node failover, hang-watchdog, and PBS resubmit each operate at a different scope.

Inner loops catch most failures; outer loops catch the rest.

Restarts, since TPC’26: native --auto-retry

The failover cascade moved out of bash and into ezpz:

ezpz launch --auto-retry --np 512 -- python -m torchtitan.train …

• Classifies each attempt → success / walltime / bad-node / stuck-pre-training; swaps a spare, re-execs. Ships in saforem2/ezpz#144.
• Broke the 20B/512N stall: the sync chain sat walltime-blocked for weeks (a PBS mid-save kill left a stale step-4500/ placeholder); auto-retry relaunch drove 4,400 → 5,400 cleanly.

What generalizes, what doesn’t

Generalizes across vendors / sites / models

• Bit-exact deterministic smoke gate after every upstream sync
• lm-eval as the ground truth for “is it actually learning?”
• Spare-node failover wrapper — same idea on Slurm
• Launcher / env autodetect — push every vendor-shaped assumption out of training code

Doesn’t generalize (needs per-(config, hardware, version) tuning)

• torch.compile decisions (and it hurts MoE on XPU)
• AC boundaries (MoE + AC + compile = grief)
• EP↔FSDP frontier
• Collective tuning (XCCL vs gloo fallbacks, NCCL env)
• Optimizer stability at the bf16 corner (LR-finder ranking ≠ sustained stability)

What it took to train AERIS

On the same machine (Aurora), but with its own separate stack, AERIS¹⁷ is a pixel-level Swin diffusion transformer for generative weather & climate prediction, the first diffusion model taken to this scale, a useful contrast point for what full-machine training demands.

• A separate codebase, but the same operational realities we hit: bf16 stability, full-machine scaling, and failure tolerance at 10k nodes.
• Its hard new problem was scaling diffusion stably at high resolution, which SWiPe’s window-parallelism addresses without inflating the global batch.

samf.sh/talks/2026/07/14

Thanks

AuroraGPT team: Venkat Vishwanath, the AI/ML Group at ALCF, collaborators across ANL.

Argonne Leadership Computing Facility: Aurora time, Sunspot staging.

Intel: Intel Max 1550 XPU + oneAPI / XCCL / IPEX support throughout.

Code & docs

• ezpz: github.com/saforem2/ezpz
• torchtitan fork (experiments/ezpz): github.com/saforem2/torchtitan
• These slides: samf.sh/talks/2026/07/14

This research used resources of the Argonne Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC02-06CH11357.

Questions?

Appendix: backup slides

Material that didn’t make the main path but is here for Q&A.

• Open questions
• Silent-correctness bugs (bf16 RMSNorm freeze · TP loss reporting)
• Post-training: CPT · SFT · GRPO
• Failover engineering deep-dive
• yeet-env tarball broadcast scaling

Open questions: the ask

• A portable bit-exact regression suite across vendors — does anyone have one?
• torch.compile at 1T scale — defensible decision tree?
• Async-checkpoint Pareto frontier: recovery time × frequency × storage cost in production
• Optimizer failure at 80B+: SophiaG (Hessian-diagonal estimate saturates) + Muon (Newton-Schulz iterations overflow bf16) both diverge — algorithmic limit or fp-precision artifact?
• MoE EP↔FSDP scaling boundaries at 16B / 64B / 100B
• Make xccl honor train_timeout_seconds so we don’t have to rely on an stdout-idle watchdog as the hang-detection ground truth

Operational reality — bad-node failover

5+ production jobs killed by bad-node failures in 2 weeks. Pattern: PBS gives us 256 / 512 nodes, one is bad, training crashes or hangs after N hours, walltime gone.

Failover wrapper: request select=N+spare (~2%, min 4). Split into active + spare pool. On crash, scrape bad nodes from log, swap a spare in, retry.

qsub -q prod -l select=522 -v NHOSTS_TRAIN=512 \
    submit_agpt_2b_aurora_venv_failover.sh

Handles 6 recurring crash modes. Does not handle silent hangs — those still need a heartbeat watchdog.

Silent-hang detection: ezpz launch --timeout / --retries

The problem. xccl on XPU silently ignores train_timeout_seconds, so a torchtitan job stuck in a hung collective sits consuming the full PBS walltime instead of aborting. Every collective hang quiets stdout (every rank blocks in the same call, nothing reaches the log) — that’s the signal we can act on.

ezpz launch --timeout 600 --retries 3 \
    python -m torchtitan.train --config-file ./config.toml

• --timeout SECONDS — kill the launched process if its stdout goes idle (not walltime) for this many consecutive seconds. Returns exit code 124 (matches GNU timeout(1)).
• --retries N — re-execute on any non-zero exit (including the watchdog’s 124) up to N times. Exponential backoff: 5s → 10s → 20s → 40s → 60s (capped).

Scope caveat. Watches only the process ezpz launch spawns directly. If qsub runs a wrapper script that internally invokes python train.py, the watchdog needs to live inside that script (or you wrap the inner call with ezpz launch too).

ezpz launch --timeout: one hang/recover cycle

Every collective hang shows up as silence on stdout — the process is “alive” by kill -0 but nothing is happening. The watchdog fires on the absence of progress, not on a heartbeat ping.

--auto-retry: bad-node failover, on tap

Allocate spares up front, swap them in on failure:

# 522 nodes allocated, train on 512, keep 10 as spares.
# Loop until success, walltime, or spare exhaustion.
ezpz launch --auto-retry --np 512 -- python -m torchtitan.train …

• Classifies each attempt’s exit → success / walltime / bad-node / stuck-pre-training
• On bad-node: scrapes the failing host from the log, swaps in a spare, re-execs
• Guards against config bugs: 2 consecutive attempts with zero step= markers → stop (don’t burn the whole spare pool on a broken run)

Ships in saforem2/ezpz#144. Same scraper as the bash-lib path; pure-Python loop on top.

Failover wrapper: caught a real silent hang in production

Job 8505298, 2026-05-23. Attempt 1 trains cleanly steps 1→37, then log goes completely silent at step 37. No traceback, no MPI error, no rank dying. Just dead.

Three new pieces had to fire in sequence on a real-world hang to prove production-readiness: --timeout=1800 watchdog · exit 124 classification distinct from PBS exit 143 · failover_swap_one_blind() when no specific bad node can be identified. They did.

Full writeup: docs/experiments/agpt/aurora/20260523-failover-silent-hang-recovery-8505298.md

Silent bug #1 — bf16 master ⇒ RMSNorm frozen

Symptom. Loss curves looked reasonable. lm-eval scores didn’t move — ARC-Easy stuck at ~0.27 (random baseline) for 17K+ steps.

Cause. training.dtype=bfloat16 → bf16 master copy. RMSNorm weights init at 1.0; bf16 ULP at scale 1.0 is ~7.8e-3. Per-step optimizer update is ~1.6e-5 — every update rounds to zero. All 25 RMSNorm tensors stayed at exactly 1.0 from step 100 → 17,400.

Why other params trained fine. Linear layers init at std≈0.02 → bf16 ULP at scale 0.02 is ~3.8e-5, same magnitude as the update. RMSNorm’s larger init scale = coarser ULP = updates lost in rounding.

Fix. Default training.dtype=float32, FSDP MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=fp32). Master is fp32; forward/backward stay bf16. Extra ~1 GB master at 2B, ~10 GB at 20B — under budget.

Silent bug #1 — the smoking gun

Lesson: loss looks like training. lm-eval is the only ground truth for “is the model actually learning?” Add a periodic eval gate.

Silent bug #2 — TP loss reported / dp_world_size

• Symptom: Step-1 loss for agpt_* (vocab=256128) should be ln(256128) ≈ 12.45 - On TP=1 we see 12.95 ✓ - On TP=2 after upstream commit 1786292d (2026-04-27): - step-1 reported as 1.07 ✗ — exactly 12.84 / 12 where dp_world_size = 12

• Cause: _dist_reduce() short-circuits on DTensor with full_tensor()

  • But the loss is Replicated on the TP mesh, and the reduction was requested over batch_mesh (orthogonal)
  • The short-circuit silently drops the cross-batch sum

• Why it survived review: Gradients + optimizer steps are correct

  • Only the loss: field that lands in stdout / W&B is wrong
  • Loss curves look “reasonable” — just 1/12 of the true value
  • Filed as pytorch/torchtitan#3204; our workaround calls loss.full_tensor() before dist_sum in ezpz/trainer.py:503-516

• Lesson: Bit-exact smoke caught this immediately; the prior 2B TP=1 baseline gave us a number to disagree with

ezpz yeet: Efficiently Running 50k Python Processes

Nodes	yeet (s)	First-step (s)	Per-node (ms)
8	69.7	29.3	8,712
16	89.7	31.6	5,606
32	89.2	20.9	2,788
64	91.2	34.6	1,425
128	110.4	30.5	862
256	132.9	37.6	519
512	174.5	44.5	341
1024	255.4	60.8	249
2048	421.4	94.8	206
4096	750.6	194.0	183

Two regimes. 8–64 nodes extract-bound (~70–91s flat, per-node cost falls 8.7s → 1.4s); ≥128 nodes broadcast-bound, each 2× in nodes adds ~1.5–1.8× wall-clock.

Full write-up: samf.sh/posts/2026/05/01

Footnotes