ezpz

ezpz is a thin layer over PyTorch distributed that removes the per-machine boilerplate: rank and topology detection, backend selection, hostfile parsing, and launcher invocation. The same script runs on Aurora, Polaris, and Perlmutter without an if machine == ... branch.

Full documentation: ezpz.cool.

The problem it solves

Moving a working PyTorch script between ALCF systems, or from ALCF to NERSC, normally means rewriting the parts that have nothing to do with the model:

Concern	Aurora / Sunspot	Polaris / Sophia	Perlmutter
Scheduler	PBS Pro	PBS Pro	SLURM
Launcher	mpiexec	mpiexec	srun
Accelerator	Intel XPU	NVIDIA	NVIDIA
Backend	xccl	nccl	nccl
Hostfile	$PBS_NODEFILE	$PBS_NODEFILE	$SLURM_NODELIST
Local rank var	PALS_LOCAL_RANKID	PMI_LOCAL_RANK	SLURM_LOCALID

ezpz detects all of it at runtime. The Supported Systems page lists the full matrix, including the hostname prefixes used for detection (x4* for Aurora, x1* for Sunspot, and so on).

It also sets the Aurora-specific environment that is easy to get wrong:

• ZE_FLAT_DEVICE_HIERARCHY=FLAT, so all 12 tiles are visible. Without it only 6 devices appear and every rank with local rank ≥ 6 fails with device index out of range.
• CCL_OP_SYNC=1, which works around an oneCCL async-completion hang that affects FSDP2 with tensor parallel degree > 1 (AuroraBugTracking#174).

Setup

module load frameworks
pip install "git+https://github.com/saforem2/ezpz"  # (1)!

1. Install from the repository. The name ezpz on PyPI belongs to an unrelated package.

Then, inside a job script or an interactive allocation:

source <(curl -LsSf https://bit.ly/ezpz-utils)
ezpz_setup_env

ezpz_setup_env loads the right modules for the current machine, activates the Python environment, and exports the machine-specific environment variables. Verify with:

ezpz doctor   # environment and dependency check
ezpz test     # distributed PyTorch smoke test

Launching

ezpz launch wraps the scheduler's launcher. It reads the hostfile, computes ranks-per-node from the detected accelerator count, and applies CPU binding:

ezpz launch python3 -m ezpz.examples.test

That command is unchanged across every supported machine; on Aurora it becomes a 12-rank-per-node mpiexec --pmi=pmix, on Perlmutter an srun. See ezpz launch for the flags, including --auto-retry for resubmitting around node failures.

Minimal PBS job on Aurora

#!/bin/bash -l
#PBS -l select=2
#PBS -l walltime=00:30:00
#PBS -q debug
#PBS -A <project>
#PBS -l filesystems=home:flare

cd "${PBS_O_WORKDIR}"
source <(curl -LsSf https://bit.ly/ezpz-utils)
ezpz_setup_env

ezpz launch python3 -m ezpz.examples.test --model small

Using it in your own script

Two calls replace the usual distributed preamble:

import ezpz

rank = ezpz.setup_torch()          # (1)!
device = ezpz.get_torch_device_type()   # "xpu" on Aurora, "cuda" elsewhere

model = MyModel().to(device)
model = ezpz.distributed.wrap_model_for_ddp(model)

1. Initializes the process group with the correct backend, binds the device, and returns the global rank. No MASTER_ADDR / MASTER_PORT plumbing.

See Distributed Training for the full API.

Features

Area	What it does	Details
Launcher	Scheduler-aware mpiexec / srun wrapper with CPU binding and --auto-retry	ezpz launch
Setup	Per-machine module loading and environment export	Quickstart
Distributed	Backend selection, rank/topology detection, DDP and FSDP2 wrapping	Distributed Training
Parallelism	FSDP2, tensor parallel, sequence parallel, 2D meshes	Parallelism
Examples	Runnable training examples with a shared model-size ladder (debug → xxxl)	Examples
Profiling	torch.profiler wrapper with automatic XPU/CUDA activity selection	Profiling
Metrics	History for per-iteration metrics, plots, and W&B / MLflow logging	History
Checkpointing	Distributed checkpoint save/restore, resumable after node failure	Checkpoint & Restart
Fault handling	Node-kill injection and automatic retry	Fault Injection
Diagnostics	ezpz doctor, ezpz test	ezpz test

Examples

Each example takes a --model size from debug (laptop-runnable) through xxxl (~10B), so the same command scales from a smoke test to a real job:

ezpz launch python3 -m ezpz.examples.test    --model small
ezpz launch python3 -m ezpz.examples.fsdp    --model large
ezpz launch python3 -m ezpz.examples.fsdp_tp --model large --tp 2   # (1)!

1. FSDP2 + tensor parallel. On Aurora this requires CCL_OP_SYNC=1, which ezpz_setup_env sets; see AuroraBugTracking#174.

The full set is listed under Examples, and ezpz benchmark runs them in sequence with a summary report.

Other systems

The same workflow applies on Polaris and Sophia (PBS, nccl) and on Perlmutter (SLURM, srun); the NERSC specifics are in the Perlmutter guide.

Getting help

Issues and questions: saforem2/ezpz. Common failure modes are collected in Troubleshooting.