Pyre에서 메모리 부족 오류를 관리하기 위한 전략

### Pyre에서 메모리 부족(OOM) 문제를 관리하기 위한 전략
* 🖥 OOM(메모리 부족 오류)은 개발자가 **Pyre 코드로 작업할 때 흔히 겪는 문제**로, 효과적인 솔루션이 필요합니다.
* 🔧 기존에는 개발자가 즉각적인 해결 방법으로 **배치 크기나 모델 크기를 줄이는 방법**을 사용했지만, 이러한 접근 방식은 뉘앙스가 부족하고 지나치게 단순할 수 있습니다.
* 💡 예를 들어, **39 및 49와 같은 인기 있는 소비자용 GPU는 약 24GB**인 반면, **A1 100은 40 또는 80GB, H100은 100GB**에 달할 수 있으므로 VRAM 제한을 이해하는 것이 중요합니다.
* 📊 모델 크기를 추정할 때는 매개변수뿐 아니라 기울기 및 최적화 상태도 고려해야 합니다. 예를 들어, **llama 7B와 같은 모델은 모든 요소를 고려할 때 최대 56GB의 메모리가 필요**할 수 있습니다.
* 📈 특히 배치 크기가 클 경우 **활성화가 메모리 사용량을 지배**할 수 있고 적절한 도구 없이는 이를 추정하는 것이 복잡할 수 있으므로 **플래시 주의와 같은 최적화**는 매우 중요합니다.
* 🔄 예를 들어 **파라미터를 4비트로 양자화**하면 모델의 크기를 **14GB에서 약 3.5GB로 줄일 수 있습니다.**
* ⚙ 완전 분할 데이터 병렬(FSDP)과 같은 기술은 **메모리 사용량을 여러 GPU에 분산**하여 VRAM 한도를 초과하지 않고도 보다 효율적으로 모델을 훈련할 수 있도록 도와줍니다.
* 🛠 ️ **CPU 오프로딩은 매개변수를 CPU로 전송**하여 메모리 부담을 완화함으로써 보다 집중적인 연산 작업을 위한 GPU 리소스를 확보할 수 있습니다.

### FSDP의 발전: FSDP 1에서 FSDP 2로 전환하기
* 🔍 FSDP 1에서 FSDP 2로의 전환은 GPU에서 텐서를 관리하는 방식에 큰 변화를 수반하며, **텐서가 고유한 ID와 메타데이터**를 유지할 수 있도록 합니다.
* 💾 FSDP 1에서는 텐서가 플랫 파라미터로 결합되어 동일한 데이터 유형과 기타 메타데이터를 공유해야 했지만, **FSDP 2에서는 각 텐서를 개별적으로 표현**할 수 있습니다.
* ⚙ FSDP 2를 구현하면 **텐서 분할 및 재배치로 인해 추가 복사본**이 발생하지만 **텐서 동일성 유지의 이점**이 이 단점보다 더 큽니다.
* 📊 초기 벤치마크 결과 FSDP 2를 사용하면 **피크 메모리 사용량은 개선**되었지만 런타임 성능은 **FSDP 1에 비해 19% 느려져** 성능 격차에 대한 추가 조사가 필요했습니다.
* 🔧 세부 추적 결과, 분산 텐서 사용으로 인한 오버헤드로 인해 **FSDP 2에서 최적화 단계가 상당히 오래 걸리는 것**으로 밝혀졌고, 인텔과 협력하여 **융합 최적화 단계를 개발하여 속도를 8배 향상**시켰습니다.
* 📉 작업 중 텐서 크기의 차이 등 다른 성능 문제도 확인되었으며, **일관된 구성을 보장하고 텐서 패킹 방법을 개선**하여 해결했습니다.
* 최종 조정을 통해 계산과 통신의 중첩을 개선하고 **FSDP 2의 고유한 기능**을 활용하여 성능과 메모리 관리를 최적화할 수 있었습니다.
* 📚 FSDP 2로의 전환은 현재 진행 중이며, **연구 커뮤니티에 정량화 및 분산형 교육 혁신**을 모색하기 위한 협력을 요청하고 있습니다.

### Pyre에서 OOM 관리 최적화: 전략 및 도구
* 🛠 메모리 부족 오류(OOM)는 **Pyre 코드에서 중요한 문제**이며, 개발자는 메모리 제약 조건을 관리하기 위한 효과적인 솔루션을 찾아야 합니다.
* 💻 Jane과 발표자를 포함한 Meta의 개발자들은 배치 크기나 모델 크기를 줄이는 것과 같은 기존 방식을 넘어 **OOM 처리를 간소화하는 도구**를 만드는 데 집중하고 있습니다.
* 📊 39 및 49와 같은 인기 있는 소비자용 GPU는 **약 24GB**를 제공하는 반면, A1 100은 **40 또는 80GB**, H100은 **약 100GB**를 제공하는 등 VRAM의 한계를 이해하는 것이 중요합니다.
* 📏 모델 크기를 추정할 때는 파라미터뿐만 아니라 메모리 요구 사항을 빠르게 증가시킬 수 있는 그라디언트 및 최적화 상태도 고려해야 합니다. 예를 들어, **Llama 7B와 같은 모델은 완전히 미세 조정할 경우 최대 56GB의 VRAM이 필요**할 수 있습니다.
* 🔍 특히 배치 크기가 클수록 **활성화가 메모리 사용량을 지배**할 수 있으므로 플래시 주의와 같은 최적화가 효율성을 위해 필수적입니다.
* 💡 양자화, 특히 **4비트 양자화**와 같은 기술을 도입하면 모델 크기를 크게 줄일 수 있지만 **그라디언트를 양자화**하면 수렴을 방해할 수 있으므로 주의해야 합니다.
* 🔄 Laura와 같은 전략을 사용하면 **선택적 파라미터 업데이트**를 통해 학습 가능한 파라미터의 작은 하위 집합에 집중하여 메모리 오버헤드를 줄일 수 있으며, **FSDP(완전 샤드 데이터 병렬)** 같은 도구는 여러 GPU에서 메모리를 관리하는 데 도움이 됩니다.
* 💾 CPU 오프로딩은 또 다른 효과적인 방법으로, **매개변수가 GPU에서 정방향 및 역방향 패스를 수행하는 동안 CPU 메모리에 상주**하도록 허용하여 GPU의 메모리 부담을 완화합니다.

### FSDP 향상: 성능 향상을 위해 FSDP 1에서 FSDP 2로 전환하기
* 🔄 이 논의는 GPU에 텐서를 배포할 때의 어려움으로 시작하여, **FSDP 1의 텐서 관리가 텐서 간 메타데이터의 강제 균일성**으로 인해 어떻게 비효율성을 초래할 수 있는지를 강조합니다.
* 반면, **FSDP 2에서는 각 텐서가 데이터 유형 및 그라데이션 요구 사항과 같은 고유한 속성을 유지**할 수 있으므로 유연성과 성능이 향상됩니다.
* 📉 그러나 **FSDP 2는 텐서 분할 및 재배치로 인해 추가 오버헤드**가 발생하여 처음에는 성능 메트릭이 느려질 수 있습니다.
* ⚙ 이 발표자는 소비자용 GPU를 사용한 벤치마킹 과정을 자세히 설명하며 **메모리 사용량이 개선되었음에도 불구하고 FSDP 2의 런타임이 이전 버전에 비해 19% 느려졌다고 밝혔습니다.**
* 🔍 옵티마이저 단계와 텐서 연산을 꼼꼼하게 추적하여 **과도한 매개변수 오버헤드와 비효율적인 커널 호출**을 비롯한 여러 성능 병목 현상을 파악합니다.
* ⚡ 이러한 차이를 해결하기 위해 퓨즈드 옵티마이저 단계와 텐서 패킹 개선과 같은 솔루션이 제안되고 있으며, 일부 유망한 결과는 이미 **특정 작업에서 8배의 속도 향상**을 달성했습니다.
* 💡 마지막으로 **FSDP 2의 기능**을 활용하여 분산 훈련이 크게 발전할 수 있는 잠재력을 강조하는 한편, **양자화 및 분산 훈련 기술의 혁신적인 적용**을 모색하기 위한 연구 커뮤니티의 협력을 촉구합니다.

메타의 개발자들은 Pyre에서 OOM을 효과적으로 관리할 수 있는 도구를 만들어 모델 훈련의 효율성을 높였습니다.

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

OOM은 Pyre 코드에서 흔히 발생하는 불편한 버그로, 배치 크기나 모델 크기를 줄임으로써 해결되는 경우가 많습니다.

인기 있는 소비자 카드는 약 24~80GB를 제공하기 때문에 VRAM 제약은 큰 문제입니다.

모델에 필요한 메모리에는 파라미터, 그라데이션, 옵티마이저가 포함되며, 이는 종종 사용 가능한 VRAM을 초과하기 때문에 배치 크기와 컨텍스트 길이가 클수록 활성화가 주요 메모리 오버헤드가 됩니다.

OOM(메모리 부족 오류)은 Pyre 코드에서 흔히 발생하는 실망스러운 문제로, 배치 크기나 모델 크기를 줄임으로써 해결되는 경우가 많습니다.

VRAM 제약은 중요한 요인으로, 인기 있는 소비자 카드에는 24~80GB의 VRAM이 탑재되어 있습니다.

모델에 필요한 총 메모리에는 파라미터, 그라데이션, 옵티마이저 크기가 포함되며, 이는 종종 사용 가능한 VRAM을 초과하여 미세 조정 및 훈련이 어렵게 만듭니다.

효율적인 파라미터 정량화 및 그라데이션 관리

파라미터를 fp32에서 int8로 양자화하는 과정에는 최대값을 찾고 그에 따라 값을 스케일링하는 과정이 포함되며, 양자화를 해제하면 이 과정이 역전됩니다.

토치 컴파일 및 환경 변수를 사용하면 효율적인 Trident 커널을 생성할 수 있으므로 커스텀 CUDA 커널을 작성할 필요가 없습니다.

4비트 양자화와 결합하여 모델 파라미터의 하위 집합을 업데이트하는 Q-Lauro 알고리즘은 모델 훈련에 효과적인 것으로 입증되었으며, 순수 PyTorch 코드를 통해 구현이 간소화되었습니다.

머신 러닝의 효율적인 매개변수 처리 및 정량화

양자화에 대한 접근 방식은 요소의 최대값을 사용하여 값을 fp32에서 int8로 스케일링하는 것입니다.

토치 컴파일은 효율적인 Trident 커널을 생성하는 데 사용되므로 커스텀 CUDA 커널을 작성할 필요가 없습니다.

모델 파라미터의 하위 집합을 업데이트하는 Q-Lauro 알고리즘은 효과적인 그라데이션 처리를 가능하게 하며 미세 조정 대회에서 널리 채택되고 있습니다.

낮은 데이터 유형 및 멀티 GPU 트레이닝을 위한 PyTorch 최적화하기

PyTorch 라이브러리인 PyR은 NF4 텐서를 지원하지 않으며 대신 IN8, FP16 및 FP8에 중점을 두고 있습니다.

PyTorch의 텐서 서브클래스를 사용하면 여러 기기에서 연산 및 데이터 유형 배포를 사용자 정의할 수 있습니다.

훈련 중 메모리 관리를 위해 데이터 병렬 처리와 파라미터 샤딩은 효과적인 전략이며, 완전 샤딩된 데이터 병렬(FSDP)은 GPU 간에 파라미터를 교환하고 메모리 사용량을 최적화하는 방법을 제공합니다.

이중 양자화 및 PyTorch의 기본 데이터 유형

이중 양자화에는 데이터를 정규화하고 더 낮은 dtype 커널을 사용하는 것이 포함되며, PyR은 IN8, FP16 및 FP8을 제공하지만 NF4는 제공하지 않습니다.

파이토치의 텐서 서브클래스를 사용하면 연산을 재정의하고 장치 간에 데이터 유형 분포를 사용자 정의할 수 있습니다.

이 프레젠테이션에서는 또한 트레이닝 중 메모리 관리에 대해 설명하고 여러 GPU를 효율적으로 사용하기 위한 기술로 완전 샤드 데이터 병렬(FSDP)을 소개합니다.

FSDP의 이해와 메모리 관리에 미치는 영향

FSDP는 모델 레이어의 래핑을 결정하고 파라미터와 레이어를 그룹화하여 메모리 사용량에 영향을 줍니다.

CPU 오프로딩은 매개변수를 CPU로 이동하여 GPU 메모리 요구 사항을 줄이는 데 사용할 수 있습니다.

FSDP 2는 텐서 병렬 처리와 여러 GPU에 텐서를 분산하는 효율성을 개선하는 것을 목표로 합니다.

FSDP(완전 샤딩 데이터 병렬)는 모델에서 레이어를 래핑하여 메모리 사용량과 매개변수 분포에 영향을 줍니다.

CPU 오프로딩은 파라미터를 보관하고 필요할 때 데이터를 공유함으로써 GPU 메모리 요구량을 줄일 수 있습니다.

FSDP는 작업하기가 까다로워 텐서 병렬 처리를 보다 효율적으로 처리하는 것을 목표로 하는 FSDP 2의 개발로 이어졌습니다.

FSDP 2는 슬라이싱과 샤딩을 통해 텐서 관리를 개선하여 각 텐서가 자신의 정체성과 메타데이터를 유지할 수 있도록 합니다.

이러한 접근 방식을 통해 텐서는 다양한 데이터 유형과 그라데이션 요구 사항을 가질 수 있으므로 학습 및 정량화 문제를 해결할 수 있습니다.

또한 FSDP 2는 메모리 결정론을 제공하며, NF4와 통합되어 보다 깔끔하고 컴포저블하게 사용할 수 있으며, 지속적인 벤치마크를 통해 FSDP 1에 비해 성능을 평가하고 있습니다.

FSDP 2는 텐서를 개별적으로 분할하여 고유한 ID와 메타데이터를 유지할 수 있도록 함으로써 텐서 분포를 개선합니다.

이 버전은 다양한 데이터 유형 및 기울기 요구 사항과 같은 일반적인 문제를 해결하여 학습에 더 많은 유연성을 제공합니다.

또한 FSDP 2는 메모리 결정론을 제공하여 예측할 수 없는 메모리 급증의 위험을 줄여주며, FSDP 1에 비해 런타임은 19% 느리지만 피크 메모리 사용량이 더 나은 것으로 나타났습니다.

최적화 도구의 새로운 트레이스는 기준선에 비해 매개변수와 단계 수가 크게 증가하여 실행 시간이 길어졌음을 보여줍니다.

텐서 서브클래스 디스패치의 사용과 출력 투영의 구성 차이가 파라미터 수 증가에 기여하는 주요 요인으로 확인되었습니다.

구성과 매개변수를 조정하기 위한 노력에도 불구하고 새로운 설정은 여전히 실행 시간은 느려졌지만 최대 메모리 사용량은 개선되었습니다.

옵티마이저의 새로운 추적은 기준선에 비해 단계와 매개변수의 수가 크게 증가했음을 보여줍니다.

파라미터가 두 배로 증가하고 텐서 서브클래스 디스패치가 사용되면서 최적화 도구의 시간이 더 오래 걸렸습니다.

구성 차이를 파악하고 일관된 파라미터로 벤치마크를 실행한 후 피크 메모리는 개선되었지만 런타임은 여전히 느려서 전진, 후진 및 최적화 단계를 집중적으로 분석했습니다.

첫 번째 차이는 C++ 디스패치 오버헤드로 인한 최적화 단계의 느린 성능으로, 인텔과 협력하여 최적화 단계를 하나의 커널로 통합함으로써 8배의 속도 향상을 달성했습니다.

두 번째 차이는 텐서 크기가 일치하지 않고 비효율적인 패킹과 관련된 것으로, 패킹 프로세스를 최적화하고 출력 투영을 정량화하여 해결했습니다.

세 번째 격차는 NF4의 정량화 제거로 인한 추가 오버헤드로 인해 발생했으며, 더 빠른 실행을 위해 Torch 컴파일 및 Triton 커널을 사용하는 등의 잠재적인 해결책이 있었습니다.

최적화 프로세스의 첫 번째 차이는 C++ 디스패치로 인한 오버헤드로 인해 최적화 단계가 3배 느려졌습니다.

이에 대한 해결책으로 인텔과 협력하여 최적화 단계를 위한 퓨즈 아톰을 만들어 속도를 8배 향상시켰습니다.

두 번째 차이는 텐서 패킹과 출력 투영의 차이로 인한 것으로, 보다 효율적인 패킹 방법을 구현하고 출력 투영을 정량화하여 해결했습니다.

이 발표자는 훈련 과정에서 CPU에서 GPU로 오프로드하는 문제와 GPU 간의 중복 통신에 대해 설명합니다.

FSDP 2는 적절한 메모리 제약과 레이어 그룹화를 통해 이러한 문제를 해결하고 동기화를 개선하는 것을 목표로 합니다.

발표자는 협업의 중요성을 강조하며 양자화 및 분산 시스템을 연구하는 연구자들이 서로 교류하며 더 깊이 탐구할 수 있도록 초대합니다.

발표자는 트레이닝 과정에서 CPU에서 GPU로 오프로드하는 문제와 GPU 간의 중복 통신에 대해 설명합니다.

FSDP 2는 메모리 제약을 보장하고 더 큰 파라미터 레이어가 공존할 수 있도록 함으로써 이러한 문제를 해결하는 것을 목표로 합니다.

발표자는 중복성을 제거하고 효율성을 개선하는 새로운 래핑 정책을 FSDP 2에서 성공적으로 구현한 사례를 강조하며 양자화 및 분산 시스템에 대한 연구 협력을 촉구합니다.