AMSP: Superdimensionamento do Treinamento de LLMs por meio de Particionamento Avançado de Estados do Modelo

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

November 1, 2023

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

cs.AI

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) têm demonstrado desempenho impressionante em diversas tarefas subsequentes. Ao treinar esses modelos, há uma tendência crescente de processar mais tokens em escalas de treinamento maiores, mas com tamanhos de modelo relativamente menores. O Otimizador de Redundância Zero (ZeRO), embora eficaz em ambientes de treinamento convencionais, enfrenta desafios de escalabilidade quando confrontado com esse paradigma emergente. Para isso, propomos um novo framework de treinamento de LLMs chamado AMSP, que realiza uma partição granular dos estados do modelo, abrangendo parâmetros (P), gradientes (G) e estados do otimizador (OS). Especificamente, o AMSP (1) constrói um espaço de partição unificado, permitindo estratégias de partição independentes para P, G e OS; (2) incorpora um particionador consciente da escala para buscar automaticamente estratégias de partição ótimas; (3) projeta um otimizador de comunicação dedicado para garantir o gerenciamento eficiente das discrepâncias de posicionamento de dados decorrentes de diferentes estratégias de partição. Nossas avaliações mostram que o AMSP alcança até 90,3% de eficiência de escalabilidade em 1024 GPUs.

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.