# Enabling Efficient Trillion Parameter Scale Training for Deep Learning Models

Tunji Ruwase, Microsoft | AI in Production 2024 | 27:36

Source: https://www.youtube.com/watch?v=JFBACDiNRzk
Channel: MLOps Community, now AAIF Live (https://www.youtube.com/@AAIFLive-x1r). Summarised by MLOps Talks.
Page: https://mlopstalks.com/talks/enabling-efficient-trillion-parameter-scale-training-for-deep-learning-models
Published: 2024-04-17
Tags: gpus, inference, open-source, training-pipelines

## TL;DR
- DeepSpeed ZeRO reduces GPU memory use by partitioning parameters, gradients, and optimizer states across devices, then offloading them to CPU memory or NVMe when needed.
- DeepSpeed Mixture of Experts increases model size without making every input use every parameter, while 3D parallelism scales the base model across tensor, expert, and data dimensions.
- DeepSpeed's data efficiency framework uses curriculum learning and token routing to reduce training cost or improve model quality at the same data cost.

## Summary
Tunji Ruwase explains how DeepSpeed addresses three limits in large-model training: memory, compute, and data. ZeRO partitions optimizer states, gradients, and parameters across GPUs instead of replicating them. ZeRO-Offload and ZeRO-Infinity move model state to CPU memory or NVMe, which lets teams train models that do not fit in GPU memory. DeepSpeed Mixture of Experts routes each input through a subset of experts, so the model can grow without proportional compute growth. DeepSpeed-TED adds tensor, expert, and data parallelism, with communication optimizations for scaling the base model. The data efficiency framework applies curriculum learning and layerwise token dropping to reduce training time or improve downstream results. Ruwase also discusses support for models up to a trillion parameters and DeepSpeed's expansion beyond CUDA to AMD, Habana Gaudi, Intel, and Apple accelerators.

## Key ideas
### Model growth has outpaced GPU memory growth
[00:46](https://www.youtube.com/watch?v=JFBACDiNRzk&t=46s)
Ruwase says language models grew from fewer than 100 million parameters in 2018 to roughly half a trillion parameters by 2022, while GPU memory grew much more slowly. He frames model scaling as a systems problem involving memory, compute, and data. A billion-parameter model trained with the Adam optimizer needs about 20 gigabytes per GPU for parameters, gradients, and optimizer state alone. That estimate excludes inputs and activations, so the practical memory requirement is higher.

### ZeRO partitions model state instead of replicating it
[03:13](https://www.youtube.com/watch?v=JFBACDiNRzk&t=193s)
ZeRO reduces GPU memory use through a sequence of partitioning stages. Stage 1 partitions optimizer state across GPUs. Stage 2 partitions gradients as well. Stage 3 partitions the parameters, so each GPU holds only a slice of the model. Ruwase says the roughly 20 bytes per parameter in the baseline falls to less than five bytes with Stage 1, less than three with Stage 2, and less than one with Stage 3. These stages let distributed training fit models that are much larger than a single GPU's memory.

### ZeRO can use CPU memory and NVMe as part of the training system
[03:36](https://www.youtube.com/watch?v=JFBACDiNRzk&t=216s)
Partitioning is one part of ZeRO. Ruwase also describes offloading model state from GPU memory. ZeRO-Offload moves optimizer state into CPU memory. ZeRO-Infinity extends this idea by placing the entire model state, including parameters, gradients, and optimizer state, in NVMe storage. He presents this as a way to train large models with fewer GPUs when the system has enough CPU memory or SSD capacity.

### Mixture of Experts grows model quality without proportional compute
[08:18](https://www.youtube.com/watch?v=JFBACDiNRzk&t=498s)
Dense models use all their parameters for every input, so increasing model size also increases computation. Mixture of Experts models route each input through only a subset of experts. In Ruwase's example, a 1.3-billion-parameter model with 128 experts has the same computation cost as the smaller dense model, while matching the quality of a 6.7-billion-parameter dense model. He says this gives similar model quality with five times less compute in the comparison he presents.

### DeepSpeed-TED combines three forms of parallelism
[11:08](https://www.youtube.com/watch?v=JFBACDiNRzk&t=668s)
Mixture of Experts does not remove every scaling problem. Ruwase says research has observed diminishing returns as the number of experts increases, so the base model still needs to scale. DeepSpeed-TED combines tensor parallelism, expert parallelism, and data parallelism. In his comparison, it scales the base model to nearly five times the size achieved with the baseline DeepSpeed Mixture of Experts setup, while limiting the comparison to 128 experts and tensor parallelism within a node.

### Communication became a major limit in expert-model scaling
[12:32](https://www.youtube.com/watch?v=JFBACDiNRzk&t=752s)
In the DeepSpeed-TED experiments, about half of iteration time came from communication. Ruwase describes duplicate token drop, which reduced all-to-all time by about 64 percent, and communication-aware activation checkpointing, which reduced all-reduce time by about 33 percent. Together, these changes improved DeepSpeed-TED performance by about 21 percent overall. Strong scaling results on the Summit supercomputer showed speedups of 20 to 29 percent over the DeepSpeed baseline as GPU count increased.

### Data efficiency can trade training cost for quality
[14:18](https://www.youtube.com/watch?v=JFBACDiNRzk&t=858s)
Ruwase says larger models also need more training data, so data becomes part of the total training cost. DeepSpeed's data efficiency framework combines a data analyzer, samplers, and routers. Curriculum learning samples examples from easier to harder rather than using a random order. Random layerwise token drop skips some tokens in middle layers with minimal effect on quality. In the GPT-2 results he presents, curriculum learning either reached similar quality in less than half the training time or improved perplexity and accuracy while using the full data set.

### DeepSpeed is expanding across models and accelerator hardware
[19:19](https://www.youtube.com/watch?v=JFBACDiNRzk&t=1159s)
Ruwase says DeepSpeed has been used to train open-source large models ranging from five billion parameters to half a trillion parameters in collaboration with NVIDIA. The project is a PyTorch-based open-source library, and he says it is integrated into many common frameworks. DeepSpeed began as a CUDA library and added support for AMD, Habana Gaudi, Intel, and Apple accelerators. In the question period, he also explains that ZeRO does not speed up single-GPU training, but it can make a model trainable on one GPU by using storage, even if the run is slower.

## Notable quotes
- "The models are growing faster than the hardware can keep up." (02:47)
- "The first is partitioning, that is, we have a way of dividing the memory state across the GPU devices rather than replicating them." (03:54)
- "The main takeaway here is that with the DeepSpeed TED with tensor parallelism we're able to scale the base model size up to close to five times compared to just using regular DeepSpeed MoE." (12:10)
- "As long as you have big enough SSD, you're good to go." (24:57)

## Tools & references mentioned
- DeepSpeed
- ZeRO
- ZeRO-Offload
- ZeRO-Infinity
- DeepSpeed Mixture of Experts
- DeepSpeed-TED
- Adam
- Megatron-Turing NLG
- GPT-2
- GPT-3
- BERT
- FlashAttention
- LoRA
- PyTorch
- NVIDIA
- Summit
- AMD
- Habana Gaudi
- Intel
- Apple
- Microsoft

## Who should watch
- You are training large Transformer models and GPU memory is limiting the model size you can fit.
- You want to compare dense models with Mixture of Experts models, including their routing and communication costs.
- Your team is considering DeepSpeed for multi-GPU training, NVMe offload, or data-efficient training.

## Editor's note

From the pack [Scaling training and compute](https://mlopstalks.com/packs/scaling-training-and-compute):

Singh connects faster training with recoverable checkpoints, while Khanna separates job placement from execution in each cluster. ZenML lets teams define Python pipeline steps and choose execution infrastructure through configuration, recording each run's inputs, outputs and code version. This connects a trained artifact to its workflow; the training system still supplies the memory, communication and checkpoint mechanisms discussed here.

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

## Related talks

- [Large Model Training and Inference with DeepSpeed](https://mlopstalks.com/talks/large-model-training-and-inference-with-deepspeed) (Samyam Rajbhandari, Microsoft DeepSpeed, 36:23)
- [Quantized LLM Training at Scale with ZeRO++](https://mlopstalks.com/talks/quantized-llm-training-at-scale-with-zero) (Guanhua Wang, Microsoft, 24:35)
- [How to Optimize Large AI Models with PyTorch](https://mlopstalks.com/talks/how-to-optimize-large-ai-models-with-pytorch) (Michael Gschwind, Meta Platforms, 57:44)
- [Efficiently Scaling and Deploying LLMs](https://mlopstalks.com/talks/efficiently-scaling-and-deploying-llms) (Hanlin Tang, MosaicML, 25:14)
- [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)