Fp8 Training From Hopper To Blackwell - Luca Wehrstedt, Meta

In this talk, the speaker discusses the implementation of FP8 training on Nvidia's latest GPUs, emphasizing its benefits and challenges. FP8 allows for significantly increased training speed due to its lower precision, providing four times the speed compared to FP32. The speaker elaborates on the nature of floating-point representation, detailing the trade-offs in precision and range. He highlights the importance of scaling data appropriately to avoid overflow and underflow, particularly during matrix multiplications. Several scaling methodologies are presented, along with considerations for optimal performance and accuracy in machine learning models. The session concludes with insights into advanced topics such as stochastic rounding and the specifics of integrating FP8 into training workflows with tools like PyTorch.

My name is Luca. I'm uh I work at Meta, though not on the PyTorch team itself, rather on the research team. And I'm here to talk to you about training models in FP8 on the latest two generations of Nvidia GPUs. My goal today is to kind of give you some a lay of the land, uh an overview of what FP8 does, what it doesn't do, what the

pitfalls are, to help you find your way around it. I won't tell you what to do exactly. I won't tell you how to do it. I won't show any code. Oddly for the venue, I won't mention PyTorch except in this slide here. It's really just like a tutorial on best practices and bad practices. With this out of the way, why would anyone want to do FP8? For

us here, the reason is speed, training speed, compute speed. And precisely because GPUs are always getting faster, each generation doubles or triples the flops, so the floating-point operations per second that the hardware can do, but they're getting even faster at lower precision. Each time you halve the size of your data type, you gain twice the speed. So, FP8 gives us four times as much speed on newer

hardware than uh FP32. And yeah, and we want to use this for when we are compute-bound, when we are stuck on the compute uh operations. There are two other reasons that someone might want to do FP8, but they're not as relevant to us here today in this room. This first reason is bandwidth. By having fewer bytes to load, you load them faster. Uh this is relevant when

you're bound on data loading, which is typically the case for inference, or for some other parts of pre-training, but not the ones we care about. This is why people one of the reasons people put the KV cache at inference in FP8. And this does not require fast, new, fancy hardware. You can do it on old GPUs. Uh it's just for storage, not for compute. One reason we

don't care at all about here today, but it's also very valid, is to store as memory. here they actually might store more memory because we need to convert tensors back and forth between low precision and high precision. So, we might actually store more data. But, again, in inference cases, this could motivate you to store the KV cache permanently in FP8. Now, with the why other than that,

see actually what FP8 is. I don't Maybe it's obvious. I want to cover it anyways. what is FP? What is float? What is floating point? It's the way we have in computers, pretty much everywhere, to represent real numbers, where we have a sign, we have a coefficient between one and two, uh which we call a mantissa. Then, we have a power of two that we multiply it

by, which we call the exponent. this is So, if mantissa and exponent had infinite precision, we could encode any real number, would be great. They don't, so we have a sampling of the real numbers. And we need to add some special odd values to make it all work out, zero and negative zero, infinities, not a number, subnormals, and a bunch of stuff that we going to sweep

under the carpet. the eight in FP8 means how many bits we use to store this data, and different floats have different number of bits that they allocate to different parts of this representation, mantissa, exponent, only one bit for the most of the cases. So, yeah, this is a trade-off. Some needs to be eight or 16 or 32, but that's where we can play with with trade-offs. the

two components, the mantissa and the exponent, they control two different aspects of a floating point, hence the One is the range. What is the smallest positive value we can encode? What is the highest positive value we can encode? How far away are they? What is maximum delta we can encode between two values? So, this is a range. The highest exponent, the more these two boundaries get pushed

out. And the second one is the precision, which is how many discrete notches we have inside this range. Like, if we go from one point in this, because again, it's a discrete sample of real numbers, to the next one, how dense are they? If we have more mantissa, we can encode smaller and smaller notches and have a finer granularity of encoding. And if anyone wonders what subnormals

are, it's not numbers that are out of the range to the bottom there. Again, very weird. I don't want to go into details, but that's where they are. Um Different floats have different ranges. So, here I put FP32 and FP8 as the two extremes. You can already see that the exponent of these 10 to the power of are widely different between the two. FP8 has a very,

very small range. And what happens if we want to convert between the above and the below? Two things could happen. If we convert a very small number, we round down to zero, so we lose all information. Anything there has just disappeared. We have a zero now, a real zero. And the other side, we overflow. Uh very very high numbers are too high for the range of FP8,

and they become infinitives, which is again, very useless. They don't encode much information. They're they're too big to be numbers. Uh so, this is a problems. There is a third thing we could do, which is is also a problem. We could round down very high numbers. In this case here, we would get a error of 10 to the power of 36. Also not always recommendable. So, we

sometimes do it, but we it's not advisable. We want to solve these problems, underflow and overflow, but we want to do so in a specific order. Top priority, avoid overflows. They really make everything go boom. If an info or non appears in the model, you stop there, you go home. You can do nothing about it. So, this is our top priority. Best effort, we want to also

avoid underflows because again we lose Stuff goes to zero when it wasn't zero, so we if you are a gradient that is not zero and we want to descend on the direction, we would like to remember that and not lose this information. So, priority one, priority two. this already If anyone here has been here for eight years or more, I haven't, so I learned this through through

colleagues. We already had this issue when we were doing FP16 and BF16 training. This was around the Volta and Ampere era. FP16 also had very small range compared to what we did before, which was FP32. And we found solutions like allowing or blocking operators from using lower precision by knowing the properties and how resilient they would be to this precision. We would scale the gradients to put

it back into the right range and avoid which is why which was all like very very error prone. We didn't like it. So, we invented BF16 which trades which is a non-standard D type if you want. Which has the same range as FP32. So, we don't lose range. We can convert without overflows, without underflows. Brilliant. And we sacrifice a lot precision. This worked brilliantly, so we learned

a thing. We learned that AI can handle low precision. It's noise basically. It's noise or quantization to add and SGD is able to handle it correctly. But, we do require or else really models don't converge. We get problems. We get infs or or NaNs. Let's see how this all applies to FP8, which is what we're here for. We cannot do BF8. If we do it, we have

all exponent, zero sign, zero mantissa. So, this solution isn't viable to us anymore. We need to take out scaling from the closet again and make it work this time, make it work well. Uh so, scaling here, that what we want to use it for, is to basically map the range of our data, the actual data that we want to train on, the more smallest value and highest

value, to the range of FP8, to the value that we can represent. To compress it down or scale it up, depending on what we need. is to avoid overflows. If we always map the highest value in the data to the highest value in FP8, we will never have an overflow. It will always be exactly, like, barely fitting in. Great. We might have some underflow because values close

to zero will get scaled down to zero, but this is our second priority, so it's fine to sacrifice it a bit. So, this is the scaling we want, and when do we want it? We want it only where we need it. So, as we said before, this is for compute optimizations. The compute happens inside matmuls. So, we want to only scale just at the very last moment

before our tensors enter a matmul. So, before we had this input and weight being BF16, we had a slow BF16 matmul. Everything was the same D type, nice and clean. Now, we want to basically add this quick, small, last-minute extra step just for the inputs only of our matmul to convert them to low precision, so that the matmul itself can kick to high precision, sorry, so to

higher speed, low precision, so to FP8 matmul. And that's the only place where we introduce FP8. All the other operations, the output of the matmul, all other operations, all the pointwises, all the reductions, we don't touch them. They're not compute bound, they won't gain anything asterisks from FP8, so we only touch the inputs to our matmuls. And what does it mean to scale? It means obtain an

FP8 tensor with our FP8 data that we have scaled down, but we need to remember those scales as well. We need to remember how we scaled down or up to scale up or down the result correspondingly and obtain the same numeric the same mathematical at least output at the end of the matmul. So, an morally is two tensors. A full size small precision FP8 tensor with the

actual data and taking along a small metadata tensor with the scales. If you use PyTorch, sorry, I broke my promise. Uh there is a class that handles it for you and makes it uh uh opaque, so don't worry. But uh this is what happen under the hood. This is the probably the biggest question of the talk. How do we scale our data? there's as many ways as

there are people in this room probably uh we can scale the whole tensor with a single big scale, so we take the global maximum of absolute maximum by the way of the whole tensor. We do it per rows, we do it per feature, per channel, per column, we do it per block, square block, rectangular block, flat block, flat column. You can envision your own uh scaling tile

scaling like quantization block. All are valid options, all have trade-offs. Uh people have tried uh very very many of them. Let's talk trade-offs. Um biggest one is accuracy. The smaller the fewer tiles if if the fewer the elements you pull together the less error you have. If you have one scale for one element, you're just doing a half precision matmul, so you're back to where you you

started. So, we want to the most fine-grained tile we can. It would limit the risk of underflows because the scale would be as small as possible, so we don't risk one very high element squeezing all the other ones in its tile down to zero. So, for accuracy, we want small tiles. For efficiency, we want large tiles. Uh because scaling up the result has a cost, which we

can amortize the bigger the tile is, which is do it once. Like thanks to linear algebra, we can only do it once per per tile. So, the more tiles we have, the more we pay a cost in descaling the result. This is only partially true on Blackwell because it's now hardware accelerated, so actually Blackwell we can have small if we follow the instructions from Nvidia, the pattern

from Nvidia, we can have very very fast descaling, and we can actually afford very small tiles. But before Blackwell on Hopper, we had to keep this into account. I'm okay with losing some of the perf that FP8 gives us to recover some of the accuracy by smaller tiles. The inverse is true for the scaling itself. Smaller tiles can be scaled faster because you can do them more

in parallel. You don't need to do a single global maximum of the whole tensor. Maybe do it in two passes. Keep everything in work memory. No, if you have small tiles, you can do it on the fly much more quickly. So, it makes the previous the pre-step the preparation step for your input faster. Here again a tradeoff. How much do you balance between making the quantization step

faster or the amount more faster. Last, we have found out that you might want to consistency between forward and backwards for your weights. It There might be risks with rounding, so underflowing a weight a parameter in the forward, so not having it contribute at all to your loss, but not having it underflow in the backward, and having it contribute to a for loss that it didn't take

part in. This can sometimes lead to better results. So, we tend to also prefer uh tiles that allow us to have the same exact overflow underflow pattern in the forward and backward, at least for the weights. There's lots to take in. There's I think there's no correct answer here. Uh you might want to go with NVIDIA record solution with your own solution. You might want to try

out a few options, but yes, there's lots of uh aspects to consider to this pretty much core question of uh how we do FP8. One last detail I think about scaling is uh how do we store this uh scaling metadata? On Hopper, we're using float 32, because why not? It's a very small tensor. We It's really like key to recovering precision once we de-scale. So, we use

float 32. On Blackwell, if we want to use the NVIDIA accelerated path, we need to use basically the exponent of a float 32, which is a 8-bit power of two weird float data type. It's basically like an integer, a uint8 that you interpret as a uh exponent. So, it's power of two uh up to eight bits. Anyways, details, but uh this also matters. Blackwell loses some precision

by having this rounded scale, but it gains it by having a much smaller uh scaling tile. One thing we've stopped doing, but I'm putting it here uh anyways, is at the very beginning, we were doing delayed scaling, where we were not actually computing the real-time based on the data tight maximum. We're using some estimates based on historical uh previous steps. This doesn't work. Uh this leads to

overflow uh much more much more easily, and overflows are bad. So, we stopped doing it. We prefer versions where we really really determine each time the actual maximum from the data. Uh going to advanced topics, this is mostly for your knowledge. I don't recommend doing it necessarily. Is you can try to limit underflows by tweaking your data bit. If you have a very very high value in

your in your scaling tile, it will squeeze all other ones to zero. But, you can rotate your uh embedding your vector your tensor a bit, and suddenly now all coordinates are more or less balanced if you do it correctly. And therefore, none of them will over will underflow. They will all be correctly encoded. Uh and this is magic because if you scale if you rotate in the

same way both operands of matmul, they just cancel out because the matmul only cares about the angle between its operands, not the actual coordinates. So, you do it to the inputs and you don't need to do anything to the outputs. Magic. How do we rotate? Uh I know of two ways. Uh one is you actually have a fully random rotation matrix, which you can sample. There's there's

tools for this. It's nice because it has really a Gaussian distribution, which is usually what we want to recover. It's a nice bell curve. It's denser in the middle. It's It's also clamped. It's a better distribution, actually. So, this is mathematically what we want. People have also been using what It's called like a Hadamard rotation, which is a very special one, which is constructed in a recursive

fractal way with only plus ones and minus ones. So, it's very easy to apply. You only add elements, no products, no complicated math. Both have downsides, in my opinion. Doing a real rotation means that you touch your bits of your um uh tensor much more, so it adds more noise. Hadamard matrices are not exactly perfect rotations. They're not uniform, so they might not give you the actual

normal distribution you want at the end. You do you. I don't recommend uh specifically either of them. I want to cover the some gotchas real quick. I don't have much time left, but let's get into it. If you remember at the beginning, there was actually two FP8 types, which have a different trade-off between a range and Nvidia, I think, plans to use the one with higher range

for gradients. We realized actually that's not that interesting. We It works actually better if you use the standard E4M3, so 4-bit exponent and 3-bit mantissa for everything. Because the scaling itself already provides enough range to recover this. we have personally stopped using E5M2. This is counterintuitive because we said that we need the range more than precision, but the scaling does it already. So, there's no need to

have even more of it inside the FP8 type itself. If you want to compare the two, this is a plot I like to use. So, one is denser in some range and then becomes sparser, one is sparser everywhere. But, we have 15 orders of magnitude where it's dense and this is usually enough for us. Now, two pitfalls on Hopper, which are not applicable to Blackwell. Hopper had

an undocumented behavior where actually Tensor Cores accumulated FP22. I don't know if it's a standard Which means that if you use what they call fast accumulation mode, in some cases, in the backward in particular, you would lose precision during accumulation. So, we recommend disabling it. It It It was a weird feature with a weird accumulation. Slow accum is not much slower, a few percent in our experience,

and it gives you much higher results. So, this is what fast accum does, if you're wondering. We recommend disabling it. Hopper also had a weird constraint on the way you needed to store your FP8 data in memory. The reduction dimension, so K in this picture in this needed to be stored contiguously in Uh which made it very painful in the forward backward because each of the tensors

that are implicated in a matmul, the input, the weight, the output gradient, they all appear each one time row-wise and one time tensor-wise. That's right. Column-wise. which means that you cannot reuse one for in both forward backward. You need to re-transpose them each time. You need to re-quantize them a second time. Very, very painful. Uh backward luckily solves this. But yeah, this was one of the main

challenges I think we had in Hopper uh because these extra copies, these extra transpositions basically ate most of the speed gains that we were hoping to obtain. I don't know if there's really a way around it, though. I want to also go back to the second reason for FP8, which is we need to transfer less memory. This actually can help us when we for the quantization step

itself, so for the pre-processing step before the matmul, because naively, if you have a BF16 baseline and we add to it a separate quantization kernel, this uh maybe doubles the time you you spend quantization cuz you need to load BF16, restore it again, reload it, and then finally store it in FP8. If you could fuse it with something like but uh torch compile, you actually it could

become faster than the baseline because you replace the expensive store in BF16 with a much cheaper store in FP8. So, quantization actually might, if done correctly through uh inductor, torch compiler, your own kernels, uh make you the whole like even the non-matmul part of your model faster than it was before. Try to do it. Uh we tried it very hard. It really paid off. So, I recommend

it. Uh I have 4 minutes, so I'm going to cover this bonus chapter cuz I think it's very nice. I only got just got into it, but I find it find it really, really cute. Uh stochastic rounding. Uh Uh, let's start with a quiz. What's behind this image? It's a grayscale image. Hard to say. If I round it to the nearest Boolean, so down to zero, up

to one. What do we see? We don't see anything again. Okay, there's a diagonal in the What if I first apply some noise and then round it to zero or one? We recover this type of image here. Which is much, much more detail. We see much more detail than we had before. And now you could guess what the mystery image was. This is called dithering for vision

people, for image people. This is stochastic rounding for floats. Concretely, what you want to do mathematically, if you're mathematically inclined, is to make sure that the at the on expectation, on average, we round a value to exactly, like mathematically, it's the value. We don't want to introduce any bias up or down. Anyway, this is very important for small values, which could risk getting underflowed each time we

quantize them. We want those to sometimes, by luck, we round it up and actually deliver some signal, keep some signal. So, this noise makes them sometimes go up a notch and not always down. You can see it as adding like a symmetrical uniform noise and then rounding to nearest, or there's another way of seeing it with rounding down. So, this is how you actually implement it. And

mathematically, there's a whole theory behind it of you want to add as many bits as the one you're canceling to really compensate for the loss of information you're introducing, and you want them to be aligned and IID and uniform and so on. In practice, we mostly approximate this for efficiency reasons. We have fewer bits, we reuse it somewhat, we reverse them. Lots of tricks to make it

a bit cheaper to to run. Still seems to kind of work. Uh, so this was all I had to say, and I'm only 2 minutes left before the end of the talk, so I'm going to break my promise again. How do you do this in PyTorch? Uh the keyword you want to look for is torch AO. Uh this is where everything that's FP8 related is put in

PyTorch. All the helpers, all the kernels, all the tensor subclasses, all the docs. So, go to torch AO and then look it up. Um if you want to see how it's to use it in your training model, torch titan has some bindings, has some examples, can show you how to deploy it, how to implement it. So, also go take a look at torch titan to learn more

about this. That's all I have. I think we have 1 minute for questions, if there are any. >> Uh I'm Andre. I'm from uh Amazon devices. Uh so, my question is uh about the FP8 uh values, and you said you considered the bigger mantissa to be more important. Have you considered fixed-point uh as the quantized value instead of uh FP8? So, like or signed integers? So, E0M7?

>> Short answer is no, because we can only use what Nvidia has hardware accelerated in tensor cores. So, we have two data types that we will run faster on the hardware. Everything else, every other data type we could imagine we could come up with, we would need to convert to something that the tensor core understands, so BF16 or even higher, and we would lose uh the performance

speedup we get. So, we are constrained by the tools we have. Um I know that in uh inference cases, for example, they do that, because they're not bound by compute, they're bound by memory uh access, memory bandwidth. So, they can afford to do compute in a slower data type, they're not fine, as long as they get the bandwidth speed they are looking for and the accuracy. But

yeah, we only looked at the those two data types as the ones we could actually efficiently run. >> So, your use case was for training? >> Yes. Yes, thank you. All right, I think we're over time, so maybe we stop here and if you want to find me around, I'll be here today.