大模型训练开销还能更小!微软推出首个FP4训练框架,训练效果与BF16相当

AIGC动态3周前发布 QbitAI
213 0 0
大模型训练开销还能更小!微软推出首个FP4训练框架,训练效果与BF16相当

 

文章摘要


【关 键 词】 FP4精度大模型训练量化策略混合精度微软研究院

微软研究院最近推出了首个FP4精度大模型训练框架,这一技术能够在保持与FP8和BF16相当的训练效果的同时,显著减少所需的存储和计算资源。该框架支持高达130亿参数规模的模型训练,并且能够处理千亿级别的训练Tokens。由于研究开展时尚未有原生支持FP4的硬件,研究者通过在FP8的TensorCore上模拟FP4来实现这一技术。

该框架采用了定制化的FP4矩阵乘法CUDA内核,通过FP16将FP4的A和B矩阵读入共享内存并进行变换,然后使用FP4完成分块矩阵乘法,最后用FP16对中间结果进行归约,得到FP16格式的输出矩阵。框架选择了E2M1的FP4格式,即用2位表示指数,1位表示尾数,外加1位符号位,总共4位,以适应主流ML加速芯片的量化计算单元设计。

研究团队对权重矩阵W和激活矩阵A采取了不同粒度的量化策略,W进行列方向量化,而A进行行方向量化,这种策略与硬件上的并行实现方式相契合,最大化FP4在矩阵乘法上的加速效果。在模型前向传播开始时,框架对每一个线性层的权重矩阵W和输入激活矩阵A同时进行FP4量化,通过查表的方式将数值四舍五入到最近的FP4离散值,并进行逐层的量化参数校准。

为了解决量化后矩阵求导导致权重矩阵梯度几乎处处为0的问题,研究者提出了一种新颖的可微分梯度估计方法。在前向计算时使用硬量化以保证计算效率,而在反向传播时,用一个连续可微的函数重新拟合量化函数,并求导得到梯度的修正项。

针对模型隐层激活分布的长尾特征和离群点问题,研究者提出了“离群点削峰和补偿”策略,通过分位数检索找出幅值最大的离群点,将它们限幅到预设阈值范围内,并构造出稀疏补偿矩阵。

此外,该框架还采用了混合精度设计,如在梯度通信时使用FP8,在优化器状态存储时选择FP16,以及其他部分如非矩阵乘操作的计算、Loss Scaling等也采用FP16,以在保证训练数值稳定性的同时降低计算和存储开销。

该框架由微软亚洲研究院和SIGMA团队打造,所有研究人员均为华人,第一作者Ruizhe Wang是中科大在读博士生,目前在MSRA实习,研究方向包括低精度量化。项目还涉及了查正军教授、程鹏、Yeyun Gong和郭百宁等多位杰出科学家。

原文和模型


【原文链接】 阅读原文 [ 1442字 | 6分钟 ]
【原文作者】 量子位
【摘要模型】 moonshot-v1-32k
【摘要评分】 ★★★☆☆

© 版权声明
“绘蛙”

相关文章

“极客训练营”

暂无评论

暂无评论...