# Handling Multi-Terabyte LLM Checkpoints

Simon Karasik, Nebius AI | MLOps Podcast | Episode 228 | 55:37
Hosted by Demetrios Brinkmann

Source: https://www.youtube.com/watch?v=6MY-IgqiTpg
Channel: MLOps Community, now AAIF Live (https://www.youtube.com/@AAIFLive-x1r). Summarised by MLOps Talks.
Page: https://mlopstalks.com/talks/handling-multi-terabyte-llm-checkpoints
Published: 2024-04-30
Tags: gpus, orchestration

## TL;DR
- Training checkpoints for a 300-billion-parameter model can reach 3.5 terabytes and must be saved across several machines.
- Checkpoint frequency is a tradeoff: saving every hour limits lost work, while each save pauses training and can slow the run.
- Storage and networking choices change between pre-training and fine-tuning because the same checkpoint is read by very different numbers of machines.

## Summary
Simon Karasik explains what changes when machine learning infrastructure moves from traditional models to large language models. He describes training a 300-billion-parameter model on more than 1,000 GPUs and the operational problems around its checkpoints. A training checkpoint includes the model parameters and optimizer state, so it is much larger than an inference model. At this scale, a checkpoint can be about 3.5 terabytes. Simon's team saves parts of the checkpoint in parallel across machines, then uses the fast GPU network to exchange those parts. They also keep checkpoints on an exponential schedule instead of retaining every hourly copy. The conversation covers why Nebius AI uses Kubernetes and Argo Workflows, how GPU failures are replaced automatically, and why distributed logs and network tests matter. Simon also contrasts pre-training with fine-tuning, where fewer machines may be needed but each machine may have to load a much larger share of the checkpoint.

## Key ideas
### Old machine learning systems can keep training models nobody owns
[08:43](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=523s)
Simon describes a production system at Yandex where inference, data processing, and training pipelines had become disconnected over time. A model could be removed from runtime while the pipeline that trained it continued consuming CPUs. His team built a project to connect those pieces, so deleting a runtime model could trigger a warning that its training process was still active. The original models were sometimes a decade old, and the people who created them had left the company. Removing them required care because nobody wanted to break a profitable production system. One cleanup effort eventually released enough capacity to save about 1,000 CPUs.

### Large language model training needs checkpoints because failures can erase days of work
[16:20](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=980s)
Simon agrees with Demetrios Brinkmann's video-game comparison: a checkpoint lets training resume from a recent point instead of restarting from the beginning. If the loss suddenly explodes, the team can roll back. Without a checkpoint, they may have to repeat days of work. His team generally saves every hour or two. Saving is expensive because training must pause for several minutes, so saving too often slows the run. Saving too rarely increases the amount of work that can be lost. The schedule reflects how much retraining the team is willing to accept after a failure.

### A training checkpoint is much larger than the model used for inference
[23:20](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=1400s)
Simon explains that an inference model on Hugging Face may use roughly two bytes per parameter, while training requires the model parameters and optimizer state. He describes this as several floating-point numbers for each parameter, with each number taking four bytes in his example. A 7-billion-parameter model may therefore take around 70 gigabytes for training, even though its inference form is much smaller. GPU memory needs are higher still because gradients and other training data remain allocated during the run. The memory requirement can reach hundreds of gigabytes before the checkpoint is written.

### Small-scale experiments and training dashboards reduce the risk of a failed large run
[21:03](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=1263s)
Before committing to a large model, Simon's team trains 1-billion- and 3-billion-parameter models. These smaller runs can finish in a day or several days on eight or 16 GPUs, which gives the team a way to check that the training setup works. Simon also refers to scaling laws, which help them estimate whether a successful small run should extend to a larger model. During training, Weights & Biases collects losses, gradients, and per-layer information. When something goes wrong, the team can inspect a particular plot instead of relying only on the final result. Their first larger test reproduced a Llama 7-billion-parameter model before they moved to a 300-billion-parameter model.

### Multi-terabyte checkpoints have to be split and written in parallel
[23:02](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=1382s)
Simon says a checkpoint for a model larger than 300 billion parameters can be about 3.5 terabytes. It cannot fit on a single virtual machine with one terabyte of memory. His team distributes the checkpoint across the machines running training. If eight machines are involved, each one writes a part, around 300 gigabytes in Simon's example, instead of one machine handling the entire file. The same pattern applies when loading. Each host first receives a small part, then the hosts exchange their pieces over the high-speed GPU-to-GPU network. This avoids sending the full checkpoint through one machine.

### Checkpoint retention needs a schedule that does not fill storage immediately
[26:50](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=1610s)
The team creates a checkpoint about every hour, but Simon says it does not make sense to keep every hourly copy. A 3-terabyte checkpoint retained at that rate could quickly consume a petabyte of disk. Instead, they keep recent checkpoints more densely and older checkpoints more sparsely, such as a checkpoint from the previous day and another from the previous week. Checkpoints also feed a background evaluation process. That process reads new checkpoints and runs validation or benchmark work without stopping the main training job. Large checkpoint files therefore affect both storage retention and the amount of work done by downstream evaluation.

### Storage that looks like a filesystem can behave badly for checkpoint operations
[31:04](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=1864s)
Simon compares cloud storage options including NFS-style file storage and S3-style object storage. A cloud virtual machine's disk is often a network disk located on separate machines, so moving data between compute machines is different from moving a local laptop disk. His team tried an S3 filesystem layer that made a bucket appear like a mounted disk. It worked until their checkpoint library moved data. S3 implements a move by copying the object and deleting the old one, which made training much slower at checkpoint scale. They chose Tractor AI for the moment because they already used it for data processing and knew its reliability and management model.

### Kubernetes and Argo Workflows coordinate a training job spread across many machines
[36:07](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=2167s)
Nebius AI uses Kubernetes to orchestrate training, even though Simon learned about Slurm only later from clients who use it. He associates Slurm with scientific and MPI-oriented backgrounds, while his team already had Kubernetes experience. They use Argo Workflows to coordinate the workers. A 100-machine training job needs all its processes to work together, and a failure on one machine should cause the job to fail in a controlled way. The Kubernetes setup also includes a GPU operator. It checks GPU health and can mark a node as unsuitable when a GPU is too hot or otherwise broken.

### Pre-training and fine-tuning put different pressure on storage and networking
[45:38](https://www.youtube.com/watch?v=6MY-IgqiTpg&t=2738s)
Fine-tuning uses far fewer GPUs than pre-training because the dataset and job are smaller, but the model checkpoint remains about the same size. Simon says a 3-terabyte checkpoint may need to be loaded by two or four machines during fine-tuning, rather than being divided across 100 machines during pre-training. Pre-training can spread small checkpoint pieces over many hosts, while fine-tuning needs a fast path from storage into fewer hosts. That changes the storage design and the network bottleneck. Simon's team also tests links between training nodes because one slow node can hold back the whole distributed job. Cloud infrastructure lets them disable a broken GPU and receive a replacement.

## Notable quotes
- Simon Karasik: "If you don't have a checkpoint you're in trouble, you need to redo like days of work." (17:25)
- Simon Karasik: "If I do checkpoints too frequently I make the training slower, but if I don't do checkpoints I'm in the risk." (18:11)
- Simon Karasik: "If one node is slow for some reason, it looks like all training is slow." (39:19)
- Simon Karasik: "Unless you're good with something simpler, you don't need super fancy deploying models on eight GPUs." (50:50)

## Tools & references mentioned
- Nebius AI
- Yandex
- PyTorch
- Hugging Face
- Llama
- Weights & Biases
- Kubernetes
- Slurm
- Argo Workflows
- Kubeflow
- S3
- NFS
- Tractor AI
- OpenAI
- GPT-2

## Who should watch
- You are building infrastructure for distributed LLM training and need to understand what checkpointing adds to the system.
- Your training jobs run on many GPUs, and failures, slow nodes, or difficult-to-debug logs are costing you compute time.
- You are choosing between object storage and shared file storage, or moving from pre-training to fine-tuning with the same large model.

## Editor's note

Simon Karasik explains that a production system at Yandex had inference, data processing, and training pipelines that became disconnected, leaving training for models no one used. ZenML keeps workflow steps, inputs, outputs, and code versions recorded for each run, so teams can trace a model back to the data and code that produced it. The same pipeline can run on different infrastructure through configuration.

Written by the MLOps Talks editors (the ZenML team), not by the speaker.

## Related talks

- [Composable Memory for GPU Optimization](https://mlopstalks.com/talks/composable-memory-for-gpu-optimization) (Bernie Wu, MemVerge, 55:19)
- [Fixing GPU Starvation in Large-Scale Distributed Training](https://mlopstalks.com/talks/fixing-gpu-starvation-in-large-scale-distributed-training) (Kashish Mittal, Uber, 52:49)
- [Efficient GPU infrastructure at LinkedIn](https://mlopstalks.com/talks/efficient-gpu-infrastructure-at-linkedin) (Animesh Singh, LinkedIn, 59:14)
- [Enabling Efficient Trillion Parameter Scale Training for Deep Learning Models](https://mlopstalks.com/talks/enabling-efficient-trillion-parameter-scale-training-for-deep-learning-models) (Tunji Ruwase, Microsoft, 27:36)
- [How We Cut LLM Latency 70% With TensorRT in Production](https://mlopstalks.com/talks/how-we-cut-llm-latency-70-with-tensorrt-in-production) (Maher Hanafi, Betterworks, 1:05:20)