Slaying OOMs with PyTorch FSDP and torchao

Yeah, so hi everyone. So we're here to talk about slaying ooms. So ooms are maybe notoriously and one of the most annoying bugs to deal with in PyTorch code. And so both Jane and I are developers on PyTorch Core at Meta. And so we wanted to sort of talk to you about a lot of the tools we've been building to make this process a bit easier to deal with. So traditionally, the way I've seen people deal with ooms is this way, which is basically, People see an Umm, they're like, okay, crap, like what do I do? And then the sort of two like nuclear options are you like either reduce your batch size or you like reduce the size of your model. You go like, oh, this thing's too big. Like let me just have something smaller. But this is like a very like coarse tool and there's like, you know, certainly a lot more finer-grained things you could do with like a bit more knowledge. So the first constraint you have is essentially like how much VRAM you have on your chip. So for example, for 3090s and 4090s, which are like very popular consumer cards, you have about 24 gigs. And then for the A100s, you have like either 40 or 80. I think the H100 is like about 100 if I recall correctly. But my point is, is that like you're always gonna be VRAM starved. And specifically for consumer cards, you know, if I were to speculate, like I would speculate that the 5090 probably also has like around 24 gigs of VRAM. And so we're always gonna be in the sort of like VRAM constrained environment. But again, as we're thinking about memory for a model, instead of just thinking about, oh, it's like something is blah gigabytes, let's be a bit more concrete about what's involved in estimating the size of a model. So you have 3 core buckets. So for example, let's say you say, oh, I'm downloading Llama 7B. 7B is referring— 7B is the number of parameters. And if each of those parameters is an FP16, then you need 2 bytes per parameter, which means that the total size of the model is about 14 gigs. And indeed, if you download Llama 7B on your disk, that's roughly how big it's going to be. But that's not enough. But if that was enough, you could just run Llama 7B on a 3090. But you can't. And how come, right? So the reason why is, well, if you're doing fine-tuning or training, you also need to basically per parameter, you have gradients. And your gradients will also be in FP16. And so you basically end up with another 14 gigs. And then finally, like a detail everyone seems to forget all the time is also the size of your optimizer. So the single most popular optimizer used in the world is Adam. And Adam's like, the amount of memory that Adam takes is twice the amount of parameters. So basically if your parameters are 14 gigs, Adam takes 28. So if you sum all of these up, you get like 14 plus 14 plus 28, which is 56 gigs, which is bigger than 40 gigs. So bigger than most GPUs that, that people have. So this is sort of the traditional, what people would refer to as a full fine-tuning. I also sort of neglected to mention activations. So activations are basically the— let's say you're running, for example, you have your weights. It takes certain times a certain input, wx. So the output of that is your activations. Activations tend to dominate the VRAM of your model at larger batch sizes and larger context lengths. So that's why optimizations like flash attention are really important. And I'm not going to talk too much about activations, like, in the beginning. But they're a bit harder to estimate. Like, they don't have as clean of a formula as the gradients and optimizers do. But there are tools that can help you estimate it, which are a bit better than just doing math that I've found can be a bit finicky and error-prone. Anyway, so the first thing you might think, like, looking at this picture, you're like, okay, like, why the heck does Adam take this much memory? I'm going to instead use another optimizer, maybe like SGD, which has no memory overhead. Sorry, can I ask a question real quick? Is there a rule of thumb about how much for the activations that you try to give yourself headroom for? I usually always estimate it. Jane, have you found a good heuristic for it?

For activations, it usually corresponds to your batch size. Estimating it is doable for things like LLaMA or transformers where you could literally sit down, do some math, and figure out the sizes of everything. But otherwise, for other models, it is— I think it requires— there's no straight-up formula like, oh, Adam is 2x params, that type of stuff.

Yeah, so for what it's worth though, the reason I showed this picture is if you sort of as you slowly, as the batch size gets to about 1,000 or the context length gets to about 3,000 or so, 99% of the memory overhead of your model is going to be activations. So I just think of it more mentally as if I go to large batch size and context length, this is the bottleneck. Otherwise, it's not. So again, back to Adam, you might think, hey, there's always a new optimizer that comes out. And people say, oh, there's this new fancy optimizer. It's more memory efficient. The problem with a lot of that work is that like Adam is basically so used and so popular because it works. And there's like tons of— there's like, like there's like close to a decade worth of like papers and people showing anecdotal results showing that it's like a great optimizer. And that hasn't been true for a lot of like newer optimizers that have been introduced. So conceptually, you might think this is the bottleneck, but it ends up being like a very poor first thing to try to make. Like replacing Adam is very, very challenging as a first step. Okay, so let's instead take a look at the parameters. So like we said, basically at Llama 7B, we have the parameters and we have 14 gigs at FP16. And you might have heard of like 4-bit quantization. And so what we're going to do is we're going to basically take every one of those weights and turn them into int4. For an int4, it's actually a sub-byte dtype. So basically every int4 is half a byte. And so roughly you get a model that's like about 3.5 gigs this way. So yeah, great. So we found basically a good approach of dealing with the parameters. Conceptually, the way this works, it's not like a to call. The way a lot of quantization kernels look like, basically, if you wanted to, for example, cast fp32 to an int8, generally, the formulas look very similar to this, which is you basically go over all of the elements, every element of your vector or your tensor. You find what the maximum value is, and then you use that to figure out how to correctly scale the values in the sort of the new int8 domain. So the formula for formulas for this ends up looking very much like this. And basically, quantizing is going from FP32 to int8, and then dequantizing is going from int8 back to FP32. So this terminology is very common in the community. Great. So you might think now, well, how do we make sure that this process is fast? And historically, you would rely on people writing custom CUDA kernels for this, but those are hard to write. And so as a good workaround, like what I've been using a lot has been torch.compile in this case where I just made a couple of like simple changes to this code where I decorated the functions with torch.compile. And then I also added this environment. Okay, I see a question by Roy, by Rui. You have to unmute yourself. All right, maybe I'll keep going then. So what we've done here is we've just decorated these functions with torch.compile. And there's also this very handy environment variable called TORCH_LOGS_OUTPUT_CODE that I use all the time. And if you enable this, you can actually code-generate an efficient Triton kernel that basically— so in this case, this is how the quantization code would work. So the point is, is that you don't really need to learn how to like write CUDA or Triton kernels in this case. So you can use it as a starting point and get efficient like kernels out of it. So this doesn't require a lot of domain expertise. See, Andres also has a question. All right, people, can't they unmute? Maybe that's the problem.

There's also questions in the Q&A. I can read some of them if that— if you want to address them now, Mark.

Yeah, let's do that. Thank you.

Okay, the first one says, how do these calculations for memory also apply to machines with multiple GPUs to get to 24 gigabytes, like two 12-gigabyte cards?

Two 12— I see. So regarding for like, if you hypothetically, you know, if NVIDIA released like a hypothetical GPU with like 200 gigs of VRAM, like the same sort of calculations would help. But like for multiple GPUs, I think like Jane's going to be talking a lot more about how this works with like sharding, model sharding. So I'm not going to be covering that right now. Else it's going to come in a few minutes.

Cool.

Cool. So let's go back to gradients. So remember, we had the parameters, then we had the gradients. The gradients is you have, like, it's, again, like, you need a gradient per model parameter. So let's say we quantize the gradients to 4-bit. Like, would this work? And the answer is it simply does not anecdotally. Like, your model will not converge because there's just, like, not information in the backward x-pass. So this is just like scientifically, if you were to just run convergence studies, you wouldn't get anywhere doing this. So that's how we know this is not like very fruitful. Okay, but there's other tricks. Like basically the main trick to make this work is LoRA. And what LoRA is, is basically— I think a lot of people might have already talked about LoRA here— but the core idea is that instead of updating the weights for all of your parameters, you pick a subset of the parameters for which you update the weights. And we call these like adapters. And so the idea now is that instead of basically quantizing the gradients directly, we make it so that it's only a very small percentage of the parameters that actually need gradients. And so this lets us shave off, like let's say in the QLoRA paper, only like 2% of the total model parameters are trainable. And will thus have gradients and optimizer state associated with them. So basically doing that plus quantizing your parameters to 4 bits gets you QLoRA. So this is exactly what the QLoRA algorithm is at a high level. So, okay, so we got this working. Anecdotally also, because a lot of the stuff is scientific in the sense that we know that QLoRA works because everyone uses it and says it works. And last year, I helped host a NeurIPS competition about fine-tuning. There was no entry that did not use QLoRA. Like, it was, like, by far, like, 99% of the meta, and all of the winners used QLoRA then. However, implementing QLoRA is kind of tricky. Like, basically, QLoRA was mostly implemented by Tim Detmers, and if you look at the file, it's about, like, 4,000 lines of CUDA code. Actually, some people here may know me by the Discord group CUDA Mode. And the way this term came about was like Tim Demers was describing his process for writing CUDA kernels, and it was basically he sort of like sits down in a room, like no music, no lights, nothing, and he just like wrote the kernels in like a single night. I personally do not have the ability to write such kernels in a single night. So for me, what has been much more accessible is basically writing those kernels in pure PyTorch and compiling them. And there's like one of the reasons why this file is very long, by the way, is I gave you this nice picture of QLoRA. But the algorithm has a lot more details. Like basically the weights aren't in int4, they're in a format called nf4, which basically mimics a normal distribution a bit better. You also can't just matrix multiply nf4 tensors, so you need to upcast them to be f16 and then do a matmul. Remember when I said that it's very important to figure out what the max is in the original tensor? Well, this makes you very sensitive to outliers, and that's why people do it in blocks. Then QLoRA also has basically scales per block, but then it also quantizes the scales in what they call double quantization. And so it's just like a lot, like basically it's just a lot of math and CUDA that you need to write to be productive. And the alternative we have now at this point is like basically this kernel, which was written by Driscoll, who's also on the PyTorch team. And what he did was in basically about like 900 lines of Python code, got the NF4 tensor from QLoRA working for FSDP without writing any custom CUDA code. So this is this is all pure Python. So you can see, for example, like here where it says like double quantization, like the full algorithm is, okay, well, we have these like kind of four tensors, then we're doing the double quantization, we're doing a normalization, and then we return. And this is all like Pythonic code, so you can like add breakpoints, you can read it, you can understand it. So again, like this is a, like, this is not like a great intro quantization algorithm to understand like the core idea is really covered here. But if you understand this, you're sort of like well on your way to understanding like more complex kernels. So make sure to just go poke around at that like when you have some free time. All right. So the other thing though is that like within PyTorch itself, like PyTorch does not have a concept of an nf4 tensor. Like PyTorch goes down to int8, it has fp16, and it recently has fp8, and that's it. But we've seen a lot of people like experiment with lower dtype kernels. They're just not doing it with PyTorch. And today they have actually a way of creating native data types. And so this uses some modern machine machinery called tensor subclasses. So it's probably the feature that PyTorch devs are the most excited about internally. And what this does is you can basically override PyTorch ops such that, for example, with nf4, the way you do a matrix multiplication over an input is you cast the weight from basically nf4 or int4 to the, basically the weight of the input, which is an FP16. And then you do like a matrix multiplication. You can customize all of this logic using like operations, using tensor subclasses in this way. And most notably, you can also define like what the semantics should be for how this data type should be distributed over multiple devices. So if you're getting QLoRA not composing with any sort of PyTorch subsystem, generally subclasses is like one way of modifying the behavior of PyTorch purely in Python. Without being a developer on the PyTorch team. So yeah, I know this was a lot, but yeah, I guess now I'm going to hand it off to Jane, who's going to talk a lot more about how we got this working over multiple devices. So let me— I'll stop sharing my screen. Cool.

There's also a question in the Q&A you can answer in the meantime about fine-tuning full-parameter Llama 3 8B on 24 gigabytes. So the question is, so no way to do that, basically?

So there's no way to fine-tune, like, yeah, floating point Llama 3.8B on 24-bit. Yes, that's going to be very challenging. It's just like fundamentally you would need a smaller model. And it's kind of why QLoRA has become such a dominant algorithm to work with this size.

I am sharing my screen. Well, I'm trying to share my screen. Ah, there it is. There it is. Okay, can people see it?

Not yet.

How about now?

There we go. There we go.

Okay, so following up with your questions, like, oh dang, if we only have one GPU, it might not be enough to fit all the things we want to do. So on the left here, I did a little illustration of what memory looks like when you're training. So it goes left to right. So in your forward, as you're training, you gain activations that you're saving for the backward. In the backward, you start using them, so they go down. but then in the backwards, you're building up your gradients for your parameters. So the grads also goes up. And at one point, the activations are all gone and the grads need to be there for the optimizer step. So, and note that optimizer step, the state is an Atom W state and it is about 2x bigger than the params. I like measured it. So that's the left side. And just huge disclaimer, there's more in your model. Like when you're doing math, you sometimes need scratch space. So there are intermediate tensors that are not in this illustration, but they will not, they won't matter as much. Okay. If people have questions on that already, please ask now because you will be seeing this little boy like a lot. No questions? Great. Okay. So let's say it's too tall at the peak here. You see it's just taller than your GPU and GPU is sad, but it's okay. GPU has a twin. So So now the question is to you. What happens, what would you do if you had 2 GPUs? How would you fit it within that? Like it's clearly dividable. So what's the way to do it? Do people have answers? Has anyone commented on the Discord? I can't see the Discord right now.

No answer in the Discord yet.

Wow, okay. I will answer my own question. So, oh, I guess it wasn't so obvious after all, but we'll start with just parallelizing your data. So as mentioned earlier, parallelizing your data will cut down on the batch size, which contributes to this activation memory. So like your params are the same, everything else is the same because they relate to the params like Mark said, but when you slash your activations in half, you get that peak to be lower on each device. Note that everything else is replicated. But let's make this harder. Let's say you have less memory than that, and even after doing that data parallelism, it was still ooming, it's still out of memory. What else can you do? And here, you can get, you can go do what Mark was mentioning, but not with quantization. You can shard your parameters in half. You can be like, I want half of my params to live on the first one and half to live on the other. And since the qurads and atom_w state correspond with the params, each of them also become halved. So now you're like, wow, I'm doing great. This is awesome. I now can go on with my life. But note that this is not the entire picture because there are some complications when you do sharding across anything, 'cause when you shard anything, you kind of, at some point, you're gonna need to bring them back. So imagine you're doing your model training at this current moment. You're running on GPU 0, you're doing your first linear, and you're like, oh crap, I only have half of my parameters. The other half of me is in that GPU over there. What am I gonna do? Well, what you're gonna do is you're gonna be like, yo, other GPU, we gotta talk. We gotta exchange some parameters. And you will do that. And so what that really looks like is for every step, every layer you run through, you're gonna need a little more memory that's just representing the layer that's currently getting processed. And FSDP will do this in a way. So yeah, this technique, it's called fully sharded data parallel. You don't need to remember that. It's okay. We're talking about FSUP the whole time anyway. But like, it will save, it will bring in the memory you need to do like a linear, like a MatMul, and once it's done, it'll be like, oh, I don't want this anymore, and then it will put that back and drop it. And it will keep doing that to make sure that you don't peak too much. All right. But you're like, how does, how do we know that it's small? Well, you don't. Well, you do, because you're the user and you get to determine what gets like what gets considered as a layer in FSDP. So this tree thing might look a little scary, but this is LLaMA 2. LLaMA 2 is a transformer decoder that kind of branches into a bunch of things, including 32 transformer decoder layers. And then those branch into a bunch of things, including attention. And then if you do LoRA, then your LoRA has linears, and it just keeps going, dot, dot, dot, dot, dot, dot. But it's a big tree and how FSDP wraps things determines what gets brought in when you need something. So if that's a little confusing, it's okay, but you can think of each of these blobs like this green blob is one layer, this big blue blob is another layer, but FSDP would wrap in this case if you're specifying linear, then it will be like, okay, I found a linear, that's my one layer, and then you'll find another linear and that's one layer, and it will iterate from bottom up to do that. And because, so in this specific wrapping policy, the linears are wrapped and the transformer decoder layers are wrapped, and then everything else gets wrapped. So each blob is its own layer. And if you're like, Jane, why are you telling me about this? This is so confusing. I am now lost. Don't worry. So the big key point here is that The blobs correspond to how big this, like, little orange dotted box is gonna be. So, like, the bigger the blobs, the more memory you're gonna need to bring in at a time, so the bigger that box is going to be. So the way you wrap can really influence the memory you use. Okay, pausing here. Questions, comments? Oh. Thanks. Okay, so let's—

We're actually getting a question around MOEs, but yeah, it's like, how does model tensor parallelism work for MOEs?

Ah, you are getting very ahead. But yeah, there's— so you're right, there's a lot more ways to parallelize, but today we're gonna focus on FSDP. And the cool thing about FSDP2, which we're gonna introduce, is that it will handle that tensor parallelism more easily than today's FSDP. We'll get into that soon. Okay, so let's say after you do all of that, you still oom. Like, what can you do now? What else is there? What is left? And that's where CPU offloading comes in. And it's nice because CPU is like your little brother on the side who's like not doing anything as you're like trying to beef up stuff, but it can hold things. So you can make it hold your parameters as you are iterating. So in this case, you're just like, let's just put the parameters in CPU, and when we need it, we will move what we need to GPU, do the all-gather, do that sharing of knowledge, sharing of data and parameters, and then move on with our merry lives. And with that, with CPU offloading, you get your memory to be much, much smaller because the biggest chunks here are now living in CPU. So note that we really want the GPU, like GPUs are accelerators. They're meant to do like beefy work. And your forward and backward are the beefy work in a model usually. So for the optimizer step, people are like, it's fine. We don't need the GPU for that. And in this case, we'll only bring in the parameters for forward backward and be like, okay, we can put that on CPU now, which means the optimizer step runs on CPU. And you also save a lot of space on your GPU if you host your whole Atom state there. So There's that. Okay. So you might be wondering, Jane, FSDP has been published for, like, a long time now. Why are you explaining this to us? What is the point? People use it. In fact, that's true. People like Answer.AI, who are wonderful, they already built out an integration for FSDP and bits and bytes params for bit to make QLoRA happen. But it's kind of annoying to work with FSDP1. They had to do a lot. And we came out with her parameter FSDP, which I will also call FSDP2 for later. And what is that? So let's start with the status quo. Like, what is it today? Let's say you, just for our toy example, you have 3 tensors that you need to distribute across your 2 GPUs. And they are these shapes. So the goal is that you want to make that exchange of, you know, when you're talking to the other GPU, that thing efficient. And Nickel, which is the software and driver stuff that does it, it requires that each GPU will give the same tensor size. So those are, that's the constraint. What does FSDP1 do today? Okay, what it does is it flattens all your tensors. So, and this is what they look like in memory. So flattening is actually pretty chill. And then it's gonna like line them up in a line and then it's just gonna like slice it in the middle. And if it's not even, it will add a padding at the end. And then because now it's just like arbitrarily sliced in the middle, it will be like, all right, this half goes to 1, this half goes to 0, I guess 0 and 1. And so you'll end up with something like this where tensor 1 and I guess like a little more than a third of T2 will live on GPU 0. And then the other half of this, including the padding, will live on GPU 1. And this is nice, but note that the way this is implemented today is that T1 and half of T2 is gonna get smooshed into one tensor, which is a big con. We'll get into that later. And same thing with T2, T3, and the padding. That gets smooshed in one tensor. And we call that tensor a flat parameter 'cause it's like so flat. And some reasons why you might already be thinking, hmm, this might not be a good idea after all is the fact that this forces T1 and T2 and T3 to all have the same dtype, to have the same, like, requires_gradness, and all the other metadata you might want your tensors to have. Okay, well, what if we tried down, like, a different path of dividing things in two? So what we're gonna do is we're gonna slice every tensor first. We're gonna cut T1 in half, T2 in half, and T3 in half. Great, except we notice that T2 needs some padding because it's not, it's 3 by 3. You can't really cut that in half. We'll do it, it's fine, we'll shard. And that way, what this will look like is every tensor gets its own representation on this GPU. And this is great. The main pro is that they keep their identity. Like if T1 was UN8 and T2 were BF16, totally fine, they can stay that way. But in the previous case, in FSDP1, you wouldn't even be able to put them together. That's just like, or you'd have to hack around it a lot. And this gets into the QLoRA stuff soon, very soon as well. Okay. There is a con to this. Because of all the slicing and rearranging you have to do, there are extra copies to FSDP2. So that's, that's a con, but the pros it gives are so much more worth it. And Just a recap of before in more clearer terms, like a flat parameter would force all three of these tensors to share all the metadata they have that a tensor can have. And in fsdb2, because they are individual separate tensors, we call them dtensors because they're not like, you know, the full tensor, they're a distributed tensor, they're smaller. They can be themselves. They get to keep their identities. They can have their own dtype, their own requires_grad. And so you're like, okay, but why? So if you think of quantization, which Mark talked about earlier, what if you wanted, you know, T1 to be uint8, T2 to be full size, like float FP32, any other size. In the first regime, unacceptable. The second regime, totally fine. No one is gonna bother you. FSDP will do that. Another thing that is very popular nowadays that LoRA is, is you don't really want to train the big stuff because that will require big grads, big optimizer step. So what if T2 is frozen? You don't want it to actually train. And T3 is, you know, the LoRA adapter that you do want to train. In that case, in your first world, you still can't have that because a tensor can only have one requires_grad. And the flat parameter here will force you to either make it requires_grad or not requires_grad, or to force you to do a lot of fancy hacking to make it work. All of these things, all these concepts that are like, oh, I wish I had this. In fsdp1 today, you would, it would be difficult. But in fsdb2, it's for free. And another thing that's really cool about fsdb2 that is its own other lecture entirely is memory determinism. So, uh, one of the major implementation changes is that now fsdb2 actually guarantees that you will have only that small little sliver of memory before, like this, this little, this little orange thing. Whereas FSDP1 actually didn't do it well enough and could cause memory spikes that are not deterministic. But yeah, for this one, you should read the blog linked here if you, if you want more details. Okay, so now that we have FSDP2 and we're like, this should be easier to use, let's do it, let's use it. And Wei did do that, we did do that. So Wei, who's another dude on the team, he wrote this PR here that puts together fsdb2 and nf4. And it works, it works, it's great. But we know, like, like, okay, like fsdb2 is cleaner, it's more composable, but the last question remains of like, can this actually replace fsdb1? Like, we would love to use it, but can you tell us that it, it is good on perf, that we won't be slower than before. And so that's what the next few slides are gonna be. All right, pausing here to see if people have questions, thoughts. If not, we're gonna go with the plan. All right, so here's the plan. The plan is I'm gonna go get some GPUs, we're gonna run some benchmarks, and then we're gonna make sure those are the same benchmark, and then once they are, we're gonna record some gaps and figure out what the gaps are and if we could make them faster. All right. So getting some GPUs, this is the easiest part of the journey. You just go to Vast AI and then you ask for, well, first you need to have money. And then you go and you're like, give me two 3090s or 4090s. And, and I got two. They are 24 gigabytes each for VRAM. There's some other details here if you care, but they're not super relevant this time. Just know that I have two. I have consumer hardware. And they are 24 each. So I ran a bunch of benchmarks on answer.ai's train.py, which allows, which is our, like, baseline, like FSDP1 and BNB, that's our baseline. And I'm using the batch size 8 as a baseline, and just so that it works. And if you're curious, you could, if you wanted to run this for yourself, the command is right here. Feel free to copy-paste that in the future. But you could just pay attention now. Okay, so I ran the same thing in the torchtune recipe. One difference in torchtune and train.py is that it uses a YAML for the config versus the command line. It's like just different. And I did have to tweak the YAML quite a bit to make sure that I was running the same config. And since, and then these were my results. So peak memory-wise, we were doing better. And for runtime though, we were like 19% slower. So someone here might be like, eh, FSTP2, we know it's stricter about memory, we know it requires extra copies, that makes sense that we're better at memory and worse at runtime, right? But no, no, no, we gotta be diligent. And very quickly, if you look at the traces, it reveals some troubling shenanigans. So here's the two traces. On the top is the baseline, on the bottom is our new trace. Can you spot the difference? Okay, there's a lot of differences, so I'll just go with my favorite one. I work on optimizers, and those little light blue things are optimizer steps, and immediately I was like, dude, the optimizer is taking so much longer. What could that be? And so this is where I get into the actual tracing stuff. I wonder if I— I wonder if I can actually show you the traces. That would be pretty cool. Okay, I'm gonna stop sharing to reshare, and then we can do that. Let's just share my entire screen. Doop. Okay, do people see traces?

Yep.

Okay, awesome. So wait, I'm gonna delete this. So on the left is our baseline, on the right is the Slowboy. So in our baseline, we're gonna go to the second step because the first step is always full of like warmup stuff and initiating stuff. So we're just gonna, we're gonna ignore that and we're gonna go here because every other step after this is much more normalized. And something you can do, I'm using Perfetto. I don't know if people are familiar with Perfetto already, But it's been pretty helpful. Yeah.

Doop.

Okay. And something that is super useful and nice in Profeto is you can like highlight a region and it will give you profiler— profiler— profiler— wow, I cannot talk today— results on here. So here you can see that there are a bunch of— it tells you what thing takes the longest. It's like the moles take 77 milliseconds and there are 768 of them. And on this side, when we do that, we're going to notice some very different numbers. So here the mole also takes the longest, but there's 1,700 of them compared to 700. And you might be like, what is going on here? But let's go look at the smallest number. In, in Optimizers, there's only one divide. You can just take my word for that. So here we know that there are 384 divs, which means that there are 384 parameters. Here, we see that there are 800, 896, which is like more than double. And so let's like go find a div. Like, where are they? And here you can just like doo doo doo. But you can already notice that everything kind of gets doubled here, whereas in here, they are just called immediately. So like this A10MOL goes directly to A10_2. This A10MOL though goes to A10MOL again. And you're like, wait, what is going on? Why is that? And this is where you learn the painful things of tensor subclass dispatch. So since we're using dtensors, it means that it goes into this MOL as a dtensor. And then dtensor is like, all right, let me do my metadata unwrapping. And now you get to go to A10MOL as a plain tensor now. So there are double the amounts. And just to spare you some math, spare you some math, it turns out that if we divide this MOL by 2 or the div by 2 or any of the things that run once, it shows that we actually are running training on more parameters than we thought. So in the left side, we're only running it for 384, and the right side, we're running 64 more parameters. Like, can people guess where this came from? I will show you what, like, war story of debugging has shown me. So in the end, I realized that this was a config difference where if you are fast at reading, you might have already spotted the difference here. And the answer is that, oh, in train.py, they do not LoRaFi the output projection, whereas in our world, we do do that. And since LoRaFying means you add 2 adapters and there are 32 layers, 32 times 2 is 64 more extra parameters to train. So yeah, that was bad. And that was a great lesson because it was like, I was not measuring apples to apples. And I needed to do some other things to make sure that we were. So the first thing is like making sure the dataset was the same, making sure that every parameter was the same, changing the wrapping policy to be the same. And after I did all of that, I ran another benchmark. And here, I mean, that was like also weeks of work, by the way. Like, it was not easy to figure out every little difference and why they were different. But after all of that, I was like, let me run it again. Maybe it will be better. But no, it was still slow. It was actually slower than before. But the peak memory was a lot better. So, however, I was still happy because even though it may feel like we took a step back, we actually made a giant leap forward on blocking the first real step, which is like, now that we have apples to apples, there are things that will match, and then we should just look at things that are different. And so that's where I could start playing my very long game of spot the difference. Um, the first thing I did though was like measure their runtime. And here you can see that the forward, backward, and optimizer were, were like the culprits, and that's, that's, that's how I knew what to focus on. Okay, so this is a big slide. I, okay, if at this point you have not looked at the traces yet, but you would like to, I sent a link to the Google Drive in the Discord, and there are the traces that you want. There's like, there are 4 of them and the 2 you care about to look at now are the ones that don't start with bad or final. The other 2, do the AnswerAI one and the TorchTune one. But I'm gonna, I already found these gaps for you and it'd be fun, if you find it fun, if you wanna learn more about traces, it'd be fun if you could find each of these yourself, like where they are located and do the same exploration I did. But because this is a presentation and I do wanna save time, we're gonna dive right into what these gaps were and where, like, how we ended up fighting them. So the very first gap is, yeah, so we're gonna go, and this is the link, by the way. Okay, the very first gap is the optimizer step was still slower, but remember how earlier I hinted that there was all this overhead here for dtensor? And the way it works is because of C++ dispatch, It just takes a long time to do that kernel call. And because there are so many calls, all of this overhead adds up to make the optimizer step 3 times slower. And also another minor detail, if you look at this, this is 64 microseconds and this is 81. The answer for that is because the parameter is not contiguous, but that's more minor. So the solution here is actually really chill. We worked with Intel. We were like, Hey, what if we just had a fused Adam, like your optimizer step, but all happening in one kernel so you can dispatch just once and get results? So this avoids all that overhead because there's just one kernel now versus 384 times however many ops there are, and it also leverages vectorization. So we go from about 1 second to 120 milliseconds, which is like an 8x speedup. Speed up. So that's one gap that is now gone. Alright, pausing here to see if people have questions.

I've been like speed answering everyone on Discord, so.

Oh, okay, nice. Okay, okay. I was like, are they lost? But no, okay.

They're very happy actually. Like people are saying they love this kind of debugging and yeah, people love it.

Okay, well, let's keep going. So the next 3 are a little boring, but we're gonna go through, and there was a lot of pain in this one actually. This second gap was crazy. So I went and the way I literally did this was the most brute force way you can imagine. Like you open your trace, you find the first forward, you find the first GPU kernel, and you're just like, do they match in size, shape, input, everything? And I would do that for you. But we are a little short on time, so I'm gonna just show you what the, the difference is here. And the first difference that was major was that the second allgather in every linear, like every layer, was somehow 5 milliseconds longer. And that was when I like needed to click on them and figure out how big it were, they were, and I, I noticed that in the, on the left side for train.py, there were just fewer things getting allgathered. Like, it was just, it was just not the same thing. And I was like, why is it not the same thing? So I did a lot of printing and I hacked around fsdb2. And what that yielded was me writing down the sizes of every tensor that got packed up in the allgather and realizing that the difference was because in our nf4 metadata where Answer AI did not pack their scalars and quantization factor. They just like, for bits and bytes, when they did their 4-bit thing, they use a dictionary. They have their own method. And this is actually, the reason they couldn't pack it is because FSDP-1 had restrictions, by the way. Like, it just wouldn't do it for them. So they needed to work around that. So that was one piece of the puzzle where we just packed everything in one go. But in the other bigger piece of the puzzle, the big, big difference, like that was just like 1 million. Bytes, whatever, doesn't matter. But the other thing was like so many more, it was like, like 12 million bytes. And that was when we realized that when we opted out of LoRA, the output projection did not get quantized in our world. So it was like 4 times the size of AnswerAI's version. And I was like, oh, why don't we do that? And then I talked to the TorchTune team and they're like, oh yeah, we should do that. And so we should do that. That's the conclusion. We should do that. The first one is intended behavior, so we don't really need to worry about that. But the second one, we should do it. And we will hit you up when that happens. So this gap, I would mark as yellow. Like, we know what to do. It's solved. Okay, third gap is realizing that there were just like more overhead. And remember what Mark was saying, how when you have NF4, you have to, you can't just like put that through your gem kernel. Like, CUDA's gonna complain. You need to dequantize, get it to the BF16. And then put it through the CUDA kernel. And it turns out that Tim Detmers, you know, brilliant guy, already wrote a really fast version of that, whereas we just get stuck with the little raggedy version that tries every op. So that's also where we are slower just because of additional overhead. But again, this is not a major problem. Solutions. We could use torch.compile. I tried it. It was not trivial. So I was like, I will look into this later when I have a month of time. And then, or when I don't have a presentation on Monday. And then the second step is to just use Triton kernels. So Dris, our coworker, already wrote them. I didn't wanna copy-paste them for the sake of this presentation, but if you wanna copy-paste them, go for it. No one's judging. They're good. They work. And so we're like, okay, we also know how to fix this one. The third one was really stupid. This one is definitely the worst gap. Where basically there were just very different ops happening before the SDPA. And this was just because we used two different ROPE algorithms. TorchTune was like, we are traditional. We are going to use the original meta algorithm. We work there. So there will be no numerical differences. And everybody else in the world is like, oh, but we want it to be faster. So it's fine. And then The real solution here is just for TorchStream to offer more of these, and that's also in progress. So yeah, but okay, let's talk about the most fun one. The most fun gap I noticed is— I realize this is a little hard to read, but on the left side, note that there are two things happening. The first, this green thing here, is the CPU offload where you go from CPU to GPU, and the second one is when you, after you've moved it to the GPUs, you like have them talk to each other. And here in train.py in FSDP1, we're like, wait, how come this is overlapped? Like, you see how this is— there's no gap here, but this one is so badly overlapped. Look at that. Look at that exposition. It's bad. And this was really frustrating because this was costing us like 10 milliseconds each step. And I was like, I wonder why this is. But actually, this is very much expected behavior. And this is part of the bigger memory constraints that FSTP2 promises you. So FSTP2 is promising that, hey, we're not only looking at all gathers, we're also going to make sure that before we bring, like, CPU to GPU, that you have the space you have. So it is guaranteeing the constraint that only 2 layers of parameters will be allowed at a time. And, and, and, and that is why on the left, because of how FSDP1 was implemented, it didn't do that very well. So you'd get random memory spikes and FSDP2, you're promised to never get that. But someone looking at this trace will be like, but this is kind of, I don't, what, what if I'm okay with non-deterministic memory? Like it's not happening now. Like maybe I can just go on with my life. But no, no, no, we have an answer for you. And the answer is the reason it's so exposed is not because FSDP is too strict. That's not the problem. The problem is that the overlap of computation and communication was too different. The computation here is super duper tiny because it corresponds to this tiny linear here. And then the, here when you're CPU offloading, you're actually trying to bring this big boy back in. So it's like the mismatch in the layer sizes was really causing this problem. So what could we do? Well, it's fine. We'll just change the wrapping policy, have bigger layers. And it's really just, hey, don't wrap those linears by themselves. Just group them in so we can just have every transformer decoder layer be its own layer. And note that this is only possible with FSDP2. You can't have the right-hand side in FSDP1. Why? Because the LoRA linears have the base weights, the base weights are quantized. They're NF4 tensors. They're gonna be small. They're UN8. Whereas your linear tensors, those are BF16 or whatever training thing you want. And in FSDP1, they can't be brought together under the same family 'cause they're gonna be forced into one big flat parameter. But in FSDP2, they can coexist. And the change is actually really easy to do. The policy is just a matter of commenting out these two lines here. And once we do that, the solution is like, they're both overlapped. It's kind of crazy. Look, like, there are, there's no more, no more exposition at all, where even in the first case, wait, even before here, in the train.py one, this allgather was exposed, also due to the same reason. That's not even true at all here. And this, this wrapping policy, this new one, is only possible because of FSDP2, which is a great point here. All right. So now we've fought all our gaps and things work. So it's your turn. You should try them out. You should, you should really, like, it doesn't have to be NF4. If you have another favorite lower precision thing, you can try it. If you want to try out more LLaMA stuff with it, you can. There are now scripts ready to do that. So yeah. One disclaimer though, we are working on this, like, this is not a forever thing. Checkpointing of just FSDP and QLoRA does not work yet. So that's just FYI, we're working on it. Sorry, you can't do that. But you can try, you can just like try and play with it, among other things. So yeah, here I would love to, I mean, Mark and I are speaking here, but this was really made possible by all these amazing people. So Driss wrote the original NF4 stuff. Andrew is the main dude who designed FSDP2. Wei ended up taking Andrew's work and making it compose with Driss's work. So he, he like amalgamated both of them. And then TorchTune people, so like Rohan and Evan, they were super helpful. They wrote the LoRA recipes and they're the foundation of which we built and showcased our work. And of course Mark, who's amazing. So Thanks, Mark.

Yeah, so I really hope this gives people some more context around what we're thinking here. We did want to showcase a useful recipe with the QLoRA and FSTP composition. But really, this is kind of like— our call to action here would really be, if you're doing interesting research at the intersection of quantization and distributed, we'd really, really love to hear from you. So if you're working with more exotic DTypes or more exotic forms of parallelism, A lot of this work should like really, really be helpful. And you know, we have all these links here that can give you some more context. I guess we'll pause here if people have any questions. Otherwise, you know, we'll probably be hanging out on Discord for a couple more hours just to see if people have any questions.