Reproduce “Towards Efficient Pre-training: Exploring FP4 Precision in Large Language Models” on OpenResearch

About this paper

Jiecheng Zhou, Ding Tang, Rong Fu, Boni Hu, Haoran Xu, Yi Wang, Zhilin Pei, Zhongling Su and 3 more

This work explores training large language models (LLMs) using 4-bit floating-point (FP4) precision, which substantially reduces computational costs. A mixed-precision strategy combined with a two-stage training schedule enables LLMs to be trained with FP4, achieving comparable performance to FP16-trained models while theoretically reducing computational costs by about 30%.

Key insights

• Different components of Transformer architectures exhibit varying sensitivities to precision reduction, necessitating a mixed-precision approach.
• A two-stage training schedule, starting with FP4 and concluding with a short FP16 recovery phase, can mitigate quantization noise and achieve comparable performance.
• Critical modules like attention (QKV projections) and weight gradients require higher precision (FP8) to maintain model accuracy and training stability.

Summary by alphaXiv.