Lightning Talk: Accelerating PyTorch Models With Torch.compile's C++ Wrapper Mode - Bin Bao, Meta

This talk discusses the acceleration of PyTorch models using the torch.compile C++ wrapper mode presented by Min Bal from Meta. The speaker begins by addressing the importance of identifying model bottlenecks, which can be categorized as compute-bound, memory-bound, or CPU-bound. They explain how optimizations through faster hardware, quantization, and kernel fusion can mitigate these issues. The core of the discussion involves the torch.compile architecture, particularly how replacing Python with C++ for orchestration can reduce CPU overhead and improve performance. The talk also delves into challenges associated with C++ compilation and the integration of custom operations within the framework. The results showcase improvements in performance across various benchmarks, while noting that increased compilation time is a trade-off for these enhancements.

My name is Min Bal, I'm from Meta, and today I'm going to talk about accelerating PyTorch models with torch.compile C++ wrapper mode. So, I will start with the motivation, and then dive a little bit deeper into the technical details, some of the challenges for this problem, and how And towards the end, I'm going to show more data, and then I'll mention the future work. So, if you're

in this talk, I I assume you all care about how to speed up your PyTorch model, right? So, obviously, in order to speed up your model, first you need to understand the bottlenecks of your model. Roughly speaking, like PyTorch models, like if they have bottlenecks, then roughly speaking, they can be categorized into three groups, right? Compute-bound, memory-bound, and CPU-bound. So, for compute-bound, it basically means like your

model uses some heavy ops that consumes a lot on your GPU. And to to optimize that, people can use like faster hardware, or you you can use like lower precision and quantization, those things. To help with the memory bottleneck, and people can use technologies like kernel fusion, which TorchInductor is really good at. And also, a lower precision can still help here because you just basically have smaller

data footprint. Right? And then, another kind is called CPU-bound, which means like your kernels are still waiting for your CPU side to finish all the work before you can launch your kernel. You uh you're kind of wasting your GPU resources waiting for the CPU side. So, I want to mention that uh which kind of bottleneck it is for your model is highly your implementation dependent, right? Like

some models, they may be like compute-bound, some may be the dynamics can change as you optimize through all those bottlenecks, right? Let's say your model is compute-bound, but then as you optimize it, then suddenly you start to see like CPU part becomes the bound. For the CPU-bound, uh one common technology called CPU CUDA graphs, and that I think there was a talk yesterday from Daniel from Nvidia.

And then, basically, it says like CUDA graphs works well when it works, but sometimes it doesn't apply. So, here it comes another solution, we call it C++ wrapper. A a quick dive into, like, you know, recap about torch.compile architecture. So, basically, we have TorchDynamo as a back front end that captures graph. And then, we use a graph and feed it into AOTAutograd, which does, you know, further

like backward capture, and then decomposition, functionalization. And then, we feed the graph to TorchInductor. So, what comes out of TorchInductor, this compiler, is basically two parts, right? Like, one is compute kernels, basically like Triton by default, is what gets executed on your GPU. And the other part is wrapper code. And by default, code is in Python, and what it does is like it orchestrates all the kernel

launches, you know, prepares things, does tensor allocation, synchronization, things like that. So, remember, it's by default it's in Python, right? And we all know Python is kind of slow if if compared to C++, so C++. The idea is quite simple here, right? In order to cut your CPU side of overhead, why don't we just replace Python with C++, right? And then, we can save the overhead from

like running your Python execution. So, for those who heard like AOTInductor before, like a quick mentioning here is AOTInductor is a solution we had for non-Python deployment. So, if you look at that, it's you know, it goes through the export path, it's not really like torch.compile path, it's a it's a different export path. a graph, and you compile it you you you compile it in a binary,

and you can use it from C++ for a non-Python environment's inference. So, what's different here is C++ wrapper is really works with torch.compile. It's still like in the this like JIT compile domain, and it gets compiled. It's not a really like ahead-of-time compilation. And then, but the good part of that about that is you can still live with Dynamo graph breaks, right? You can just compile like

each sub graph into C++. So, obviously, it works with torch.compile means like it works inference both. But they two thing really shares the C++ code back end generation part. So, how do you enable that? Well, it's quite simple, there's a config for Inductor, you just set C++ wrapper to true, or you can use with this environment environment variable. Actually, you will need a C++ compiler on your

host. let me dive a little bit deeper into some of the challenges, right? Like, with C++ wrapper, suddenly now you start to see the separation between compile time and run time, right? With Python, you don't Uh with C++ wrapper, like now we have to compile your C++ code at the compile time, and then at the run time, you kind of run what's being compiled. Uh for the

major part of the wrapper code is kind of straightforward, right? But then, what gets tricky is for the kernel parts. Remember I said the kernels are generated as Triton. Triton's basically in Python. So, how do we use that in your C++ wrapper, right? So, there are two ways to solve this problem. One is say, I kind of eagerly compile my Triton kernel at the compile Right? At

that time, uh there's like no real execution, what you can do is it's based on your the the the input to your Triton kernels, you can generate the some random inputs based on the size, based on the stride, state, D type, all those things, and use that to eagerly trigger Triton compilation auto tuning. That's one solution, but there's problem with that solution is sometimes, let's say, your

tensor may be an integer tensor, which is going to be used as an indexing into your other tensors, right? Randomly generating inputs for for those might cause IMA issues. sometimes like real inputs gets you better like tuning results. So, another way in against like, you know, comparing to the eager compilation is we call it lazy compilation. So, basically, at the C++ compile time, C++ compile time, you

don't really compile the Triton kernel first. What you do is you just generate a bunch of those stop functions calls, and then at the run time, the first time you see the code, you're going to trigger actually Python compilations into those kernels. So, you at the compile time, you just basically can still only compile C++ code. At the run time, you you rely on the calling back

uh trigger those Triton compilation and auto tuning. So, that's what we call like lazy compilation. Yeah, there's a model once features like, you know, TMA on those kernels. Similarly, we can do something like this, you know, we we basically generate those stop function calls then trigger corresponding like, you know, Triton code into at the run time. I'm going to skip that just sorry. Also, for torch.compile, there's

uh the compile PyTorch compiler team tried really hard to make sure it composes PyTorch features, right? One of them is custom ops. then how do we handle that in this context, right? How do we do that in with C++ wrapper? there are two kinds of ways we can register custom ops in PyTorch. One is more like the tradi- you write your kernel in C++, you call the,

you know, registration APIs and kind of register your your ops. And that is actually a easier path for the C++ wrapper because under the hood, you still have the C++ code. Right? So, all we need to do is at the compile code gen time, we generate code and trying to do the lookup at the run time. There's another group of custom ops basically registered through Python. And

then, at that, you don't really have C++ implementation for those code. So, how do we call that at a compile time? Sorry, at a run time, right? Uh with the C++ wrapper, what we do is we generate a, you know, corre- and it's we can call back into Python at the run time. So, we still rely on Python to do the op lookup, and then call the

ops at run time. Uh I mentioned CUDA graphs at the beginning, so it's another technique people can reduce CPU overhead. So, how compared to each other, right? So, basically, my view of this is they they they compose. Like, they can help each other. It's not a, you know, one against the because like CUDA graphs, sometimes it works, sometimes it doesn't, then for the CUDA graphs, like let's

say your model has breaks. Dynamo graph breaks, then you have 10 sub graphs. Right? Maybe five of them can use CUDA graphs, five of them cannot. Then well, you can easily apply the C++ wrapper on the rest of the CUDA graph non-CUDA graphs. So, they should really like help So, let's jump into the result part. This is measured on three benchmarks suites that we commonly use in

a lot of like torch compile like papers or online like nightly [clears throat] dashboard. So, basically hugging face team and and torch bench. Measured on three different kind of GPUs and they all use the same kind of CPU. the the Let me explain the number. So, basically it's a speed up over the default inductor, which is a Python wrapper. Right? And then H100 on This is a

geo mean of all the like benchmark suites in that in that suite. Right? And then hugging face we see like 4% improvement and and teams in one and torch bench like 8%. You may think the number is not really that large, but uh it really depends on your Like teams and teams to be tends to be like smaller speed up because they use more the original models.

They use more like compute heavy apps like convolution. Right? Remember at the beginning like if you if your model is a compute bound, then CPU overhead may not be that much. But if you're like on the other hand, if you're like there tends to be like smaller It's not heavily like compute bound. CPU bound tends to be more like obvious. So, that's why you see torch bench

has a larger bar over there. And what's interesting is about like if we move on to the B200, right? Because faster. And then remember that I said that if you want to reduce your compute bound, one way is to to bump up to a you know faster hardware. So, that's exactly what's happening here. Like if you have B200, suddenly you start to see more CPU bound here.

That's why you see you know larger B200. And then at the last but not the least is we can do something similar on the MI300 side, right? So, which means this technology is really not specific to any GPU. It's really across platform. Yeah, I'm going to skip this and cuz And quickly mention like you know It's not all like all good. It's we we are trading something

for this good performance, right? One thing we're trading for is a compilation time. Because we have to use C++ compiler. And then there are techniques we can be used to to reduce these uh compilation you know, precompile headers and and code splitting not doing like inlining or lower optimization level and things like this. So, there's a bar here that basically shows like you know, each techniques can

help us to further reduce the compilation time. So, roughly speaking at this stage is is you know, 1.3x to to 2x slow down comparing to default inductor in terms of compilation time. Yes. Now it comes to this end. Basically, it's available already today on nightly. And I have linked some IPRs. IPRs if you want to dive into some technical details, you can take a look So, for

future work uh Like one thing is we have to make it composable with more non-default optimizations. The coverage is just not there yet. And then also we need to further reduce compilation time. Okay.