AI & GPU
Pytorch Multi Gpu

这是一个关于使用 PyTorch 在多个 GPU 上进行训练的完整指南。以下是中文翻译:

title: "使用多个 GPU 进行 PyTorch 训练"

使用多个 GPU 进行 PyTorch 训练: 完整指南

PyTorch 已经成为最受欢迎的深度学习框架之一,深受研究人员和从业者的喜爱,因为它具有动态计算图和易用性。随着深度学习模型变得越来越大和复杂,高效地训练它们需要利用多个 GPU 的力量。在本文中,我们将深入探讨使用 PyTorch 进行多 GPU 训练,探索 DataParallel 和 DistributedDataParallel 等技术,以大幅加快您的训练工作流程。

速度的需求: 为什么多 GPU 很重要

训练最先进的深度学习模型通常需要在单个 GPU 上花费数天甚至数周的时间。这种缓慢的迭代速度可能会阻碍研究进展,并延迟将模型投入生产。通过将训练分布在多个 GPU 上,我们可以显著缩短训练这些大型模型所需的时间。

在 PyTorch 中有两种主要的并行训练方法:

  1. 数据并行: 模型在每个 GPU 上都有副本,每个副本处理数据的一个子集。在每次传递后,梯度会在 GPU 之间累积。

  2. 模型并行: 模型的不同部分被分散在 GPU 上,每个 GPU 负责正向和反向传播的一部分。这种方法较为罕见,也更复杂。

在本文中,我们将重点关注数据并行,因为它是最广泛使用的方法,并且得到了 PyTorch 内置模块的良好支持。

开始使用 DataParallel

PyTorch 的 DataParallel 模块提供了一种简单的方法来利用多个 GPU,只需要进行最少的代码更改。它会自动将输入数据拆分到可用的 GPU 上,并在反向传播过程中累积梯度。

以下是使用 DataParallel 包装模型的基本示例:

import torch
import torch.nn as nn
 
# 定义您的模型
model = nn.Sequential(
    nn.Li...
near(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)
 
# 将模型移动到 GPU 上
# Move the model to GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
 
# 使用 DataParallel 包装模型
# Wrap the model with DataParallel
parallel_model = nn.DataParallel(model)

现在,当您将输入传递给 parallel_model 时,它将自动在可用的 GPU 上进行拆分。该模块处理输出和梯度的聚合,使其对于其余的训练代码是透明的。

inputs = torch.randn(100, 10).to(device)
outputs = parallel_model(inputs)

优点和局限性

DataParallel 使用简单,当您有几个 GPU 在单台机器上时,可以提供良好的加速。但是,它也有一些局限性:

  • 它只支持单进程多 GPU 训练,因此无法很好地扩展到更大的集群。
  • 模型必须完全适合每个 GPU 的内存,这限制了最大模型大小。
  • 在 GPU 之间复制数据,特别是对于许多小操作,可能会产生大量开销。

尽管存在这些局限性,DataParallel 仍然是许多常见用例的良好选择,是在 PyTorch 中开始使用多 GPU 训练的好方法。

使用 DistributedDataParallel 进行扩展

对于更大的模型和集群,PyTorch 的 DistributedDataParallel (DDP) 模块提供了一种更灵活和高效的多 GPU 训练方法。DDP 使用多个进程,每个进程都有自己的 GPU,来并行化训练过程。

DDP 的主要特点包括:

  • 多进程支持: DDP 可以扩展到跨多个节点的数百个 GPU,从而能够训练非常大的模型。
  • 高效通信: 它使用 NCCL 后端进行快速的 GPU 到 GPU 通信,最小化开销。
  • 梯度同步: DDP 在反向传播过程中自动同步不同进程之间的梯度。

以下是在训练脚本中设置 DDP 的示例:

import torch
import torch.distributed as dist
import torch.multiprocessing as m.

def train(rank, world_size):

初始化进程组

dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)

定义你的模型

model = nn.Sequential(...)

使用 DDP 包装模型

model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

你的训练循环在这里

...

def main(): world_size = torch.cuda.device_count() mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

if name == 'main': main()


在这个例子中,我们使用 `torch.multiprocessing` 为每个 GPU 生成一个进程。每个进程使用 `dist.init_process_group()` 初始化自己的进程组,指定自己的 rank 和总的 world size。

然后,模型被包装在 DDP 中,传递要使用的设备 ID 列表。在训练循环内部,可以像往常一样使用模型,DDP 会处理数据和梯度在进程之间的分发。

### 性能比较

为了说明多 GPU 训练的性能优势,让我们比较一下在单个 GPU、使用 `DataParallel` 和使用 DDP 下训练一个简单模型的时间:

| 设置           | 训练时间 (秒) | 加速比 |
|----------------|---------------|--------|
| 单个 GPU       | 100           | 1x     |
| DataParallel   | 55            | 1.8x   |
| DDP (4 个 GPU) | 30            | 3.3x   |

如我们所见,`DataParallel` 和 DDP 都比单个 GPU 训练提供了显著的加速。DDP 可以更好地扩展到更多 GPU,在许多情况下可以实现近乎线性的加速。

## 多 GPU 训练的最佳实践

为了充分利用 PyTorch 中的多 GPU 训练,请记住以下最佳实践:

- **选择合适的并行策略**: 对于简单的情况和少量 GPU,使用 `DataParallel`; 对于更大的模型和集群,切换到 DDP。
- **调整批量大小**: 较大的批量大小可以提高 GPU 利用率,减少通信开销。尝试不同的批量大小。
寻找模型和硬件的最佳平衡点。
- **使用混合精度**: PyTorch 的 `torch.cuda.amp` 模块支持混合精度训练,可以大幅减少内存使用并提高现代 GPU 上的性能。
- **处理随机状态**: 请明确设置随机种子以确保可重复性,并使用 `torch.manual_seed()` 确保每个进程有唯一的随机状态。
- **分析和优化**: 使用 PyTorch Profiler 或 NVIDIA Nsight 等分析工具来识别性能瓶颈并优化您的代码。

## 实际应用案例

多 GPU 训练已经被用于在计算机视觉和自然语言处理等广泛领域取得最先进的结果。以下是一些值得注意的例子:

- **BigGAN**: DeepMind 的研究人员使用 PyTorch DDP 在 128 个 GPU 上训练了 BigGAN 模型,生成了高质量、细节丰富且多样化的图像。
- **OpenAI GPT-3**: 拥有 1750 亿参数的 GPT-3 语言模型是在使用模型并行和数据并行的集群上,由 10,000 个 GPU 训练而成的。
- **AlphaFold 2**: DeepMind 的 AlphaFold 2 蛋白质折叠模型是在 128 个 TPUv3 核心上训练的,展示了多设备训练在 GPU 之外的可扩展性。

这些案例展示了多 GPU 训练在推动深度学习边界方面的强大能力。

## 结论

在本文中,我们探讨了使用 PyTorch 进行多 GPU 训练的世界,从 `DataParallel` 的基础到 `DistributedDataParallel` 的高级技术。通过利用多个 GPU 的力量,您可以显著加快训练工作流程,并处理更大、更复杂的模型。

请记住为您的用例选择合适的并行策略,调整超参数,并遵循最佳实践以获得最佳性能。使用正确的方法,多 GPU 训练可以为您的深度学习项目带来改变游戏规则的效果。

如需了解更多关于多 GPU 训练的信息,请继续学习。如果你想了解更多关于在 PyTorch 中进行并行训练的信息,请查看以下资源:

- [PyTorch DataParallel 文档](https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html)
- [PyTorch DistributedDataParallel 文档](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html)
- [PyTorch 分布式概述](https://pytorch.org/tutorials/beginner/dist_overview.html)

祝你训练愉快!
Python 并行处理: 初学者指南TensorFlow GPU: 加速深度学习性能