FP8-LM: Training FP8 Large Language Models
Authors: name: Houwen Peng; name: Kan Wu; name: Yixuan Wei; name: Guoshuai Zhao; name: Yuxiang Yang; name: Ze Liu; name: Yifan Xiong; name: Ziyue Yang; name: Bolin Ni; name: Jingcheng Hu; name: Ruihang Li; name: Miaosen Zhang; name: Chen Li; name: Jia Ning; name: Ruizhe Wang; name: Zheng Zhang; name: Shuguang Liu; name: Joe Chau; name: Han Hu; name: Peng Cheng
Published: 2023-10-27 17:59:51

Abstract: 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}.