AMSP: Escalado Masivo del Entrenamiento de LLM mediante Partición Avanzada de Estados del 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

Resumen

Los Modelos de Lenguaje a Gran Escala (LLMs, por sus siglas en inglés) han demostrado un rendimiento impresionante en diversas tareas posteriores. Al entrenar estos modelos, existe una creciente tendencia a procesar más tokens en escalas de entrenamiento más grandes, pero con tamaños de modelo relativamente más pequeños. El Optimizador de Redundancia Cero (ZeRO), aunque efectivo en entornos de entrenamiento convencionales, enfrenta desafíos de escalabilidad cuando se enfrenta a este paradigma emergente. Con este fin, proponemos un nuevo marco de entrenamiento para LLMs llamado AMSP, que realiza una partición granular de los estados del modelo, incluyendo parámetros (P), gradientes (G) y estados del optimizador (OS). Específicamente, AMSP (1) construye un espacio de partición unificado, permitiendo estrategias de partición independientes para P, G y OS; (2) incorpora un particionador consciente de la escala para buscar automáticamente estrategias de partición óptimas; y (3) diseña un optimizador de comunicación dedicado para garantizar una gestión eficiente de las discrepancias en la ubicación de datos que surgen de diversas estrategias de partición. Nuestras evaluaciones muestran que AMSP alcanza una eficiencia de escalabilidad de hasta el 90.3% en 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.