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、正確的損失縮放和硬體相容性的最佳實踐的採用將幫助您充分發揮混合精度的強大功能。隨著模型不斷增長，並且增長沒有上限，混合精度訓練將仍然是擴充套件模型的重要工具。

列印頁面