Lightning Talk: Graph Based Pipeline Parallelism - Sanket Purandare, Meta & Simon Fan, Meta PyTorch

This talk introduces graph-based pipeline parallelism, a novel approach developed by the Meta PyTorch team for enhancing model training efficiency. The speaker explains how traditional pipeline parallelism creates idle time on devices, referred to as pipeline bubbles. To tackle this issue, various scheduling techniques, including smaller microbatches and operation reordering, are outlined. Advanced methods such as zero-bubble pipeline parallelism and dual pipe scheduling are discussed, focusing on their ability to overlap forward and backward computations to optimize performance. The speaker highlights the complexity of manual backward writing in PyTorch, leading to the proposal of a graph-based solution. This new method promises significant flexibility and customization in pipeline scheduling, leveraging a compiler to manage backward passes and threading automatically. Finally, experimental results showcase impressive speedups on multiple GPUs, validating the effectiveness of the proposed approach.

Bonjour, hey. Uh, my name is Simon and uh, this is my colleague Sanket and uh, we're from the Meta PyTorch team. Today we're here to present you graph-based pipeline parallelism, a new solution to author the most complex pipeline schedules. So to recap, pipeline parallelism shards your model on the model depth dimension. Now this is different than sharding on the model dimensions as in theory it lets you

keep larger blocks for your compute so that you can run them more efficiently. Uh, but in practice the new problem introduced is now you have pipeline bubbles. So if you were to naively pipeline your model so that you can see that in gray is most of the time we actually spent uh, the devices are actually idle. Uh, and so schedules were developed to counteract this point. Uh,

so you'll see schedules introducing smaller microbatches, increasing the number of microbatches, uh, and reordering the operations so that the later devices can run sooner in order to uh, reduce the Uh, and more advanced techniques like introduced by zero bubble pipeline parallels it is uh, the decomposition of the backwards. So we exploit the fact that the backward pass uh, doesn't need to run in full for the next

stage to run. So you only need the backwards with respect to the inputs uh, so for the next stage to start and then you can delay the backward with respect to weights uh, until later in the pipeline to fill another gap. Uh, and when the original authors uh, published this uh, they included code examples forked from Megatron uh, where they basically had to manually write their entire

backwards for all of the model layers which they wanted to do this decomposition on. Uh, and later on the distributed pipelining package introduced uh, new abstractions to help author this but even today we're still seeing issues with composability with compile and activation checkpointing. Uh, and even more recently more advanced schedules have come out. So Dual Pipe which was designed specifically to hide the MoE communications from expert

parallel uh, uses the microbatch uh, the forward of a micro batch and overlaps it with the backwards of another micro and in the operator level schedule is very fine grain and follows a specific sequence. Now, authoring this is still very difficult today in PyTorch and we don't really have a good way for it because you need to coordinate between multiple Python threads and it's usually if you

and there's more complexity involved if you want to go with a different order than what you had written in your original model code. And so this generally requires the full manual backward writing. And so we summarize the pain points of today's authoring. You have limited controls over the backwards which generally means that you will write your manual backwards. You have to handle streams and threads management and

if you're not careful you will hit a lot of composability issues with compile and So we introduce graph based pipeline parallel where if you are willing to put up with the upfront cost of model tracing we promise you the full control of your backwards provided by the compiler. We'll promise you that you don't need to manage your threads and streams manually but you can if you want

and we'll show you that this is compiled first where you can get out of the box performance with inductor. And most importantly we show that it is highly customizable. These are graph passes that you can apply as you wish and customize to your desired pipeline schedule. So briefly go over the flow of how we do this as you would do the same as you usually do which

is you split the model into stages, you parallelize it, you compile the forward and backwards and then for graph based pipeline parallel you would then apply graph transforms to extract the graphs that are needed for your particular schedule and you would integrate those into your runtime. So let's build some intuition. Suppose you have a MLP forward here's a sample code. The compiler would provide you with the

full backward graph for this trace code. So in blue you have the inputs to this graph and in red you have the outputs. And we show that backward decomposition is actually as simple as choosing which outputs you want to keep and discarding all the nodes that are not involved in this competition path. So, here you would have a graph for computing only the gradient with respect to

inputs. And then similarly, you can do the same in order to get a graph that only computes the gradient with respect to weights. So, this suffices for the zero-bubble schedule. For the dual pipe schedule, the overlap forward we use a slightly different technique that we'll call multiplexing. So, in multiplexing, we'll combine these two graphs into a disjoint graph. And this disjoint graph has the same inputs and

outputs as the original two graphs. But now now you have nodes that are independent from each other in this graph, so you can interleave them and reorder them as you wish. And this allows you for powerful combinations of expressing which the the specific overlap we need and we show that it's sufficient to do dual pipe. So, now I'll be handing it off to Sanket who will explain

how these graph transforms are being used in our pipeline runtime. Thanks, Simon. So, now I'll introduce the schedule IR and this IR can be used to describe any schedule that is used in pipelining today. So, we have a bunch of actions like execute forward, execute backward, or execute backward only with respect to input and so on. And using these set of canonical actions which are indexed by

the stage and microbatch respectively as you can see on the screen, we can describe any pipeline schedule today. So, now this is here is an example of a compute-only pipeline parallel. And then we add PP comes to it. So, essentially you're able to describe any schedule with this IR. And this IR is currently upstreamed in torch distributed pipelining and most schedules can be described using it. So,

now as Simon described, the entire workflow goes as follows. You have the whole model and then you split the your NN module into different stages. Each stage you apply your favorite SPMD and then finally you obtain a joint forward and backward graph. Next uh you can apply your activation checkpointing or uh communication fusion also called as bucketing passes to obtain the optimized joint forward and backward graph

corresponding to each stage in the pipeline. Then you partition your uh joint graph into forward and backward The next optimization to think about is you have FSDP collectives which shard and unsharded your parameters at stage boundaries and you need to do them only once. Because once you have unsharded your parameters, you essentially want to execute all the micro batches with it and then reduce at the end.

So we take the forward and backward graphs and split them into two components. The forward graph with no FSDP collectives and the unsharded graph and similarly for backward graph, the backward with no FSDP collectives and the reduce grad graph. Finally, as Simon had described, we split the backward graph into two parts. The backward that computes the gradient with respect to input and the one that And for

dual pipe-like schedules, you have a joint action which executes the forward and backward together. So you take two different graphs from stages, one corresponding to forward and backward. And if first you produce a concatenated graph and then you apply a graph pass that does the forward backward overlapping. So in summary, you have these seven graphs corresponding to each action in the schedule IR. Now what happens at

runtime is uh we maintain the following states. We have graph callables. You have per micro batch caches and send and receive buffers. And for each stage we maintain the sharded and unsharded parameters and grads. And let's say you encounter an action. So the pipeline run time will choose the right graph corresponding to the action, use the right micro batch, and use the send receive buffers and store

the outputs in the cache. So this is how the run time operates. Finally, we did some experimentation on 64 GPUs of eight comprising of eight nodes on H100s and with a smaller version of DeepSeek model. And when we compare the performance of compile and eager and since we are graph based first and always composable with different techniques, we get out of the box compiler performance and you

can see for dual pipeline schedules we see a 70% speed up over eager. For zero bubble we are able to see 60% and for interleave 1F1B around 50%. Next, if we want to strip away the compile benefits and only see what our graph passes offer, here is the breakdown. Where by using zero bubble we are getting 7% speed up over 1F1B and then further by using dual

pipe schedule to overlap MOE comms, we are getting around 7 17% speed up. So we had promised you that we'll also provide flexibility for scheduling. And what you can do is just annotate your user code with different annotations and then these annotations get forwarded to your FX graph nodes. And once you have these annotations for your graph nodes, you can schedule them however you want. So, another

thing which you can do is during your warm-up and cool-down stages, you usually don't have the forwards and backwards, but you have only one of them. So, in this case, you can do a fine-grained annotation. And with this fine-grained annotation, you can use the forward, chunk it into two, uh let's say nano batches, and then essentially define a schedule that overlaps different chunks in forward and overlaps

different chunks in backward. So, by using this fine-grained scheduling, we give you the user full their pipeline schedule, and which can be just written as a graph pass. Thanks.