Enabling State-of-the-art Asynchronous Execution in Torch.compile With CUDA Streams - Michael Lazos

This talk discusses the integration of state-of-the-art asynchronous execution in Torch Compile using CUDA streams, presented by Michael Lazos, a software engineer at Meta. The speaker explains the concept of streams as different execution queues that allow concurrent kernel execution for better compute overlap and memory transfer hiding. They delve into the architecture of Torch Compile, particularly how Torch Dynamo tracks streams and how the changes made to AOT Autograd enable synchronized backward passes while preserving stream orderings. The talk further explores the design challenges encountered during implementation, use cases such as microbatch overlapping and activation offloading, and the resulting improvements in performance and memory efficiency for transformer models in PyTorch. Finally, the speaker highlights the upcoming release of these features in version 2.12 of the framework.

Hi everybody. Um, so we're going to go ahead and uh get started. And so uh today I'll be talking about my work uh enabling state-of-the-art uh asynchronous execution in Torch Compile uh with CUDA streams. Um, just quick introduction. I'm Michael Lazos uh and I'm a software engineer at Meta. So first off I'll go over a quick agenda. So first we'll talk about uh what are streams, um

what are they used for, and um then we'll talk about the overall architecture of the system and how I integrate it into Torch Compile. And then we'll go through a bunch of design challenges we had when uh implementing this work. And then we'll talk about some use cases, uh namely microbatch overlapping and activation offloading. And then we'll cover some results and if we have any time left

uh we can do some Q&A. So first off I'll I'll give some background here. So what are streams? Um, so streams you can think of as uh uh different execution queues. And so they can run concurrently and so you can run multiple kernels at the same time. And uh typical use case for this is to do comp compute overlapping and uh memory data transfer hiding. And um

whenever you do concurrent execution you usually need some method of uh performing synchronization. And so um this is no different and events are used to do the synchronization between different streams. And uh finally one thing that's uh less known about streams is that they can be used to synchronize uh across devices as well. So you can use events um to synchronize uh across different GPUs on the

same host. And so here's just a a quick example usage to give some detail. Um, so here you can have GPU zero and GPU one um and kernel A and kernel B will run on two different streams. And um the way this works is kernel A will start and then there'll be an event record which captures the work in progress on that stream. And then you can

then wait on that event which will block until the work that was recorded uh completes. And then kernel B will uh begin execution uh after uh kernel A has completed. And so at a high level, here's um the entire torch compile system and the changes we had to make to implement the stream handling. So at the top here is torch dynamo and um this is the the

graph capture mechanism in torch compile. And um the changes we had to do there were to uh symbolically track the current stream as your Python code executes. We had to annotate uh the FX graph nodes as they're uh captured. And we also had to populate a stream and event table um which uh we'll get more into in the next slide. And then in the middle here, we

have AOT autograd. Um so this uh phase is used to generate the backwards graph and to do a few different uh compiler passes. And um the changes we had to make there is we had to make the backward pass uh synchronized. And we also had to preserve uh the different stream orderings and um ensure that nodes are never reordered uh across event boundaries. And then uh finally

we end with torch inductor, which has uh which does the codegen and scheduling. So this handles codegenning the different stream assignments and emitting the uh event record and waits. And then finally we also have to um restrict uh fusions across streams and events. And um um and then we also had to uh prevent memory use across streams as And so now we'll um this is a a

high-level overview. We'll dive deeper into each of these, so um to to give more detail cuz this is a a uh just a a bird's-eye view. And so first off is how do we track the streams in torch dynamo? So this piece is um pretty simple for the most part. Um the user code here on the left has uh some kernels running on different streams. The way

you do this in eager is you have these context managers like with uh cuda.stream s1, which we defined up here. And then um these kernels are all are all running on different streams. The the most recent uh stream context manager entered is the uh one that uh the the kernel will run on. And so in Dynamo we just uh match these stack semantics with a a symbolic

stack. And you can see here the state of the stack as we execute this code. And um on the right there is the FX graph that's traced. And um the stream assignments of each of the nodes uh that are traced. And uh the next piece here is uh how do we actually support uh streams and events as as graph inputs? So here you'll have um kernel A

and kernel B. Like if we're compiling just this region without these two streams, we need to pass these in uh as arguments to the graph. Um and doing that in in torch compile would actually be very difficult because today AOT autograd uh only supports tensors as arguments. And the reason that is is because uh it only supports basically objects that you can compute gradients for. And so

changing that would require a large refactor of AOT autograd. And so to get around this when torch Dynamo sees uh this code it will rewrite the the bytecode of the the Python function to call uh the trace And so um before we call the traced graph we pass in a uh or we have a global object table that we write to um with some bytecode right before

we call the graph. And that's uh this piece here. And so each stream and event will have an index into that table. And then in the FX graph um you'll just read from that table instead of having to explicitly pass the streams or events as inputs. And that's the pieces over here where you set um where you see get user object by index and uh set current

stream. So you're basically reading the the um stream and event object from the table, setting the current stream, and then and running the the kernel on that And um you can see on the right here that this is kind of verbose, and it'd be kind of, uh, weird to see this in your graphs all the time. Um, and it'd be kind of hard to read and track

the the different streams. And so, um, the way we actually represent, uh, the different stream ops in the IR is we kind of hide this lookup behind, um, these these different stream ops that we, uh, defined. And so, these are kind of the familiar ops that I talked about earlier where we have record event and wait event, um, and wait stream, and, um, and so, these will

basically internally do the get user object by index lookup. And then, instead of doing a set stream for each node in the graph, um, we'll store the, uh, stream index on, uh, the metadata of each graph node. And so, the bottom here shows, uh, an example FX graph with all of the different, um, stream indices assigned to the nodes and the, uh, event uh, or the different,

uh, ops that we use there. yeah, like I said, the the record event and wait event, those indices, uh, correspond to a stream and an event And so, now we'll, uh, go into the AOT Autograd piece. Um, so, the trickiest part here is we, uh, need to generate the backward pass, and we want to match the same parallelization strategy that you use in the forward pass. So,

whatever the user did, uh, in the forward pass, uh, we want the backward uh, the backward pass nodes that correspond to the forward pass kernels to have the same, uh, parallelization semantics. So, they should have the same streams. And, uh, the reason we want to do this was cuz that just matches eager. So, we want to compile to be as faithful as possible to eager. And so,

uh, in this example here, um, basically, what we would do is, uh, we'd go when when we, um, trace the forward pass out, when AOT Autograd is generating the backward pass, um, we just propagate those stream indices from the that I showed in the previous slide. Each node has an assigned stream. We just propagate that to the backward pass, so the backward nodes will also have the

same stream. Um, and then, uh, we do a a pass over uh, all those nodes and um, we assign or we add synchronization. And so, the way we do that is we just look at a kernel, and if any of its arguments are on a different stream than the kernel itself, we insert the synchronization. And then, uh, the one last wrinkle here what happens for nodes that

don't have an associated forward pass node. And so, in this graph here, um, forward A and forward B both touch W. Um, they both use so, when when you're going to uh, generate the backward pass for this, um, you're going to have a gradient accumulation node. And that node doesn't have any analog in the forward pass. And so, the way we, um, handle this case is the

same as eager once again, where we, uh, have the gradient accumulation node take the the stream of the first user of it. So, in this case, uh, grad and A is the first user of of grad accum. And, um, so, it'll be on that stream. And then, uh, we still follow the same rules before the synchronization, where if grad and B, it's on a different stream than

grad accum, so we add a record event and a wait event in it. And this shows the details of the, um, uh, of the synchronizations. Uh, and so, next the the other thing we had to handle is AOT Autograd will often, uh, reorder nodes uh, to do perform different optimizations. One example of this is is comp compute overlap. if you look at this code on the left

here, um, this e.record is going to record kernel 1, 2, and 3. And wait will, uh, the e.wait here will wait until those complete before it does kernel 4. Um, and uh, this is really important cuz if you if AOT Autograd reordered kernel 4 before the wait, then, uh, the program would no longer be correct. There'd be a race condition. And so, um, the way we handle

this is, uh, we basically add fake dependencies. So, we look, uh, here is this this op called control depths Um, this allows you to add arbitrary inputs to another op. And so we utilize this to we do an analysis pass to find the tensors that are used that basically cross the record and weight boundary here and that they are used on a different stream. So we know

A, B and C are used on a different stream S1 here. And so we want to pass through A, B and C through this new control depths wrapped record and depths wrapped weight. And that will basically allow you to have the extra dependencies that will prevent the reordering. So on the right we get A, B and C from kernel 1 2 3 and then they'll pass through

the event record and these are just equivalent to these arguments we're just passing them through and creating an alias. And the same thing happens here with the event weight. And so now any graph pass that we have in AOT autograd will respect those dependencies cuz it's we're just making the dependencies explicit in the graph. The next piece is how we handle functionalization and input mutation. So I'm

going to give some background here first to add some detail before we talk about the actual problem we're going to need to address. But for for the background here on the left we have this user code which has an in-place op here mul. And let's say X is a graph input. So this would be an input mutation. What functionalization does is it takes this mul and makes

it out of place. So instead of doing the in-place op where we mutate X directly we'll just create a copy which you can see here. and run the rest of the graph. And then in order to be correct the caller of this will still want to see the mutation after the compiled region returns. So after this function returns X should should be updated. And so the way

we do this in torch compiler is we add this copy underscore epilogue to actually perform the mutation update. And the reason we do this is that functional graphs are way easier to optimize. You don't need to worry about aliasing and you don't need to worry when you're reordering ops whether you're going to violate move an op where it might be before mutation occurred or after. And so

this basically makes writing the the compiler passes a lot easier. and since the caller still needs to see the mutation we we still have to do this copy underscore at the end. Um so this piece I should set off some alarm bells cuz we moved a copy underscore we moved the update from the front of the graph to the back. And in the previous slide we showed

um that we don't want to move things across event boundaries. And so um this shows the actual issue that will happen when we do that. So the user has a record here any event outside the compiled region waiting on this will expect X to have the new value. And the issue here is if we move the the the update for X past the record the the the

e.wait here will trigger too early and this function will see old old version of X. we kind of just cop out here and we analyze and find this exact scenario and we throw a compiler error. So you can actually restructure this code the user will have to restructure to basically manually functionalize this. basically create a copy and use the copy down here instead. Um there's cases where

this can still work where you don't need to do that all of the dependencies of this record if all the weights are within the compiled region then we're okay. So things will still work and we won't need to actually do that. So, that's where we opt to do a compiler error here um because uh if we wanted to handle this, we would have to do some uh

pretty large refactoring of AOT Autograd to handle uh a a non a non-functionalized graph. And then um the final piece in the um overall architecture is how we uh generate streams uh with TorchInductor. And so, um here you can see uh without streams, Inductor will just perform all these kernel fusions. Um you can take the ReLU mul and the add and uh generate a fused kernel of

all of those. Um the issue is if your um if your kernels are operating on different streams, uh you shouldn't be able to do that fusion cuz the user has indicated they want to uh parallelize the execution of those two kernels. And so, the main piece we had to do in Inductor was to just analyze uh which nodes had um different uh what the stream assignments of

each node and if kernels were supposed to run different streams, we shouldn't fuse them. And then uh one piece that I don't show in this slide is um uh buffer reuse. We also want to make sure that if any um any tensor uh should only be reused on the on the stream that it was uh created And then kind of uh along those lines, uh we'll talk

about cross-stream memory safety. this is kind of a case that's more subtle for users of streams. Um you have to always ensure that when you allocate a stream, the uh all uses of that stream on side streams are finished before you do another allocation. And what I mean by that is uh in this example, uh we'll have this user code. Um it'll use uh this kernel B

on a side And if we uh delete X, you can see that we haven't done any synchronization here. So, kernel B might still be using it. And um we'll then reallocate here, and this might reuse X. So, the the way the CUDA caching allocator works is it has a dedicated memory pool for each uh stream. And so, it tries to be as as efficient as possible. So,

once uh there's no references to X left, it'll just uh try and reuse that memory. And the correct way to handle this is to add a synchronize and then then delete. So, that way you know that any um uses of X uh in this case, you know kernel B was running on a side stream, we're going to wait on S1 to make sure that stream is finished

before we perform the way we handle this uh in in torch uh is that we need to make sure that uh we respect uh when a user calls del X. Um so, in the usual case, in the non-cross stream case, when um we see a del X, um we actually just ignore it because torch dynamo and uh torch compiler in general um is um knows the exact

lifetime of every tensor. So, you know the last usage where the last usage is in the graph is, and so you know when you can call del. Um for the multi-stream case, um cuz of the previous example I showed you, you don't you want to respect where the user called del, cuz it might be at a later it might be at a later time uh where the

other stream is finished with it. here we call um uh the synchronize here and and call del X and then do the reallocation. And the way we handle this in dynamo is we and we find the specific case, and we insert an op in the graph which does the synchronization. So, that's called a sync dealloc op. And um we end up with a timeline there on the

right, which is actually correct. Um this new allocation will only happen after that stream is finished with it. Um and then the now we'll talk about applications. Um So, the first application I want to talk about is uh microbatch overlap. So in this um scenario the user might have a bunch of compute ops with some com ops mixed in. And so this is really common in the

distributed workloads where you'll have an all reduce or an all gather. we really want to try and overlap these to get better latency. So you can see on if we overlap the communication and the we can get better latency. And so you can utilize a side stream to do that communication while the computer is happening to get better GPU utilization. Yeah, we can you can also do

this you could do this automatically to your model but we'll also respect if the user manually does this. And then another application this can be useful for is activation offloading. And so in this case it's kind of similar but instead of communication you'll use you'll be doing memory transfers during the the compute. And so this will allow you to overlap the memory transfer with the compute. And

the reason you want to do that is instead of the typical case where you may save a tensor for your backward pass from the forward pass. Um you'll be able to get lower peak memory by transferring the tensor to the CPU in the meantime. in this example you can see that we try to fully offload the activations of layer zero layer one with the computer the subsequent

layer. And then when you do the backward pass you prefetch the different activations right before they're used. And so this allows you to lower um to lower the peak memory with out any any runtime degradation ideally. So you'd have to for for that to happen your model needs to have enough compute to um basically be able to hide the the memory transfers. And so, here's some uh

results uh from Um so, I ran this on a on a transformer model with uh multiple layers. Um and here we just have the the model dim, uh the number of layers, the batch size, and the sequence length. And um this is just a um typical like LLM uh architecture setup uh with uh different transformer layers. And uh here you can see the scaling that I talked

about. So, on the right here, uh we have the potential mem uh peak memory savings. And so, um with the small model, you can see it's only like 10% and you have about 9% runtime overhead. And so, um in that case, it might not be worth it. You'd rather be using this uh where you can completely hide the runtime overhead and get pure memory savings. And so,

you can see that uh the scaling that I talked about happens as the model gets bigger, the runtime overhead goes down um because you have more compute to hide the the memory transfers. And then uh the memory savings actually go up because you can you have larger activations to hide. Uh and then uh to end on this uh we can you can use uh streams today. Uh

it's going to come out in the next uh release in in 2.12. And um it uh you can run your code as is today with uh with your stream code in eager um other than the one caveat that I mentioned before with the uh input Um and you can see the the eager streams documentation at that link to uh run it. And um that's all I have.

So, uh I'm ready for any questions. >> Um, I could say uh, the question is what about observability when you have different streams and events synchronizing. I don't know if torch profiler will show those events today, but I know that you'll at least see the the kernel stall. You'll see that the launches won't happen until later. So, if you're able to trace the kernel back to the

source code, then you'll which which I the torch profiler does support, you can see which kernels are are waiting. Yeah. Um, in that case it wouldn't get cleaned up if you've reassigned X um, cuz there's still be a reference that would exist, right? Or It's not a weight, it's a temporary Yeah, but when you call del, it wouldn't deallocate is my point cuz there would still be

a the reference count would go up when you assign the alias. You're creating an alias basically when you do X assign Y or Y assign X. And Yeah, you're creating an alias. So, when you del X um, the actual tensor memory wouldn't get deallocated. Yeah. Oh, oh, okay. I think I see what you're saying. I think it would still work because Dynamo would look at we would

look at the the Y and we'd know that that Y was multi-stream because we look we use uh um the version the version counter. I love it. But yeah, we can talk about that after. Yeah. Um yeah, today um I called this CUDA streams, but uh there is an uh a device-agnostic torch. stream, which this also supports, so you can do it on AMD as well. Um

and um yeah, I could I'm going to publish a blog post on this in the coming weeks, so I can try other devices as well to show uh some results. Yeah. Oh. that's definitely could be possible. We wouldn't allow uh synchronization across different hardwares. You would have to stick to one type of hardware. Um so you would have to if if you have mixed hardware, you'd have

to use uh the the AMD version of streams. I mean, you could use the torch an stream that I'm talking about, and that'll handle the different devices. if you change it uh in the future, I think your concern is valid. I I don't think there's any guarantee that the performance wouldn't uh change. The balance would You'd have to tune it again to balance where you want to

do the parallelization. if it's multi-node, you wouldn't use uh the same kind of events across nodes. You'd have to use actual communication. Um for a single node with eight GPUs on um I wouldn't be nervous. I would expect it to work. Yeah. Um there shouldn't be and if there is, it's probably a bug. So, yeah. Um mostly uh I'd say when PyTorch was developed, it was originally

for CUDA and so it's uh historical. Yeah. If we could have made an agnostic one from the beginning, I think we would have. Yeah. Mhm. Um so, when we would get this graph uh in Torch Inductor would actually fuse that away usually. It would fuse it into another kernel. Yeah. Well, it's just that Yeah, well, that's true actually. It wouldn't it wouldn't do this fusion, but also

this Yeah, I guess we if if it was possible to do the fusion over the record, you're correct that it would actually work again. that that yeah, that could that could potentially fix it, but Yeah, you're right that right now with the event record we wouldn't allow this anyway. Yeah. Um that's an interesting idea. I I think Maybe we could I could actually maybe loosen this restriction

even more now. If I could do some analysis to figure out, oh, maybe this could actually just fuse and there's no other there's no update on X after the record. So, that could loosen this restriction a little bit actually. Yeah. Um I guess what do you mean by the default stream? Like setting it to different things or don't know if you're waiting for everything on every other

stream. at the beginning you would have to wait for kernels I guess in this code example, you would have to wait for any kernels on the default stream to finish. So, sometimes people do insert a record and a wait to like they'd put a wait here and then put a wait on the side stream as well. And that allows you to make sure that the kernels kind

of start at the same time. So, I've seen that before, but I don't know if there's any other uh uh sequential uh piece there. Yeah. Okay, I think that's it. That's it. Thanks everybody for coming. Yeah.