Optimizing Error Recovery for Cost-Ef... Radostin Stoyanov, Andrey Velichkevich & Viktória Spišáková

This talk addresses the optimization of error recovery for cost-efficient distributed AI model training using Kubernetes. The speaker presents a real-world scenario in which a protein structure prediction job running on a GPU failed after 34 days due to overheating, resulting in significant computation loss. The discussion highlights the challenges presented by distributed workloads, particularly in relation to synchronization issues and the necessity of gang scheduling. The speaker categorizes common failure types into infrastructure, AI engine, and user errors, emphasizing the importance of recovery mechanisms such as model and infrastructure checkpoints. The focus then shifts to transparent checkpointing and its implementation challenges, including the need for efficient GPU state management and the integration of compression techniques to improve restore times. The speaker indicates ongoing work to enhance checkpointing functionalities within the Kubernetes ecosystem, particularly using the Kubeflow SDK. Finally, the talk covers future enhancements aimed at further optimizing distributed AI training processes.

Hello everyone. Hi. Um, today we are talking we are going to talk about our work on optimizing error recovery for cost-efficient distributed AI model training with Kubernetes. And this is joint work with um, Andrej Vilicic and Viktorija Pisak. As well as our collaborators Adrian Reber and Peter Hunt. And my supervisors uh, Rodrigo Bruno and Luis Armando. So, Vicky. So, this is a real scenario from Czech National

Research Infrastructure. There was a protein structure prediction job that has been running uh, with a single GPU on a signal on a single node. After 34 days of running, the node just overheated, the job got stuck, and reboot was necessary. This means that the job um, was just lost. It didn't implement any sort of recovery mechanism, it didn't implement any checkpoint mechanism. So, we lost 34 days

of compute. Now, this was only one GPU and one node, but imagine this situation scaled to much bigger sizes such as eight GPUs, 64 GPUs, 500 GPUs, or even more than 1,000 GPUs. With so many GPUs, these consequences are much more um, urgent. So, distributed workloads pose different challenges to the systems. Um, especially distributed workloads that consist of uh, parallel of parallel uh, frequently communicating tasks. These

tasks really want to avoid any synchronization issues. Moreover, these tasks also involve require gang scheduling. Uh, gang scheduling means that either all tasks are scheduled and need to start or no task can start. Um, gang scheduling is a very expensive operation. So, once we manage to schedule a gang, we also should, uh, make sure that it keeps on running. However, workload failures happen due to various reasons,

and if any single task that, uh, is, uh, a part of a gang experiences any failure, it may, um, make the job stall or even die. This means that a distributed job is only as reliable as the least reliable infrastructure component. We will focus on the, um, distributed trainings in this talk. So, I will talk about distributed training workload failures. There have been multiple studies that already

tried to describe the most common reasons for the uh, for the failures of the distributed trainings, and they came up with a simple categorization. So, the first category are infrastructure failures, and the problems such as network issues, slow file system mounts, or NVLink or InfiniBand errors fall into this category. These errors usually exhibit, uh, later during the workload run time, and they tend to exhibit as timeouts.

The second category are AI engine failures. So, problems such as CUDA failures, uh, CUDA initialization failures, CUDA version mismatches, or out-of-memory problems fall into this category. This, these failures tend to exhibit a bit sooner into the workload run time. The third category are user errors, and the user errors include normal human caused errors, such as misconfigurations, permission errors, or wrong semantics. These errors also tend to exhibit

quite sooner into the workload lifetime. Regardless regardless of what error type happened, if it is scaled to distributed trainings or distributed workloads, the consequences can be pretty bad. Um some studies report on these consequences. For example, one study reports that during 54 days of training, a job experienced more than 460 interruptions, and nearly 80% of those were unexpected interruptions due to hardware errors. Another study reports that

from 3 to 23 hours the mean time between failures is between 3 to 23 hours on older GPUs. And these failures together pile up to a few million dollars of of lost compute costs. Of course, there are some methods to prevent these huge costs, and we know two methods at the moment. One method are model checkpoints. This is applicable to training workloads, and it is basically error

recovery in user code. Model checkpoints require restarting training jobs from the from the checkpoints, and rerunning application code. This might involve potentially large initialization overheads, and also frequent model checkpointing leads to GPU idle times. Another method is infrastructure checkpoints, which stands for transparent system level system level error recovery. Infrastructure checkpoints enable checkpointing of jobs without any user code changes, and it they also support transparent job migrations

within existing schedulers. These two methods are different, and we focus on the second one, so on infrastructure infrastructure applied to the distributed domain. Hi. Thanks so much, Vicky. Um as Victoria mentioned, we were mainly focusing on transparent checkpointing. And um over the past few years, we explored how we can synchronize the checkpointing across multiple nodes in Um so, how we can coordinate distributed snapshots. And last year,

we were exploring how we can perform GPU checkpointing. And um essentially, we are building on that focusing on how we can enable fault tolerance for training jobs. And um the GPU checkpoint functionality is pretty straightforward. It sounds complicated, but essentially, we have four main mechanisms that the GPU driver exposes to user space. It allows us to essentially lock the GPU computation, meaning that um we essentially pause

what is running on the GPU. And then, the checkpointing operation essentially moves the GPU state to host memory. And this allows us to perform essentially to save a memory dump, to save all the memory pages of the application. And uh in the same way, we can restore it back back to host memory, and back to the GPU. And the unlock operation will essentially resume the GPU computation.

one thing that we have been as exploring is how to essentially make this easier for deployment in Kubernetes. And the main challenge with existing with the existing approach is that it requires a separate binary tool called CUDA checkpoint, which um essentially, we execute every time for every process, every thread. And um the CUDA API The checkpoint restore APIs have been made publicly available um public, so we

can essentially integrate this with the CUDA plugin. And this not only allows us to make it easier to deploy, but also eliminates the overheads of fork exec and then parsing the output of the external binary. So, this is the first change we were exploring in terms of improving the checkpoint restore functionality. Um, one of the challenges that we noticed when evaluating, for example, inference workloads and training

workloads is that essentially uh, writing all the GPU state to disk is um, very slow. It's essentially the bottleneck and we are trying to find a way of optimizing and um, over the past few years in the CRIU community, we were exploring how we can integrate the essentially a compression mechanism with within the checkpoint restore pipeline. So, the idea is that we can compress memory pages before

writing them to disk and essentially reduce the amount of data that we have to read and write uh, to disk. So, this not only um, is more efficient, but it also allows us to avoid IO overhead because um, the existing approach for compression is just to save the checkpoint data, then to uh, create a compressed copy of this and this essentially doubles the storage requirements. Um, so,

we have been working towards integrating this and one thing that um, we sort of discussed in the community is since we're compressing every memory page, we can also identify uh, which memory pages don't compress very well, but also which memory pages contain zero bytes. So, instead of um, saving the um, incompressible data to uh, in the checkpoint, we can essentially skip this step and during restore, we

can essentially eliminate the overhead of decompression. And for zero bytes, we just um, have a marker that indicates that um, this is zero bytes that will be restored. Um, so, we have these additional optimizations that allow us to uh improve the overhead of compression. And what we can see in the results is um for example with VOM inference we we can reduce significantly the size of the

checkpoint. Um, but essentially the main thing is that um we can improve the restore time. So essentially by reading less data from disk, uh you can essentially resume faster. And we we see this also for um training jobs, but um here since um most of the weights and the training state um um is not redundant data, essentially compressing this the compression ratio is is lower. You can

see that we we gain some benefit of um the get between 10 and 20% of reduction in terms of uh checkpoint size. So this was essentially the next challenge we were trying to solve. So um the problem uh with um essentially checkpointing training jobs is that um they can be very large. Like if you have uh H100 or uh B200 GPU, you have to save uh large

amount of GPU state. And the way we do this with um uh for example with CPU workloads is by essentially keeping track of which memory has changed since the last checkpoint. And this is implemented with essentially um a mechanism called soft dirty bit. It's um page table entry in the Linux kernel that keeps track of which memory pages have been modified. And then CRIU essentially uh saves

only the memory um since the last checkpoint. The problem um when we checkpoint the GPU state into host memory, this creates new memory locations every time. So we cannot identify what memory has changed. So, we can still use this mechanism for CPU state, but for GPU state, we had a different solution. So, what we decided to do is to compute a checksum for the GPU memory, and

then essentially keep these check checksums as part of the checkpoint. And then, during the next iteration, we can essentially use these checksums to identify the memory that has changed, then save only this memory. And what we can see is that, for example, when we do fine-tuning with quantized models, then the frozen the frozen states of the frozen weights of the model don't don't change, so we can

skip large amount of state when saving the checkpoint. that with different quantizations, we significant results. Like, for example, checkpoints by 90% the size of checkpoints by 90% by using this deduplication mechanism. But, it doesn't work very well when we essentially full training with full with BF16, for example, when we um use high precision. this is where the um essentially this this best optimizations that we were able

to create so far. And now, we were essentially looking into how this could be used with um cloud-native AI systems. Or, in this case, we have multiple stacks in these of in the soft in the software stack. So, for example, we focused on enabling distributed checkpointing functionality at a low level, essentially the Linux kernel, the GPU driver, and then we were exploring how to enable this also

in Kubernetes. We have a checkpointing API that is available today. Um but data scientists use um usually Python and the Kubeflow SDK to interact with with Kubernetes, and we were exploring how we uh checkpointing functionality across the stack. the essentially what we have um identified is essentially the Python SDK interface allows uh data scientists to create the uh essentially the the train jobs uh the train job

objects and to essentially manage the tasks running in Kubernetes, and Andrei has a very good talk describing how this system is implemented, and I would recommend to look at this, but essentially it automates most of the management tasks and uh creating pods out of um essentially the uh recipes that are used in Kubeflow trainer. And what we identified as main use case is the suspend resume functionality,

where we essentially create a checkpoint when uh jobs are being suspended. So, um essentially we extended the train job controller to um essentially trigger the checkpointing mechanism when the suspend um functionality is when suspend is set to true, and this allows us to uh for example, um create a checkpoint when the user manually suspends and resumes a job, but also when uh scheduling jobs with different priorities

and low priority jobs are being preempted. And I have a short demo just to demonstrate how this works, how this looks like. So, in this case, we um create two um uh train jobs that communicate. So, this is a um a simple example of distributed training. So, this is um training GPT-2 uh model using PyTorch um with Q uh Kubeflow um what we do now is we

trigger the suspend mechanism. And this uh starts creating a checkpoint checkpoints of the pods. And once the checkpoints are created, um the Kubeflow trainer will uh perform the normal suspend, which will terminate the And um essentially, these checkpoints are saved to local disk um as essentially tar archives, which are then converted to OCI images that can be used during restore. So, once the uh once the pods

terminate, uh we can essentially use these checkpoints. One interesting and important thing here is the IP address. So, um the way we make this work at the moment is um by using the same IP address we used during checkpoint and Um changing the IP address requires different ways of um essentially handling this, but um one of the main challenge here is when we have nickel communication between

um between the train jobs. So, we have to abort all the nickel communications before the checkpoint and then to rebuild them after restore. So, this is kind of similar to how PyTorch will perform distributed checkpoint restore, but essentially, we do this at low level with Kube. um present one of the use cases that we are exploring as future work is uh recently introduced functionality in Git Flow

trainer allows us to keep track of the progress of training jobs and we were exploring how we can integrate this checkpointing mechanism to periodically create snapshots and the idea here is to overhead in terms of both performance and storage when we save checkpoints by keeping track of our sequential thanks so much for listening and I will be happy to answer any questions. Thanks. Hi, thank you for

the the great talk. I'm very interested in two different use cases. One you mentioned earlier in the first talk this day about warm warm starts. And I wonder how it relates to like to this how what is the readiness of that and related to this one like how ready is this because I I don't get if it's already available through through the trainer or is it something

we have to wait or contribute like I didn't get the right like maturity that we are right now but thank you it's very interesting. Yeah, um so just to answer the first question um so this is more like research work that we have been doing. Um so we are working both on improving the GPU checkpointing functionality at Krio. Um so in in in essentially in this presentation

we kind of have three main contributions like three main new things that we're presenting. The first is um essentially uh self um essentially the CUDA plugin is improved by um containing the integrating the um communication with the GPU driver essentially using the CUDA APIs. We also we introduced the deduplication method. So, this is something new like I mean it's very active area of research to be honest.

It it evolves very quickly, but we also collaborate with the Nvidia team on the GPU checkpointing functionality and we keep track of new features being released. One of the main challenge that people have essentially is the iterative checkpoints. So, this deduplication mechanism that we were presenting is kind of trying to solve this problem. So, memory has changed on the GPUs something people cannot do today. And also

checkpointing distributed workloads is and handling the nickel communication is what we are trying to solve with this talk. So, this is essentially Well, we have patches. We have to discuss them in and I will get approvals and integrate this and with Andre we have roadmap item for Kubeflow trainer where we want to introduce this functionality, but we have to also talk with the Kubernetes community in terms

of how this will be integrated with job sets. And yeah, essentially it's a long process, but the idea of this talk is to present what we can do today and what are the solutions that we have. I also want to mention that if you're interested in contributing, please join our weekly meetings and either in the checkpoint restore working group or in the batch working group or the

trainer call. And this will be very helpful to learn more about the use cases that you have and better support this in the open source community.