Skip to content

📝 Ready-to-go Examples⚓︎

This section contains ready-to-run examples demonstrating various features of the ezpz library.

New to ezpz? Start with ezpz.examples.test for a minimal DDP training loop, then move to ezpz.examples.fsdp when you're ready to shard model parameters. The remaining examples layer on additional features — Vision Transformers, tensor parallelism, Hugging Face integration — so pick whichever matches your workload.

Prerequisites⚓︎

  • ezpz installed: pip install ezpz or pip install git+https://github.com/saforem2/ezpz
  • PyTorch (auto-detected at runtime)
  • MPI for multi-GPU training (mpi4py + MPICH or OpenMPI). Single-GPU and CPU work without MPI.

Running Examples⚓︎

All examples use the same launch pattern:

ezpz launch python3 -m ezpz.examples.<name> [args]

For example, to run the test example:

ezpz launch python3 -m ezpz.examples.test

On a laptop this uses a local mpirun; inside a PBS or SLURM job, the scheduler is auto-detected.

Picking an Example⚓︎

Links Example When to use Level
· · test Starting point — simplest DDP training loop Beginner
· · fsdp Model too large for one GPU, or memory-efficient training Beginner
· · vit Vision Transformer with FSDP + optional torch.compile Intermediate
· · fsdp_tp Very large models needing 2D parallelism (FSDP + TP) Advanced
· · diffusion Diffusion model training with FSDP Intermediate
· · hf Fine-tune causal LM with Accelerate + FSDP Intermediate
· · hf_trainer Using HF Trainer with ezpz's launcher Beginner
· · inference Distributed HF inference (benchmark / generate / eval modes) Intermediate

Fine-tuning or training an LLM?

See the Fine-Tuning & Training LLMs guide for a task-oriented walkthrough of the three paths (hf_trainer, hf, fsdp_tp) — which to pick, larger-model targets, resource requirements, and how to scale.

Model size ladder⚓︎

All 5 training examples (test, fsdp, vit, diffusion, fsdp_tp) expose the same --model {debug,s,m,l,xl,xxl,xxxl} preset ladder targeting consistent parameter counts. Architectures differ; the ladder positions are aligned.

Preset Target test fsdp vit diffusion fsdp_tp
debug smoke ~110K ~10K ~50K ~50K ~10K
s (small) ~100M 107M 76M 87M 101M 125M
m (medium) ~250M 248M 227M 204M 274M 246M
l (large) ~500M 449M 605M 632M 500M 495M
xl ~1B 858M 1.21B 1.21B 939M 1.21B
xxl ~5B 3.43B 4.84B 5.44B 5.50B 5.93B
xxxl ~10B 9.88B 9.68B 9.67B 11.4B 11.34B

Long-form aliases (small, medium, large, xlarge, extra-large, xxlarge, extra-extra-large, xxxlarge, extra-extra-extra-large) map onto the same short-name canonical presets.

Breaking change as of this release: --model small (and other long-form aliases) now resolve to the production-scale ladder above. Old toy-scale small/medium/large (~250K–1M params) are gone. Use --model debug for the laptop-runnable smoke-test path.

For Llama-specific configs, fsdp_tp additionally exposes the torchtitan-flavored agpt-2b (1.99B params) and agpt-20b (20.74B params) presets, which reproduce the AuroraGPT registry exactly.