AMSP: Hochskalierung des LLM-Trainings durch fortschrittliche Partitionierung von Modellzuständen

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

November 1, 2023

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben beeindruckende Leistungen bei verschiedenen Downstream-Aufgaben gezeigt. Bei der Ausbildung dieser Modelle besteht eine zunehmende Tendenz, mehr Tokens auf größeren Trainingsskalen zu verarbeiten, jedoch mit relativ kleineren Modellgrößen. Der Zero Redundancy Optimizer (ZeRO), obwohl effektiv in konventionellen Trainingsumgebungen, hat Schwierigkeiten mit Skalierungsproblemen, wenn er mit diesem aufkommenden Paradigma konfrontiert wird. Zu diesem Zweck schlagen wir ein neuartiges LLM-Trainingsframework AMSP vor, das eine granulare Partitionierung der Modellzustände vornimmt, einschließlich Parameter (P), Gradienten (G) und Optimiererzustände (OS). Insbesondere (1) erstellt AMSP einen einheitlichen Partitionierungsraum, der unabhängige Partitionierungsstrategien für P, G und OS ermöglicht; (2) integriert einen skalenbewussten Partitioner, um autonom nach optimalen Partitionierungsstrategien zu suchen; (3) entwirft einen speziellen Kommunikationsoptimierer, um die effiziente Verwaltung von Datenplatzierungsdiskrepanzen zu gewährleisten, die sich aus verschiedenen Partitionierungsstrategien ergeben. Unsere Auswertungen zeigen, dass AMSP eine Skalierungseffizienz von bis zu 90,3 % über 1024 GPUs erreicht.

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.