Breaking the Memory Wall: A Study of I/O Patterns and GPU Memory Utilization for Hybrid CPU-GPU Offloaded Optimizers

TL;DR

The paper tackles the memory wall in training massive transformers when optimizer state is offloaded to host memory using DeepSpeed. It empirically characterizes fine-grained I/O patterns and GPU memory usage across iteration phases to reveal bottlenecks and optimization opportunities. Key findings show that blocking device-host transfers, CPU-bound updates, and PCIe contention limit performance, but there are actionable opportunities to improve overlap by buffering gradients during the final backward pass and by dynamically adjusting offload during updates. These insights pave the way for flexible hybrid CPU-GPU offloading that reduces memory footprints while mitigating cost and performance penalties in transformer and LLM training.

Abstract

Transformers and 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 slow and often takes in the order of weeks or months. Thanks to 3D model parallelism (data, pipeline, and tensor-level parallelism), the training can scale to a large number of GPUs, which reduces the duration of the training but dramatically increases the cost. Even when a large number of GPUs are available, the aggregated GPU memory is often not enough to hold the full training state (optimizer state, model parameters, and gradients). To compensate, state-of-the-art approaches offload the optimizer state at least partially to the host memory and perform hybrid CPU-GPU computations. Such flexible solutions dramatically reduce the GPU memory utilization, which makes it feasible to run the training on a smaller number of GPUs at the cost of performance penalty. Unfortunately, the challenges and bottlenecks of adopting this strategy are not sufficiently studied by state-of-the-art, which results in poor management of the combined host-GPU memory and poor overlapping between data movements and computations. In this paper, we aim to fill this gap by characterizing the behavior of offloaded training using the DeepSpeed runtime. Specifically, we study the GPU memory utilization over time during each iteration, the activity on the PCIe related to transfers between the host memory and the GPU memory, and the relationship between resource utilization and the steps involved in each iteration. Thanks to this study, we reveal opportunities for future improvements of offloading solutions, which enable greater flexibility to optimize the cost-performance trade-off in the context of transformer and LLM training.

Figures (5)

• Figure 1: Data, pipeline, and tensor parallel runtime training. Compute node configuration consisting of 4 A100-40GB GPUs.
• Figure 2: Subgroup sharding of model parameters, gradients, and activations on GPUs.
• Figure 3: Breakdown of different training stages for 13B model running on 1 GPU with TP=1, DP=1, micro-batch=1, gradient accumulation=2, ZeRO stage=3. The optimizer is completely offloaded to the CPU.
• Figure 4: Memory and PCIe utilization statistics of single iteration when training 13B model.
• Figure 5: Breakdown of tensor sizes transferred between device and host (Figure \ref{['fig:d2h-tensors']} and Figure \ref{['fig:h2d-tensors']}). PCIe transfer throughput during various subfunctions of the update phase (Figure \ref{['fig:breakdown-subfunctions']}).