Automatic Mixed Precision (AMP) - Gradient Scaling (梯度缩放)

Automatic Mixed Precision {AMP} - Gradient Scaling {梯度缩放}1. Gradient Scaling (梯度缩放)2. 混合精度原理与计算过程2.1. float32 与 float16 数据格式2.2. 混合精度原理2.3. 混合精度计算过程2.4. 特定计算操作对计算精度的要求ReferencesAutomatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html自动混合精度包 - torch.amphttps://docs.pytorch.ac.cn/docs/stable/amp.htmlAutomatic Mixed Precisionhttps://docs.pytorch.org/tutorials/recipes/recipes/amp_recipe.html自动混合精度https://pytorch-cn.com/tutorials/recipes/recipes/amp_recipe.htmlAutomatic Mixed Precision exampleshttps://docs.pytorch.org/docs/main/notes/amp_examples.html混合精度原理与计算过程 (AMP)https://www.hiascend.com/document/detail/zh/Pytorch/60RC1/ptmoddevg/trainingmigrguide/PT_LMTMOG_0083.html3.6. 自动混合精度训练https://fleet-x.readthedocs.io/en/latest/paddle_fleet_rst/collective/collective_performance/amp.html自动混合精度训练 (AMP)https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/performance_improving/amp_cn.htmlAscend Extension for PyTorchhttps://www.hiascend.com/document/detail/zh/Pytorch/730/index/index.html混合精度https://www.mindspore.cn/tutorials/experts/zh-CN/r1.10/others/mixed_precision.html混合精度训练是指在训练的过程中同时使用单精度 (FP32) 和半精度 (FP16)。1. Gradient Scaling (梯度缩放)If the forward pass for a particular op hasfloat16inputs, the backward pass for that op will producefloat16gradients.Gradient values with small magnitudes may not be representable infloat16. These values will flush to zero (“underflow”), so the update for the corresponding parameters will be lost.如果某个操作的正向传播具有float16输入则该操作的反向传播将产生float16梯度。幅度较小的梯度值可能无法在float16中表示。这些值将变为零 (下溢)因此对应参数的更新将丢失。To prevent underflow, “gradient scaling” multiplies the network’s loss(es) by a scale factor and invokes a backward pass on the scaled loss(es).Gradients flowing backward through the network are then scaled by the same factor. In other words, gradient values have a larger magnitude, so they don’t flush to zero.为防止下溢梯度缩放 (gradient scaling) 将网络的损失乘以一个比例因子并在缩放后的损失上调用反向传播。然后流经网络的梯度将乘以相同的因子。换句话说梯度值具有更大的幅度因此不会变为零。Each parameter’s gradient (.gradattribute) should be unscaled before the optimizer updates the parameters, so the scale factor does not interfere with the learning rate.在优化器更新参数之前应取消缩放每个参数的梯度 (.grad属性)以免比例因子干扰学习率。NoteAMP/fp16 may not work for every model! For example, most bf16-pretrained models cannot operate in the fp16 numerical range of max 65504 and will cause gradients to overflow instead of underflow. In this case, the scale factor may decrease under 1 as an attempt to bring gradients to a number representable in the fp16 dynamic range. While one may expect the scale to always be above 1, our GradScaler does NOT make this guarantee to maintain performance. If you encounter NaNs in your loss or gradients when running with AMP/fp16, verify your model is compatible.AMP/fp16 可能并非适用于所有模型例如大多数 bf16 预训练模型无法在 max 65504 的 fp16 数值范围内运行并且会导致梯度溢出而不是下溢。在这种情况下比例因子可能会小于 1以尝试将梯度带入 fp16 动态范围内可表示的数字。虽然您可能期望比例始终大于 1但我们的 GradScaler 不保证这一点以保持性能。如果您在使用 AMP/fp16 时遇到损失或梯度中的 NaN请验证您的模型是否兼容。https://github.com/pytorch/pytorch/blob/v2.10.0/torch/cuda/amp/grad_scaler.pyclass torch.cuda.amp.GradScaler(init_scale65536.0, growth_factor2.0, backoff_factor0.5, growth_interval2000, enabledTrue)Seetorch.amp.GradScaler.torch.cuda.amp.GradScaler(args...)is deprecated. Please usetorch.amp.GradScaler(cuda, args...)instead.https://github.com/pytorch/pytorch/blob/v2.10.0/torch/cpu/amp/grad_scaler.pyclass torch.cpu.amp.GradScaler(init_scale65536.0, growth_factor2.0, backoff_factor0.5, growth_interval2000, enabledTrue)Seetorch.amp.GradScaler.torch.cpu.amp.GradScaler(args...)is deprecated. Please usetorch.amp.GradScaler(cpu, args...)instead.2. 混合精度原理与计算过程大语言模型的成功证明了增大神经网络的参数规模能够提升模型性能但同时也增大了对加速器内存、算力及通信传输带宽的要求。大量实验已经证明可以使用半精度浮点数训练大模型也不会对模型性能带来显著影响相反低精度计算作为正则化的一部分反而能够给模型泛化能力带来好处。但目前低精度训练对模型的统计学影响也并不那么清晰所以整个训练过程单纯使用低精度浮点运算非常具有挑战性。在此情况下依赖实践经验的混合精度技术成了一个现实的选择。2.1. float32 与 float16 数据格式https://www.hiascend.com/document/detail/zh/Pytorch/60RC1/ptmoddevg/trainingmigrguide/PT_LMTMOG_0083.htmlfloat16是指采用 2 字节 (16 位) 进行编码存储的一种数据类型float32是指采用 4 字节 (32 位) 存储的数据类型。以 FP16 为例1 个符号位 sign 表示正负符号5 个比特位表示指数 exponent全 0 和全 1 有特殊用途因此二进制范围为 00001~1111010 个比特位表示分数 fraction。假设S SS表示 Sign 的十进制值E EE表示 Exponent 的十进制值F FF表示 Fraction 的十进制值。规格化的 FP16 的值为Value ( − 1 ) S × 2 E − 15 × ( 1 F 2 10 ) \text{Value} (-1)^{S} \times 2^{E - 15} \times (1 \frac{F}{2^{10}})Value(−1)S×2E−15×(1210F​)规格化的 FP32 的值为Value ( − 1 ) S × 2 E − 127 × ( 1 F 2 23 ) \text{Value} (-1)^{S} \times 2^{E - 127} \times (1 \frac{F}{2^{23}})Value(−1)S×2E−127×(1223F​)规格化的 BF16 的值为Value ( − 1 ) S × 2 E − 127 × ( 1 F 2 7 ) \text{Value} (-1)^{S} \times 2^{E - 127} \times (1 \frac{F}{2^{7}})Value(−1)S×2E−127×(127F​)使用float16代替float32计算有如下优点内存占用减少float16内存占用比float32更小。在模型训练时可以设置更大的 batch size。加速计算float16训练速度可以是float32的 2~8 倍。使用float16代替float32也可能带来如下问题数值溢出由于float16的值区间比float32的值区间小很多所以在计算过程中很容易出现上溢出 (Overflow 65504 6550465504) 和下溢出 (Underflow 6 × 10 − 8 6 \times 10^{-8}6×10−8)溢出之后就会出现 NaN 的问题。在深度学习中由于激活函数的梯度往往要比权重梯度小更易出现下溢出的情况。舍入误差舍入误差指的是当梯度过小小于当前区间内的最小间隔时该次梯度更新可能会失败。float16相对于bfloat16有更高的精度但是表示范围更小。2.2. 混合精度原理为了避免float16表示范围小引起的浮点上溢和下溢可以通过在训练过程中使用动态损失缩放 (Loss Scale) 来解决使用混合精度加速使用float16进行乘法和存储只使用float32进行加法操作避免累加误差。使用损失缩放 (Loss Scale)反向传播前将损失变化手动增大2 N 2^N2N倍因此反向传播时得到的中间变量激活函数梯度则不会溢出。反向传播后将权重梯度缩2 N 2^N2N倍恢复正常值。使用amp.autocast和amp.GradScaler来实现amp.autocast混合精度加速AMP 会使用amp.autocast自动为算子选择合适的数值精度用户不需要手动对模型参数的数值类型进行转换。该操作可以作为 Python 上下文管理器 (Context Manager) 和装饰器来使用用来指定脚本中某个区域或者某些函数按照 AMP 运行。混合精度是先将float32的模型的参数拷贝一份并转换成float16。AMP 规定的float16算子 (例如 Convolution、Matmul) 以float16计算。AMP 规定的float32算子 (例如 BatchNormalize、Softmax) 在输入和输出时精度为float16计算时为float32在反向传播时依然是混合精度计算得到数值精度为float16的梯度。优化器的操作是利用float16的梯度对float32的参数进行更新。amp.GradScaler损失缩放 (Loss Scale)AMP 使用amp.GradScaler提供梯度缩放 (scale) 操作缓解反向传播时float16的梯度数值溢出问题。在优化器更新参数前AMP 会自动对梯度缩放进行还原 (unscale)对用于模型优化的超参数不会有任何影响。2.3. 混合精度计算过程https://www.hiascend.com/document/detail/zh/Pytorch/60RC3/ptmoddevg/trainingmigrguide/LMaccuracy_0020.htmlhttps://www.hiascend.com/document/detail/zh/Pytorch/730/ptmoddevg/trainingmigrguide/PT_LMTMOG_0095.htmlScale Up 阶段前向计算后反向传播前将得到的损失值增大2 K 2^K2K倍。Scale Down 阶段反向传播后将权重梯度缩小2 K 2^K2K倍恢复 FP32 值进行存储。保留一份 FP32 格式的权重主备份同时优化器状态也应以 FP32 格式存储。初始化缩放系数 (Scale Factor)。对每一个训练 Step3.1. 将权重复制一份到 FP16 格式。3.2. 使 FP16 格式的权重和激活值进行前向传播。3.3. 将最终的 Loss 乘以缩放系数 S。3.4. 使用 FP16 格式进行后向传播包括权重、激活值及其对应的梯度。3.5. 若检测到权重梯度中出现 Inf 或 NaN3.5.1. 减小 S 值。3.5.2. 跳过当前权重更新步骤重新开始下一个训练步骤。3.6. 将权重梯度乘以 1/S。3.7. 梯度累积或者梯度累积足够步后使用 FP32 更新主权重。3.8. 如果之前 N 步都没有看到 Inf 或者 NaN增加 S 值。3.9. 在上述过程中有几处计算必须要以 FP32 完成。主权重的更新因为累加能够导致精度误差积累所以必须要以 FP32 计算。缩放因子必须是 FP32 类型甚至 1/S 的计算要将 S 转成双精度数求倒数再转回到 FP32。使用 BF16 格式的半精度数时因为 BF16 有更大的表示范围所以一般无需使用 Loss 缩放机制。但是 BF16 数值精度比 FP16 更差所以在步骤 3 的第 7 点做梯度累积的时候需要使用 FP32否则有可能会因为梯度累积误差导致模型不收敛。另外 BF16 比 FP16 多 15% 的运行时内存主要原因在于梯度累积时需要转 FP32。2.4. 特定计算操作对计算精度的要求https://www.hiascend.com/document/detail/zh/Pytorch/60RC3/ptmoddevg/trainingmigrguide/LMaccuracy_0021.html在大规模深度神经网络的设计与训练过程中针对不同计算操作选择适宜的精度类型这一决策既要基于深厚的数学原理分析又需丰富的实践经验支持。Softmax 函数Softmax 函数中有指数运算容易发生上溢出和下溢出并且分母有累加操作目前主流的算子库只使用 FP32 来实现 Softmax。Cross Entropy 损失计算Cross Entropy 用来计算交叉熵损失和很多损失函数一样需要使用 FP32 计算以保持数值稳定性所以算子实现只接受 FP32 输入。矩阵乘法 (Matmul) Matmul 矩阵乘法一般来说可以使用半精度数但在大语言模型的 Attention 层Attention 分数与 V 张量的矩阵乘法需要用 FP32以便维持数值稳定性避免上溢减小累积错误。归一化函数归一化函数如 LayerNorm、BatchNorm、RMSNorm 需要用 FP32 计算。LayerNorm 要计算方差有累加操作如果用 FP16则计算过程中容易发生上溢和比较大的误差从而影响最后的收敛。BatchNorm 与 LayerNorm 类似也需要将输入张量 Upcast 到 FP32 计算。RMSNorm 是 LayerNorm 的扩展在开源 LLaMA 模型中使用PyTorch 目前并没有提供 RMSNorm 的实现。如果开发者自己实现需要注意将输入张量强制 Upcast 到 FP32 计算结果可以转回半精度数主要原因也是它需要计算 L2 均值涉及累加操作低精度运算容易累积误差。卷积操作与池化卷积操作属于逐单元矩阵乘法对精度不敏感可以使用 BF16不会对模型收敛性产生影响。和卷积类似的还有池化同样可以使用 BF16。ReLU 和 GeLU 属于逐点运算同样对精度不敏感。循环神经网络 (RNNs) RNN 一般对精度比较敏感 。LSTMCell 对精度不敏感可以用 BF16。模型起始与结束层通常大语言模型的第一层和最后一层对精度敏感比如 GPT 的 Embedding 层 (词嵌入和位置嵌入层) 需要使用 FP32。一般的输出层即便不是 Softmax计算和输出结果也需要是 FP32。通信算子涉及到梯度或者激活值累加的通信算子例如 All-Reduce、Reduce-Scatter都需要把输入张量 Upcast 到 FP32 以保证数值稳定性。References[1] Yongqiang Cheng (程永强), https://yongqiang.blog.csdn.net/[2]Automatic Mixed Precision package - torch.amp, https://docs.pytorch.org/docs/stable/amp.html