Scaling PyTorch on Intel XPU

Field notes from AI for Science on Aurora

Interactive version: live diagrams + linkssamf.sh/talks/2026/10/21

Aurora at ALCF

ALCF	Argonne Leadership Computing Facility at Argonne National Lab: DOE open science on Aurora + Polaris
Aurora	10,624 nodes · 127,488 XPUs · system paper
This talk	Running PyTorch on Aurora: backend fixes, training, and recovery

Talk outline

Scientific AI on Aurora

I co-led MProt-DPO and AERIS, and co-lead AuroraGPT.

ezpz: same training program on Intel, NVIDIA, AMD.²

PyTorch XPU and Aurora timeline

PyTorch XPU backend components

XPU integration issues

The CUDA path worked. The XPU path inherited defaults that needed fixing.

From Megatron-DeepSpeed to torchtitan

The fork trained the 2B model to 7.77T tokens. Each PyTorch upgrade required another rebase and patch audit.

One training program across vendors

# example.py
import torch
import ezpz
rank = ezpz.setup_torch()
world_size = ezpz.get_world_size()
dev = ezpz.get_torch_device()      # {xpu, cuda, cpu, ...}
# FSDP2 by default — accepts DeviceMesh
model = ezpz.wrap_model(torch.nn.Linear(128, 10).to(dev))
print(f"Hello from {rank}/{world_size}!")

$ ezpz launch python3 example.py
Hello from rank 0/24!
Hello from rank 1/24!
# ... etc

# your own script — no changes needed
ezpz launch python3 your_training_script.py
# full HF model + dataset via ezpz examples
ezpz launch python3 -m ezpz.examples.fsdp_tp \
    --dataset=allenai/olmo-mix-1124
    --model=meta-models/Muse-Spark-Open-816B-A42B \
    --local-batch-size=1
    # ...etc

Benchmark the training step

1. Fix the build, model, batch size, precision, and node count.
2. Run the same training step on each backend.
3. Measure step time, tokens/s, collective time, and memory.
4. Increase node count. Rerun after kernel or collective changes.

An op test can pass while the distributed training step hangs.

Check loss and throughput at each node count. Test timeout and recovery behavior separately.⁴

Changes in the torchtitan training stack

ezpz stayed in place. Parallelism, collectives, and checkpoints moved out of the fork.

Testing at every layer

The streaming loader contributes ~0.05 loss variation from shuffle order. The 50-step check uses a ±0.10 tolerance.

AuroraGPT: production training progress

Shown tokens / raw loss2B MDS reference7.771T / 2.0332B 256N stage 14.674T / 2.6522B 512N stage 14.674T / 2.6872B 256N stage 26.085T / 2.4692B 512N stage 26.927T / 2.481Solid: 256 nodes · dashed: 512 nodes. Training mixes differ.

Stage 2: 6.927T cumulative tokens on 512 nodes · 6.085T on 256 nodes.

Detecting and recovering from stalled jobs

Job 8505298: step 37, then silence. No traceback. Watchdog at 30 min (exit 124), wrapper swapped a node, relaunched. Attempt 2 reached step 296.

_set_pg_timeout has no xpu branch: timeout is stored, never enforced.⁷

Failure modes at ~120k XPUs

Test node exclusion, checkpoint reload, and scheduler resubmission before a long run.⁸

Process, node, and allocation restarts

Remaining XPU gaps

Run the training step, save a checkpoint, interrupt the job, and resume it.

Bringing up the next accelerator

1. Test compute and collectives together.
2. Use the same training program on each machine.
3. Compare loss against the reference run.
4. Measure full training steps at scale.
5. Interrupt, restart, and check the resumed loss.

Links and further reading

Further reading: ezpz portability · Launching 50k processes · Data and training at ALCF

Acknowledgements

This work used resources of the Argonne Leadership Computing Facility, a DOE Office of Science user facility at Argonne National Laboratory, supported under contract DE-AC02-06CH11357.

Aurora was built in partnership with HPE and Intel: Intel Max 1550 XPU (PVC), oneAPI / oneCCL / XCCL, and the upstream torch.xpu work tracked in Act A. Thanks to the ALCF operations and support teams, the PyTorch and torchtitan maintainers, and the science teams behind GenSLM, MProt-DPO, AERIS, and AuroraGPT.

Aurora allocation programs

Anyone (!) can apply for time on any of the ALCF systems, including our AI Testbed!

For a DD proposal, find me after the talk: [email protected] / @saforem2.

Next at ALCF: Solstice and Equinox

Announced with NVIDIA and Oracle:

The hardware, fabrics, launchers, networks, and file systems keep changing.

The same training code still needs to launch, communicate, checkpoint, and recover at scale.

Backup: inside an Aurora node

6 GPUs × 2 tiles = 12 torch.xpu devices in FLAT mode. GPU Max 1550 (PVC), 768 GiB GPU HBM per node.

Session feedback

Scan to share feedback with us and the organizers.