AMSP: Reducing Communication Overhead of ZeRO for Efficient LLM Training

TL;DR

AMSP tackles the high communication cost of ZeRO in large-scale LLM training by introducing flexible, per-component sharding for Parameters, Gradients, and Optimizer States (P, G, OS) across three strategies: Full-Replica, Full-Sharding, and Partial-Sharding. It formalizes a dependency-aware optimization to minimize inter-GPU communication under GPU-memory constraints and implements an execution engine that aggressively overlaps communication with computation and placement on leaf-spine networks. The approach yields substantial gains, achieving up to 52% Model FLOPs Utilization on 1024 GPUs for LLaMA-13B and LLaMA-7B, and improves training throughput by 1.4–12.7x over MiCS and ZeRO++ baselines. These results demonstrate AMSP’s practical impact for efficient distributed LLM training at scale, including successful training of InternLM on thousands of GPUs and a comprehensive micro-benchmark-based understanding of communication overlap benefits.

Abstract

Training large language models (LLMs) encounters challenges in GPU memory consumption due to the high memory requirements of model states. The widely used Zero Redundancy Optimizer (ZeRO) addresses this issue through strategic sharding but introduces communication challenges at scale. To tackle this problem, we propose AMSP, a system designed to optimize ZeRO for scalable LLM training. AMSP incorporates three flexible sharding strategies: Full-Replica, Full-Sharding, and Partial-Sharding, and allows each component within the model states (Parameters, Gradients, Optimizer States) to independently choose a sharding strategy as well as the device mesh. We conduct a thorough analysis of communication costs, formulating an optimization problem to discover the optimal sharding strategy. Additionally, AMSP optimizes distributed LLM training by efficiently overlapping communication with computation. Evaluations demonstrate up to 52\% Model FLOPs Utilization (MFU) when training the LLaMA-based model on 1024 GPUs, resulting in a 1.56 times improvement in training throughput compared to newly proposed systems like MiCS and ZeRO++.

Figures (20)

• Figure 1: Overview of GPU memory allocation for model states with different strategies. ZeRO-1 and ZeRO-3 significantly reduce memory consumption for model states compared to standard data parallelism. MiCS and ZeRO++ are proposed to mitigate communication overhead, particularly cross-node communication time, in comparison to the ZeRO approach.
• Figure 2: Micro-benchmark of training LLaMA-7B across a scale of GPUs, ranging from 8 to 512, while maintaining a global batch size of 4M tokens. The micro-batch size $B$ is consistently set to $1$ in all tests. Panel (a) illustrates the GPU memory consumption of model states. Panel (b) depicts the time taken for forward and backward computations. Panel (c) presents the latency of three communication operations with a fixed message size of 256MB.
• Figure 3: Performance evaluation of collective communication operations using NCCL. The assessment is conducted with varying message sizes (in bytes). GPU nodes are linked using 4 Mellanox Infiniband HDR NICs (800 Gbps bandwidth in total). The notation $8 \! \times \! A$ GPUs indicates that the tests were conducted on $A$ nodes, with each node housing 8 NVIDIA Ampere GPUs (A800) connected by NVLINK.
• Figure 4: Optimizing model states sharding through the dependency rule. In this instance, when $s_p=s_g=2$, there's no need to set $s_{os}=1$ as it would store redundant optimized states, incurring additional communication costs.
• Figure 5: Analysis of inserted collective communication operations when individually sharding parameters, gradients, and optimizer states.