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.
2
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.
3
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.
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
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
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
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.
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
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
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
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.
"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
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.
