大语言模型的工程技巧（二）—

大语言模型的工程技巧（二）——混合精度训练

news2024/10/7 17:31:28

一、概述

在人工智能领域，反向传播算法（计算参数梯度的算法）是非常重要的。而在进行反向传播计算时，必须将经过膨胀的计算图存储在内存中（如果使用GPU运算，那么将存储在GPU的专用内存中）。然而，这种存储量相当庞大，在整个计算图的存储结构中，数值存储占据了最大的比例。这些数值包括各个节点的计算结果（来自向前传播的输出），以及相应的梯度（这些梯度是来自反向传播的结果）。虽然梯度累积技术可以通过分解计算图来限制计算图的膨胀，从而降低内存的使用，但面对庞大的模型时，即便是单个数据点的计算图，其所需的内存都是巨大的。例如，大语言模型的参数数量可能高达数十亿甚至上百亿。

二、什么是混合精度训练？

为了解决这个具有挑战性的问题，需要采取额外的优化策略来降低内存的使用。在深入探讨这些策略之前，我们需要更详细地了解数字在计算机中的存储方式。一般而言，数值计算结果使用32位浮点数（需要4字节来存储，使用32位的二进制的方式表示）存储。这种存储方式被称为单精度浮点数。那么，如果使用16位二进制数表示一个数值，会产生什么影响呢？

这种方法的好处之一是能够立即减少所需的存储空间，同时提升计算速度。然而，这种方法也存在一个明显的缺陷，即能够表示的数值范围受限。为了便于讨论，下面以能够表示的最小正数为例。使用16位浮点数，能够表示的最小正数是 $2^{-24}$ （相比之下，32位浮点数能够表示的最小正数为 $2^{-149}$ ）。当实际的数值小于这个阈值时，计算机会错误地将其视作0，这就是浮点数下溢（Underflow）。

为了尽可能地减少这类错误的发生，可以混合精度训练（Mixed Precision Training）算法，顾名思义，它是指在模型训练过程中使用不同的数值精度来处理不同部分的计算。

三、算法细节

这一算法包含两个主要部分。

精度分层处理：在这种训练中，模型本身（模型参数）依然使用32位浮点数进行存储，参数更新过程也使用32位浮点数。在模型的向前传播和反向传播过程中，转而使用16位浮点数进行计算。具体情况如图1所示。

图1

引入比例因子（Scale Factor）：在数学上，要防止浮点数下溢是相当容易的，只需要将模型损失乘以一个较大的常数n，该常数也被称为比例因子。根据链式法则，这将导致所有节点的梯度都增大n倍。这种方法确保了梯度落入16位浮点数表示的范围，从而解决浮点数下溢问题。在使用这些梯度进行参数更新时，需要将引入的缩放移除，也就是将梯度除以n。将这个过程与精度分层处理相结合，如图2所示。

图2

混合精度训练方法的优势在于，在保持适当的模型表示能力的同时，显著降低了内存开锁。通过将高精度的32位浮点数与16位浮点数的计算相结合，在不牺牲模型性能的前提下，显著减少内存需求，使计算机能够处理更大规模的模型和数据集。

四、代码实现

在实际应用中，PyTorch已经提供了相应的封装函数，分别是torch.cuda.amp.autocast和torch.cuda.amp.GradScaler。其中autocast实现的是第一部分——精度分层处理；GradScaler实现的是第二部分——引入比例因子。借助这两个工具，在优化算法中使用混合精度训练就变得很容易了。示意代码如下：

# 常规的模型训练实现
for epoch in range(0): 
    for input, target in zip(data, targets):
        # 启动混合精度训练
        with torch.autocast(device_type=device, dtype=torch.float16):
            output = net(input)
            loss = loss_fn(output, target)

        # 在触发反向传播之前，启动缩放因子
        scaler.scale(loss).backward()

        # 更新模型参数
        scaler.step(opt)
        scaler.update()

        opt.zero_grad()