AMSP: Super-Scalabilità dell'Addestramento di LLM tramite Partizione Avanzata degli Stati del Modello

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

November 1, 2023

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) hanno dimostrato prestazioni impressionanti in vari task downstream. Durante l'addestramento di questi modelli, si sta osservando una crescente tendenza a elaborare più token su scale di addestramento più ampie, ma con dimensioni del modello relativamente più ridotte. L'ottimizzatore Zero Redundancy (ZeRO), sebbene efficace negli ambienti di addestramento convenzionali, incontra difficoltà di scalabilità quando si confronta con questo paradigma emergente. A tal fine, proponiamo un nuovo framework di addestramento per LLM chiamato AMSP, che adotta una partizione granulare degli stati del modello, comprendendo parametri (P), gradienti (G) e stati dell'ottimizzatore (OS). Nello specifico, AMSP (1) costruisce uno spazio di partizione unificato, consentendo strategie di partizione indipendenti per P, G e OS; (2) incorpora un partizionatore consapevole della scala per cercare autonomamente le strategie di partizione ottimali; (3) progetta un ottimizzatore di comunicazione dedicato per garantire una gestione efficace delle discrepanze di posizionamento dei dati derivanti da diverse strategie di partizione. Le nostre valutazioni mostrano che AMSP raggiunge un'efficienza di scalabilità fino al 90,3% su 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.