torch.compile and Diffusers: A Hands-On Guide to Peak Performance - Sayak Paul, Hugging Face

Key Takeaways

Hi folks, uh thanks for attending this session. My name is Shyuk. Uh I work as a research engineer at HuggingFace. I work on image generation, video generation and also a little bit of kernels work and that requires me to all to ma work with a lot of torch compile stuff which is exactly what I'm going to present today. It's not going to be a general uh introduction to torch compile or anything like that. My talk will focus more on how diffusion based workflows uh integrate with torch comp. I also want to get an understanding of how many of you work with image and video generation and related stuff mostly LLMs you work with image videos and so on. But anyway, the fun thing is the takeaways that I'm going to be sharing today. Uh these are model agnostic. Some specific to diffusion, but most of them will will be fairly agnostic uh in nature. So with that in mind, let me actually start. And just as a disclaimer, this is not my sole work. uh it took a village to get torch compile uh to get us the want uh from diffusion specific workflows. So we collaborated with animesh from meta and I also had my colleague Benjamin who also helped us quite a bit and this is roughly going to be the agenda. uh since you folks do not image and video generation a lot I will try to give you the context of the library that I work on at hugging face the diffusers library and then I'll also spend some uh time discussing what torch compile is and how that integrates uh within the diffusers library and then this the second point the meat of this talk where we will discuss how we can set up our modeling implementations up for success to get the maximum benefits out of torch compile and I'll also present some exciting diffusion specific workflows uh that work well with torch compile and what what are need to make uh in order to ensure that I'll also talk about setting up tests because no one really write test cases for torch compile except us I'm very proud about it so I'm going to spend some time discussing that topic as well and hopefully there will be some time for Q&A but I'll be around uh so if you feel free to grab me and ask questions uh so yeah let's start so diffusers uh it's a library that's primary maintained at hugging face it's an open source python library for state-of-the-art image video audio and text stuff late uh diffusion based text generation ising up so we also provide support for diffusion based text generation as well uh We prioritize performance and accessibility. Uh and it's also PyTorch native. Uh we used to support Jax as well. Uh but then we realized it's not worth it probably. Um yeah, let me also give you uh of the library. If you uh how many of you have worked with the transformers library? I expected more. Uh but the flavor almost remains the same. like you have the same from pre-trained API you have the same torch DT type component transformers they changed it to DT type my bad uh have to do the device placement explicitly you start with a prompt a cat holding a sign that says hello world and voila it's only a few lines of code since I care care about code indentation I have written it like that but if you do not care about all of that you could have fit it in probably five lines but anyway this you go from a prompt all the way to this cute cat. So yeah, now torch compile. Here's the big elephant in the room. It's a just in time compiler. It's not ahead of time. Uh uh you can compile your models or functions through torch compile and if it's a torch.n.module, you can also call the compile. Uh it's known for its uh user friendliness and it's also known for delivering speedups right out of the gate. if your model is implemented in the right way which is what I'm going to discuss and it deeply integrates uh into diffusers which is also something uh I am going to be talking about now here are some sharp bits although torch compile is known for delivering excellent speed ups right out of the gate it also has its sharp bits and yeah usage is easy as I mentioned but getting the maximal benefits because you are spending ing dollars utilizing those H100, GB10s and so on. You want to squeeze every penny out of those investments, right? So getting the maximal benefits out of torch compile is not always that obvious because your modeling implementation can have graph breaks. It can have recompilation issues. It can also have other issues postings CPU overheads and so on. So all of these issues can get in the way of maximum torch compile benefits. So you would want to be definitely careful uh about these things and of course no one likes to spend time when the model is compiling. It can take ages. We are we are not that patient time can also be pretty high. And the good news is I'm going to be sharing some tricks to reduce that timing uh time down at least by 5x. So hold on. Um I'm going to be sharing my slides so you do not have to click pictures so we can focus on the discussion. Um and also just to give you of the complexity involved in diffusion workflows. Diffusion models are not single models. Um this is what it takes to go from this text prompt all the way to this cute cat. So it's like a pipeline interconnected between different components. So let's let's unpack that a little bit. you have text. Now, we need some way to, you know, get some internal representation out of this text prompt, right? So, you of course would want to use text encoders. Uh, modern diffusion models often use multiple text encoders. I'm going to be uh you uh taking the flux model as my example. It uses two text encoders. Another one is T5XL in order to get some internal representations out of your prompt. And then once we have the text embeddings, you start with some noisy latence. The details are not important here. Just just just internalize uh the schematics. You start with noisy latence and step because diffusion is an iterative process. You have to be aware of the time steps and then you feed all of that back into a diffusion backbone which is almost always transformer based and then you invoke uh the diffusion transformer in an iterative manner. You get the refined latency that back into a decoder and you get the actual RGB values. So from getting uh from this text prompt all the way to this image it takes a couple of models. It's not just a single model like language or vision language models right so you need to be aware of this facts like when I say actually a pipeline it's more justified to call it a pipeline because it's not a single model and let's also try to see who is the big elephant in the room it's the diffusion transformer that takes the most amount of compute is the slowest model in the room all the other components coders and the decoder It takes only a tiny bit of compute like all the other components are just one shot but the diffusion transformer is invoked multiple times and it's unlike GANs. In GANs it's one shot but they are almost dead. Diffusion is the uh new hot kid. So but it uh it's the most compute intense uh component in the pipeline. So this is where we will be spending all of our time in. This is exactly the component that we will try to compile and along the way we will see what it takes to set it up for success with torch compile. So that's the what's coming in this talk. Now yeah let's let's take the big elephant. So this is where we will be focusing all of our efforts. So yeah now I want to get into what's most important in this talk. How can you set set your graph for success so that it delivers almost maximal gains with torch compile? So quick gains eliminate graph breaks as much as possible. And some common patterns include numpy ops in between your forward implementations. Please do not do that. It's not that hard to eliminate numpy ops and hard to eliminate those item calls because if you have data dependent calls within your forward or whatever methods torch compile is going to break. So get rid of those dot item calls and also these kinds of control flow patterns because the compiler does not know what x dots sum is going to be during run. So try to write your implementations in a way that avoids all all of these patterns. And how do you know that your model does not have a graph break? Well, instead of compiling without full graph equals to, you compile it with full graph equals true and that's it. Immediately once you in first call, if there are graph breaks, graph breaks it, it's going to complain. So use compile with full graph equals true so that you know in your mind that my model is yeah my my model does not have it have any graph breaks recompilation triggers. So would be mismatches in shapes device and dtypes because the torch compiler specializes for certain shapes uh and it puts guards around them. So as the minute those guards are vi violated uh it triggers uh a recompilation. So get rid of recompilation issues as much as possible know that there's recompilation. Well instead of just calling your model you put that call under a with context uh and you do it with error on recompiles and that's it. So at this point in time you know if your model has full uh graph breaks if your model is going to trigger any compilations. So eliminate those and you are already 50% out of the way to get maximum torch compile benefits. So that's that's good news. That's already good news. Now I want to spend a bit more time around eliminating recompilation triggers because these obvious uh so shapes are a big thing because maybe you want a different resolution image. Maybe you do not always want a square image. Maybe you want a rectangular image. Maybe you compiled your model for 1024 by 1024 resolution. Now you want to scale it up to 2048. That's obviously going to cause recompilation if you are not careful enough. Right? So what you do is you mark that tensor region with mark dynamic and that should just work. And if you do not know which tensors are are going to be dynamic, you can specify dynamic equals true. Although that's recommended and instead of using plain Python scalers, wrap it around torch tensors because let's say your function has something like x and c and it compiled for that specific constant and if that constant changes, it's going to trigger recompilation. So don't do that. Just wrap it around torch tens. That's it. Um sometimes recompilations are unavoidable. uh if you have worked with offloading like you offload certain model parameters to CPU and you onload them as you are doing computations to save up on memory those will obviously cause device mismatches and in those cases you can't really compilations so those are probably acceptable if you are prioritizing memory savings right and if you know your recompilation limits beforehand maybe specify the recompilation limit so that you are absolutely sure that it's not going to trigger recompilation more than twice or whatever that number is. Now, if you want to know more about these kinds of triggers and how to optimize your torch compile guards, I can recommend this reading. It's from the Dynamo uh team itself, Anim. So yeah um my recommendation is when you are writing your model use torch compile already like you for your training to complete you do not have to wait for post- training to complete as soon as you are done with your implementation throw some dummy inputs at your modeling implementation with torch compile and see what are the obvious issues that you can eliminate because that kind of forces you to write write your model in a way that with torch compile and it's not a big deal and if you are having to spend too much time in it maybe it's a feedback for the torch compile team. Uh so yeah report it and of course we do not want to wait for my model compilation to complete because it takes ages right regional compilation the idea is you identify the most compute inensive subm modules within your models like for example a model that's transformer based the transformer subm modules will be the most comput inensive almost always right and All of those blocks usually the same compute graph but their parameters change. So you can actually compile that compute graph and reuse it. So that's the idea behind uh regional compilation. So only compile those and in code it it's extremely simil like extremely simple. So yeah you an idea of how regional compilation looks like for a simple transformer based model. Um yeah use regional compilation. Now you are not convinced. You in your head you are probably asking does it work? Like does it give me the speed up? Does it actually reduce the compilation time? I I don't like to blabber much. I'll give you the numbers to do the talking. Delivers the same inference speed up. cuts down the compilation time about 7x. So yeah, now we are talking believe in numbers. So yeah, you have your numbers. Yeah, hopefully I was able to convince you that you should use regional compilation as much as you can and yeah, maybe impress others. Oh yeah. Um so in diffusers I wanted to spend some time around discussing how we integrate this regional compilation thing in diffusers. So we have this method called compile repeated blocks which exactly does what it says. So instead of compiling the transformer with the compile method with full graph equals true repeated blocks method with full graph equals true and in order for this work uh we expose this attribute called repeated blocks. So for the flux model which is a fairly modern and state-of-the-art model from Black Forest Labs, we just specify which are the repeated blocks like which names and that's it. So yeah and huge thanks to Animesh from the PyTorch team for contributing this directly to diffusers. Uh it was really great working with him and he's the one that shipped uh this method and championed the whole you know paradigm of regional compilation. So yeah, now in the context of diffusion, there still can be, you know, little quirks that can get in the way of maximum torch compile benefits. For example, device to host transfer where your GPU is sitting idle because your CPU is not fast enough and there could be, you know, CPU overheads as well is I hate to be the bearer of bad news, you have to profile. So and profiling can cannot always be trivial uh because you need to know what exactly those places uh that you should look for in a torch profile trace and it can be pretty complicated. So this is why we uh where you where we walk you through an entire torch profile trace of you know modern models like flux and we tell you what to look for so that you can eliminate these issues uh and benefit from torch compile even further. Now I wanted to also discuss the popular features uh from the diffusion world integrate with torch compile and this is offloading and quantization. So typically modern image and video generation models can take 60 GBs or 70 GBs but we are GPU poors uh at our hearts. So we don't cannot afford that uh that kind of VRAM. So we usually offload our parameters but also quantize. Offloading alone uh sometimes is not enough. So we also have to quantize. Now when you are offloading and when you are quantizing you essentially can also uh incur recompile uh and I think that's fine but the point is compilation can still be beneficial. So uh and I'm going to present numbers to make my point so bear with me. So what you have to do in order to make all of this work at least as far as offloading is concerned you have to disable uh theing and offloading region within your implementation with this uh disabled decorator and that should take care of its compatibility with torch compile. And when you are working with quantization in order to make it work with like in order to make it compatible with torch comp register the quantization ops uh so that the compiler can identify them during runtime and once it's done once you are basically combining offloading with quantization you are of course reducing the footprint that's the goal we want to have like we are from 34 gigs all the way down to 12 gigs. Um you are of course getting benefits of regional compilation. You can also use regional compilation here but you also get get inference benefits like without compilation you have about 12 seconds and with compilation you have about 9.8 It's a 20% 20% decent speed up I would say and it's it's free like why would you not do it if it's free like it's just one line of code change and it's done but the thing is we do the heavy lifting for you like we identify which region in the implementation should be marked with torch compiler disable uh care of the fake registration ops and so on. So yeah all of that heavy lifting is already done for you you just have to be able to use them. So yeah now Laura I think you must have worked with Laura if even if you have worked with LLMs Laura adapters to customize your LLMs and so on uh in the diffusion world as well because you can get models to follow particular culture particular ethnicities uh and so on. So and Laura if we take if we try to approach it from the first principles when you inject or when you load an adapter insideally changes the architecture right it fundamentally changes the parameter space and so on and not all Lora uh come with same ranks like maybe you loaded a lura with the rank 200 maybe then you loaded another lura with 32 rank and so on. So since these ranks can vary from one another, recompilation issues are bound to happen uh like this is obvious. We try to solve it uh by following a set of constants. For example, we we want to load the Lora with the maximum rank first and the subsequent Loras are within that rank and if they target the same set of modules or a subset of modules, we won't be incurring such recompilation triggers because because the compiled graph will have already recognized the kinds of modules it should target and the kind of weight shapes it should specialize on. Yeah, these are like the minimal set of constraints that we impose in case you are hot swapping different Lora adapters and the code also looks fairly simple because we do the heavy lifting for you already. So you basically load the first Lora adapter with the maximum rank and then you compile it. You are free to do whatever hot swapping you want to do. So yeah, that's the idea. And does it work? Boy, does it work. It does. Hell yeah it does. 2.5x uh 2.05x sorry it does 2x is a lot of speed up. So it was the benchmark was conducted on flux two adapters with the maximum rank of 32 but it can scale. Uh so yeah it does work with compile like you get free benefits but in order to make it work you had to as I as I discussed you had to take care of some changes. So yeah, of course we care about tests. All the features that I just discussed are all thoroughly tested. We test for graph bras. We test for dynamic shapes. We test for offloading and compilation. We test for quantization and compilation. We also test for loa hot swapping because as someone untested features are not features. So yeah, here is your complete test suit. So yeah, but tests are nice. We do not want to be causing any benchmark regressions because compile is supposed to deliver you speed ups. Now in order to optimize our want to cause benchmark regressions like we do not want the model to get slowed down, right? So we have a regular benchmarking workflow that lets us spot any improvements or regressions because if there's a if there's a decrease in the throughput we should worry we should get concerned we have all of that set up and that's how we try to harden the oper operalization of torch compile within diffusers. So yeah, all the all the stuff that I presented is actually uh from a blog post which is already public. But here are some takeaways for you as well for the portability and less framework overhead. You can also uh AOT compile uh most of this stuff that works and yeah the Laura related compilation stuff is here in this blog post and yeah you can scan the QR if you want to see the slides. Think I have finished two minutes earlier. Oh yeah I have a minute. So I can take a quick question and get the stage ready for the next speaker or you can ask. Yeah, go ahead. Blackwell. Um, uh, my benchmarks were from H100, it's equally well to Blackwell, like if you, especially when you have larger B sizes. It also works on RTX 4090 as well. Um, yeah, thank you. Feel free to grab me for any questions outside. So, thanks.