AMSP: Super-Scaling LLM Training via Advanced Model States Partitioning

Qiaoling Chen; Qinghao Hu; Zhisheng Ye; Guoteng Wang; Peng Sun; Yonggang Wen; Tianwei Zhang

返回文献探索

Paper · arXiv 2311.00257

AMSP: Super-Scaling LLM Training via Advanced Model States Partitioning

Qiaoling Chen, Qinghao Hu, Zhisheng Ye, Guoteng Wang, Peng Sun, Yonggang Wen, Tianwei Zhang

9 upvotesNovember 1, 2023arXiv 预印本

AI 摘要

A novel LLM training framework, AMSP, addresses scaling challenges for large-scale token processing using model state partitioning and a scale-aware partitioner.

Large Language Models (LLMs)Zero Redundancy Optimizer (ZeRO)model statesparametersgradientoptimizer statespartitioning spacescale-aware partitionercommunication optimizerscaling efficiency

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.