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++.
@article{arxiv.2311.00257,
title = {AMSP: Reducing Communication Overhead of ZeRO for Efficient LLM Training},
author = {Qiaoling Chen and Qinghao Hu and Guoteng Wang and Yingtong Xiong and Ting Huang and Xun Chen and Yang Gao and Hang Yan and Yonggang Wen and Tianwei Zhang and Peng Sun},
journal= {arXiv preprint arXiv:2311.00257},
year = {2024}
}