Skip to content

🍋 ezpz

Write once, run anywhere.

ezpz makes distributed PyTorch launches portable across NVIDIA, AMD, Intel, MPS, and CPU—with zero-code changes and guardrails for HPC schedulers.

It provides a:

  • 🧰 CLI: ezpz that provides utilities for launching distributed jobs
  • 🐍 Python library ezpz for writing hardware-agnostic, distributed PyTorch code

  • 📝 Pre-built examples:

    All of which:

    • Use modern distributed PyTorch features (FSDP, TP, HF Trainer)
    • Can be run anywhere (e.g. NVIDIA, AMD, Intel, MPS, CPU)

Checkout the 📘 Docs for more information!

🐣 Getting Started

  1. Setup Python environment:
    To use ezpz, we first need a Python environment (preferably virtual) that has torch and mpi4py installed.

    • Already have one? Skip to (2.) below!

    • Otherwise, we can use the provided src/ezpz/bin/utils.sh2 to setup our environment:

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

      [Optional] Note: This is technically optional, but recommended.
      Especially if you happen to be running behind a job scheduler (e.g. PBS/Slurm) at any of {ALCF, OLCF, NERSC}, this will automatically load the appropriate modules and use these to bootstrap a virtual environment.
      However, if you already have a Python environment with {torch, mpi4py} installed and would prefer to use that, skip directly to (2.) installing ezpz below

  2. Install ezpz1:

    1
    uv pip install "git+https://github.com/saforem2/ezpz"

    • Need PyTorch or mpi4py?

      If you don't already have PyTorch or mpi4py installed, you can specify these as additional dependencies:

      1
      uv pip install --no-cache --link-mode=copy "git+https://github.com/saforem2/ezpz[torch,mpi]"

    • ... or try without installing!

      If you already have a Python environment with {torch, mpi4py} installed, you can try ezpz without installing it:

       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
      # pip install uv first, if needed
uv run --with "git+https://github.com/saforem2/ezpz" ezpz doctor

TMPDIR=$(pwd) uv run --with "git+https://github.com/saforem2/ezpz" \
    --python=$(which python3) \
    ezpz test

TMPDIR=$(pwd) uv run --with "git+https://github.com/saforem2/ezpz" \
    --python=$(which python3) \
    ezpz launch \
        python3 -m ezpz.examples.fsdp_tp

  3. Distributed Smoke Test:

    Train simple MLP on MNIST with PyTorch + DDP:

    1
    ezpz test

    See: [📑 ezpz test | W&B Report] for sample output and details of metric tracking.

🐍 Python Library

At its core, ezpz is a Python library designed to make writing distributed PyTorch code easy and portable across different hardware backends.

See 🐍 Python Library for more information.

✨ Features

  • See 🚀 Quickstart for a detailed walk-through of ezpz features.

  • 🪄 Automatic:

    • Accelerator detection: ezpz.get_torch_device(),
      across {cuda, xpu, mps, cpu}
    • Distributed initialization: ezpz.setup_torch(), to pick the right device + backend combo
    • Metric handling and utilities for {tracking, recording, plotting}: ezpz.History() with Weights & Biases support
    • Integration with native job scheduler(s) (PBS, Slurm)
      • with safe fall-backs when no scheduler is detected
    • Single-process logging with filtering for distributed runs

📓 Examples

👀 See Examples for ready-to-go examples that can be used as templates or starting points for your own distributed PyTorch workloads!

🧰 ezpz: CLI Toolbox

Once installed, ezpz provides a CLI with a few useful utilities to help with distributed launches and environment validation.

Explicitly, these are:

1
2
3
ezpz doctor  # environment validation and health-check
ezpz test    # distributed smoke test
ezpz launch  # general purpose, scheduler-aware launching

To see the list of available commands, run:

1
ezpz --help

🧰 CLI Toolbox

Checkout 🧰 CLI for additional information.

🩺 ezpz doctor

Health-check your environment and ensure that ezpz is installed correctly

1
2
ezpz doctor
ezpz doctor --json   # machine-friendly output for CI

Checks MPI, scheduler detection, Torch import + accelerators, and wandb readiness, returning non-zero on errors.

See: 🩺 Doctor for more information.

✅ ezpz test

Run the bundled test suite (great for first-time validation):

1
ezpz test

Or, try without installing:

1
2
3
4
TMPDIR=$(pwd) uv run \
    --python=$(which python3) \
    --with "git+https://github.com/saforem2/ezpz" \
    ezpz test

See ✅ Test for more information.

🚀 ezpz launch

Single entry point for distributed jobs.

ezpz detects PBS/Slurm automatically and falls back to mpirun, forwarding useful environment variables so your script behaves the same on laptops and clusters.

Add your own args to any command (--config, --batch-size, etc.) and ezpz will propagate them through the detected launcher.

Use the provided

1
ezpz launch <launch flags> -- <cmd> <cmd flags>

to automatically launch <cmd> across all available3 accelerators.

Use it to launch:

  • Arbitrary command(s):

    1
    ezpz launch hostname

  • Arbitrary Python string:

    1
    ezpz launch python3 -c 'import ezpz; ezpz.setup_torch()'

  • One of the ready-to-go examples:

    1
2
    ezpz launch python3 -m ezpz.test_dist --profile
ezpz launch -n 8 -- python3 -m ezpz.examples.fsdp_tp --tp 4

  • Your own distributed training script:

    1
    ezpz launch -n 16 -ppn 8 -- python3 -m your_app.train --config configs/your_config.yaml

    to launch your_app.train across 16 processes, 8 per node.

See 🚀 Launch for more information.

📝 Ready-to-go Examples

See 📝 Examples for complete example scripts covering:

  1. Train MLP with DDP on MNIST
  2. Train CNN with FSDP on MNIST
  3. Train ViT with FSDP on MNIST
  4. Train Transformer with FSDP and TP on HF Datasets
  5. Train Diffusion LLM with FSDP on HF Datasets
  6. Train or Fine-Tune an LLM with FSDP and HF Trainer on HF Datasets

⚙️ Environment Variables

Additional configuration can be done through environment variables, including:

  1. The colorized logging output can be toggled via the NO_COLOR environment var, e.g. to turn off colors:

    1
    NO_COLOR=1 ezpz launch python3 -m your_app.train

  2. Forcing a specific torch device (useful on GPU hosts when you want CPU-only):

    1
    TORCH_DEVICE=cpu ezpz test

  3. Changing the plot marker used in the text-based plots:

    1
2
3
4
    # highest resolution, may not be supported in all terminals
EZPZ_TPLOT_MARKER="braille" ezpz launch python3 -m your_app.train
# next-best resolution, more widely supported
EZPZ_TPLOT_MARKER="fhd" ezpz launch python3 -m your_app.train

➕ More Information

  1. If you don't have uv installed, you can install it via:

    1
    pip install uv

    See the uv documentation for more details. ↩

  2. The https://bit.ly/ezpz-utils URL is just a short link for convenience that actually points to https://raw.githubusercontent.com/saforem2/ezpz/main/src/ezpz/bin/utils.sh ↩

  3. By default, this will detect if we're running behind a job scheduler (e.g. PBS or Slurm). If so, we automatically determine the specifics of the currently active job; explicitly, this will determine:

    1. The number of available nodes
    2. How many GPUs are present on each of these nodes
    3. How many GPUs we have total

    It will then use this information to automatically construct the appropriate {mpiexec, srun} command to launch, and finally, execute the launch cmd. ↩