并行训练实践
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。 ...