FP8-LM: Entrenamiento de Modelos de Lenguaje Grande en FP8

FP8-LM: Training FP8 Large Language Models

October 27, 2023

Autores: Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, Yuxiang Yang, Ze Liu, Yifan Xiong, Ziyue Yang, Bolin Ni, Jingcheng Hu, Ruihang Li, Miaosen Zhang, Chen Li, Jia Ning, Ruizhe Wang, Zheng Zhang, Shuguang Liu, Joe Chau, Han Hu, Peng Cheng

cs.AI

Resumen

En este artículo, exploramos los formatos de datos de baja precisión FP8 para el entrenamiento eficiente de modelos de lenguaje grandes (LLMs). Nuestra idea clave es que la mayoría de las variables, como los gradientes y los estados del optimizador, en el entrenamiento de LLMs pueden emplear formatos de datos de baja precisión sin comprometer la precisión del modelo y sin requerir cambios en los hiperparámetros. Específicamente, proponemos un nuevo marco de precisión mixta automática FP8 para el entrenamiento de LLMs. Este marco ofrece tres niveles de utilización de FP8 para optimizar la precisión mixta y el entrenamiento paralelo distribuido para LLMs. Incorpora gradualmente gradientes de 8 bits, estados del optimizador y aprendizaje distribuido de manera incremental. Los resultados experimentales muestran que, durante el entrenamiento del modelo GPT-175B en la plataforma GPU H100, nuestro marco de entrenamiento de precisión mixta FP8 no solo logró una notable reducción del 42% en el uso real de memoria, sino que también funcionó un 64% más rápido que el marco BF16 ampliamente adoptado (es decir, Megatron-LM), superando la velocidad de Nvidia Transformer Engine en un 17%. Esto reduce en gran medida los costos de entrenamiento para modelos base grandes. Además, nuestra metodología de entrenamiento de precisión mixta FP8 es genérica. Puede aplicarse sin problemas a otras tareas, como el ajuste de instrucciones de LLMs y el aprendizaje por refuerzo con retroalimentación humana, ofreciendo ahorros en los gastos de ajuste fino. Nuestro marco de entrenamiento de baja precisión FP8 es de código abierto en {https://github.com/Azure/MS-AMP}{aka.ms/MS.AMP}.

English

In this paper, we explore FP8 low-bit data formats for efficient training of large language models (LLMs). Our key insight is that most variables, such as gradients and optimizer states, in LLM training can employ low-precision data formats without compromising model accuracy and requiring no changes to hyper-parameters. Specifically, we propose a new FP8 automatic mixed-precision framework for training LLMs. This framework offers three levels of FP8 utilization to streamline mixed-precision and distributed parallel training for LLMs. It gradually incorporates 8-bit gradients, optimizer states, and distributed learning in an incremental manner. Experiment results show that, during the training of GPT-175B model on H100 GPU platform, our FP8 mixed-precision training framework not only achieved a remarkable 42% reduction in real memory usage but also ran 64% faster than the widely adopted BF16 framework (i.e., Megatron-LM), surpassing the speed of Nvidia Transformer Engine by 17%. This largely reduces the training costs for large foundation models. Furthermore, our FP8 mixed-precision training methodology is generic. It can be seamlessly applied to other tasks such as LLM instruction tuning and reinforcement learning with human feedback, offering savings in fine-tuning expenses. Our FP8 low-precision training framework is open-sourced at {https://github.com/Azure/MS-AMP}{aka.ms/MS.AMP}.