Strategies for Managing Out of Memory Errors in Pyre

Developers at Meta created tools to effectively manage OOMs in Pyre, enhancing model training efficiency.

### Strategies for Managing Out of Memory (OOM) Issues in Pyre
* 🖥 OOMs, or **Out of Memory errors**, are common challenges developers face when working with Pyre code, prompting the need for **effective solutions**.
* 🔧 Traditionally, developers have resorted to **reducing batch sizes or model sizes** as immediate fixes, but these approaches lack nuance and can be **overly simplistic**.
* 💡 Understanding VRAM limitations is crucial; for instance, popular consumer GPUs like the 39s and 49s have around **24GB**, while A1 100s have **40 or 80GB**, and H100s can reach **100GB**.
* 📊 When estimating model size, it's essential to consider not just the parameters but also the **gradients and optimizer states**; for example, a model like llama 7B can require **up to 56GB of memory** when accounting for all factors.
* 📈 Optimizations such as flash attention are vital as **activations can dominate memory usage**, particularly with larger batch sizes, and estimating these can be **complex without proper tools**.
* 🔄 New strategies like quantization can significantly reduce memory requirements; for example, quantizing parameters to **four bits** can shrink a model's size from **14GB to approximately 3.5GB**.
* ⚙ Techniques such as Fully Sharded Data Parallel (FSDP) help **distribute memory usage across multiple GPUs**, allowing for **more efficient model training** without exceeding VRAM limits.
* 🛠 ️ CPU offloading can also alleviate memory pressure by **transferring parameters to the CPU**, thereby **freeing up GPU resources for more intensive computational tasks**.

### Advancements in FSDP: Transitioning from FSDP 1 to FSDP 2
* 🔍 The transition from FSDP 1 to FSDP 2 involves a **significant change** in how tensors are managed across GPUs, allowing tensors to **maintain their unique identities and metadata**.
* 💾 In FSDP 1, tensors were combined into a **flat parameter**, which forced them to **share the same data type and other metadata**, while FSDP 2 allows for **individual representation of each tensor**.
* ⚙ The implementation of FSDP 2 introduces extra copies due to the slicing and rearranging of tensors, but the benefits of **maintaining tensor identity** outweigh this drawback.
* 📊 Initial benchmarks showed that while peak memory usage improved with FSDP 2, runtime performance was **19% slower** compared to FSDP 1, prompting further investigation into the performance gaps.
* 🔧 Detailed tracing revealed that the optimizer steps were taking **significantly longer** in FSDP 2 due to overhead from using distributed tensors, which led to collaboration with Intel to develop a **fused optimizer step**, resulting in an **8X speed-up**.
* 📉 Other performance issues were identified, including differences in tensor sizes during operations, which were resolved by **ensuring consistent configurations** and **improving tensor packing methods**.
* The final adjustments allowed for **better overlap of computation and communication**, leveraging the unique capabilities of FSDP 2 to **optimize performance and memory management**.
* 📚 The transition to FSDP 2 is ongoing, with a call to the research community for collaboration on exploring **quantization and distributed training innovations**.

### Optimizing OOM Management in Pyre: Strategies and Tools
* 🛠 OOMs, or **Out of Memory errors**, are a significant challenge in Pyre code, prompting developers to seek effective solutions for **managing memory constraints**.
* 💻 Developers at Meta, including Jane and the speaker, are focused on creating tools to **simplify the handling of OOMs**, moving beyond traditional methods like **reducing batch sizes or model sizes**.
* 📊 Understanding VRAM limitations is crucial, with popular consumer GPUs like the 39s and 49s offering **around 24 GB**, while A1 100s provide either **40 or 80 GB**, and H100s around **100 GB**.
* 📏 When estimating model sizes, it's essential to consider not just the parameters, but also **gradients and optimizer states**, which can quickly escalate memory requirements; for instance, a model like Llama 7B can require **up to 56 GB of VRAM** when fully fine-tuned.
* 🔍 Activations can dominate memory usage, especially with larger batch sizes, making optimizations like **flash attention vital for efficiency**.
* 💡 The introduction of techniques like quantization, specifically **four-bit quantization**, can significantly reduce model size, but care must be taken with gradients, as quantizing them may **hinder convergence**.
* 🔄 Strategies such as Laura allow for selective parameter updates, **reducing memory overhead by focusing on a small subset of trainable parameters**, while tools like FSDP (Fully Sharded Data Parallel) help **manage memory across multiple GPUs**.
* 💾 CPU offloading is another effective method, allowing parameters to reside in **CPU memory** while performing forward and backward passes on the **GPU**, thus **alleviating memory pressure** on the GPU.

### Enhancing FSDP: Transitioning from FSDP 1 to FSDP 2 for Improved Performance
* 🔄 The discussion begins with the challenges of distributing tensors across GPUs, highlighting how tensor management in FSDP 1 can lead to **inefficiencies due to forced uniformity in metadata across tensors**.
* In contrast, FSDP 2 allows each tensor to maintain its **unique properties**, such as **data type and gradient requirements**, which enhances **flexibility and performance**.
* 📉 However, FSDP 2 introduces additional overhead due to the **slicing and rearranging of tensors**, which can initially result in **slower performance metrics**.
* ⚙ The speaker details the benchmarking process using consumer-grade GPUs, revealing a **19% slowdown in runtime** for FSDP 2 compared to its predecessor, despite improvements in **memory usage**.
* 🔍 Through meticulous tracing of optimizer steps and tensor operations, the speaker identifies several performance bottlenecks, including **excessive parameter overhead** and **inefficient kernel calls**.
* ⚡ Solutions such as fused optimizer steps and better tensor packing are proposed to address these gaps, with some promising results already yielding an **8X speed-up** in certain operations.
* 💡 The final takeaway emphasizes the **potential for significant advancements in distributed training** by leveraging the capabilities of FSDP 2, while also inviting collaboration from the research community to explore **innovative applications of quantization and distributed training techniques**.

Yeah, so hi everyone. Um, so we're here to talk about slaying OOMs. So OOMs are like maybe like the notoriously and like one of the most annoying bugs to deal with in Pyre code. Uh, and so both like Jane and I are like developers on Pyre Core at Meta, and so we wanted to sort of talk to you about like a lot of the tools we've been building to make this process like a bit easier to deal with. Um, so traditionally the way I've seen people deal with OOMs is this way, like which is basically people seeing OOMs, they're like, Okay, crap, like what do I do? And then the sort of two 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 is too big, like that means just have something smaller. Um, but this is like a very coarse tool and there's like, you know, certainly a lot more finer grain things you could do with like a bit more

a bit more knowledge. So, um, the first constraint you have is essentially like how much VRAM you have on your chip. So, for example, for 39s and 49s, which are like very popular consumer cars, you have about 24 gigs. And then for the A1 100s, you have like either 40 or 80. I think the H100 is like about 100, if I recall correctly. Uh, but my point is, is that like you're always going to be VRAM starved, and specifically for consumer cards. Uh, you know, if I were to speculate, like I would speculate that the 590 probably also has like around 24 gigs of VRAM. And so, we're always going to be in the sort of like VRAM constraint environment. Um, but like again, like as we're thinking about like memory for a model, uh, instead of just sort of thinking about like, oh, it's like something is blah gigabytes, let's sort of like be a bit more concrete about like what's involved when estimating the size of a model. Um, so you have like

three core buckets. So, for example, let's say, uh, you say, Oh, like I'm downloading llama 7B. The 7B is referring to the number of parameters. Uh, and if each of those parameters is in fp16, then you need two bytes per parameter, which means that the total size of the model is about like 14 gigs. And indeed, if you download like llama 7B on your disc, like that's roughly how big it's going to be. But that's not enough. Like, but if that was enough, you know, you could just like run llama 7B on a 390, but you can't. And how come, right? So, the reason why is like, well, like if you're doing fine-tuning or training, you also need to, uh, 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 Adams like Bas. 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 people have. So this is sort of the traditional what people would refer to as a full fine-tune. I also sort of like neglected to mention activations. So activations are basically the like let's say you're running like for example you have like your weights. It takes certain times a certain inputs WX. So like the output of that is your activations. Activations tend to dominate the VRAM of your model at larger batch sizes and larger context length. So that's why optimizations like flash attention are really important. And I'm not going to talk too much about activations like for 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 found can be a bit finicky in error 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? Like, I'm going to instead use another Optimizer, maybe like SGD, which has like no memory overhead. Sorry, a question. Is there a rule of thumb about how much, like, you know, for the activations that you kind of try to give yourself headroom for? Um, I usually always estimate it like, like Jane, have you found like a good heuristic for it? Like, uh, for activations, it usually corresponds to your batch size. Uh, estimating it is doable for things like Llama or Transformers where you could like literally sit down, do some math, and figure out the sizes of everything. But otherwise, for other models, it is, I think

it requires like there's no like straight-up formula like oh atom is 2x per RS that type of stuff yeah so for what it's worth though like I the reason I showed this picture is like if you sort of as you slowly as the batch size gets to about like a thousand or the context length gets to about like 3,000 or so like 99% of the memory overhead of your model is going to be activations so like I just think of it more mentally as if I go like the large bze in context length this is the bottleneck otherwise it's it's not um so again like like back to Adam like you might think like hey like there's there's always a new optimizer that comes out and people say oh like there's this new fantasy 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 papers and people showing anecdotal results showing that it's like a great optimizer, and that hasn't been true for a lot of newer optimizers that have been introduced. So, conceptually, you might think this is the bottleneck, but it ends up being a very poor first thing to try to make. Replacing items is very, very challenging as a first step. Okay, so let's instead take a look at the parameters. So, like we said, basically Alama 7B, we have the parameters, and we have 14 gigs at fp16. And you might have heard of four-bit quantization U. And so what we're going to do is we're going to basically take every one of those weights and turn them into in four. For an in four, it's actually a sub-byte type. So basically, every four is half a byte, and so roughly you get a model that's like about three and a half gigs this way. So yeah, great. So we we

we found like basically like a good approach of dealing with the parameters. Conceptually, the way this works, it's not like a two-call. The way a lot of quantization kernels look like, like basically if you wanted to, for example, cast fp32 to 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 de-quantizing is going from int8 back to fp32. So this terminology is very common in the community. Great. So you might think now like, well, like how do we make sure that this process is fast, and historically you would rely on people writing custom CUDA kernels for

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 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. Uh, you have to mute yourself. All right, maybe I'll keep going then. Um, so what we've done here is like 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 to generate an efficient Trident kernel that basically... So, in this case, this is like how the quantization code would work. So, the point is that you don't really need to learn how to write Cuda or Trident kernels in this case, so you can use it as a starting point and get efficient kernels out of it. So, this doesn't require a lot of domain.

expertise. See, Andre also has a question. All right, people can they un unmute. Maybe that's the problem. There's 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 GBT like 2 12 gigabyte cards 2 12. I see. So so regarding for like like 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 charting model sharting, so I'm not going to be covering that right now. Else is going to come in a few minutes. Cool. Um, cool. So let's go back to gradients. So so remember we had the

parameters, then we have the gradients. The gradients are you have like it's again like you need a gradient per model parameter. So let's say we quantize the gradients to four bits. 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 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 Laura. And what Laura is, is basically I think a lot of people might have already talked about Laura 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. Like 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 Kora 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 bit gets you Q lauro. So, this is exactly what the Q algorithm is at a high level. So, great, so we got this working. Anecdotally also, like, because a lot of the stuff is scientific in the sense that, like, we know that Ur works because everyone uses it and says it works. And last year, I helped host a Europe's competition about fine-tuning. There was no entry that did not use qor. Like, it was by far, like, 99% of the meta, and all of the winners used Cur. Then, however, implementing Cura is kind of tricky. Like, basically, Cur was mostly implemented by Tim Demmer's, and if you

look at the file. It's about like 4,000 lines of CUDA code. Um, actually, some people here may know me by the Discord group CUDA mode. And the way this term came about was like Tim Demmer 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 um is basically writing those kernels in pure PyTorch and compiling them. And then like one of the reasons why this file is very long, by the way, is I gave you this nice picture of qora, but the algorithm has a lot more details. Like basically the weights aren't an 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 bf16 and then do a matmo. You 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 Kora 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 Cura that you need to write to be productive. And the alternative we have now at this point is like basically this colel, which was written by bris cus, also on the py team. And what he did was in basically about like 900 lines of python code. U Got the N4 tensor from Kora work for fstp without writing any custom cud code. So this is all pure python. Um, 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 N4 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. Um, so again, like this is a, like this is not like a great intro quantization algorithm to understand like the core idea, you know, was 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 f around at that like when you have some free time. All right, so the other thing though is that like within PyR itself, like PyR does not have a concept of an NF4 tensor. Like PyR goes down to IN8, 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, with lower dtype kernels. They're just not doing it with PyTorch. And today they have actually a way of creating like native data types. And so this uses some modern machine machinery called tensor.

subclasses, so it's probably the feature that PyTorch developers are the most excited about internally. What this does is you can basically override PyTorch operations such that, for example, with NF4, the way you do a matrix multiplication over an input is you cast the weight from basically AN4 or N4 to the basically the weight of the input, which is FP16, and then you do like a matrix multiplication. You can customize all of this logic using operations using tensor subclasses in this way. Most notably, you can also define what the semantics should be for how this data type should be distributed over multiple devices. So, if you're getting Kora 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 like now I'm going to like 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 38b 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 Lama 38b on 24. Yes, that's going to be like very CH. It's just like fundamentally, you would need a smaller model, and it's kind of why like Q has become such a dominant algorithm to work with 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. Now we can. 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've did a little illustration of what memory looks like.

when you're training, so it goes left to right. So, in your forward as your 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 backward, you're building up your gradients for your parameters, so the grads also go up. And at one point, um, the activations are all gone, and the grads need to be there for the optimizer step. So, and note that Optimizer step, the state isn't adom W State, and it is about 2x bigger than the pams. I like measured it, um, so that's the left side. And just a 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, uh, 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 uh 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 two 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 answers 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 prams are the same, everything else is the same because they relate to the prams, 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 hard. Let's say you have less memory than that, and even after doing that data.

parallelism. It was still oozing. It's still out of memory. What else can you do? And here, um, 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 cads and Adam W State correspond with the params, each of them also becomes 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's some complications when you do sharding across anything. Because when you shard anything, you kind of, at some point, you're going to need to bring them back. So imagine you're doing your model training at this current moment. You're running on GPU zero. You're doing your first linear and you're like, Oh crap, I only have half of my parameters. The other half of me isn't that GPU over there. What am I gonna do? Uh, well, what

you're gonna do is you're gonna be like your other GPU. We got to talk. We got to 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 going to 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 can, you don't need to remember that. It's okay. We're talking about FSDP the whole time anyway. But like, it will save, it will bring in the memory you need to do like a linear, like a map mole. 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? How does, how do we know that it's small? Well, you don't. Well, you

do 'cause you're the user, and you get to determine what gets like what gets considered Ed as a layer in fsdp. So, this tree thing might look a little scary, but this 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 Laura, then your Laura has linear, and it just keeps going, do do dot do do 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, 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 kind of iter iterate from bottom up to

do that and because so in this specific wrapping policy, the linear are wrapped, and the Transformer decoder layers are wrapped, um, 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 going to be. So, like, the bigger the blobs, the more memory you're going to 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, thanks. Okay, we're actually getting a question around Moes, but uh, yeah, it's like, How does model tense rep parallelism work for Ovis? Ah, you are getting very ahead, but uh, yeah, there's so you're right. There's a lot more ways to parallelize, but today we're going to focus on fsdp, and the cool thing about fsdp 2, which we're going to 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, who are wonderful, they already built out an integration for FSDP and bits and bytes, PAMs for Bit, to make Kora happen. But it's kind of annoying to work with FSDP. One, they had to do a lot, and we came out with her.

parameter fstp, which I will also call fsp2 for later. And what is that? So let's start with the status quo, like what is it today. Uh, let's say you just for a toy example, you have three tensors that you need to distribute across your two 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. Um, it requires that each GPU will give the same tensor size. So those are, that's the constraint. What does fsdp 1 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 going to like line them up in a, in a line, and then it's just going to like slice it in the middle. Uh, and if it's not even, it will add a padding at the end. And then because now it's just like

arbitrary sliced in the middle. It will be like, all right, this half goes to one, this half goes to zero, oh, I guess zero and one. And so you'll end up with something like this, where tensor one and I guess like a little more than a third of T2 will live on GPU zero, and then the other half of this, including the padding, will live on GPU one. And this is nice, but note that the way this is implemented today is that T1 and half of T2 is going to 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 because it's like so flat. And some reasons why you might already be thinking, Hey, 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 um.

uh uh like requires gradness and all the other metadata you might want your tensors to have. Okay, well, what if we try down like a different path of dividing things in two? So, what we're going to do is we're going to slice every tensor. First, we're going to 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 3x3. 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 the 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 fsdp 1, 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 Kora.

stuff soon, uh, 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 fsdp 2, so that's a con. But the pros it gives are so much more worth it. And just a recap of before, uh, 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 fdb2, because they are individual separate tensors, we call them D tensors because they're not like, you know, the full tensor. They're a distributed tensor. They're smaller, uh, 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 have you wanted, you know, T1 to be un8, T, uh, T2 to be full size like float fp32, any other size in the first regime, unacceptable. The second regime

totally fine. No one is going to bother you. FSDP will do that. Another thing that is very popular nowadays that Laura 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, that you don't want it to actually train, and T3 is, you know, the Laura 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 parameter here will force you to either make it requires grad or not requires grad, or 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 FSDP 1 today, you would, it would be difficult, but in FSDP 2, it's for free. And another thing that's really cool about FSDP 2 that is its own other lecture entirely is memory determinism. So, uh, one of the major implementation changes is that now

fsdp 2 actually guarantees that you will have only that small little sliver of memory before, like this little orange thing, whereas fsdp 1 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 links here if you want more details. Okay, so now that we have fsdp 2 and we're like, This should be easier to use. Let's do it. Let's use it. And Wey did do that. We did do that. So Wey, who's another dude on the team, he wrote this PR here that puts together fsdp 2 and nf4, and it works. It works. It's great. But we know, like, Okay, like, fsdp 2 is cleaner. It's more composable. But the last question remains of, like, Can this actually replace fsdp 1? Like, we would love to use it, but can you tell us that it is good on perf, that we won't be slower than before?

and so that's what the next few slides are going to be. All right, pausing here to see if people have questions, thoughts. If not, we're going to go with the plan. All right, so here's the plan. The plan is I'm going to go get some GPUs, we're going to run some benchmarks, and then we're going to make sure those are the same benchmarks. And then once they are, we're going to 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, well, first you need to have money. And then you go and you're like, give me two 3090s or 4090s. 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 train.py, which allows, which is our like Baseline like fsdp 1 and BMB, that's our Baseline, and I'm using the batch size eight as a Baseline, and just so 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 torch tune recipe. One difference in torch tune and train.py is that it uses a yaml for the config versus the command line. It's just different. And I did have to tweak the yaml quite a bit to make sure that I was running the same config. 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, Ah, fsp2, 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 got to 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 uh, so this is where I get into the actual tracing stuff. H. I wonder if I can actually show you the traces. That would be pretty cool. Okay, I'm going to stop sharing to reshare, and then we can do that. Let's just share my entire screen. 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 slow boy. So, in our Baseline, we're going to go to the Second Step because the first step is always full of warm-up stuff and initiating stuff. So, we're just gonna ignore that, and we're going to go here because every other step after this is much more normalized and something you can do. Um, I'm using Perfetto. I don't know if people are familiar with Perfetto already, but it's been pretty helpful. Yeah, okay. And something that is super useful and nice in Perfetto is you can highlight a region, and it will give you profiler results on here. So, here you can see that there are a lot 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 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 896, which is like more than double. And so let's go find a div, like where are they? And here, you can just like, but you can already notice that everything kind of gets doubled here, whereas in here, they are just called immediately. So like this A10 mole goes directly to A102. This A10 mole though goes to A10 mole again, and you 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 DET tensor, it means that it goes into this mole as a det tensor, and then detor is like all

right, let me do my metadata unwrapping, and now you get to go to atmo as a plain tensor now. So, there are double the amount, and just to spare you some math, um, spare you some math, uh, it turns out that if we divide this mole by two, or the div by two, 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. 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, uh, 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 glorify the output projection, whereas in our world, we do do that. And since glorifying means you add two adapters, and there are 32 layers, 32 * 2

is 64 more extra parameters to train um so yeah that was bad and that was a great lesson because it was like I was not measuring Apple 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 data set 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 like it was not easy to figure out every little difference and why they were different uh but after all of that I was like let me run it again maybe it'll 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 unblocking 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 the culprits. And that's how I knew what to focus on. Okay, so this is a big slide. Um, I okay, if at this point you have not looked at the traces yet but you would like to, uh, I sent a link to the Google Drive in the Discord, and there are the traces that you want. There's like, there are four of them, and the two you care about to look at now are the ones that don't start with bad or final. The other two do, the answer AI one and the torch tune one. But I'm going to, I already found these gaps for you, and it'd be fun if you find it fun. If you want to 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 want to save time, um, we're going to 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 GNA 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 Det tenser, um, 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 three times slower. And also, another Minor Detail, if you look at this, this is 64 M uh micros seconds, and this is 81. The answer for that is because the parameter is not continuous, 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 atom 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 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 three are a little boring, but we're going to go through, and 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, 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. Um, and I would do that for you, um, but we are a little short on time, so I'm going to just show you what the differences here. And the first difference that was major was that the second a in every linear, like every layer, was somehow 5 milliseconds longer. And that was when I needed to click on them and figure out how big they were. And I noticed that in the on the left side for train.pi, there were just fewer things getting all gathered. 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 fsdb 2. And when 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 scalers and quantization Factor. They just like for bits and bites when they did their four-bit thing. They used a dictionary. They had their own method. And this is actually the reason they couldn't pack it. It's because fsdp 1 had restrictions, by the way. Like it just wouldn't do it for them. So they didn't 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 12 million bytes. And that was when we realized that when we opted out of Laura, the output projection did not get quantized in our world. So it was like four times the size of answer AI version. And I was like, Oh, why don't we do that? And then I talked to the torch tune 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. Um, 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, uh, 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 is going to complain. You need to de-quantize, get it to the BF16, and then put it through the CUDA kernel. And it turns out that Tim Demmer's, 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, or when I don't have a presentation on Monday. And then, uh, the second step is to just use Triton kernels. So, Dris, our coworker, already wrote them. I didn't want to copy-paste them for the sake of this presentation, but if you want to 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, this one is definitely the worst gap. What? Where basically there were just very different ops happening before the uh sdpa. And this was just because we used two different rope algorithms. Torch tune 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 Torch to offer more 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 into the GPUs, you like have them talk to each other. And here in train.py in FSDP 1, 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 FSTP 2 promises you. So, FSDP 2 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 two layers of parameters will be allowed at a time. And that is why on the left, because of how FSP1 was implemented, it didn't do that very well. So, you'd get random memory spikes. And FSP2, you're promised to never get that. But someone looking at this tray will be like, But this is kind of... I don't... 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 uh the here when you're CPU offloading. You're actually trying to bring this big boy back in. So, 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 it. It's really just hey don't wrap those Linear by themselves. Just group them in so we can just have every uh Transformer decoder layer be its own layer. And note that this is only possible with fsp2. You can't you can't have the the right hand side and ofp1. Why? Because the L the Lura linear have the base weights. The base weights are quantized. They're nf4 tensors. They're going to be small. They're un8. Whereas your linear tensors those are bf16 or whatever training thing you want. And in fsb1 they can't be brought together under the same family.

because they're going to be forced into one big flat parameter, but in FSDB 2, 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, um, 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 Exposition at all, uh, where even in the first case. Wait, even before here in the train. Pi one. This all gather 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 FS DP2, which is a great point here. All right, so now we fought all our gaps and things work, um, so it's your turn. You should try them out. You should, you should really like it. It doesn't have to be N4. If you have another favorite lower precision thing, you can try it. If you want to try out, uh, more, more llama stuff with it, you

can there are now scripts ready to do that. Uh, so yeah, one disclaimer though, we are working on this. Like, this is not a forever thing. Checkpointing of just fsdp and Cur 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 J wrote the original N4 stuff. Andrew is the main dude who designed fsdpt. Wey ended up taking Andrew's work and making it compos with dr's work. So he like Amal amalg both of them and then torch tune people. So like Rohan and Evan, they were super helpful. They wrote the Laura recipes and they're the foundation of which we built and showcased our work. And of course, Mark, who's amazing. So thanks, Mark. Um, yeah, so like I really hope this gives people some more context.

around what we're thinking here. Like, we did want to showcase a useful recipe with the Cur and fstp composition, but really this is kind of like, you know, 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 D types or more exotic forms of Paralis, a lot of this work should be really, really helpful. And you know, we have all these links here that can give you some more context. U, 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, if people have any questions.

Slaying OOMs with PyTorch FSDP and torchao

Hamel Husain

OOMs are a common and frustrating bug in Pyre code, often addressed by reducing batch size or model size.

VRAM constraints are a major challenge, with popular consumer cards offering around 24 to 80 GB.

The memory required for a model includes parameters, gradients, and optimizers, often exceeding available VRAM, making activations the dominant memory overhead at larger batch sizes and context lengths.

OOMs (Out of Memory errors) are a common and frustrating issue in Pyre code, often addressed by reducing batch size or model size.

VRAM constraints are a significant factor, with popular consumer cards having 24 to 80 GB of VRAM.

The total memory required for a model includes parameters, gradients, and optimizer size, often exceeding the available VRAM, making fine-tuning and training challenging.

Efficient Parameter Quantization and Gradient Management

The process of quantizing parameters from fp32 to int8 involves finding the maximum value and scaling the values accordingly, with de-quantization reversing the process.

Using torch compile and environment variables can generate efficient Trident kernels, eliminating the need for custom CUDA kernel writing.

The Q-Lauro algorithm, which updates a subset of model parameters, combined with 4-bit quantization, has proven effective in training models, with its implementation being simplified through pure PyTorch code.

Efficient Parameter Handling and Quantization in Machine Learning

The approach to quantization involves scaling values from fp32 to int8 using the maximum value of the elements.

Torch compile is used to generate efficient Trident kernels, eliminating the need for custom CUDA kernel writing.

The Q-Lauro algorithm, which updates a subset of model parameters, allows for effective gradient handling and has been widely adopted in fine-tuning competitions.

Optimizing PyTorch for Lower Data Types and Multi-GPU Training

PyR, the PyTorch library, does not support NF4 tensors, focusing instead on IN8, FP16, and FP8.

Tensor subclasses in PyTorch allow for the customization of operations and data type distribution across multiple devices.

To manage memory during training, data parallelism and parameter sharding are effective strategies, with Fully Sharded Data Parallel (FSDP) offering a way to exchange parameters between GPUs and optimize memory usage.

Double Quantization and PyTorch's Native Data Types

Double quantization involves normalizing data and using lower dtype kernels, with PyR offering IN8, FP16, and FP8 but not NF4.

PyTorch's tensor subclasses allow for overriding operations and customizing data type distribution across devices.

The presentation also discusses memory management during training and introduces fully sharded data parallel (FSDP) as a technique to efficiently use multiple GPUs.

Understanding FSDP and Its Impact on Memory Management

FSDP determines the wrapping of model layers, influencing memory usage by grouping parameters and layers.

CPU offloading can be used to reduce GPU memory requirements by moving parameters to the CPU.

FSDP 2 aims to improve tensor parallelism and efficiency in distributing tensors across multiple GPUs.

FSDP (Fully Sharded Data Parallel) wraps layers in a model, influencing memory usage and parameter distribution.

CPU offloading can reduce GPU memory requirements by holding parameters and sharing data when needed.

FSDP has been challenging to work with, leading to the development of FSDP 2, which aims to handle tensor parallelism more efficiently.

The Benefits and Implementation of FSDP 2

FSDP 2 improves tensor management by slicing and sharding, allowing each tensor to maintain its identity and metadata.

This approach enables tensors to have different data types and gradient requirements, addressing challenges in training and quantization.

FSDP 2 also offers memory determinism and has been integrated with NF4 for a cleaner, more composable usage, with ongoing benchmarks to evaluate its performance compared to FSDP 1.

Advancements in FSDP: From Version 1 to 2

FSDP 2 improves tensor distribution by slicing them individually, allowing them to maintain their unique identities and metadata.

This version addresses common challenges, such as varying data types and gradient requirements, offering more flexibility in training.

Additionally, FSDP 2 provides memory determinism, reducing the risk of unpredictable memory spikes, and has shown better peak memory usage, though it is 19% slower in runtime compared to FSDP 1.

The new trace for the optimizer shows a significant increase in the number of parameters and steps compared to the baseline, leading to longer execution times.

The use of tensor subclass dispatch and a configuration difference in output projection are identified as key factors contributing to the increased parameter count.

Despite efforts to align the configurations and parameters, the new setup still resulted in slower runtime, but with improved peak memory usage.

Analyzing Performance Traces in Optimization

The new trace for the optimizer shows a significant increase in the number of steps and parameters compared to the baseline.

The doubling of parameters and the use of tensor subclass dispatch are causing the optimizer to take longer.

After identifying the config difference and running a benchmark with consistent parameters, the peak memory improved, but the runtime remained slow, leading to a focused analysis of the forward, backward, and optimizer steps.

The first gap was the optimizer step's slow performance due to C++ dispatch overhead, which was resolved by working with Intel to fuse the optimizer step into one kernel, resulting in an 8X speed-up.

The second gap involved mismatched tensor sizes and inefficient packing, which was addressed by optimizing the packing process and quantizing the output projection.

The third gap was caused by additional overhead from de-quantizing NF4, with potential solutions including using Torch compile and Triton kernels for faster execution.

Identifying and Resolving Optimization Gaps

The first gap in the optimization process was caused by overhead from C++ dispatch, making the optimizer step three times slower.

The solution involved working with Intel to create a fused atom for the optimizer step, resulting in an 8X speed-up.

The second gap was due to differences in tensor packing and output projection, which were resolved by implementing a more efficient packing method and quantizing the output projection.

The speaker discusses the challenges of CPU to GPU offloading and the overlapping communication between GPUs in the training process.

FSDP 2 aims to address these issues by ensuring proper memory constraints and grouping layers, leading to better synchronization.

The speaker emphasizes the importance of collaboration and invites researchers working on quantization and distributed systems to connect for further exploration.

FSDP 2 aims to address these issues by guaranteeing memory constraints and allowing larger parameter layers to coexist.

The speaker highlights the successful implementation of a new wrapping policy in FSDP 2, which eliminates the overlap and improves efficiency, while also calling for collaboration in research on quantization and distributed systems.