PyTorch 单机多GPU 训练方法与原理整理

这里整理一些PyTorch单机多核训练的方法和简单原理，目的是既能在写代码时知道怎么用，又能从原理上知道大致是怎么回事儿。如果只是炼丹，有时候确实没时间和精力深挖太多实现原理，但又希望能理解简单逻辑。

PyTorch单机多核训练方案有两种：一种是利用nn.DataParallel实现，实现简单，不涉及多进程；另一种是用torch.nn.parallel.DistributedDataParallel和torch.utils.data.distributed.DistributedSampler结合多进程实现。第二种方式效率更高，但是实现起来稍难，第二种方式同时支持多节点分布式实现。方案二的效率要比方案一高，即使是在单运算节点上。

为方便理解，这里用一个简单的CNN模型训练MNIST手写数据集，相关代码：

• model.py：定义一个简单的CNN网络
• data.py：MNIST训练集和数据集准备
• single_gpu_train.py：单GPU训练代码

方案一

核心在于使用nn.DataParallel将模型wrap一下，代码其他地方不需要做任何更改:

model = nn.DataParallel(model)

为方便说明，我们假设模型输入为(32, input_dim)，这里的 32 表示batch_size，模型输出为(32, output_dim)，使用 4 个GPU训练。nn.DataParallel起到的作用是将这 32 个样本拆成 4 份，发送给 4 个GPU 分别做 forward，然后生成 4 个大小为(8, output_dim)的输出，然后再将这 4 个输出都收集到cuda:0上并合并成(32, output_dim)。

可以看出，nn.DataParallel没有改变模型的输入输出，因此其他部分的代码不需要做任何更改，非常方便。但弊端是，后续的loss计算只会在cuda:0上进行，没法并行，因此会导致负载不均衡的问题。

如果把loss放在模型里计算的话，则可以缓解上述负载不均衡的问题，示意代码如下：

class Net:
    def __init__(self,...):
        # code
    
    def forward(self, inputs, labels=None)
        # outputs = fct(inputs)
        # loss_fct = ...
        if labels is not None:
            loss = loss_fct(outputs, labels)  # 在训练模型时直接将labels传入模型，在forward过程中计算loss
            return loss
        else:
            return outputs

按照我们上面提到的模型并行逻辑，在每个GPU上会计算出一个loss，这些loss会被收集到cuda:0上并合并成长度为 4 的张量。这个时候在做backward的之前，必须对将这个loss张量合并成一个标量，一般直接取mean就可以。这在Pytorch官方文档nn.DataParallel函数中有提到：

When module returns a scalar (i.e., 0-dimensional tensor) in forward(), this wrapper will return a vector of length equal to number of devices used in data parallelism, containing the result from each device.

这部分的例子可以参考：data_parallel_train.py

方案二

方案二被成为分布式数据并行(distributed data parallel)，是通过多进程实现的，相比与方案一要复杂很多。可以从以下几个方面理解：

1. 从一开始就会启动多个进程(进程数等于GPU数)，每个进程独享一个GPU，每个进程都会独立地执行代码。这意味着每个进程都独立地初始化模型、训练，当然，在每次迭代过程中会通过进程间通信共享梯度，整合梯度，然后独立地更新参数。

2. 每个进程都会初始化一份训练数据集，当然它们会使用数据集中的不同记录做训练，这相当于同样的模型喂进去不同的数据做训练，也就是所谓的数据并行。这是通过torch.utils.data.distributed.DistributedSampler函数实现的，不过逻辑上也不难想到，只要做一下数据partition，不同进程拿到不同的parition就可以了，官方有一个简单的demo，感兴趣的可以看一下代码实现：Distributed Training

3. 进程通过local_rank变量来标识自己，local_rank为0的为master，其他是slave。这个变量是torch.distributed包帮我们创建的，使用方法如下：

   import argparse  # 必须引入 argparse 包
   parser = argparse.ArgumentParser()
   parser.add_argument("--local_rank", type=int, default=-1)
   args = parser.parse_args()
   

   必须以如下方式运行代码：

   python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 train.py
   

   这样的话，torch.distributed.launch就以命令行参数的方式将args.local_rank变量注入到每个进程中，每个进程得到的变量值都不相同。比如使用 4 个GPU的话，则 4 个进程获得的args.local_rank值分别为0、1、2、3。

   上述命令行参数nproc_per_node表示每个节点需要创建多少个进程(使用几个GPU就创建几个)；nnodes表示使用几个节点，因为我们是做单机多核训练，所以设为1。

4. 因为每个进程都会初始化一份模型，为保证模型初始化过程中生成的随机权重相同，需要设置随机种子。方法如下：

   def set_seed(seed):
       random.seed(seed)
       np.random.seed(seed)
       torch.manual_seed(seed)
       torch.cuda.manual_seed_all(seed)
   

使用方法通过如下示意代码展示：

from torch.utils.data.distributed import DistributedSampler  # 负责分布式dataloader创建，也就是实现上面提到的partition。

# 负责创建 args.local_rank 变量，并接受 torch.distributed.launch 注入的值
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()

# 每个进程根据自己的local_rank设置应该使用的GPU
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)

# 初始化分布式环境，主要用来帮助进程间通信
torch.distributed.init_process_group(backend='nccl')

# 固定随机种子
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# 初始化模型
model = Net()
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 只 master 进程做 logging，否则输出会很乱
if args.local_rank == 0:
    tb_writer = SummaryWriter(comment='ddp-training')

# 分布式数据集
train_sampler = DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)  # 注意这里的batch_size是每个GPU上的batch_size

# 分布式模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)

详细代码参考：ddp_train.py

DDP 有用技巧

官方推荐使用方案二（DDP），因此这里收集了一些在使用 DDP 过程中的技巧。

torch.distributed.barrier

在阅读 huggingface/transformers 的源码时，比如 examples/run_ner.py 中会看到如下代码：

    # 加载预训练模型和分词器
    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()  # 确保只有分布式训练中的第一个进程下载模型和词汇表
    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                          num_labels=num_labels,
                                          cache_dir=args.cache_dir if args.cache_dir else None)
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
                                                do_lower_case=args.do_lower_case,
                                                cache_dir=args.cache_dir if args.cache_dir else None)
    model = model_class.from_pretrained(args.model_name_or_path,
                                        from_tf=bool(".ckpt" in args.model_name_or_path),
                                        config=config,
                                        cache_dir=args.cache_dir if args.cache_dir else None)

    if args.local_rank == 0:
        torch.distributed.barrier()  # 确保只有分布式训练中的第一个进程下载模型和词汇表

上述代码的目的是加载预训练模型并将其读入内存。如果让四个进程各自都下载一次显然是不合理的，那么如何才能实现只让一个进程下载呢？这时就可以使用 barrier 函数。当从属进程（local_rank != 0）运行到第一个 if 语句时，会被 barrier 阻塞，只能等待；而主进程可以继续往下执行，完成模型的下载和读入内存的操作。然而，在第二个 if 语句处再次遇到 barrier 时，主进程是否会被阻塞呢？答案是不会，因为主进程和从属进程已经集合在一起了（即 barrier 被解除），于是所有进程都可以继续向下执行。当然，此时各进程的执行进度不同：主进程已经完成了模型的读取，所以会从第二个 if 继续执行；而从属进程尚未进行模型读取，只会从第一个 if 开始执行。

由此可见，barrier 类似于一个路障，会将进程拦住，直到所有进程都集合齐后才会放行。这种机制适用于以下场景：只有一个进程负责下载数据，其他进程可以直接使用已下载好的文件；或者只有一个进程负责数据预处理，其他进程则使用预处理并缓存好的数据等。

模型保存

模型的保存与加载方式与单 GPU 情况有所不同。在这里需要将参数以 CPU 方式保存到存储中，因为如果直接保存 GPU 上的参数，.pth 文件中会记录这些参数所属的 GPU 编号，导致在加载模型时会尝试将参数加载到相应的 GPU 上，从而引发错误。例如：

RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

模型保存的方式是一致的，但需要注意的是，在方案二中存在多个进程同时运行，因此可能会有多个模型被保存到存储中。如果使用共享存储，则需注意文件名冲突的问题。通常情况下，只需在 rank 0 进程上保存模型参数即可，因为所有进程的模型参数是同步的。

torch.save(model.module.cpu().state_dict(), "model.pth")

模型的加载：

param = torch.load("model.pth")

以下是 huggingface/transformers 代码中使用的模型保存代码：

if torch.distributed.get_rank() == 0:
    model_to_save = model.module if hasattr(model, "module") else model  # 处理分布式/并行训练
    model_to_save.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

同一台机器上运行多个 DDP 任务

假设希望在一台拥有 4 个 GPU 的计算机上运行两个 DDP 任务，每个任务使用两个 GPU，很可能会遇到如下错误：

RuntimeError: Address already in use
RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1544081127912/work/torch/lib/c10d/ProcessGroupNCCL.cpp:260, unhandled system error

出现该错误的原因是两个 DDP 任务的通信地址发生了冲突。此时需要显式地为每个任务设置不同的通信地址。

在 torch.distributed.launch 中指定不同的 master_addr 和 master_port

# 第一个任务
export CUDA_VISIBLE_DEVICES="0,1"
python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29501 train.py

# 第二个任务
export CUDA_VISIBLE_DEVICES="2,3"
python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.2 --master_port=29502 train.py

参考资料

PyTorch 多 GPU 训练——单计算节点——你需要的一切

使用 PyTorch 编写分布式应用程序

PyTorch 多 GPU 训练总结（DataParallel 的使用）

PyTorch 单机多 GPU 训练快速上手指南

本指南基于 pytorch-multi-gpu-training 项目，帮助开发者快速掌握 PyTorch 在单机环境下的多 GPU 训练方案。主要介绍两种核心方法：简单易用的 DataParallel (DP) 和高效推荐的 DistributedDataParallel (DDP)。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 Windows (需配置相应 CUDA 环境)
• GPU: 2 块或以上 NVIDIA GPU，已安装兼容的 NVIDIA 驱动
• CUDA: 已安装与 PyTorch 版本匹配的 CUDA Toolkit

前置依赖

• Python 3.6+
• PyTorch (建议 1.4+)
• torchvision
• tensorboard (可选，用于日志可视化)

国内加速安装推荐： 使用清华源或阿里源可显著提升下载速度。

# 使用清华源安装 PyTorch (示例：CUDA 11.3)
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目主要为代码示例集合，无需额外安装工具包。只需确保上述依赖已就绪，并克隆或下载项目代码即可。

# 假设项目代码已在本地
cd pytorch-multi-gpu-training
# 确认目录包含 model.py, data.py, single_gpu_train.py 等文件
ls

基本使用

PyTorch 单机多 GPU 训练主要有两种方案，请根据需求选择。

方案一：DataParallel (DP)

特点：实现极简，无需多进程，适合快速验证或小规模实验。但在单节点上效率略低于 DDP，且存在主卡负载不均问题。

核心代码： 只需在模型定义后包裹一层 nn.DataParallel。

import torch.nn as nn
from model import Net # 假设这是你定义的模型

# 初始化模型
model = Net()

# 【关键步骤】包裹 DataParallel
# 自动将模型复制到所有可见 GPU，并将输入数据拆分
model = nn.DataParallel(model)

# 后续训练代码 (loss 计算、backward、optimizer step) 无需更改
# 注意：Loss 计算默认在主卡 (cuda:0) 进行

运行方式： 直接像单 GPU 一样运行脚本，PyTorch 会自动识别所有可用 GPU。

python single_gpu_train.py
# 或者如果你修改了代码使用了 DP
python data_parallel_train.py

---

方案二：DistributedDataParallel (DDP) [官方推荐]

特点：基于多进程实现，每个进程独占一个 GPU，通信效率高，支持多节点扩展。是生产环境和大规模训练的首选。

核心步骤：

1. 代码修改： 需要在代码中加入进程初始化、参数解析、Sampler 设置及模型包裹逻辑。

   import argparse
   import torch
   import torch.distributed as dist
   from torch.nn.parallel import DistributedDataParallel as DDP
   from torch.utils.data.distributed import DistributedSampler
   from model import Net
   from data import train_dataset
   
   # 1. 解析 local_rank 参数 (由启动命令注入)
   parser = argparse.ArgumentParser()
   parser.add_argument("--local_rank", type=int, default=-1)
   args = parser.parse_args()
   
   # 2. 设置当前进程使用的 GPU
   torch.cuda.set_device(args.local_rank)
   device = torch.device('cuda', args.local_rank)
   
   # 3. 初始化分布式环境 (后端推荐 nccl)
   dist.init_process_group(backend='nccl')
   
   # 4. 固定随机种子 (保证多进程模型初始化一致)
   seed = 42
   torch.manual_seed(seed)
   torch.cuda.manual_seed_all(seed)
   
   # 5. 准备数据 (使用 DistributedSampler 切分数据)
   # 注意：DataLoader 中不要设置 shuffle=True，由 Sampler 控制
   train_sampler = DistributedSampler(train_dataset)
   train_loader = torch.utils.data.DataLoader(
       train_dataset, 
       sampler=train_sampler, 
       batch_size=64  # 这里的 batch_size 是单个 GPU 上的数量
   )
   
   # 6. 初始化模型并移动到设备
   model = Net().to(device)
   
   # 7. 【关键步骤】包裹 DistributedDataParallel
   model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank)
   
   # 8. 定义优化器和损失函数
   criterion = torch.nn.CrossEntropyLoss()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
   
   # 9. 训练循环 (仅在主进程打印日志，避免输出混乱)
   if args.local_rank == 0:
       print("Start training...")
   
   for epoch in range(10):
       train_sampler.set_epoch(epoch) # 每个 epoch 打乱数据顺序
       for inputs, labels in train_loader:
           inputs, labels = inputs.to(device), labels.to(device)
           
           optimizer.zero_grad()
           outputs = model(inputs)
           loss = criterion(outputs, labels)
           loss.backward()
           optimizer.step()
           
           if args.local_rank == 0 and step % 100 == 0:
               print(f"Epoch {epoch}, Loss: {loss.item()}")
   

2. 运行方式： 必须使用 torch.distributed.launch 模块启动脚本，指定 GPU 数量。

   # 使用 4 个 GPU 进行训练
   python -m torch.distributed.launch --nproc_per_node=4 ddp_train.py
   

   注：若需在同一台机器运行多个 DDP 任务，需通过 --master_port 区分端口，并配合 CUDA_VISIBLE_DEVICES 指定显卡，详见原文技巧部分。

模型保存特别提示 (针对 DDP)

在 DDP 模式下保存模型时，务必保存 model.module 的状态字典，以避免加载时出现 GPU 索引错误。

# 仅在主进程保存
if args.local_rank == 0:
    # 正确做法：访问 .module 属性获取原始模型
    torch.save(model.module.cpu().state_dict(), "model_ddp.pth")