AI中的混合精度(AMP)

文摘   2024-09-23 16:03   新加坡  

本文涉及到的详细测试代码和测试步骤放置于:

https://github.com/xinyuwei-david/david-share.git下的:Deep-Learning/

AMP,本文中不再赘述代码实现。欢迎给repo点亮Star,您的点赞是作者持续创作的动力。


一、AMP的介绍

AMP(Automatic Mixed Precision,自动混合精度)确实指的是在训练过程中同时使用单精度(FP32)和半精度(FP16)的混合训练方法。

Automatic Mixed Precision (AMP) 是一种混合使用 32 位 (float32) 、 16 位 、8位 数据类型的方法。具体来说:

  • float32 (32 位):用于需要更大动态范围的操作,例如归约操作。

  • 16 位和8位:用于可以在较低精度下更快执行的操作,例如线性层和卷积层。


具体来说,AMP的主要特点包括:

 

  1. 混合精度计算

  • 在训练过程中,某些计算(如矩阵乘法)使用半精度(FP16)进行,以提高计算速度和减少内存占用。

  • 其他对精度要求较高的计算(如归约操作、损失计算)仍然使用单精度(FP32)进行,以确保计算的准确性。

  • 损失缩放

    • 由于FP16的精度较低,直接使用FP16进行梯度计算可能会导致梯度消失问题。为了解决这个问题,AMP使用损失缩放技术,将损失值放大,以确保梯度在FP16表示范围内。

  • 自动处理

    • AMP库(如PyTorch的torch.cuda.amp)可以自动处理大多数混合精度训练的细节,包括损失缩放、权重复制等。用户只需添加几行代码即可实现混合精度训练。

    使用AMP的好处:

     

    • 提高训练速度:由于FP16计算速度更快,使用AMP可以显著提高模型训练的速度。

    • 减少内存占用:FP16占用的内存只有FP32的一半,因此可以在同样的GPU内存下训练更大的模型或使用更大的批量大小。


    二、AI训练中常见的操作与精度

    “规约”这个词在中文中可以拆分为“规”和“约”两个部分来理解:

    • “规”在这里可以理解为“规则”或“规范”。在规约操作中,它指的是按照某种规则或方法来处理数据。例如,求和、求平均值、找最大值等操作都有各自的规则。

    • “约”在这里可以理解为“约简”或“简化”。在规约操作中,它指的是将大量的数据简化为一个或几个有用的结果。例如,把一组数字求和得到一个总和,把一组数字求平均值得到一个平均值等。

    结合起来:

    “规约”就是按照一定的规则(规)对数据进行处理,从而将数据简化(约)为一个或几个有用的结果的过程。


    归约操作:

    归约操作通常需要更高的精度来避免数值误差累积和满足动态范围需求。以下是一些常见的归约操作及其精度选择:

    1. 求和 (Sum)

    • 通常使用FP32,因为累加多个数值时,低精度可能会导致数值误差累积。

  • 求平均值 (Mean)

    • 通常使用FP32,因为计算平均值涉及求和操作,低精度可能会导致数值误差累积。

  • 最大值/最小值 (Max/Min)

    • 通常使用FP32,因为需要准确找到最大值或最小值。

  • 归一化 (Normalization)

    • 通常使用FP32,因为归一化操作需要精确的数值计算。

  • Softmax

    • 通常使用FP32,因为Softmax需要将数值转换为概率分布,低精度可能会导致数值不稳定。

  • 方差 (Variance)

    • 通常使用FP32,因为计算方差涉及求和和平方操作,低精度可能会导致数值误差累积。

  • 标准差 (Standard Deviation)

    • 通常使用FP32,因为计算标准差涉及方差计算,低精度可能会导致数值误差累积。

  • L2范数 (L2 Norm)

    • 通常使用FP32,因为计算L2范数涉及平方和求和操作,低精度可能会导致数值误差累积。

  • L1范数 (L1 Norm)

    • 通常使用FP32,因为计算L1范数涉及求和操作,低精度可能会导致数值误差累积。

    非归约操作:

    非归约操作对精度的要求较低,可以使用较低的精度来提高计算速度和减少内存占用。以下是一些常见的非归约操作及其精度选择:

    1. 矩阵乘法 (Matrix Multiplication)

    • 通常使用FP16或FP8,因为这些操作对精度的要求较低,可以在较低精度下更快地执行。

  • 卷积运算 (Convolution)

    • 通常使用FP16或FP8,因为这些操作对精度的要求较低,可以在较低精度下更快地执行。

  • 激活函数 (Activation Functions)

    • 通常使用FP16或FP8,因为这些操作对精度的要求较低,可以在较低精度下更快地执行。

  • 反向传播 (Backpropagation)

    • 通常使用FP16或FP8,因为这些操作对精度的要求较低,可以在较低精度下更快地执行。


    三、为什么不能强制指定16位

    直接选择FP16或BF16进行训练也是一种常见的做法,尤其是在硬件支持的情况下。每种精度格式都有其优缺点,选择哪种精度取决于具体的应用场景和硬件支持情况。

    FP16(半精度浮点数)

    • 优点

      • 内存占用少:FP16占用的内存只有FP32的一半,可以在同样的GPU内存下训练更大的模型或使用更大的批量大小。

      • 计算速度快:FP16的计算速度通常比FP32快,特别是在支持FP16的硬件(如Nvidia的Tensor Cores)上。

    • 缺点

      • 精度问题:FP16的精度较低,可能会导致梯度消失或溢出问题,特别是在训练深层神经网络时。

    BF16(Brain Floating Point)

    • 优点

      • 更好的数值稳定性:BF16保留了与FP32相同的指数范围,但尾数位数较少,因此在数值稳定性上比FP16更好。

      • 硬件支持:许多现代硬件(如Google的TPU和Nvidia的A100 GPU)都支持BF16。

    • 缺点

      • 内存占用:虽然BF16的内存占用比FP32少,但比FP16多。

     
    尽管直接使用FP16或BF16进行训练是可行的,但AMP(自动混合精度)提供了一些额外的优势:

    • 自动处理精度问题:AMP可以自动处理FP16的精度问题,如梯度消失和溢出,通过损失缩放等技术来确保训练的稳定性。

    • 灵活性:AMP可以根据操作的类型自动选择使用FP16或FP32,从而在保持数值稳定性的同时最大化性能。

    总结

     

    • 直接使用FP16或BF16:适用于硬件支持且对数值稳定性要求不高的场景。

    • 使用AMP:适用于需要在性能和数值稳定性之间取得平衡的场景,特别是当你不想手动处理精度问题时。

    四、强制指定16位可能带来问题

    想象你在做一个蛋糕,蛋糕的味道和质量取决于你使用的各种配料的比例和数量。这里,配料的比例就像是神经网络训练中的梯度。

    FP32 vs FP16

    • FP32:就像是用精确的厨房秤来称量每一种配料。你可以精确到克,这样你就能确保每次做出来的蛋糕味道都一样。

    • FP16:就像是用一个不太精确的量杯来量配料。量杯只能精确到大约的数量,所以每次做出来的蛋糕味道可能会有些不同。

    梯度爆炸和梯度消失

    在神经网络训练中,梯度是用来更新模型参数的。如果梯度太大或太小,都会影响训练效果。

    梯度爆炸

    • 比喻:想象你在做蛋糕时,突然把一大堆糖倒进了面糊里。结果,蛋糕会变得非常甜,甚至无法食用。

    • FP16:由于FP16的精度较低,当梯度值很大时,FP16可能无法准确表示这些大值,导致梯度变得更大,最终导致梯度爆炸。

    • FP32:FP32有更高的精度,可以更准确地表示大梯度值,避免梯度爆炸。

    梯度消失

    • 比喻:想象你在做蛋糕时,只加了一点点糖,结果蛋糕几乎没有甜味。

    • FP16:由于FP16的精度较低,当梯度值很小时,FP16可能无法准确表示这些小值,导致梯度变得更小,最终导致梯度消失。

    • FP32:FP32有更高的精度,可以更准确地表示小梯度值,避免梯度消失。

    为什么FP16会出现这些问题,而FP32不会

    • 精度:FP16的精度较低,只有16位,其中5位用于指数,10位用于尾数。这意味着FP16能表示的数值范围和精度都比FP32小得多。

    • 数值范围:FP32有32位,其中8位用于指数,23位用于尾数。这使得FP32能表示的数值范围和精度都比FP16大得多。

    解决方法

    • 损失缩放:为了避免FP16的这些问题,可以使用损失缩放技术。通过将损失值放大,可以确保梯度在FP16的表示范围内,从而避免梯度消失和爆炸的问题。

    • 混合精度训练:使用AMP(自动混合精度)可以自动处理这些问题,通过在需要高精度的地方使用FP32,在其他地方使用FP16,从而在性能和稳定性之间取得平衡。


    FP16被比喻为“不太精确的量杯”是因为它的表示范围和精度较低。让我们深入了解一下:

    精度和表示范围

    • FP16(半精度浮点数)

      • 位数:16位

      • 指数位:5位

      • 尾数位:10位

      • 表示范围:大约 (10^{-5}) 到 (10^{5})

      • 精度:由于尾数位只有10位,FP16只能表示有限的精度,这意味着它在表示非常小或非常大的数值时会有较大的误差。

    • FP32(单精度浮点数)

      • 位数:32位

      • 指数位:8位

      • 尾数位:23位

      • 表示范围:大约 (10^{-38}) 到 (10^{38})

      • 精度:由于尾数位有23位,FP32可以表示更高的精度,能够更准确地表示非常小或非常大的数值。

    为什么FP16是不太精确的量杯

    • 有限的尾数位:FP16只有10位尾数,这意味着它在表示数值时只能保留有限的有效数字。就像一个不太精确的量杯,它只能大致量出配料的数量,而不能精确到克。

    • 数值误差:由于精度较低,FP16在进行计算时容易产生数值误差。这些误差在深度学习训练中会累积,导致梯度爆炸或梯度消失的问题。

    • 表示范围有限:FP16的指数位只有5位,这限制了它能表示的数值范围。对于非常大的或非常小的数值,FP16可能无法准确表示,导致计算结果不准确。

    比喻解释

    • FP16:就像一个只能大致量出配料数量的量杯,可能会导致蛋糕的味道每次都有些不同,因为配料的比例不够精确。

    • FP32:就像一个精确的厨房秤,可以精确到克,确保每次做出来的蛋糕味道都一样,因为配料的比例非常精确。

    解决方法

    • 损失缩放:通过将损失值放大,可以确保梯度在FP16的表示范围内,从而避免梯度消失和爆炸的问题。

    • 混合精度训练:使用AMP(自动混合精度)可以自动处理这些问题,通过在需要高精度的地方使用FP32,在其他地方使用FP16,从而在性能和稳定性之间取得平衡。


    五、微软AMP

    MS-AMP 是微软开发的用于深度学习的自动混合精度包。

    特点:

    • 支持 O1 优化:对权重和权重梯度应用 FP8,并支持 FP8 通信。

    • 支持 O2 优化:支持两种优化器(Adam 和 AdamW)的 FP8。

    • 支持 O3 优化:支持分布式并行训练和 ZeRO 优化器的 FP8,这对于训练大规模模型至关重要。

    • 提供四个应用 MS-AMP 的训练示例:Swin-Transformer、DeiT、RoBERTa 和 GPT-3。

      与 Transformer Engine 相比,MS-AMP 具有以下优势:

    • 通过访问一个字节而不是半精度或单精度,加速内存受限的操作。

    • 减少训练模型的内存需求,使得可以训练更大的模型。

    • 通过传输低精度梯度,加速分布式模型的通信。

    • 通过更大的小批量数据,减少大规模语言模型的训练时间。

    下表展示的不涉及归约操作。计算 (GEMM) 是指通用矩阵乘法(General Matrix Multiplication)。

    下面表格展示的是归约操作:


    https://azure.github.io/MS-AMP/docs/introduction/

    https://azure.github.io/MS-AMP/docs/user-tutorial/optimization-level

    大魏分享
    https://github.com/davidsajare/david-share.git
     最新文章