并行训练实践

Thu, 24 Jul 2025 00:00:00 +0000

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。

未完成 on Ganko Space

并行训练实践

Accelerate

DataParallel (DP)

代码示例