精度及其作用：BFloat16训练

训练一个参数量达到数十亿甚至数万亿的专家混合模型（Mixture of Experts）对现代硬件构成了严峻挑战。权重、激活和梯度的庞大规模会占用大量GPU内存，并对计算吞吐量提出高要求。尽管分布式训练策略可以分割模型，但优化数据本身的数值格式提供了一种辅助且有效的方法来提升效率。正因如此，低精度数据类型，特别是BFloat16，变得必不可少。

在大多数框架中，对所有计算使用完整的32位浮点精度（FP32）是默认设置，它提供了宽广的动态范围和高精度。然而，每个FP32参数需要4字节存储空间。对于一个拥有5000亿参数的模型，仅权重就将占用2太字节内存，这使其无法在任何单个加速器上运行。

解决方案是采用低精度格式。通过减少用于表示每个数字的比特数，我们可以大幅降低内存占用，并在硬件支持下加速计算。

浮点数值类型

过去，FP32的主要替代方案是FP16（半精度），它使用16位。尽管FP16成功地将内存消耗减半，但它有一个显著缺点：动态范围非常有限。其较小的指数分配使其在训练期间容易出现数值不稳定。梯度可能非常小或非常大，容易变为零（下溢）或无穷大（上溢），从而使训练过程不稳定或中止。

这促成了BFloat16（BF16），即“脑浮点”格式的出现，这是一种专为深度学习工作负载设计的格式。与FP16一样，它也只使用16位。然而，它做出了不同的权衡。BF16为指数分配了与FP32相同数量的比特，从而保持了其宽广的动态范围。这代价是尾数（或小数部分）的比特数减少，而尾数负责精度。

下图显示了这三种格式的位分配。

FP32、BF16和FP16位结构的比较。BF16保留了FP32的8个指数位，确保了相似的动态范围，而FP16牺牲了指数位以获得更多的尾数（精度）位。

对于深度学习来说，BF16的宽动态范围远比高精度重要。神经网络对权重和激活中的噪声及较低精度表现出很强的适应性。通过防止FP16常见的下溢和上溢问题，BF16提供了一个更稳定的训练环境，通常可以近乎直接替代FP32。

使用BFloat16的混合精度训练

虽然可以简单地将整个模型转换为BF16，但一种更有效的方法，即混合精度训练，是标准做法。这种方法将BF16的速度和内存优势与FP32在训练循环重要部分的稳定性结合起来。

使用BF16的典型混合精度工作流程如下：

1. 在FP32中保留主权重： 模型的权重主副本以完整的FP32精度存储。这作为数据的权威来源，确保小的梯度更新不会因BF16的较低精度而丢失。
2. 将权重转换为BF16进行计算： 对于每个训练步骤，FP32主权重都会被转换为BF16。
3. 在BF16中执行前向和反向传播： 前向和反向传播中的所有矩阵乘法及其他计算都使用BF16权重和激活执行。现代GPU拥有专用硬件，例如NVIDIA的Tensor Cores，可以显著加速BF16操作。
4. 更新FP32主权重： 所得的梯度（为BF16格式）随后用于更新FP32主权重副本。

标准混合精度训练步骤中的数据流。计算在BF16中加速，而权重更新在FP32中执行以保持稳定性。

这种方法兼具两者优点：16位计算的内存和速度优势，以及32位权重更新的数值稳定性。对于MoE模型而言，这不仅是一种优化；更是一项实现途径。将权重和激活的内存占用减半，使得训练更大、专家数量更多的模型变得可行。

在PyTorch中的实际实现

像PyTorch这样的深度学习框架提供了简单的上下文管理器来自动化混合精度训练。如果您的硬件支持BF16（例如NVIDIA A100或H100系列GPU），启用它异常简单。

以下是使用torch.autocast的典型训练循环示例。

import torch

# 确保您的模型和数据位于支持BF16的设备上
device = "cuda" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "cpu"
model = MyMoEModel().to(device)
optimizer = torch.optim.AdamW(model.parameters())
data = torch.randn(64, 1024, device=device) # 示例数据

# torch.autocast 上下文管理器会自动处理
# 将符合条件的操作转换为指定的dtype。
with torch.autocast(device_type=device, dtype=torch.bfloat16):
    # 模型的正向传播以BF16运行
    output, aux_loss = model(data)

    # 损失计算也可以在autocast上下文内部进行
    main_loss = loss_fn(output, target)
    total_loss = main_loss + aux_loss

# 梯度是基于BF16正向传播计算的
# .backward() 调用发生在autocast上下文之外
total_loss.backward()

# 优化器更新主FP32权重
optimizer.step()
optimizer.zero_grad()

请注意其简便性。autocast上下文管理器自动处理操作到BF16的转换。当使用FP16训练时，需要一个名为GradScaler的额外组件来缩放损失，以防止梯度下溢。由于BF16拥有更大的动态范围，这个损失缩放步骤通常不需要，这进一步简化了训练代码，并减少了一个需要调整的超参数。这种固有的稳定性使得BF16成为训练MoE等庞大而复杂模型的首选。