Scaling Deep Learning Training with MPMD Pipeline Parallelism

Anxhelo Xhebraj, Sean Lee, Hanfeng Chen, Vinod Grover
12/18/2024
cs.DCcs.LGcs.PL

No confident code implementation found for this paper. 1 possible match can be shown below.

Abstract

We present JaxPP, a system for efficiently scaling the training of large deep learning models with flexible pipeline parallelism. We introduce a seamless programming model that allows implementing user-defined pipeline schedules for gradient accumulation. JaxPP automatically distributes tasks, corresponding to pipeline stages, over a cluster of nodes and automatically infers the communication among them. We implement a MPMD runtime for asynchronous execution of SPMD tasks. The pipeline parallelism implementation of JaxPP improves hardware utilization by up to $1.11\times$ with respect to the best performing SPMD configuration.

View on arXivView PDF

Code Implementations

No confident code match yet

We couldn't find an author-owned or strongly-evidenced community implementation for this paper. 1 weaker match is hidden by default — verify before relying on them.

No code implementations found yet.

Know of an implementation? Let us know in the comments below!

Similar Papers

Cite this paper

@article{xhebraj2024scaling,
  title  = {Scaling Deep Learning Training with MPMD Pipeline Parallelism},
  author = {Anxhelo Xhebraj and Sean Lee and Hanfeng Chen and Vinod Grover},
  year   = {2024},
  eprint = {2412.14374},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2412.14374v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2412.14374.svg)](https://codeofpaper.com/paper/2412.14374)

Discussion