DeepSpeed - 混合精度训练

混合精度训练是深度学习中的一种革命性方法，其中模型的训练速度更快、效率更高。这种方法混合使用 16 位浮点算法，有时使用 32 位浮点算法，以在模型精度和最大硬件效率之间取得良好的平衡。微软的 DeepSpeed 库允许通过减少内存和计算时间轻松扩展大型模型。

什么是混合精度训练?

混合精度训练对大多数计算使用较低精度算法，并保留更高精度。在这种情况下，FP32 至关重要。主要目标是降低计算成本、加快训练速度并节省内存使用量。

浮点格式

以下是浮点格式 −

FP32 单精度 − 深度学习中常用的 32 位浮点格式。
FP16 半精度 − 一种 16 位浮点格式，其计算速度比常规浮点快得多。
BF16 (BFloat16) − FP16 的一种变体，具有更宽的指数范围，并且进一步支持更可靠的训练。

使用 FP16/BF16 和 FP32 的训练模型大大减少了训练时间。它通常发生在 GPU 和 TPU 上的大规模模型训练中。

DeepSpeed FP16 和 BF16

DeepSpeed 原生支持 FP16 和 BF16 混合精度训练模式。这将允许开发人员扩展深度学习模型而不影响其性能和准确性。它看起来是这样的。

DeepSpeed FP16 混合精度训练

您需要做的就是稍微修改您的配置，以便将 fp16 添加到那里。这是一个初始化 FP16 混合精度训练的示例配置文件:

{
   "train_batch_size": 64,
   "gradient_accumulation_steps": 4,
   "fp16": {
      "enabled": true,
      "loss_scale": 0,
      "loss_scale_window": 1000,
      "hysteresis": 2,
      "min_loss_scale": 1
   }
}

内存效率 − 几乎将 GPU 内存占用量减半。

训练速度 − 使用 16 位精度加速。

BF16 混合精度训练

当 FP16 精度导致模型不稳定时，BF16 或 BFloat16 非常方便。DeepSpeed 在 AMD/NVidia 的 GPU 和 Google TPU 上原生支持 BF16。要使用 BF16 通过 DeepSpeed 进行训练，您还需要以以下形式更新配置:

{
   "train_batch_size": 64,
   "gradient_accumulation_steps": 4,
   "bf16": {
      "enabled": true
   }
}

Python 示例 (BF16)

以下示例演示了 BF16 混合精度训练的用法−

import deepspeed
def model_engine(model, optimizer, config):
    model, optimizer, _, _ = deepspeed.initialize(
        model=model,
        optimizer=optimizer,
        config=config
    )
    return model, optimizer

# 示例模型和优化器
model = YourModel() # 使用您的模型
optimizer = YourOptimizer(model.parameters())

# 为 BF16 加载 DeepSpeed 配置
ds_config = "deepspeed_bf16_config.json" # 您的 DeepSpeed 配置的路径

# 使用 BF16 初始化 DeepSpeed 模型
model_engine, optimizer = model_engine(model, optimizer, ds_config)

# 训练您的模型
for batch in dataloader:
    outputs = model_engine(batch)
    loss = criterion(outputs, targets)
    model_engine.backward(loss)
    model_engine.step()

稳定训练 − BF16 确保稳定性，尤其是对于大型模型。

高效训练 − 内存和计算效率接近 FP16。

混合精度训练的优势

以下是混合精度训练的主要优势 −

更少的内存使用 − 它使用 16 位精度进行大多数计算，32 位精度的内存使用量是其一半。这将允许在更大的模型或更大的批量大小上进行训练，而无需增加硬件要求。
加速 − 硬件加速器(例如 GPU 或 TPU)可以比标准(32 位)浮点数快几个数量级地评估低精度计算。这是一个巨大的加速，特别是对于大型模型而言。
无精度损失 − 混合精度可确保对精度最敏感的计算 - 例如，梯度累积确实以 32 位精度运行，因此即使在其他地方谨慎使用，模型的精度也会得到保留。

混合精度训练的挑战

以下是混合精度训练中的一些挑战 −

数值稳定性 − 以较低精度进行训练可能会导致数值稳定性的丧失，尤其是在 FP16 中。这可能会导致梯度下溢或上溢，从而导致优化期间收敛性不佳。
精度损失 −在某些模型中，以混合精度运行时的性能可能会受到影响，因此必须在不同精度级别进行管理。
硬件兼容性 − 并非所有硬件都支持混合精度训练。因此，在开始训练之前，在使用混合精度策略时，请确保您的硬件设计为支持 FP16 或 BF16 精度。支持 FP16 和 BF16 的一些硬件包括 Nvidia 的 Tensor Core、Google 的 TPU 等。

混合精度训练的最佳实践

以下是有效实施混合精度训练的一些最佳实践 −

1.合适的硬件

只有针对 FP16 或 BF16 计算进行了优化的硬件(例如 Nvidia 的 Tensor Cores 或 Google 的 TPU)才能充分使用混合精度。

2. 自动混合精度 (AMP)

自动混合精度库 − DeepSpeed 和 PyTorch 支持对混合精度训练进行最少的代码更改。只需启用 AMP，它就会让框架自动为您在 FP16/32 或 BF16/32 中的不同精度之间进行动态切换。

import torch
from torch.cuda import amp

# 初始化 amp autocast 和 GradScaler
autocast = amp.autocast
GradScaler = amp.GradScaler

# 创建 GradScaler
scaler = GradScaler()

for data in dataloader:
    optimizer.zero_grad()
    with autocast():
        outputs = model(data)
        loss = criterion(outputs, target)

    # 缩放损失和反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

稳定高效的训练 − AMP 确保在 FP16/32 中正确执行操作，消除梯度下溢等。

3. 损失缩放跟踪和稳定

DeepSpeed 和 PyTorch 自动提供损失缩放。训练时会自动调整比例，以避免数值不稳定。

{
   "fp16": {
      "enabled": true,
      "loss_scale": 0,  // Automatic loss scaling
      "loss_scale_window": 1000,
      "hysteresis": 2,
      "min_loss_scale": 1
   }
}

更准确的模型 − 损失缩放有助于避免梯度消失，从而使模型以稳定的方式收敛。

4. 内存和速度分析

分析模型以跟踪使用混合精度训练节省的内存和加速数量。使用 PyTorch 的 torch.profiler 等工具来监控以下指标 −

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
) as profiler:
    for step, batch in enumerate(dataloader):
        outputs = model(batch)
        loss = criterion(outputs, target)
        loss.backward()
    profiler.step()

print(profiler.key_averages().table(sort_by="cuda_time_total"))

优化的内存和速度 − 分析有助于确保重新体现混合精度训练的真正优势。

总结

虽然使用 DeepSpeed 进行混合精度训练在加速模型训练、节省内存和实现准确性方面非常出色，但当您利用 FP16 或 BF16 等格式时，您现在可以以显著降低的计算成本处理大量模型和数据集。虽然并不适合胆小的人，但采用 AMP 方面的最佳实践、适当的损失缩放和硬件兼容性将帮助您挖掘混合精度的真正威力。混合精度训练仍将是扩展模型的重要工具，因为模型只会继续变大，而增长没有上限。

DeepSpeed 教程

DeepSpeed 有用资源