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

# Sample model and optimizer
model = YourModel()  # Use your model
optimizer = YourOptimizer(model.parameters())

# Load DeepSpeed config for BF16
ds_config = "deepspeed_bf16_config.json"  # path to your DeepSpeed config

# Initialize DeepSpeed model with BF16
model_engine, optimizer = model_engine(model, optimizer, ds_config)

# Train your model
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 Core 或 Google 的 TPU。

2. 自动混合精度 (AMP)

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

import torch
from torch.cuda import amp

# Initialize amp autocast and GradScaler
autocast = amp.autocast
GradScaler = amp.GradScaler

# Create a GradScaler
scaler = GradScaler()

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

    # Scale the loss and backward pass
    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、正确的损失缩放和硬件兼容性的最佳实践的采用将帮助您充分发挥混合精度的强大功能。随着模型不断增长，并且增长没有上限，混合精度训练将仍然是扩展模型的重要工具。

打印页面
上一页
下一页
广告