Reproduce “Scaling Deep Learning Training with MPMD Pipeline Parallelism” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Anxhelo Xhebraj, Sean Lee, Hanfeng Chen, Vinod Grover
JaxPP, developed by NVIDIA, enhances JAX's SPMD programming model by introducing a task-based, Multiple-Program Multiple-Data (MPMD) runtime for flexible, user-defined pipeline parallelism in large deep learning models. This system achieved a 44.6% speedup over JAX SPMD pipeline parallelism and a 1.11x throughput improvement compared to JAX FSDP on GPT-3 175B, demonstrating 92.87% weak scaling efficiency on up to 1024 GPUs.
Key insights
- MPMD pipeline schedules, such as Interleaved 1F1B, significantly reduce memory requirements and activation rematerialization costs compared to SPMD-based GPipe schedules.
- Asynchronous Point-to-Point (P2P) communication efficiently overlaps computation and communication, minimizing overhead in distributed training environments.
- Optimizing gradient accumulation for models with weight sharing by transforming the summation of partial gradients reduces inter-actor communication, especially for large embedding tables.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.