Pytorch并行分布式训练四（DDP）

最新推荐文章于 2026-04-24 09:03:03 发布

原创

最新推荐文章于 2026-04-24 09:03:03 发布 · 3.1k 阅读

标签

#pytorch #人工智能

本文介绍了PyTorch中的DistributedDataParallel（DDP）模块，它是多GPU训练的主要方法，支持单机多卡和多机多卡场景。相较于DataParallel，DDP在多GPU环境下表现更优。文章通过对比DP和DDP，展示了DDP的启动方式，包括launch和mp模块。并提供了DDP的使用示例，包括数据加载、模型保存与加载、模型并行的结合。强调了DDP在保存模型时应确保所有进程完成后再加载，以及如何与模型并行配合使用。

Getting Started with Distributed Data Parallelhttps://pytorch.org/tutorials/intermediate/ddp_tutorial.html

DDP是目前Pytorch推荐的多GPU训练方法，它支持单机多卡，多机多卡等情况。目前pytorch对DDP技术中进程的启动方式又分为launch脚本启动，和mp模启动。就目前而言，pytorch更推荐采用mp的方法，但launch的方法使用的也很广，所以下面一并做一些介绍。

1.DataParallel 和 DistributedDataParallel之间的比较

首先，DP是单进程多线程，只可以在单机中工作。DDP是多进程的，在单机或多机情况下都可以工作。
尽管在单机工作的情况下，DP也要比DDP更慢。这是由于DP受到了GIL锁、每次前向传播复制模型、和额外的输入散射输出收集、等因素的限制。
从之前的文章可以知道，当你的模型太大以至于单卡无法加载时，你必须借助“模型并行”来解决。DDP至此结合“模型并行”，但DP却不支持。注意：当DDP结合“模型并行时”，每一个DDP进程将使用“模型并行”，所有的进程间总体使用“数据并行”。

2.一个Launch示例

此方法借助于torch.distributed.launch模块，使用得交广泛，但现在pytorch官方更推荐于mp方法（见下节）

import os

import argparse

import torch

import torch.distributed as dist

import torchvision

import torch.nn as nn

import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

   

class ToyModel(nn.Module):

    def __init__(self):

        super(ToyModel, self).__init__()

        self.net1 = nn.Linear(10, 10)

        self.relu = nn.ReLU()

        self.net2 = nn.Linear(10, 5)

    def forward(self, x):

        return self.net2(self.relu(self.net1(x)))

   

def setup(rank, world_size):

    # 注：设置默认使用的GPU设备，当遇到CUDA error: an illegal memory access was encountered

    #     请检查这块是否有设置

    torch.cuda.set_device(rank)

   

    os.environ['MASTER_ADDR'] = 'localhost'

    os.environ['MASTER_PORT'] = '12355'

   

    # 注：第一个参数选择后端，nccl后端是单机多卡情况下的推荐，比gloo快很多。

    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():

    dist.destroy_process_group()

def demo_basic(rank):

    print(f"Running basic DDP example on rank {rank}.")

   

    trainset = torchvision.datasets.CIFAR10(...)

    # 注：数据需要用torch.utils.data.distributed.DistributedSampler处理

    train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)

    trainloader = torch.utils.data.DataLoader(trainset,

                                              batch_s