微软FP4训练框架:软件性能新突破!
微软FP4训练框架:软件性能新突破!
微软研究院近日发布首个FP4精度的大模型训练框架,这一突破性成果有望进一步降低大模型训练的资源需求。在相同超参数设置下,该框架能够达到与FP8以及BF16相当的训练效果,这意味着所需的存储和计算资源可以更少。
技术突破:FP4模拟训练实现高性能
研究团队通过在FP8 TensorCore上模拟FP4训练,成功实现了这一突破。实验表明,使用该方法训练的模型规模最高可达130亿参数,训练Tokens数量达到千亿级别。值得注意的是,由于研究开展时尚未有原生支持FP4的硬件,因此当前的实现是通过在FP8的TensorCore上进行模拟。如果未来能够使用真正的FP4硬件,预计效果还将进一步提升。
为了实现FP4精度训练,研究团队采用了多项创新技术:
定制化FP4矩阵乘法(GeMM)CUDA内核:通过在内核中使用FP16读取和转换FP4矩阵,完成分块矩阵乘法,最后再将中间结果归约为FP16格式输出。
E2M1数据格式:采用2位表示指数、1位表示尾数的FP4格式,加上1位符号位,总共4位。这种格式设计充分考虑了当前主流ML加速芯片的量化计算单元特点。
差异化量化策略:对权重矩阵W和激活矩阵A采取不同粒度的量化策略,W采用列方向量化,A采用行方向量化,以最大化发挥硬件并行优势。
可微分梯度估计方法:在前向计算时使用硬量化保证效率,反向传播时通过连续可微函数拟合量化函数,确保参数更新的准确性。
离群点削峰补偿策略:针对隐层激活分布的长尾特征,通过分位数检索方法识别并限幅离群点,同时构造稀疏补偿矩阵,确保计算精度。
此外,该框架还采用了混合精度设计,在梯度通信时采用FP8,在优化器状态存储时选择FP16,其他计算环节则采用FP16,以在保证训练稳定性的同时降低计算和存储开销。
团队背景:全华人团队的创新力量
该框架由微软亚洲研究院和SIGMA团队共同开发,所有研究人员均为华人。第一作者为中科大科研部部长查正军教授,通讯作者为MSRA高级首席研究经理程鹏和首席研究经理Yeyun Gong。团队成员还包括MSRA杰出科学家、常务副院长郭百宁等。
行业影响:推动AI技术普及应用
这一突破不仅展示了微软在AI基础设施领域的持续创新实力,也为全球AI研究社区提供了新的技术参考,有望进一步推动AI技术的普及和应用。
未来展望:FP4技术的前景
随着AI模型规模的不断扩大,对计算资源的需求也在持续增长。FP4训练框架的出现,为解决这一挑战提供了新的思路。未来,随着支持FP4的硬件逐步普及,这一技术有望在更多场景中得到应用,进一步降低AI训练的成本和门槛。