特定张量计算操作对计算精度的要求

在大规模深度神经网络的设计与训练过程中，针对不同计算操作选择适宜的精度类型至关重要，这一决策既要基于深厚的数学原理分析，又需丰富的实践经验支持。本节内容总结自大量论文中的实践经验，介绍影响模型性能的关键计算操作对精度的具体要求。

• Softmax函数：softmax函数中有指数运算，容易发生上溢出和下溢出，并且分母有累加操作，目前主流的算子库只使用FP32来实现softmax。
• Cross Entropy损失计算：cross_entropy用来计算交叉熵损失，和很多损失函数一样，需要使用FP32计算以保持数值稳定性，所以算子实现只接受FP32输入。
• 矩阵乘法（Matmul）：matmul矩阵乘法一般来说可以使用半精度数，但在大语言模型的attention层，attention分数与V张量的矩阵乘法需要用FP32，以便维持数值稳定性，避免上溢，减小累积错误。
• 归一化函数：归一化函数如layer_norm、Batch_norm、RMSNorm需要用FP32计算。layer_norm要计算方差，有累加操作，如果用FP16，则计算过程中容易发生上溢和比较大的误差，从而影响最后的收敛。Batch_norm与layer_norm类似，也需要将输入张量upcast到FP32计算。RMSNorm是layer_norm的扩展，在开源LLaMA模型中使用，PyTorch目前并没有提供RMSNorm的实现。如果开发者自己实现，需要注意将输入张量强制upcast到FP32计算，结果可以转回半精度数，主要原因也是它需要计算L2均值，涉及累加操作，低精度运算容易累积误差。
• 卷积操作与池化：卷积操作属于逐单元矩阵乘法，对精度不敏感，可以使用BF16，不会对模型收敛性产生影响。和卷积类似的还有池化，同样可以使用BF16。ReLU和GeLU属于逐点运算，同样对精度不敏感。
• 循环神经网络（RNNs）：RNN一般对精度比较敏感 。LSTM-cell对精度不敏感，可以用BF16。
• 模型起始与结束层：通常大语言模型的第一层和最后一层对精度敏感，比如GPT的embedding层（词嵌入和位置嵌入层）需要使用FP32。一般的输出层即便不是softmax，计算和输出结果也需要是FP32。
• 通信算子：涉及到梯度或者激活值累加的通信算子如all_reduce、reduce_scatter，都需要把输入张量upcast到FP32以保证数值稳定性。
• 批量大小与精度的关系：大规模网络中，batch越大，对数值精度越敏感，所以当batch加大引起收敛问题时，要优先考虑某些运算是否有数值稳定性问题。
• 浮点精度问题：在GLM-130B等大型模型训练中，浮点精度问题引起的模型Loss起飞常常找不到明确的原因，有些会自动恢复，有些会有GNorm起飞的前兆随之有Loss起飞甚至NaN Loss。此时一种有效策略是跳过异常步的数据或者调整超参数。
• BF16训练：用BF16训练深度网络，如果发现不收敛现象，应该尝试使用权重衰减技术，例如设置PyTorch adam优化器的weight_decay超参数（0.1）。
• 大语言模型预训练中的混合精度：在大语言模型预训练过程中，对于未知的混合精度引起的数值稳定性问题，GLM-130B训练过程中一个可以借鉴的经验是将embedding层的梯度收缩到原来的0.1，或者对embedding层做特别的梯度裁剪。