• 教程 >
  • 使用 DDP 进行多 GPU 训练
快捷方式

入门 || 什么是 DDP || 单节点多 GPU 训练 || 容错 || 多节点训练 || minGPT 训练

使用 DDP 进行多 GPU 训练

创建于:2022 年 9 月 27 日 | 最后更新:2024 年 11 月 03 日 | 最后验证:未验证

作者:Suraj Subramanian

你将学到什么
  • 如何通过 DDP 将单 GPU 训练脚本迁移到多 GPU

  • 设置分布式进程组

  • 在分布式设置中保存和加载模型

GitHub 上查看本教程使用的代码

前提条件
  • DDP 工作原理的高级概述

  • 一台具有多个 GPU 的机器(本教程使用 AWS p3.8xlarge 实例)

  • 安装 CUDA 的 PyTorch

观看下方视频或在 youtube 上同步学习。

前一篇教程中,我们对 DDP 的工作原理有了一个高级概述;现在我们来看如何在代码中使用 DDP。在本教程中,我们将从一个单 GPU 训练脚本开始,并将其迁移到单节点上运行在 4 个 GPU 上。在此过程中,我们将讨论分布式训练中的重要概念,并在代码中实现它们。

注意

如果你的模型包含任何 BatchNorm 层,则需要将其转换为 SyncBatchNorm,以便在副本之间同步 BatchNorm 层的运行统计信息。

使用辅助函数 torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) 将模型中的所有 BatchNorm 层转换为 SyncBatchNorm

single_gpu.pymultigpu.py 的区别

这些是你在单 GPU 训练脚本中通常需要进行的更改,以启用 DDP。

导入

  • torch.multiprocessing 是 PyTorch 对 Python 原生多进程的封装

  • 分布式进程组包含所有可以相互通信和同步的进程。

import torch
import torch.nn.functional as F
from utils import MyTrainDataset

import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
import os

构建进程组

  • 首先,在初始化进程组之前,调用 set_device,它为每个进程设置默认 GPU。这对于防止 GPU:0 上的挂起或过度内存占用非常重要。

  • 进程组可以通过 TCP(默认)或共享文件系统进行初始化。请阅读更多关于进程组初始化的信息

  • init_process_group 初始化分布式进程组。

  • 阅读更多关于选择 DDP 后端的信息

def ddp_setup(rank: int, world_size: int):
   """
   Args:
       rank: Unique identifier of each process
      world_size: Total number of processes
   """
   os.environ["MASTER_ADDR"] = "localhost"
   os.environ["MASTER_PORT"] = "12355"
   torch.cuda.set_device(rank)
   init_process_group(backend="nccl", rank=rank, world_size=world_size)

构建 DDP 模型

self.model = DDP(model, device_ids=[gpu_id])

分布输入数据

  • DistributedSampler 将输入数据分块分发到所有分布式进程。

  • The DataLoader 结合了数据集和

    采样器,并为给定数据集提供一个可迭代对象。

  • 每个进程将接收一个包含 32 个样本的输入批次;有效批次大小是 32 * nprocs,在使用 4 个 GPU 时为 128。

train_data = torch.utils.data.DataLoader(
    dataset=train_dataset,
    batch_size=32,
    shuffle=False,  # We don't shuffle
    sampler=DistributedSampler(train_dataset), # Use the Distributed Sampler here.
)
  • 在每个 epoch 开始时调用 DistributedSampler 上的 set_epoch() 方法是必要的,以确保在多个 epoch 中正确地进行洗牌。否则,每个 epoch 将使用相同的顺序。

def _run_epoch(self, epoch):
    b_sz = len(next(iter(self.train_data))[0])
    self.train_data.sampler.set_epoch(epoch)   # call this additional line at every epoch
    for source, targets in self.train_data:
      ...
      self._run_batch(source, targets)

保存模型检查点

  • 我们只需要从一个进程保存模型检查点。如果没有这个条件,每个进程都会保存一份相同的模型副本。阅读更多关于使用 DDP 保存和加载模型的信息此处

- ckp = self.model.state_dict()
+ ckp = self.model.module.state_dict()
...
...
- if epoch % self.save_every == 0:
+ if self.gpu_id == 0 and epoch % self.save_every == 0:
  self._save_checkpoint(epoch)

警告

Collective calls 是在所有分布式进程上运行的函数,用于将某些状态或值收集到特定进程。集体调用要求所有 rank 都运行集体代码。在此示例中,_save_checkpoint 不应包含任何集体调用,因为它仅在 rank:0 进程上运行。如果你需要进行任何集体调用,应在 if self.gpu_id == 0 检查之前进行。

运行分布式训练任务

  • 包含新的参数 rank(替换 device)和 world_size

  • 调用 mp.spawn 时,rank 由 DDP 自动分配。

  • world_size 是整个训练任务中的进程数。对于 GPU 训练,这对应于使用的 GPU 数量,并且每个进程在一个专用的 GPU 上工作。

- def main(device, total_epochs, save_every):
+ def main(rank, world_size, total_epochs, save_every):
+  ddp_setup(rank, world_size)
   dataset, model, optimizer = load_train_objs()
   train_data = prepare_dataloader(dataset, batch_size=32)
-  trainer = Trainer(model, train_data, optimizer, device, save_every)
+  trainer = Trainer(model, train_data, optimizer, rank, save_every)
   trainer.train(total_epochs)
+  destroy_process_group()

if __name__ == "__main__":
   import sys
   total_epochs = int(sys.argv[1])
   save_every = int(sys.argv[2])
-  device = 0      # shorthand for cuda:0
-  main(device, total_epochs, save_every)
+  world_size = torch.cuda.device_count()
+  mp.spawn(main, args=(world_size, total_epochs, save_every,), nprocs=world_size)

代码示例如下

延伸阅读

文档

访问 PyTorch 全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源