Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

TL;DR

Deep Optimizer States is designed and implemented, a novel technique to split the LLM into subgroups whose update phase is scheduled on either the CPU or the GPU based on the proposed performance model that addresses the trade-off between data movement cost, acceleration on the GPUs vs the CPUs, and competition for shared resources.

Abstract

Transformers and large language models~(LLMs) have seen rapid adoption in all domains. Their sizes have exploded to hundreds of billions of parameters and keep increasing. Under these circumstances, the training of transformers is very expensive and often hits a ``memory wall'', i.e., even when using 3D parallelism (pipeline, tensor, data) and aggregating the memory of many GPUs, it is still not enough to hold the necessary data structures (model parameters, optimizer state, gradients, activations) in GPU memory. To compensate, state-of-the-art approaches offload the optimizer state, at least partially, to the host memory and perform hybrid CPU-GPU computations. However, the management of the combined host-GPU memory is often suboptimal and results in poor overlapping between data movements and computations. This leads to missed opportunities to simultaneously leverage the interconnect bandwidth and computational capabilities of CPUs and GPUs. In this paper, we leverage a key observation that the interleaving of the forward, backward and update phases generate fluctuations in the GPU memory utilization, which can be exploited to dynamically move a part of the optimizer state between the host and the GPU memory at each iteration. To this end, we design and implement \proj, a novel technique to split the LLM into subgroups, whose update phase is scheduled on either the CPU or the GPU based on our proposed performance model that addresses the trade-off between data movement cost, acceleration on the GPUs vs the CPUs, and competition for shared resources. We integrate our approach with DeepSpeed and demonstrate 2.5$\times$ faster iterations over state-of-the-art approaches using extensive experiments.

Figures (17)

• Figure 1: Model parallelism techniques with optimizer state completely offloaded to the host memory: (a) Conventional pipeline and tensor parallelism for a model with 4 layers; (b) DeepSpeed's ZeRO-3 hybrid data and model parallelism; (c) Zoom on the model and CPU-offloaded optimizers of a single data-parallel rank; and subgroup sharding of parameters on each rank. Similar to the sharding of FP16 parameters into 4 subgroups ($SG 1 \dots SG 4$), GPU-resident FP16 gradients and FP16 activations and host-resident FP32 parameters, FP32 gradients, FP32 momentum, and FP32 are sharded in 4 distinct subgroups.
• Figure 2: Iteration time for different models when scaling subgroup sizes.
• Figure 3: GPU memory util. without (top) and with (bottom) activation checkpoint.
• Figure 4: PCIe link util. at different training phases for a 20B parameters model.
• Figure 5: Working of optimizer update step with different approaches for 8 subgroups per GPU (2 subgroups statically residing on GPU). Our approach illustrates an example where 33% of the updates are scheduled on the GPU.