Lightning Talk: The $100K GPU Mystery: Why Your AI Training Dies at 99% - Michael Ifeanyi, Google

This talk addresses the issue of GPU memory allocation failures during AI training jobs, particularly when monitoring tools like Nvidia SMI indicate sufficient free memory. The speaker, Michael Ifeanyi, a technical solutions engineer at Google, explains that the actual problem lies in memory fragmentation, which can prevent large tensors from being allocated despite having an overall surplus of free memory. He illustrates this with an analogy of a parking lot, where contiguous space is necessary for large objects. The session then outlines practical solutions for managing GPU memory, such as monitoring fragmentation with the torch.cuda.memory_stats API, reducing memory pressure through smaller batch sizes or gradient checkpointing, and employing mixed precision calculations. Finally, he emphasizes the importance of understanding the limitations of Nvidia SMI and the need for continuous monitoring to ensure training jobs can proceed smoothly.

Good evening everyone. Have you ever wondered what the reason could possibly be when your AI training job crashes at 99% and at the same time your GPU monitoring tells you that you do have available memory. Now how's that possible? My name is Michael Ifeanyi. I'm a technical solutions engineer at Google. Let's dive in. So this will be the road map to our talking points today. We'll address

the problem, look at the contiguous memory root cause and take some solutions and key takeaways. Now if you look at this slide, this is actually a little bit suspicious. Now the reason why I say that is because Nvidia system management interface tells us that everything's fine. Now that's the monitoring for GPU. Everything's great. You do have 6.71 GB free and there's plenty of room for your job.

Now but the problem here is that we do have an allocation for the tensor and that is being rejected. So you if you take a closer look at the CUDA out of memory error of saying what happened, you do see that there is um an attempt to allocate 7.54 GB but it tells you that it's out of memory. Now the real story is that the Nvidia SMI

basically is just reporting the total free memory but it's not telling you that you have the contiguous free memory and that's the distinction of the root of this entire problem. Now you say what is contiguous memory? I'll give an analogy here like a parking lot where you have like like a school bus. So even if you have like five slots empty, the school bus cannot park in

the different five parking lots at the same time. Now that's the same thing with GPU memory. When PyTorch needs to allocate a large tensor, it needs one solid block. So at the top here you do have 10 GB free all in one chunk and you also do 7 GB allocated successfully. Now on the bottom here you do have the fragmentation that's appearing. Now you can see that

we do have the the blue the blue blocks representing active data in use and then we have fragmented memories on the green spots. Now if you sum this up together, it's still the same 10 GB free memory which Nvidia SMI is going to reveal to you. However, you're going to get a different result because the data is actually fragmented in different places and it's unusable and that

results to your job crashing. I'm going to pull out the test that I did on the Tesla T4 GPU just to expatiate more on the report the results. So on the left side you still see Nvidia SMI telling us that you have free plenty of room to carry out your training job and on the right hand you see PyTorch telling you that hey, you do have peak

fragmentation that's showing up in your training job. if you look at the bottom, this is actually the GPU memory layout and where do we have the blue blocks that are currently in use and you have the fragmentation blocks and then you do have the free block of GPU 6.71 GB. However, the revolution says that it's available but it's not at the same time because fragmentation has already

started eating into the GPU. How do we prevent this? There's like three steps that you could use at least for a start to prevent this. First of all, you want to monitor the right metrics. Now you want to utilize the torch.cuda.memory_stats API to track fragmentation. That's by taking away the reserved taking away the reserved from the allocation and then the remainder should give you the ratio to

track consider that as an important metric that would reveal the amount of fragmentation that's in your training then you also want to reduce the memory pressure. There are different ways that you can try to reduce the memory pressure. You could also use smaller batches where you if you had like 64, you could reduce that to 16 but you'll be trading that for speed. You could also enable

gradient checkpointing. That's kind of like memory optimization during training to reduce the memory usage during training. This would also trade your speed for more memory and this is slower. So you can also use mixed precision whereby you are either using the floating point 16 bit or the brain floating point 16 bit as well. So basically what this does is that you are involving the GPU to do

most of the calculation in 16 bits as opposed to 32 bits. Now the third level of solution that you could actually utilize in your setup is to deploy a demon set to extract this memory fragmentation. You could also export that to an observability tool like Prometheus where you you're being triggered where the threshold is being met. That's another fancy way of doing it and then there's the

option of draining the nodes. Now when you drain these nodes, keep in mind that you want to also implement checkpointing in your code and the reason why we're doing this is because when you restart those nodes, you want the job to continue from where it paused before you drain the node. Now there's three key takeaways here for today. First of all, we have an understanding that Nvidia

SMI does have a blind spot. It only shows the free memory and not the total contiguous free memory blocks. Fragmentation appears in an accumulated phase during and that prevents allocation even with apparent free memory. And third, you always want to monitor for fragmentation so that it doesn't creep up on your job. Thank you very much.