AMSP : Surdimensionnement de l'entraînement des LLM via un partitionnement avancé des états du modèle

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

Résumé

Les modèles de langage de grande taille (LLMs) ont démontré des performances impressionnantes sur diverses tâches en aval. Lors de l'entraînement de ces modèles, il y a une tendance croissante à traiter plus de tokens sur des échelles d'entraînement plus larges, mais avec des tailles de modèles relativement plus petites. L'optimiseur Zero Redundancy (ZeRO), bien qu'efficace dans les environnements d'entraînement conventionnels, rencontre des difficultés de mise à l'échelle face à ce nouveau paradigme. Pour répondre à cela, nous proposons un nouveau cadre d'entraînement pour les LLMs, appelé AMSP, qui effectue une partition granulaire des états du modèle, incluant les paramètres (P), le gradient (G) et les états de l'optimiseur (OS). Plus précisément, AMSP (1) construit un espace de partition unifié, permettant des stratégies de partition indépendantes pour P, G et OS ; (2) intègre un partitionneur sensible à l'échelle pour rechercher de manière autonome les stratégies de partition optimales ; (3) conçoit un optimiseur de communication dédié pour assurer une gestion efficace des écarts de placement de données résultant de diverses stratégies de partition. Nos évaluations montrent qu'AMSP atteint jusqu'à 90,3 % d'efficacité de mise à l'échelle sur 1024 GPU.

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.