Production Pre-Training at Scale:

The Good, the Bad, and the Restarts

Lessons from AuroraGPT

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

2026-06-03

Outline

Notes

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

  • across {Intel, NVIDIA, AMD} hardware?
  • while also mitigating hardware and system failures ?

• Tension between:

  • long running continuous (weeks/months+) pre-training jobs
  • typical facility scheduling policies (e.g. 8-12 hours per day)

• In addition to:

  • node failures
  • file system instabilities
  • (silent!) HW failures
  • frequent checkpointing and restarting to mitigate

The stack in one slide

Current stack:

• 🍋 saforem2/ezpz (+ ezpz.cool)
• 🧠 saforem2/torchtitan@ezpz: FSDP · TP · PP · EP · MoE
• 📚 zhenghh04/blendcorpus: weighted blending across datasets

Old stack (reference):

• 🪦 argonne-lcf/Megatron-DeepSpeed:
  • AuroraGPT-2B reference (~7.77T tokens)
  • pre-torchtitan

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 first: the only option at the time

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

• at scale
• on Intel XPU

Supporting context:

• PyTorch FSDP1 had Intel XPU gaps — collectives, AC patterns, optimizer-state sharding
• torchtitan existed as a research project but not yet on XPU / not yet usable for our config
• blendcorpus had a stable MDS integration; the data pipeline was already a known quantity
• The fork was ours — every Aurora-specific patch could land same-day

MDS was the pragmatic choice. By mid-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

2B reference: training loss

2B reference + torchtitan overlay

Why we moved to torchtitan

The trade-off we accepted: living on a fast-moving upstream (pytorch/torchtitan main) instead of a tagged release. That trade is the “fork tax”, the rest of this talk.

2B loss: MDS full trajectory vs torchtitan

2B eval: MDS reference vs torchtitan

20B eval: all-production overlay (2B + 20B)

The Restarts: At Scale, Failure is the Default

Production pre-training is long, large, and routinely interrupted:

• 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⁹

At our scale, the question isn’t if training breaks — it’s how cheaply you recover.

“The Restarts”: three layers of recovery

Bad-node failover, hang-watchdog, and PBS resubmit each operate at a different scope. Together they turn “lost a 12 h walltime” into “lost a 10 min sub-step.”

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

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

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
• AC boundaries (MoE + AC + compile = grief)
• EP↔FSDP frontier
• Collective tuning (XCCL vs gloo fallbacks, NCCL env)

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: sam.onl/posts/2026/05/01

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 + Muon both diverge — algorithmic or numerical?
• 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

Thanks

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

Argonne Leadership Computing Facility: Aurora time, Sunspot staging.

Code & docs

• ezpz: github.com/saforem2/ezpz
• torchtitan fork (experiments/ezpz): github.com/saforem2/torchtitan
• These slides: sam.onl/talks/2026/06/03

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.

Silent-correctness bugs — the two bugs that almost killed the v1 production runs. Loss looked fine, training kept going, and only an external eval gate (or a bit-exact smoke comparison) flagged that something was wrong.

• Bug #1: bf16 master ⇒ RMSNorm frozen — every update rounds to zero at scale 1.0
• Bug #2: TP loss reported as true_loss / dp_world_size — looked reasonable on plots, caught only by smoke comparison

blendcorpus — the data layer (deterministic shard indexing, weight-stable resumption; pain: upstream shard churn, dataloader skew across backends).

Operational reality — bad-node failover — the production table + qsub wrapper that swaps spares in on crash.

Cross-version churn — APIs shift faster than vendor support catches up. Example: DeviceMesh + AC + TP=2 + torch 2.13 → AOT-autograd assertion on every 80B config; pin torch 2.10 until the upstream fix lands.

Optimizer + LR-finder — AdamW · Muon · SophiaG sweep across {2B, 20B, 80B} on Aurora / Sunspot / Polaris; speedrun rounds with Mano + QK-Norm; yeet-env tarball broadcast scaling.

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.

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

Smith 2015 / Gugger: exponentially ramp LR over ~10% of training, record EMA-smoothed loss, pick LR at the steepest descent (or blow-up point / 10 as a conservative default).

LR-finder — 2B sweep + optimal LR per config

Findings: AdamW most LR-tolerant; SophiaG needs ~10× lower LR; SophiaG + Muon both break at 80B (bf16 overflow in Newton-Schulz / Hessian estimate on 9216-dim matrices).

Optimizer speedrun — Mano LR sweep, WSM, Schedule-Free

Headline numbers (Round 3, 10B tokens / 8N / GBS=384):

Rank	Config	Loss	TPS/GPU
1	AdamW	2.711	7,354
2	AdamW + QK-Norm	2.720	7,480
3	Mano + QK-Norm	2.854	7,346
4	Mano	2.875	7,429
5	Muon	DNF	—

AdamW wins at GBS=384. Mano lands ~0.16 behind — LR finder was tuned at GBS=48 and the manifold update under-shoots at large batch.

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

blendcorpus — the boring layer that breaks first

What it does

• Weighted blending across N scientific corpora; deterministic shard indexing so a given (global_step, dp_rank) always sees the same bytes
• Weight-stable resumption: changing per-corpus weights mid-run doesn’t re-shuffle history

Where it bites

• Upstream shard churn — corpus owners re-emit shards (dedup, re-tokenization). Same logical corpus, different bytes → silent history drift unless you pin shard SHAs in the manifest.
• Cross-backend dataloader skew — num_workers, prefetch, NUMA pinning behave differently on XPU vs CUDA hosts. Same seed, different arrival order on the GPU → not reproducible across vendors.

Smoke test catches this: bit-exact first-step loss requires the dataloader to deliver identical batches across runs, vendors, and resumptions.

Footnotes