jia-zhuang/pytorch-multi-gpu-training：整理 pytorch 单机多 GPU 训练方法与原理

技术解读

以下内容由 GitRead 中文社区编辑团队基于项目 README 与官方资料整理，结合中文开发者的常见使用场景补充技术选型建议。具体以项目官方文档为准。

深度解析 jia-zhuang/pytorch-multi-gpu-training 项目

项目背景与痛点

在现今的深度学习领域，训练神经网络模型效率的提升是一个核心问题。尤其是在大规模数据集和复杂模型结构面前，单GPU的计算能力往往难以满足高效训练的需求。此时，单机多GPU训练便成为了一种重要的解决方案，能够显著加速模型的训练过程。然而，多GPU环境下的训练并不如单GPU那么简单，它需要面对诸如数据同步、梯度聚合、通信延迟等多种挑战。

项目“jia-zhuang/pytorch-multi-gpu-training”应运而生，它提供了多GPU训练的详细方法和技巧，涵盖了从基础知识到实现技术的全方位指导。通过深入剖析项目的内容，我们可以一窥多GPU训练背后的实现机制和遇到的难点。

核心技术揭秘

技术架构与实现原理

此项目采用了 PyTorch 提供的多 GPU 支持机制来实现高效的单机多 GPU 训练。核心在于数据并行（Data Parallel）模式，通过 torch.nn.DataParallel 将模型复制到每个可用 GPU 上，同时在每个 GPU 上分配独立的数据子集来并发执行前向和后向传播。关键的几个技术点包括：

数据划分：将输入数据均匀地分配到各个GPU上，确保每个GPU的利用率最大化。
梯度同步：通过梯度聚合机制将各个GPU计算的梯度加权平均后同步到同一架构，保证参数的一致性。

graph TD A[输入数据] --> B[模型复制与数据划分] B --> C{"批处理"} C -->|"GPU 1"| D[前向传播] C -->|"GPU 2"| E[前向传播] C -->|"GPU 3"| F[前向传播] C -->|"GPU 4"| G[前向传播] H[梯度反向传播] --> I[梯度同步与聚合] I --> J{"反向参数同步"} J --> K[模型更新]

核心代码片段

import torch
import torch.nn as nn

class MultiGPUNetwork(nn.Module):
    def __init__(self):
        super(MultiGPUNetwork, self).__init__()
        ## 网络定义
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3),
            nn.ReLU(),
            nn.Flatten()
        )
    
    def forward(self, x):
        return self.model(x)

## 数据并行模型实例
device_ids = [0, 1, 2, 3]  # 这里根据实际GPU数量调整
model = nn.DataParallel(MultiGPUNetwork(), device_ids=device_ids)
model.cuda()

功能亮点与差异

与其它同类项目相比，此项目有几个显著的优势：

详细的理论说明：在官网说明文档中提供了全面详细的理论介绍和实现举例，对于初学者尤其友好。
实用性与高效性： 采用的数据并行方法计算性能得到有效释放，同时能够做到较好地同步各GPU之间的梯度信息，确保了训练过程中的一致性和平滑度。
灵活性：你可以根据具体的硬件配置修改参数，或者在多台机器之间扩展来实现更好的性能优化。

应用场景与落地建议

场景

适用于需要处理大规模数据、大型模型的业务场景，如图像识别、自然语言处理等任务。

注意事项

需要根据实际运行环境（GPU 数量、内存、网络带宽等）配置参数。
在多GPU环境进行生产部署前，建议进行全面的性能测试以确保模型在多种条件下的稳定性和速率。
检查模型在Edge Case上的表现，并考虑加入更多的容错机制。

综合评价

优点：提供了全面的理论讲解和实践代码，能明显加快模型训练速度。
缺点：灵活性有待进一步提高，特别是在数据分割和梯度收集的具体实现上可能会遇到不同的痛点。优化空间依然很大，具体应用时可能需要针对具体需求进行定制化调整。

此项目为多GPU训练提供了一个非常有价值的实验和实践空间，帮助开发者更好地理解其背后的技术原理及实现细节，提供了优化训练过程的实用指导。

关于

pytorch multi-gpu-training

863 Stars
88 Forks
4 Watching

简要分析

热度分

1043

价值分

544

活跃状态

活跃

GitHub

语言占比

Python

贡献者

jia-zhuang

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)，是通过多进程实现的，相比与方案一要复杂很多。可以从以下几个方面理解：

从一开始就会启动多个进程(进程数等于GPU数)，每个进程独享一个GPU，每个进程都会独立地执行代码。这意味着每个进程都独立地初始化模型、训练，当然，在每次迭代过程中会通过进程间通信共享梯度，整合梯度，然后独立地更新参数。
每个进程都会初始化一份训练数据集，当然它们会使用数据集中的不同记录做训练，这相当于同样的模型喂进去不同的数据做训练，也就是所谓的数据并行。这是通过torch.utils.data.distributed.DistributedSampler函数实现的，不过逻辑上也不难想到，只要做一下数据partition，不同进程拿到不同的parition就可以了，官方有一个简单的demo，感兴趣的可以看一下代码实现：Distributed Training
进程通过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。
因为每个进程都会初始化一份模型，为保证模型初始化过程中生成的随机权重相同，需要设置随机种子。方法如下：
```
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会看到一下代码：

    # Load pretrained model and tokenizer
    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

    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()  # Make sure only the first process in distributed training will download model & vocab

上述代码要实现预训练模型的下载和读入内存，如果4个进程都分别下载一遍显然是不合理的，那如何才能实现只让一个进程下载呢？这个时候就可以使用barrier函数。当slave进程(local_rank!=0)运行到第一个if时就被barrier住了，只能等着，但master进程可以往下运行完成模型的下载和读入内存，但在第二个if语句时遇到barrier，那会不会被barrier住呢？答案是不会，因为master进程和slave进程集合在一起了(barrier)，barrier会被解除，这样大家都往下执行。当然这时大家执行的进度不同，master进程已经执行过模型读入，所以从第二个if往下执行，而slave进程尚未执行模型读入，只会从第一个if往下执行。

可以看到barrier类似一个路障，进程会被拦住，直到所有进程都集合齐了才放行。适合这样的场景：只一个进程下载，其他进程可以使用下载好的文件；只一个进程预处理数据，其他进程使用预处理且cache好的数据等。

模型保存

模型的保存与加载，与单GPU的方式有所不同。这里通通将参数以cpu的方式save进存储, 因为如果是保存的GPU上参数，pth文件中会记录参数属于的GPU号，则加载时会加载到相应的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.

模型保存都是一致的，不过时刻记住方案二中你有多个进程在同时跑，所以会保存多个模型到存储上，如果使用共享存储就要注意文件名的问题，当然一般只在rank0进程上保存参数即可，因为所有进程的模型参数是同步的。

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  # Take care of distributed/parallel training
    model_to_save.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

同一台机器上跑多个 ddp task

假设想在一台有4核GPU的电脑上跑两个ddp task，每个task使用两个核，很可能会需要如下错误：

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 task通讯地址冲突，这时候需要显示地设置每个task的地址

specifying a different master_addr and master_port in torch.distributed.launch

# 第一个task
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

# 第二个task
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训练-单运算节点-All you need

WRITING DISTRIBUTED APPLICATIONS WITH PYTORCH

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

暂无评论

pytorch-multi-gpu-training

技术解读

深度解析 jia-zhuang/pytorch-multi-gpu-training 项目

项目背景与痛点

核心技术揭秘

技术架构与实现原理

核心代码片段

功能亮点与差异

应用场景与落地建议

场景

注意事项

综合评价

相关仓库

关于

简要分析

语言占比

贡献者

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

方案一

方案二

ddp有用的技巧

torch.distributed.barrier

模型保存

同一台机器上跑多个 ddp task

参考

评论