Accelerate

Accelerate 是 Hugging Face 开发的深度学习多卡并行训练库，构建在 torch_xla 和 torch.distributed 之上。它提供了一个简单的 API，将与多 GPU、TPU、混合精度训练相关的样板代码抽离了出来，让用户可以在不同设备上轻松地进行分布式训练和混合精度训练，而无需重复编写大量样板代码。

另外，Accelerate 还提供了很多性能优化的功能，使得大规模训练和推理变得简单、高效且适应性强。

在学习如何使用 Accelerate 之前，首先回顾一下并行训练的其他方法。

DataParallel (DP)

DP 是最常见的并行训练方式，通过将数据拆分（split）到各个 workers 中（每个 worker 拥有完整模型）进行并行计算，以解决 Batch Size 过大的问题。因为求导以及加和都是线性的，所以数据并行在数学上是等价的。

DataParallel 是 PyTorch 中最容易实现的并行方案，只需要增加一行代码：model = nn.DataParallel(model)。

代码示例

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

# 数据集的长度为 100，batch size 为 32
# 模型是一个简单的 fc 层，输入长度是 5，输出是 2
input_size, output_size = 5, 2
batch_size, data_size = 32, 100

model = Model(input_size, output_size)

if torch.cuda.device_count() > 1:
    print(f"Detected {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

rand_loader = DataLoader(
    dataset=RandomDataset(input_size, data_size),
    batch_size=batch_size, 
    shuffle=True
)

for data in rand_loader:
    input = data.to(device)
    output = model(input)

在上述代码中，batch_size=32。由于使用了 DataParallel，在有 2 个 GPU 时，一个 batch 被划分成了 2 份，即 tensor.split(16)，分别送往两个 GPU。

此外，在第一次调用 model.to(device) 时，模型被加载到了第一个 GPU 设备上；而在第一次调用 output = model(input)（即第一次执行前向传播）时，模型被复制到了其余 GPU 上。

DataParallel 迭代过程

  1. 初始化：主 GPU 初始化模型，然后复制到每个 GPU 上。
  2. 分发：主 GPU 读取数据批次，将数据切分成多个子 batch 分发给不同的 GPU。
  3. 前向计算：在每个 GPU 上独立完成前向计算，输出结果被收集（Gather）到主 GPU 上计算 Loss。
  4. 梯度计算：Loss 被分发（Scatter）到每个 GPU，各自通过反向传播（BP）计算梯度。
  5. 同步更新：每个 GPU 的梯度被聚合（Reduce）到主 GPU，然后在主 GPU 上更新模型权重。
  6. 广播：在下一次迭代前，主 GPU 将更新后的参数广播（Broadcast）到其它 GPU。

局限性 (Parameter Server 架构)

  • 通信开销大：采用 PS 架构，每个 worker 都要与参数服务器频繁通信。
  • 负载不均衡：主 GPU 承担了聚合与分发的重任，负载过大，导致 GPU 利用率不足。
  • 扩展性差：仅支持单机多卡模式，无法实现多机多卡训练。

DistributedDataParallel (DDP)

在 DDP 中，不再有“主 GPU”的概念，每个 GPU 都在执行相同的任务。推理、Loss 计算、梯度计算等流程在各个 GPU 上独立并行完成。

DDP 并行训练的核心在于模型间的梯度同步。这是通过 All-Reduce 通信操作实现的，保证每个 GPU 都能得到完全相同的梯度。此外，每个进程拥有独立的优化器。由于初始状态一致且梯度更新同步，因此每轮迭代后不同进程间的模型参数保持完全一致，保证了数学上的等价性。

更多细节见文章：HF 高效训练技术[https://blog.csdn.net/qq_56591814/article/details/134099476?spm=1001.2014.3001.5501]

Accelerate 的使用

Accelerate 可以直接集成进 PyTorch，使用过程非常直观。

1. 实例化 Accelerator

首先，实例化一个 Accelerator 类。它会自动初始化分布式训练环境，无需手动判断是单 GPU、单机多卡还是多机多卡。

from accelerate import Accelerator

accelerator = Accelerator()
  • 可以通过传入 cpu=True 或 fp16=True 强制指定环境。
  • 不再需要手动调用 .to(device) 或 .cuda()，accelerator 会处理好设备放置。如果需要显式获取设备，请使用 accelerator.device。

2. 准备对象 (Prepare)

使用 prepare() 方法转换 PyTorch 对象。这是 API 的核心，支持四种主要类型：

  • Models (torch.nn.Module)
  • Optimizers (torch.optim.Optimizer)
  • DataLoaders (torch.utils.data.DataLoader)
  • Schedulers (可选)
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)