[Distributed w/ TorchTitan] Training with Zero-Bubble Pipeline Parallelism

username: H-Huang
cooked: <p>with <a href="https://github.com/H-Huang" rel="noopener nofollow ugc">Howard Huang</a>, <a href="https://github.com/wconstab" rel="noopener nofollow ugc">Will Constable</a>, <a href="https://github.com/kwen2501" rel="noopener nofollow ugc">Ke Wen</a>, <a href="https://github.com/soulitzer" rel="noopener nofollow ugc">Jeffrey Wan</a>, <a href="https://github.com/haocizhang" rel="noopener nofollow ugc">Haoci Zhang</a>, <a href="https://github.com/donglimm" rel="noopener nofollow ugc">Dong Li</a>, <a href="https://github.com/vivien-chu" rel="noopener nofollow ugc">Weiwei Chu</a></p>
<h1><a name="p-461795-tldr-1" class="anchor" href="#p-461795-tldr-1"></a>TL;DR</h1>
<p>In this post, we’ll dive into a few key innovations in <code>torch.distributed.pipelining</code> that make it easier to apply pipeline parallelism, including zero-bubble schedules, to your models. And we’ll highlight an end to end example of training LLMs with <code>torch.distributed.pipelining</code> composed together with FSDP and Tensor Parallelism in TorchTitan, and share learnings that helped improve composability and clean up the model.</p>
<ul>
<li>Support for manual model partitioning or compiler-based graph splitting</li>
<li>Declarative schedule format (IR) decouples schedule writing from executor implementation, and provides a library of schedule generators.</li>
<li>Autograd support for automatic zero-bubble avoids the need for a custom autograd.function and manually partitioned backwards</li>
</ul>
<h1><a name="p-461795-distributed-with-torchtitan-series-2" class="anchor" href="#p-461795-distributed-with-torchtitan-series-2"></a>Distributed with TorchTitan Series</h1>
<p>The GitHub repository <a href="https://github.com/pytorch/torchtitan" rel="noopener nofollow ugc">TorchTitan</a> is a proof of concept for large-scale LLM training using native PyTorch, designed to be easy to understand, use, and extend for different training purposes, supporting multi-dimensional parallelisms with modular components. In this series of topics, we introduce the latest PyTorch features for distributed training enabled in Torchtitan.</p>
<ul>
<li>Topic 1: <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323">Enabling Float8 All-Gather in FSDP2 </a></li>
<li>Topic 2: <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487">Introducing Async Tensor Parallelism in PyTorch</a></li>
<li>Topic 3: <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250">Optimizing Checkpointing Efficiency with PyTorch DCP</a></li>
<li>→ Topic 4: Training with Zero-Bubble Pipeline Parallelism</li>
</ul>
<h1><a name="p-461795-introducing-pipeline-parallelism-in-pytorch-3" class="anchor" href="#p-461795-introducing-pipeline-parallelism-in-pytorch-3"></a>Introducing Pipeline Parallelism in PyTorch</h1>
<p>Pipeline Parallelism is notoriously difficult to apply to a model, and also more bottlenecked by bubble overhead, compared to data parallelism. However, recently with the invention of the <a href="https://arxiv.org/pdf/2401.10241" rel="noopener nofollow ugc">zero-bubble schedule</a>, pipeline parallelism is more competitive, and it’s an important ingredient for large scale training (and for distributed inference as well).</p>
<p><code>torch.distributed.pipelining</code> APIs make it easier to apply pipeline parallelism to your model and experiment with new schedule formats. For more detailed API documentation, see our <a href="https://pytorch.org/docs/main/distributed.pipelining.html" rel="noopener nofollow ugc">docs page</a>. For more usage guidelines, see our <a href="https://pytorch.org/tutorials/intermediate/pipelining_tutorial.html" rel="noopener nofollow ugc">tutorial</a> and <a href="https://github.com/pytorch/torchtitan" rel="noopener nofollow ugc">TorchTitan</a>.</p>
<h2><a name="p-461795-model-partitioning-schedule-definition-and-execution-4" class="anchor" href="#p-461795-model-partitioning-schedule-definition-and-execution-4"></a>Model Partitioning, Schedule Definition, and Execution</h2>
<p>Pipeline parallelism in PyTorch is built upon three fundamental abstractions: model partitioning, schedule definition, and execution. These are embodied in the classes of <code>PipelineStage</code>, which handles model partitioning, <code>PipelineSchedule</code> which defines the schedule based on a configuration of stage counts, ranks, or microbatches, and <code>PipelineScheduleRuntime</code> which manages execution.</p>
<p><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXcVT4MSaXd-j8aYiYTB4NQQiE5nswjWm-AUyq022mvawFXLveOoGcJHLstAWjlRb2HrA5VvxHy_ygxLCry4Wu2U-QCeE5ovr52oYSrITgIufuFK9uIxlJRviPj7socbjkdpMVBi?key=MJ_pH9FsYYWRWqmAS2OMtQ" alt="" role="presentation" width="690" height="273"></p>
<p><code>PipelineStage</code> provides methods for shape inference, ensuring that the dimensions of tensors are correctly determined as they traverse through different stages. When these stages communicate, the tensor shapes they exchange are contingent upon the input microbatch’s shape. During execution, stages engage in metadata exchange to guarantee the precise transfer of data, accompanied by validation of expected shapes. This process is crucial to prevent any potential hangs, as it ensures that all stages are in agreement regarding tensor shapes.</p>
<p><code>PipelineStage</code> also offers the flexibility of both <a href="https://pytorch.org/docs/main/distributed.pipelining.html#option-1-splitting-a-model-manually" rel="noopener nofollow ugc">manual</a> and automatic splitting. Users can manually partition models using an <code>nn.Module</code> class-based wrapper, or they can opt for <a href="https://pytorch.org/docs/main/distributed.pipelining.html#option-2-splitting-a-model-automatically" rel="noopener nofollow ugc">compiler-based splitting</a> through torch.export. This dual approach empowers users with greater adaptability, enabling them to develop utility functions that can partition a model in various ways to suit their specific needs. Compiler-based splitting, while simpler for some users, can be harder to compose with other parallelisms, so we generally recommend manual splitting.</p>
<p>The pipeline stage abstraction was chosen to more easily integrate with the execution engine; however, this integration is not mandatory. The components of our solution are designed to be orthogonal, allowing them to be adopted independently. This modularity means that while we provide a comprehensive execution engine for pipeline schedules, users are free to utilize our schedule intermediate representation (IR) to craft their own execution strategies, bypassing our API constructs if desired. Details about the IR are covered in the next section.</p>
<h2><a name="p-461795-schedule-authoring-via-ir-5" class="anchor" href="#p-461795-schedule-authoring-via-ir-5"></a>Schedule Authoring via IR</h2>
<p>One of the challenges of implementing or prototyping a new pipeline schedule is that it’s tricky to write send/recv communications in a correct way in SPMD-style code. Concretely, pipeline schedules tend to involve lots of loops and conditionals for deciding which action to issue, and also require send/recv operations to be well overlapped, which can result in code that’s fragile and can lead to hangs that are difficult to debug. That’s why we introduced a new IR-based schedule format, allowing you to break the schedule writing process down as illustrated in the following example for a trivial zero-bubble style schedule:</p>
<ol>
<li>Generate a simple “compute only” schedule trace for each rank (programmatically). Here, <code>1F0</code> means run stage 1’s 0th microbatch Forward operation. Stages are globally numbered.</li>
</ol>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXfoGGzvJl5GXl5R4lY4JSVzhNL4gHsDfZiH4YfVY7CXOCp3ntdDpy8G2esjFXLhZpcpMcEAchvXtQad00dWFTGsh4HG8LI_RSUoSR04Cw3aBFtw7X4qC2cPRzrHJ262Li9lPbtF?key=MJ_pH9FsYYWRWqmAS2OMtQ" alt="" width="340" height="194" role="presentation"></strong></p>
<ol start="2">
<li>Apply compiler passes to insert communication operations where needed and overlap them</li>
</ol>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7bfa8b50353350c249268a08e5222821b8e5e13e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7bfa8b50353350c249268a08e5222821b8e5e13e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bfa8b50353350c249268a08e5222821b8e5e13e_2_282x250.jpeg" alt="image" data-base62-sha1="hGLw4AnXzhBmFfHxI3t8q2Mv4Rg" width="282" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bfa8b50353350c249268a08e5222821b8e5e13e_2_282x250.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bfa8b50353350c249268a08e5222821b8e5e13e_2_423x375.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7bfa8b50353350c249268a08e5222821b8e5e13e_2_564x500.jpeg 2x" data-dominant-color="E4E5E6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">642×568 69.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<ol start="3">
<li>Execute the IR on an execution backend</li>
</ol>
<p>These steps are all easy to inspect and test, since the IR is serializable to CSV format or can be printed or rendered in chrometrace. That also means that if you want to make a small tweak to your schedule to see how it behaves, you’re free to directly edit the trace and run it without first figuring out how to update your generator code.</p>
<p>We also provide <a href="https://github.com/pytorch/pytorch/blob/88154024b300552aeaa9c11b9d77839ffa04bce0/torch/distributed/pipelining/schedules.py#L2387" rel="noopener nofollow ugc">tools to simulate</a> the execution of the IR just for correctness (to flush out potential ordering issues or invalid schedules) before finally executing on hardware. During development, these tools helped quickly identify bugs.</p>
<p>The IR format is also easy to extend- for example, we have exploratory work integrating with FSDP’s zero-3 where IR is used to signal the unshard and reshard operations for data parallelism.</p>
<p>You may also notice that there are separate operations in the IR for ‘FULL_BACKWARD’ (B) and ‘BACKWARD_INPUT’ (I) / ‘BACKWARD_WEIGHT’ (W). This allows for generating schedules like the zero-bubble schedule, but in order to actually execute these schedules, you’ll need a way to separate these components out of the ‘backward pass’. That’s why torch.distributed.pipelining adds support for automated backwards partitioning.</p>
<h2><a name="p-461795-autograd-support-for-zero-bubble-6" class="anchor" href="#p-461795-autograd-support-for-zero-bubble-6"></a>Autograd Support for Zero Bubble</h2>
<p>The zero-bubble paper’s key insight was recognizing that only part of the backward pass, which computes input gradients, is on the critical path to the next pipeline stage. The rest, which computes weight gradients, can be computed separately to fill in pipeline bubbles.</p>
<p>In order to run a zero-bubble schedule, you need to first partition the backward operation. In the original zero-bubble work, this is done manually which can be difficult (<a href="https://github.com/sail-sg/zero-bubble-megatron-deepspeed/blob/5a36a7c9b7b4b8ad76015973c2eb722d9631e53f/light/zbpp_light/deepspeed_engine.py#L29" rel="noopener nofollow ugc">example</a>), but may be preferred by some users. To support manual cases, we introduce <code>stage_backward_input</code> and <code>stage_backward_weight</code> <a href="https://github.com/pytorch/pytorch/pull/132691" rel="noopener nofollow ugc">methods</a>. For other users who want more of the simplicity of pytorch and autograd, we developed an automatic partitioning feature.</p>
<p><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXfAgirnVBayRvk3HU0KiN2ni68r6GTBaiRRSJVizbw6EEN82XMw1jgdswZsxjfHnEBHyYjTMSTIRBt3sGZo8V_6i_eMzxFzxTJYRMmHCQODZYwBHHXLEOm8RL4Ylyk0WfS_oCNxQ63Tu2Tb5zMzoGZP9w?key=MJ_pH9FsYYWRWqmAS2OMtQ" alt="" role="presentation" width="690" height="429"></p>
<p>Most PyTorch users are familiar with the concept of writing only the forward pass and getting the backward pass for free. This is made possible by the Autograd system, which captures a graph of backward operations as forward is run. By simply traversing and partitioning this graph, we can automatically generate functions for computing the input gradient and weight gradient of a particular pipeline stage. Critically, we save shared values used by the weight gradient computation so they do not need to be recomputed.</p>
<p><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXenw694yUaPVBP4bCGxKwLW3oTGagXn_SCLPOLtewovyV1x-W1juWc7-cabuAF4ostPUdOtvHtY8eJNsWEbZHrioTt0wSzdf6F1LjXgzVc6LOX4VgxixKd3mhgl4ebabHlm34fQ?key=MJ_pH9FsYYWRWqmAS2OMtQ" alt="" role="presentation" width="690" height="436"></p>
<p>Additionally, we have verified that memory usage is correctly released as the schedule is executed, as illustrated in the memory profile below:</p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXcPZYW3Dg_NHkZmeBGa591MuMWrtgwJJWDcuSEHYlsqkx6tLZX9VoqdUVWQdZhibXtWR0e2TakZD0s3NEaK2RUlsFn3U3KzvHsHimzNEQxLUBdwzG02_vcxzm0P7MNyOfU2GW8d?key=MJ_pH9FsYYWRWqmAS2OMtQ" alt="" width="488" height="384" role="presentation"></strong></p>
<p>To allow better modularity, users can opt to use automatic partitioning if they want, or write their own backward or weight grad update function. They can register their custom weight grad update function to be called during execution.</p>
<p>Currently torch.compile and zero-bubble are not compatible, but we are exploring how to enable them to work together.</p>
<h1><a name="p-461795-learnings-from-torchtitan-7" class="anchor" href="#p-461795-learnings-from-torchtitan-7"></a>Learnings from TorchTitan</h1>
<h2><a name="p-461795-setting-up-the-pipeline-friendly-training-loop-8" class="anchor" href="#p-461795-setting-up-the-pipeline-friendly-training-loop-8"></a>Setting up the Pipeline-Friendly Training Loop</h2>
<p>One challenge we faced when pipelining torchtitan was keeping train.py clean while accommodating the changes needed for pipeline parallelism. When pipelining, there can be multiple (virtual) stages per pipeline rank, each represented by one nn.Module and having its own Optimizer and LR scheduler. We found it was possible to refactor the training loop code in several ways to make pipelining changes easier, and minimize the number of branches or duplicate codepaths for pipeline vs non-pipeline execution. The key changes we landed on are summarized below:</p>
<ul>
<li>Keep a list of ‘model chunks’ and ‘optimizers’, rather than single objects, helping unify the checkpointer calls</li>
<li>Apply ‘SPMD parallelisms’ and torch.compile inside a ‘per-model-chunk’ helper</li>
<li>Factor the loss function as a helper that can be used inside the pipeline executor or called in the non-pipelined train loop</li>
</ul>
<h2><a name="p-461795-changes-inside-the-model-9" class="anchor" href="#p-461795-changes-inside-the-model-9"></a>Changes Inside the Model</h2>
<p>To make manual partitioning smooth, we introduced a couple of <a href="https://github.com/pytorch/torchtitan/blob/main/docs/composability.md#simplifying-the-top-level-model-forward" rel="noopener nofollow ugc">changes</a> to the model code:</p>
<ul>
<li>Switching .layers from an nn.ModuleList to nn.ModuleDict made it so that fully-qualified-names (FQNs) did not change after partitioning, which ensured smooth checkpointing</li>
<li>Writing model.forward() in a way that <a href="https://github.com/pytorch/torchtitan/blob/4f7f8831b4adbe7e6dad00f74e188e34fe2a584a/torchtitan/models/llama/model.py#L439" rel="noopener nofollow ugc">tolerated ‘None’ layers</a> let us partition simply by <a href="https://github.com/pytorch/torchtitan/blob/4f7f8831b4adbe7e6dad00f74e188e34fe2a584a/torchtitan/parallelisms/pipeline_llama.py#L80" rel="noopener nofollow ugc">deleting unused layers from a full model</a> copy for each stage. (e.g. Stage 1+ deletes the input embedding, and .forward() code handles this transparently).</li>
</ul>
<h2><a name="p-461795-model-initialization-and-rng-configuration-10" class="anchor" href="#p-461795-model-initialization-and-rng-configuration-10"></a>Model Initialization and RNG Configuration</h2>
<p>Model initialization is another area that needs to be changed for pipelining. The two key changes were (1) updating init_weights to be pipeline-stage friendly, and (2) configuring the RNG properly for pipelining.</p>
<p>Torchtitan uses meta-device initialization since it avoids initializing ‘the whole model’ on one device, which is critical for larger models that can’t fit in device memory. To implement meta initialization, you must write an <a href="https://github.com/pytorch/torchtitan/blob/4f7f8831b4adbe7e6dad00f74e188e34fe2a584a/torchtitan/models/llama/model.py#L381" rel="noopener nofollow ugc">‘init_weights’ function</a> that can be called after the model is sharded and materialized. For init_weights compatibility, we just needed to make sure it tolerated the missing layers which we deleted when applying pipelining.</p>
<p>To ensure that initialization results in unique values per stage, we must <a href="https://github.com/pytorch/torchtitan/blob/4f7f8831b4adbe7e6dad00f74e188e34fe2a584a/torchtitan/utils.py#L89" rel="noopener nofollow ugc">configure RNG seeds appropriately</a>. Since torchtitan uses DTensor based SPMD parallelisms, we mainly need to ensure that each SPMD world starts out with one shared RNG seed, because <a href="https://github.com/pytorch/pytorch/blob/88154024b300552aeaa9c11b9d77839ffa04bce0/torch/distributed/tensor/_random.py#L96" rel="noopener nofollow ugc">DTensor manages RNG offsets itself</a> to ensure different initialization of different model shards. With pipeline parallelism, we need one unique seed per pipeline rank, since each pipeline rank corresponds to a ‘SPMD world’.</p>
<h1><a name="p-461795-conclusion-11" class="anchor" href="#p-461795-conclusion-11"></a>Conclusion</h1>
<p>To demonstrate the composability of PP with that native Pytorch distributed APIs, we ran experiments on TorchTitan on a LLaMa3 405B model using 3D parallelism (FSDP + TP + PP). We also combined this with new techniques such as torch.compile, <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487">Async TP</a> and <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323">Float8 Allgather</a> and ran various job sizes ranging from 8 GPUs to 512 GPUs. We include the data of these runs in the <a href="https://arxiv.org/pdf/2410.06511" rel="noopener nofollow ugc">TorchTitan paper</a> and show a select few below:</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>GPUs</th>
<th>Schedule</th>
<th>Memory</th>
<th>Words per second (WPS)</th>
</tr>
</thead>
<tbody>
<tr>
<td>256</td>
<td>1F1B</td>
<td>87.20GiB</td>
<td>103</td>
</tr>
<tr>
<td>256</td>
<td>Interleaved 1F1B</td>
<td>89.53GiB</td>
<td>109</td>
</tr>
<tr>
<td>512</td>
<td>1F1B</td>
<td>65.79GiB</td>
<td>96</td>
</tr>
<tr>
<td>512</td>
<td>Interleaved 1F1B</td>
<td>68.44GiB</td>
<td>124</td>
</tr>
<tr>
<td>512</td>
<td>ZeroBubble (without compile)</td>
<td>82.86GiB</td>
<td>125</td>
</tr>
</tbody>
</table>
</div><p>In summary, we hope our work on <code>torch.distributed.pipelining</code> makes pipeline parallelism more accessible and efficient for large-model training and inference. By leveraging features such as model partitioning, a declarative schedule format, and automated backwards partitioning, we are able to provide a wide range of complex schedules as defaults to the user. The integration with TorchTitan further enhances the composability of different parallelism strategies, providing a robust example for production settings. As we look to the future, ongoing work will focus on automatic optimizations in the schedule runtime, hardening the IR format, and upstreaming different schedules. We welcome any contributions and suggestions from the community on this work.</p>