Lightning Talk: In-Cluster Distributed Checkpointing: Optimizing Training... - G. Kroiz & S. Mishra

Key Takeaways

Welcome everyone. Um I'm Sor. I work on PyTorch at Meta and I'm joined here by Gersonen um who comes from Google. And together we're going to talk about some of the scaling bottlenecks for fall tolerance and various solutions um that we developed to resolve those. So in a training job even though we take checkpoints periodically and we restore the training progress uh in the event of a failure uh there's still a non-trivial amount of GPU wasted GPU hours. Um one reason is the checkpointing overhead itself. Um so it we may not want to take a lot of checkpoints. Um the second is the unsafe training progress which is the red block here that you see and it can be as bad as the checkpointing interval itself and it has to be multiplied by the uh number of GPUs in the job. So for large scale jobs we want to minimize this checkpointing interval as much as possible which means that we may want to take as many checkpoints as possible. So you can already see the trade-off here. So more visually if we uh plot the wasted GPU hours or wasted training time against the checkpointing frequency the curve looks something like this. So if we are checkpointing too frequently um then the wasted training time goes up because of the checkpointing overhead and if we are checkpointing too infrequently then again wasted training time goes up because of the larger window of unsafe training progress. So for large jobs we want to take checkpoints very frequently but at the same time we want to reduce the checkpointing overhead to near zero. So what are the different um overheads during checkpointing. So I am taking the example of the asynchronous checkpointing here uh because I'm assuming anybody who cares about training efficiency uh they must already be using asynchronous checkpointing uh in there uh except staging all other steps in the checkpointing like the metadata preparation planning IO upload all those things happen asynchronously now during staging we offload the data from GPU to the host memory and that's why uh there's memory overhead uh overhead we allocate and deallocate the memory for every checkoint point save and also it's pageable memory so page falls also happen leading to inefficiency and again this is a blocking operation for GPUs the next step is the planning step uh where we figure out how different data types need to be serialized and saved in storage uh and we also prepare the metadata so there's metadata preparation cost here uh but the main bottleneck here is the collective um so rank zero coordinates across all these different ranks uh in the job and it gets uh their local metadata. Then it does validations, dduplications, uh creates the global metadata which describes the entire checkpoint. Um and this collective u what we have observed is become uh it becomes quadratically more expensive um with respect to the job size. And in the IO step uh we observe uh that the GPU utilization uh gets impacted because of the uh global interpreter lock uh contention between the trainer threads and checkpointing threads. So on the right you can see the plot of GPU utilization and the red line in the middle uh demarcates the time when staging has completed and after that the GPU utilization should come back to the previous baseline but you can see that it's sub-optimal uh it's less than expected and finally uh the checkpointing can only happen as fast as the storage bandwidth allows right so there are storage bandwidth constraints as well so we developed various optimizations for staging uh we uh now do pin memory memory based staging. So the memory is page locked. So we don't run into uh page falls. Uh also another interesting idea here is that during the forward step of the training weights do not change. So we overlap the uh staging step with the forward step. So effectively the staging uh becomes partially or completely asynchronous. For the planning step uh we introduce the rank local checkpointing where every rank uh saves and restores its own data and metadata. Another interesting idea is that across different versions of the checkpoint the metadata that describes the checkpoint that does not change. So we can cache the metadata in all these runtime objects and we can amortize the metadata preparation cost across different checkpoint saves. Finally, for IO optimizations, uh we built the process-based checkpointing which completely eliminated the gill contention issues uh that I described earlier. And for storage bandwidth constraints, we built in-cluster uh checkpointing where we stage and replicate the checkpoint within the cluster itself. To provide more technical details, I'll invite my friend Garrison here u who will talk more about it. Thank you. One of the reasons we were excited to collaborate with the torch GCP team was so that we could bring some of the latest and greatest checkpointing optimizations to our customers on Google Cloud. I want to provide some more color in what it's like to train on cloud and how to measure efficiency. Cloud customers want to maximize the return on investment, right? And care a lot about raw performance metrics such as MFU, tokens per second, and step time. These are great metrics for understanding the overall model's efficiency within a short period of time. As you have larger training jobs, due to the sheer distributed scale of these workloads, customers will run into unexpected failures and uh fault tolerant issues. Due to the highly coupled nature of AI workloads, when one node goes down, the entire training job stops making any progress. And this really emphasizes the overall importance or the importance of overall performance. Right? We don't only care about performance at a single step time, but the performance over the course of the entire training job. We define a metric goodput for this. And to understand goodput, there's a visual on the right. The bar represents the total time spent training where 80% of the time is spent training at peak performance. The remaining 20% of the time is time wasted due to unexpected interruptions, failures and stragglers and things of this nature. We could define goodput as the percent of time spent on productive training or 80%. The remaining time is considered badput 1 minus 20 or 1 - 80% which is 20%. Goodput is important to customers as it has go to market implications and could save millions of dollars. Now tying back in checkpoint, why is checkpointing so important to goodput? There are a few reasons here. During stable training, every checkpoint you save has some degree of saving overhead and this results in slower step times and ultimately longer time to converge. Furthermore, during these unexpected interruptions, you have to reload your state back into HBM. This is loading overhead, which is ultimately time not spent training. And furthermore, you have to revert state to an old checkpoint. Any progress not saved is considered lost and the time spent trained on these lost steps is considered waste in progress. The async checkpointing optimizations that Sarab went over are great for minimizing the checkpoint badput but alone are not good enough due to existing storage IO bottlenecks especially within a cloud environment where your storage is in a persistent location. So we collaborate together to develop an in-cluster checkpointing solution for our customers in Google cloud and aims to minimize checkpointing badput. We define incluster checkpointing as local checkpoints being able to save and save and uh load checkpoints locally without a persistent storage. Here we see just the cluster itself in which you could manage all of your checkpoints. The obvious uh benefit of this is the read and write speeds are much faster since you could take advantage of RAM disk. Now you could also tie this in with a more traditional persistent checkpointing uh based solution to have a multi-tered checkpointing solution. Right? You save your checkpoints frequently, frequently locally, and infrequently to persistent storage in case of any catastrophic clusterwide failures. When you add on the torch DCP async optimizations that Sar went over, you're able to push saving overhead to near zero, which allows you to fully utilize and take advantage of the ability to save checkpoints frequently locally. Now, I want to provide a bit of an illustrative example on how incluster checkpointing works. Here we have an oversimplified workload uh of a customer's training job where we have four nodes, two data replicas and we have one hot spare. You can notice the hotspare does not have checkpointing state because checkpointing state is or the the training state is updated frequently. But now let's say one node goes down, right? To resume training as quickly as possible, customers will reserve a portion of their capacity such that when this node goes down, you could swap it back into place. However, you can see node E just swapped in place doesn't have the checkpoint state. Furthermore, in place reset is not guaranteed. And what is meant by this is that as you swap in and out nodes, your overall network topology changes. And this dictates how you define your distributive collectives and groups. And so because of this, in this example, node A and node C swap positions in our network topology and thus need to transfer state. So with incluster checkpointing in our solution, we're able to replicate state during loading time, right? We see node A and C both already in the cluster swapping state and we also see node A giving its state to node E such that node E has the need to checkpoints and could resume training after all this swapping which takes advantage of the uh nickel collectives. We're able to then resume training. Now I want to like I hope that built some intuition on how incluster checkpointing works. But we also have some imperial core results. Here's a plot showing the checkpointing impact on badput. Right? On the y-axis we plot the overall badput percentage incurred from checkpointing and on the x-axis we show the number of interruptions a day. We can see that as you introduce each of these async checkpointing optimizations that s went over, we're able to reduce our badput percentage. Furthermore, with the addition of in-cluster checkpointing, we're able to even further push our the bad push percentage closer to zero. I think one of the highle takeaways here is that incluster checkpointing as well as these uh async optimizations are important and that with your depending on how frequently you run into these interruptions, you may need a various uh degree of sophistication. And in the most extreme of cases with all these optimizations, you're able to reduce badput incurred from checkpointing by over 50%. That concludes our talk. Uh we're happy to take any questions afterwards and if you're interested in learning more, we have several links attached to the slide blog post we worked on together and also like Google Cloud's approach to resiliency and checkpointing. Thank you.