AMSP: Супермасштабирование обучения больших языковых моделей с помощью расширенного разделения состояний модели

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

November 1, 2023

Авторы: Qiaoling Chen, Qinghao Hu, Zhisheng Ye, Guoteng Wang, Peng Sun, Yonggang Wen, Tianwei Zhang

cs.AI

Аннотация

Крупные языковые модели (LLM) продемонстрировали впечатляющие результаты в решении различных задач. При обучении таких моделей наблюдается растущая тенденция к обработке большего количества токенов на более крупных масштабах обучения, но с относительно меньшими размерами моделей. Оптимизатор Zero Redundancy Optimizer (ZeRO), хотя и эффективен в традиционных средах обучения, сталкивается с проблемами масштабирования в условиях этой новой парадигмы. В связи с этим мы предлагаем новый фреймворк для обучения LLM под названием AMSP, который осуществляет детальное разделение состояний модели, включая параметры (P), градиенты (G) и состояния оптимизатора (OS). В частности, AMSP (1) создает единое пространство разделения, позволяющее независимо выбирать стратегии разделения для P, G и OS; (2) включает масштабно-ориентированный разделитель для автономного поиска оптимальных стратегий разделения; (3) разрабатывает специализированный оптимизатор коммуникаций для эффективного управления расхождениями в размещении данных, возникающими из-за различных стратегий разделения. Наши оценки показывают, что AMSP достигает до 90,3% эффективности масштабирования на 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.