AMSP: Super-Scaling LLM Training via Advanced Model States Partitioning
Authors: name: Qiaoling Chen; name: Qinghao Hu; name: Zhisheng Ye; name: Guoteng Wang; name: Peng Sun; name: Yonggang Wen; name: Tianwei Zhang
Published: 2023-11-01 03:14:48

Abstract: Large Language Models (LLMs) have demonstrated impressive performance across
various downstream tasks. When training these models, there is a growing
inclination to process more tokens on larger training scales but with
relatively smaller model sizes. Zero Redundancy Optimizer (ZeRO), although
effective in conventional training environments, grapples with scaling
challenges when confronted with this emerging paradigm. To this end, we propose
a novel LLM training framework AMSP, which undertakes a granular partitioning
of model states, encompassing parameters (P), gradient (G), and optimizer
states (OS). Specifically, AMSP(1) builds a unified partitioning space,
enabling independent partitioning strategies for P, G, and OS; (2)
incorporates a scale-aware partitioner to autonomously search for optimal
partitioning strategies: (3) designs a dedicated communication optimizer to
ensure proficient management of data placement discrepancies arising from
diverse partitioning strategies. Our evaluations show that AMSP achieves up to
90.3% scaling efficiency across 1024 GPUs.