AMSP: Super-Schaalvergroting van LLM-training via Geavanceerde Modelstaatspartitionering

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

November 1, 2023

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

cs.AI

Samenvatting

Large Language Models (LLMs) hebben indrukwekkende prestaties laten zien bij diverse downstream taken. Bij het trainen van deze modellen is er een groeiende neiging om meer tokens te verwerken op grotere trainingsschalen, maar met relatief kleinere modelgroottes. De Zero Redundancy Optimizer (ZeRO), hoewel effectief in conventionele trainingsomgevingen, worstelt met schaalbaarheidsuitdagingen wanneer deze wordt geconfronteerd met dit opkomende paradigma. Daarom stellen we een nieuw LLM-trainingsframework voor, genaamd AMSP, dat een gedetailleerde partitionering van modelstatussen uitvoert, inclusief parameters (P), gradienten (G) en optimizerstatussen (OS). Specifiek doet AMSP het volgende: (1) het bouwt een uniforme partitioneringsruimte, waardoor onafhankelijke partitioneringsstrategieën voor P, G en OS mogelijk worden; (2) het integreert een schaalbewuste partitioneerder om autonoom te zoeken naar optimale partitioneringsstrategieën; (3) het ontwerpt een toegewijde communicatie-optimizer om ervoor te zorgen dat verschillen in dataplaatsing, ontstaan door diverse partitioneringsstrategieën, efficiënt worden beheerd. Onze evaluaties tonen aan dat AMSP een schaalbaarheidsefficiëntie van tot wel 90,3% bereikt over 1024 GPU's.

English

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.