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

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 generates 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 Deep Optimizer States, 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.

References in corpus (14)

• Llama 2: Open Foundation and Fine-Tuned Chat Models
• Mixed Precision Training
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
• ZeRO-Offload: Democratizing Billion-Scale Model Training
• XPipe: Efficient Pipeline Model Parallelism for Multi-GPU DNN Training
• DataStates-LLM: Lazy Asynchronous Checkpointing for Large Language Models
• DAPPLE: A Pipelined Data Parallel Approach for Training Large Models
• M6-10T: A Sharing-Delinking Paradigm for Efficient Multi-Trillion Parameter Pretraining
• Whale: Efficient Giant Model Training over Heterogeneous GPUs
• DeepSpeed4Science Initiative: Enabling Large-Scale Scientific Discovery through Sophisticated AI System Technologies
• Breaking the Memory Wall: A Study of I/O Patterns and GPU Memory Utilization for Hybrid CPU-GPU Offloaded Optimizers
• ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours

Cited by in corpus (2)

• MLP-Offload: Multi-Level, Multi-Path Offloading for LLM Pre-training to Break the GPU Memory Wall
• Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment