torch.nn.parallel.DistributedDataParallel使用

还好我不在意

已于 2023-05-17 16:21:04 修改

阅读量3.1k

点赞数 2

文章标签：深度学习 pytorch 神经网络

于 2023-05-11 22:50:33 首次发布

本文链接：https://blog.csdn.net/m0_46294481/article/details/130608283

版权

1. torch.nn.parallel.DistributedDataParallel简单介绍

Implements distributed data parallelism that is based on torch.distributed package at the module level.

这个container通过在每个模型副本上同步梯度来提供数据并行性。要同步的设备是由输入process_group指定的，默认情况下它是整个world。请注意，DistributedDataParallel不会在参与的GPU之间对输入进行分块或以其他方式进行分片；用户负责定义如何做到这一点，例如通过使用DistributedSampler。

DataParallel与DistributedDataParallel的比较：

`torch.nn.DataParallel`

DataParallel包以最低的编码障碍实现了单机多 GPU 并行处理。它只需要一行更改应用代码。需要注意的是，尽管DataParallel非常易于使用，但通常无法提供最佳性能。这是因为DataParallel的实现会在每个正向传播中复制该模型，并且其单进程多线程并行性自然会遭受 GIL 争用。

`torch.nn.parallel.DistributedDataParallel`

与DataParallel相比，DistributedDataParallel还需要设置一个步骤，即调用init_process_group。 DDP 使用多进程并行性，因此在模型副本之间没有 GIL 争用。此外，该模型是在 DDP 构建时而不是在每个正向传播时广播的，这也有助于加快训练速度。 DDP 附带了几种性能优化技术。

2. torch.distributed.init_process_group()介绍

Creation of this class requires that torch.distributed to be already initialized, by calling torch.distributed.init_process_group().

#创建此类（torch.nn.parallel.DistributedDataParallel）需要通过调用torch.distributed.init_process_group（）来初始化torch.distributed。

torch.distributed.init_process_group(backend=None, init_method=None, timeout=datetime.timedelta(seconds=1800), world_size=- 1, rank=- 1, store=None, group_name='', pg_options=None)

There are 2 main ways to initialize a process group:

Specify store, rank, and world_size explicitly.
Specify init_method (a URL string) which indicates where/how to discover peers. Optionally specify rank and world_size, or encode all required parameters in the URL and omit them.

If neither is specified, init_method is assumed to be “env://”.

参数：(process - 进程）

backend（str或Backend，可选）– 要使用的后端。根据构建时的配置，有效值包括mpi、gloo、nccl和ucc。如果没有提供后端，那么将创建gloo和nccl后端，请参阅下面的说明，了解如何管理多个后端。该参数可以用小写字符串（例如，“gloo”）表示，也可以通过Backend属性（例如，Backend.GLOO）访问。如果在每台机器上使用 nccl 后端进行多进程计算，则每个进程必须对其使用的每个GPU都具有独占访问权，因为进程之间共享GPU可能会导致死锁。ucc后端是实验性的。
init_method（str，可选）– 指定如何初始化进程组的URL。如果未指定init_method或store，则默认值为“env://”。与store相互排斥。
world_size（int，可选）– 参与工作的进程数。如果指定了store，则为必需。
rank（int，可选）– 当前进程的rank（它应该是一个介于0和world_size-1之间的数字）。如果指定了store，则为必需。
store（Store，可选）– 所有工作人员都可以访问的键/值存储，用于交换连接/地址信息。与init_method互斥。
timeout（timedelta，可选）– 针对进程组执行的操作超时。默认值等于30分钟。这适用于gloo后端。对于nccl，只有当环境变量NCCL_BLOCKING_WAIT或NCCL_ASYNC_ERROR_HANDLING设置为1时，这才适用。当设置了NCCL_BLOCKING_WAIT时，这是进程阻塞并等待集体完成的持续时间，如果时间超出，则会抛出一个异常。当设置了NCCL_ASYNC_ERROR_HANDLING时，这是一个持续时间，在该持续时间之后，集合将被异步中止并且该进程将崩溃。NCCL_BLOCKING_WAIT将向用户提供可以捕获和处理的错误信息，但由于其阻塞性，它有性能开销。另一方面，NCCL_ASYNC_ERROR_HANDLING的性能开销很小，但会因出错而导致进程崩溃。这样做是因为CUDA执行是异步的，如果出现失败的异步NCCL操作，后续的CUDA操作可能会在损坏的数据上运行，因此不能安全地继续执行用户代码。只应设置这两个环境变量中的一个。对于ucc，类似于NCCL，支持阻塞等待。然而，异步错误处理是不同的，因为使用UCC使用进度线程而不是看门狗线程
group_name（str，可选，已弃用）– 组名称。
pg_options（ProcessGroupOptions，可选）– 进程组选项，指定在构建特定进程组期间需要传入哪些附加选项。到目前为止，我们唯一支持的选项是：对于nccl后端是ProcessGroupNCCL.Options，可以指定is_high_priority_stream以便在有计算内核等待时，nccl后端可以拾取高优先级的cuda流。

简单介绍它的参数使用：

backend：一般来说，使用NCCL后端进行分布式GPU训练，使用Gloo后端进行分布式CPU训练。

world_size：可以简单理解为GPU卡的数量，比如你想在两台机器上使用9张卡运行你的模型，则需要将world_size设置为9。

在接下来的部分，我将详细介绍第二种初始化方式

首先介绍init_method：

Currently three initialization methods are supported:

TCP initialization

import torch.distributed as dist

# Use address of one of the machines
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',
                        rank=args.rank, world_size=4)

10.1.1.20是地址，23456为端口号。

Shared file-system initialization

import torch.distributed as dist

# rank should always be specified
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        world_size=4, rank=args.rank)

该初始化方法使用一个文件系统（该文件系统在一个组中的所有机器上都是共享的和可见的）以及所需的world_size。URL应以file://开头，并在共享文件系统上包含不存在的文件（在现有目录中）的路径。如果该文件不存在，文件系统初始化将自动创建该文件，但不会删除该文件。因此，您有责任确保在对同一文件路径/名称进行下一次init_process_group（）调用之前清除该文件。

此方法将始终创建文件，并在程序结束时尽最大努力清理和删除文件。换句话说，每次使用file init方法进行初始化都需要一个全新的空文件才能成功初始化。如果再次使用上一次初始化所使用的相同文件（碰巧没有被清理），这种意外的行为通常会导致死锁和故障。因此，即使这种方法会尽最大努力清理文件，但如果自动删除碰巧不成功，您有责任确保在训练结束时删除该文件，以防止下次再次使用同一文件。如果计划对同一文件名多次调用init_process_group（），这一点尤为重要。换句话说，如果文件没有被删除/清理，并且您对该文件再次调用init_process_group（），则会出现故障。这里的经验法则是，确保每次调用init_process_group（）时文件不存在或为空。

此方法假设文件系统支持使用fcntl进行锁定——大多数本地系统和NFS都支持它。

Environment variable initialization

import os
import torch.distributed as dist

# Set the environment variables
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12345'
os.environ['WORLD_SIZE'] = '4'
os.environ['RANK'] = '0'

# Initialize the process group
dist.init_process_group("gloo")

这种方法将从环境变量中读取配置，从而可以完全自定义信息的获取方式。要设置的变量包括：

MASTER_PORT - 必需；必须是rank为0的机器上的可用端口
MASTER_ADDR - 必需（除非rank为0）；rank 0 node的地址
WORLD_SIZE - 必需；可以在这里设置，也可以在对init函数的调用中设置
rank - 必需；可以在这里设置，也可以在对init函数的调用中设置

rank为0的机器将用于设置所有连接。
这是默认方法，意味着不必指定init_method（或者可以是env://）。

3. torch.nn.parallel.DistributedDataParallel详细介绍

要在具有N个GPU的主机上使用DistributedDataParallel，您应该生成N个进程，确保每个进程只在从0到N-1的单个GPU上工作。这可以通过为每个进程设置CUDA_VISIBLE_DEVICES来完成，也可以通过调用torch.cuda.set_device()（不建议使用该函数）。

import torch
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
# 必须在你使用PyTorch的代码之前运行，否则它不会生效。

print(torch.cuda.device_count())    
# 4

a = torch.rand((3, 5)).to(torch.device("cuda:4"))   
# RuntimeError: CUDA error: invalid device ordinal

os.environ["CUDA_VISIBLE_DEVICES"] = "3,2,0,1"
model = torch.nn.DataParallel(model, device_ids=[0,2,3]).cuda()

实际显卡编号----->运算显卡编号
    3     ----->     0（主卡）
    2     ----->     1
    0     ----->     2
    1     ----->     3

其中i是从0到N-1。在每个进程中，您应该参考以下内容来构建此模块：

torch.distributed.init_process_group(
    backend='nccl', world_size=N, init_method='...'
)
model = DistributedDataParallel(model, device_ids=[i], output_device=i)

为了在每个node上生成多个进程，您可以使用torch.distributed.launch或torch.multiprocessing.spawn。node（节点）指的是计算集群中的单个计算机或服务器。

torch.multiprocessing.spawn(fn, args=(), nprocs=1, join=True, daemon=False, start_method='spawn')

生成 nprocs 个进程来运行带有参数 args 的函数 fn。

如果其中一个进程以非零退出状态退出，则会终止其余进程，并引发包含终止原因的异常。如果在子进程中捕获到异常，则该异常会被forward，并将其traceback（回溯）包含在父进程中引发的异常中。

参数：

fn (function) – 函数被调用为生成进程的入口点。必须在模块的顶层定义此函数，以便对其进行pickle和spawn。这是多进程所提出的要求。
该函数被称为 fn(i, *args)，其中i是进程索引，args是参数的传递元组。
args (tuple) – Arguments passed to fn.
nprocs (int) – Number of processes to spawn.
join (bool) – Perform a blocking join on all processes.
daemon (bool) – The spawned processes’ daemon flag. If set to True, daemonic processes will be created.
start_method (str) – (deprecated)

返回值：

None if join is True, ProcessContext if join is False

Launch utility

torch.distributed软件包还在torch.distributed.launch中提供了一个启动实用程序。此助手实用程序可用于为每个节点启动多个进程，以进行分布式训练。

torch.distributed.launch是一个在每个训练节点上生成多个分布式训练进程的模块。

这个模块将被弃用，取而代之的是torchrun。

torchrun提供了作为torch.distributed.launch的功能的超集，并提供了以下附加功能：

通过重新启动所有工作程序，可以优雅地处理工作程序故障。
工作程序 RANK和WORLD_SIZE是自动分配的。
允许节点数量在最小和最大size之间变化（弹性）。

torchrun是setup.py中entry_points配置中声明的主模块torch.distributed.run的python控制台脚本（console script）。它相当于调用python -m torch.distributed.run。

如何从从torch.distributed.launch过渡到torchrun，请参考PyTorch的torchrun文档

使用

Single-node multi-worker

torchrun
    --standalone
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

Stacked single-node multi-worker

要在同一主机上运行单个节点、多个工作程序的多个实例（单独的作业），我们需要确保每个实例（作业）都设置在不同的端口上，以避免端口冲突（或者更糟的是，两个作业合并为一个作业）。要做到这一点，您必须使用--rdzv-backend=c10d运行，并通过设置--rdzv-endpoint=localhost:$PORT_k指定不同的端口。对于--nodes=1，让torchrun自动选择一个空闲的随机端口通常很方便，而不是手动为每次运行分配不同的端口。

torchrun
    --rdzv-backend=c10d
    --rdzv-endpoint=localhost:0
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

Fault tolerant (fixed sized number of workers, no elasticity, tolerates 3 failures)

torchrun
    --nnodes=$NUM_NODES
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

HOST_NODE_ADDR，格式为<host>[:<port>]（例如node1.example.com:29400），指定C10d 会合后端（rendezvous backend）的实例化和托管节点及其端口。它可以是训练集群中的任何节点，但理想情况下，您应该选择一个具有高带宽的节点。

如果未指定端口号，则HOST_NODE_ADDR默认为29400。

Elastic (min=1, max=4, tolerates up to 3 membership changes or failures)

torchrun
    --nnodes=1:4
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

HOST_NODE_ADDR，格式为<host>[:<port>]（例如node1.example.com:29400），指定C10d 会合后端的实例化和托管节点及其端口。它可以是训练集群中的任何节点，但理想情况下，您应该选择一个具有高带宽的节点。

如果未指定端口号，则HOST_NODE_ADDR默认为29400。

Note on rendezvous backend

对于多节点训练，您需要指定：

--rdzv-id: 唯一的作业id（由参与作业的所有节点共享）
--rdzv-backend: torch.distributed.elastic.rendezvous.RendezvousHandler的实现
--rdzv-endpoint: 会合后端正在运行的端点；通常采用host:port.的形式。

目前，c10d（推荐）、etcd-v2和etcd（传统）会合后端都支持开箱即用。要使用etcd-v2或etcd，请设置一个启用v2 api的etcd服务器（例如--enable-v2）。

etcd-v2和etcd会合使用etcd API v2。您必须在etcd服务器上启用v2 API。我们的测试使用etcd v3.4.3。

对于基于etcd的会合，我们建议使用etcd-v2而不是etcd，这在功能上是等效的，但使用了修改后的实现。etcd处于维护模式，将在未来版本中删除。

Definitions（worker group-工作组，job-作业）

Node - 物理实例或容器；映射到作业管理器所工作的单位。
Worker - A worker in the context of distributed training.
WorkerGroup - The set of workers that execute the same function (e.g. trainers).
LocalWorkerGroup - A subset of the workers in the worker group running on the same node.
RANK - The rank of the worker within a worker group.
WORLD_SIZE - The total number of workers in a worker group.
LOCAL_RANK - The rank of the worker within a local worker group.
LOCAL_WORLD_SIZE - The size of the local worker group.
rdzv_id - 用户定义的id，用于唯一标识作业的工作组。每个节点都使用此id作为特定工作组的成员加入。
rdzv_backend - rendezvous的后端（例如c10d）。这通常是一个强一致的键值存储。
rdzv_endpoint - 会合后端端点；通常采用<host>:<port>的形式。
一台节点运行 LOCAL_WORLD_SIZE 个worker，这些worker组成一个 LocalWorkerGroup。作业中所有节点的 LocalWorkerGroup 的并集组成 WorkerGroup。

Environment Variables

您可以在脚本中使用以下环境变量：

LOCAL_RANK - The local rank.
RANK - The global rank.
GROUP_RANK - The rank of the worker group. A number between 0 and max_nnodes. When running a single worker group per node, this is the rank of the node.
ROLE_RANK - The rank of the worker across all the workers that have the same role. The role of the worker is specified in the WorkerSpec.
LOCAL_WORLD_SIZE - The local world size (e.g. number of workers running locally); equals to --nproc-per-node specified on torchrun.
WORLD_SIZE - The world size (total number of workers in the job).
ROLE_WORLD_SIZE - The total number of workers that was launched with the same role specified in WorkerSpec.
MASTER_ADDR - The FQDN of the host that is running worker with rank 0; used to initialize the Torch Distributed backend.
MASTER_PORT - The port on the MASTER_ADDR that can be used to host the C10d TCP store.
TORCHELASTIC_RESTART_COUNT - The number of worker group restarts so far.
TORCHELASTIC_MAX_RESTARTS - The configured maximum number of restarts.
TORCHELASTIC_RUN_ID - Equal to the rendezvous run_id (e.g. unique job id).
PYTHON_EXEC - System executable override. If provided, the python user script will use the value of PYTHON_EXEC as executable. The sys.executable is used by default.

torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False, gradient_as_bucket_view=False, static_graph=False)

参数:

module（Module）– 要并行化的模块
device_ids（list of python：int or torch.device）– CUDA设备。1）对于单设备模块，device_ids只能包含一个设备id，它表示与该进程对应的输入模块所在的唯一CUDA设备。或者，device_ids也可以是None。2）对于多设备模块和CPU模块，device_ids必须为None。
当两种情况下的device_ids均为None时，前向传递的输入数据和实际模块都必须放置在正确的设备上。（默认值：None）
output_device（int或torch.device）– 单设备CUDA模块输出的设备位置。对于多设备模块和CPU模块，它必须为None，并且模块本身指示输出位置。（默认值：单设备模块的device_ids[0]）
broadcast_buffers（bool）– 在forward函数开始时启用模块同步（广播）缓冲区的标志。（默认值：True）
process_group – 用于分布式数据all-reduction的进程组。如果“None”，则将使用默认进程组，该组由torch.distributed.init_process_group（）创建。（默认值：None）#all-reduction是一种分布式计算中的通信操作，其目的是将多个进程或节点上的张量数据进行合并，得到一个最终的结果。
bucket_cap_mb – DistributedDataParallel会将参数存储到多个bucket中，以便将每个bucket的backward计算和梯度reduction过程潜在重叠，从而加速训练。bucket_cap_mb控制以兆字节（mb）为单位的bucket大小。（默认值：25）
find_unused_parameters（bool）– 遍历包含在封装模块的forward函数的返回值中的所有张量的autograd图。作为这个图的一部分没有收到梯度的参数会被预先标记为准备好可以被reduce。此外，可能已经在封装模块的forward函数中使用但不是loss计算的一部分并且因此也不会接收梯度的参数被预先标记为准备reduce。（默认值：False）
check_reduction – 此参数已被弃用。
gradient_as_bucket_view（bool）– 设置为True时，gradient将是指向allreduce通信bucket的不同偏移量的view。这可以减少峰值内存使用，其中保存的内存大小将等于总梯度大小。此外，它避免了在梯度和allreduce通信bucket之间进行复制的开销。当gradient是view时，不能在gradient上调用detach_（）。如果遇到这样的错误，请参考torch/optim/optimizer.py中的zero_grad（）函数作为解决方案进行修复。请注意，梯度将是第一次迭代后的view，因此应在第一次迭代之后检查峰值内存节省。
static_graph（bool）- 当设置为True时，DDP知道训练图是静态的。静态图意味着1）在整个训练循环中，已使用和未使用的参数集不会改变；在这种情况下，用户是否将find_unused_parameters设置为True并不重要。2）在整个训练循环中，图的训练方式不会改变（这意味着没有依赖于迭代的控制流）。当static_graph设置为True时，DDP将支持过去无法支持的情况：1）可重入backward。2）多次激活检查点。3）当模型具有未使用的参数时激活检查点。4）有些模型参数在forward函数之外。5）当存在未使用的参数时，可能会提高性能，因为当static_graph设置为True时，DDP不会在每次迭代中搜索图来检测未使用参数。要检查是否可以将static_graph设置为True，一种方法是在之前的模型训练结束时检查ddp日志记录数据，如果ddp_logging_data.get（“can_set_staticgraph”）==True，则大多数情况下也可以将static_graph设置为True。

例如：

model_DDP = torch.nn.parallel.DistributedDataParallel(model)
# Training loop
...
ddp_logging_data = model_DDP._get_ddp_logging_data()
static_graph = ddp_logging_data.get("can_set_static_graph")

4. torch.utils.data.distributed.DistributedSampler使用示例torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)
将数据加载限制到数据集的子集的采样器。
它与 torch.nn.parallel.DistributedDataParallel结合使用特别有用。在这种情况下，每个进程都可以传递一个DistributedSampler实例作为DataLoader采样器，并加载其专有的原始数据集的子集。

注意：

该数据集被假定为大小不变，并且它的任何实例总是以相同的顺序返回相同的元素。

参数：

dataset – Dataset used for sampling.
num_replicas (int, optional) – 参与分布式训练的进程数量。默认情况下，从当前分布式group中检索world_size。
rank (int, optional) – 当前进程在num_replicas中的rank。默认情况下，从当前分布式group中检索rank。
shuffle (bool, optional) – If True (default), sampler will shuffle the indices.
seed (int, optional) – random seed used to shuffle the sampler if shuffle=True. This number should be identical across all processes in the distributed group. Default: 0.
drop_last (bool, optional) – 如果为True，则采样器将丢弃数据的尾部，使其能够在replicas数量上被均匀分割。如果为False，采样器将添加额外的索引，以使数据在replicas之间可以被均匀分割。默认值：False。

警告：

在分布式模式中，在创建DataLoader迭代器之前，需要在每个epoch的开头调用set_epoch（）方法，以使shuffle在多个epoch之间正常工作。否则，将始终使用相同的顺序。

sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None),
                    sampler=sampler)
for epoch in range(start_epoch, n_epochs):
    if is_distributed:
        sampler.set_epoch(epoch)
    train(loader)

5. torch.nn.parallel.DistributedDataParallel使用示例

官方文档中的简单代码示例：

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP


os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,4,5"

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

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 demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()

def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             )

def demo_checkpoint(rank, world_size):
    print(f"Running DDP checkpoint example on rank {rank}.")
    setup(rank, world_size)

    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    CHECKPOINT_PATH = tempfile.gettempdir() + "/model.checkpoint"
    if rank == 0:
        # All processes should see same parameters as they all start from same
        # random parameters and gradients are synchronized in backward passes.
        # Therefore, saving it in one process is sufficient.
        torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)

    # Use a barrier() to make sure that process 1 loads the model after process
    # 0 saves it.
    dist.barrier()
    # configure map_location properly
    map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
    ddp_model.load_state_dict(
        torch.load(CHECKPOINT_PATH, map_location=map_location))

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn = nn.MSELoss()
    loss_fn(outputs, labels).backward()
    optimizer.step()

    # Not necessary to use a dist.barrier() to guard the file deletion below
    # as the AllReduce ops in the backward pass of DDP already served as
    # a synchronization.

    if rank == 0:
        os.remove(CHECKPOINT_PATH)

    cleanup()

class ToyMpModel(nn.Module):
    def __init__(self, dev0, dev1):
        super(ToyMpModel, self).__init__()
        self.dev0 = dev0
        self.dev1 = dev1
        self.net1 = torch.nn.Linear(10, 10).to(dev0)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to(dev1)

    def forward(self, x):
        x = x.to(self.dev0)
        x = self.relu(self.net1(x))
        x = x.to(self.dev1)
        return self.net2(x)



def demo_model_parallel(rank, world_size):
    print(f"Running DDP with model parallel example on rank {rank}.")
    setup(rank, world_size)

    # setup mp_model and devices for this process
    dev0 = rank * 2
    dev1 = rank * 2 + 1
    mp_model = ToyMpModel(dev0, dev1)
    ddp_mp_model = DDP(mp_model)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_mp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    # outputs will be on dev1
    outputs = ddp_mp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(dev1)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()

if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    if n_gpus < 5:
      print(f"Requires at least 5 GPUs to run, but got {n_gpus}.")
    else:
      run_demo(demo_basic, 5)
      run_demo(demo_checkpoint, 5)
      run_demo(demo_model_parallel, 2)

/home/root1/miniconda3/envs/TSW/bin/python /home/root1/data/TSW/DistributedDataParallel/main.py 
Running basic DDP example on rank 3.
Running basic DDP example on rank 1.
Running basic DDP example on rank 2.
Running basic DDP example on rank 4.
Running basic DDP example on rank 0.
Running DDP checkpoint example on rank 0.
Running DDP checkpoint example on rank 2.
Running DDP checkpoint example on rank 1.
Running DDP checkpoint example on rank 4.
Running DDP checkpoint example on rank 3.
Running DDP with model parallel example on rank 1.
Running DDP with model parallel example on rank 0.

进程已结束,退出代码0

单机多卡示例：

import torch
import torch.distributed as dist
import os
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torchvision.models as models
import torch.nn as nn
import torch.multiprocessing as mp
from tqdm import tqdm
import logging

from torch.nn.parallel import DistributedDataParallel as DDP

os.environ["CUDA_VISIBLE_DEVICES"] = "0,2,3,5"

# 定义数据变换
transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# 加载训练数据集
train_dataset = datasets.CIFAR10(root="cifar-10-batches-py", train=True, download=True,
                                 transform=transform)

# 加载resnet50模型
resmodel = models.resnet50()


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


def demo_basic(rank, world_size):
    # print(f"Running basic DDP example on rank {rank}.")
    # 在多进程中print函数可能会失效，建议使用logging模块来打印信息

    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)
    logger = logging.getLogger(__name__)

    setup(rank, world_size)

    # 定义数据加载器
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, sampler=train_sampler)

    # create model and move it to GPU with id rank
    model = resmodel.to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # 进行训练
    for epoch in range(5):
        train_sampler.set_epoch(epoch)
        loss_sum = 0
        for batch_idx, (data, target) in enumerate(tqdm(train_loader, desc=f"rank:{rank}")):
            data, target = data.to(rank), target.to(rank)
            optimizer.zero_grad()
            output = ddp_model(data)
            loss = loss_fn(output, target)
            loss_sum += loss
            loss.backward()
            optimizer.step()
        logger.info(f"rank: {rank}, epoch: {epoch}, loss: {loss_sum}")

    cleanup()


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             )


if __name__ == "__main__":
    run_demo(demo_basic, 4)
    print(len(train_dataset))
    print(len(train_dataset) / 64)
    print(len(train_dataset) / 64 / 4)

/home/root1/miniconda3/envs/TSW/bin/python /home/root1/data/TSW/DistributedDataParallel/test.py 
Files already downloaded and verified
Files already downloaded and verified
Files already downloaded and verified
Files already downloaded and verifiedFiles already downloaded and verified

2023-05-15 18:11:41,636 - INFO - Added key: store_based_barrier_key:1 to store for rank: 3
2023-05-15 18:11:41,641 - INFO - Added key: store_based_barrier_key:1 to store for rank: 2
2023-05-15 18:11:41,645 - INFO - Added key: store_based_barrier_key:1 to store for rank: 0
2023-05-15 18:11:41,646 - INFO - Added key: store_based_barrier_key:1 to store for rank: 1
2023-05-15 18:11:41,646 - INFO - Rank 1: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
2023-05-15 18:11:41,647 - INFO - Rank 3: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
2023-05-15 18:11:41,651 - INFO - Rank 2: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
2023-05-15 18:11:41,656 - INFO - Rank 0: Completed store-based barrier for key:store_based_barrier_key:1 with 4 nodes.
rank:1:   1%|          | 1/196 [00:16<52:32, 16.17s/it]2023-05-15 18:12:13,023 - INFO - Reducer buckets have been rebuilt in this iteration.
2023-05-15 18:12:13,042 - INFO - Reducer buckets have been rebuilt in this iteration.
2023-05-15 18:12:13,042 - INFO - Reducer buckets have been rebuilt in this iteration.
2023-05-15 18:12:13,047 - INFO - Reducer buckets have been rebuilt in this iteration.
rank:0: 100%|██████████| 196/196 [00:41<00:00,  4.77it/s]
rank:3: 100%|██████████| 196/196 [00:41<00:00,  4.78it/s]
rank:2: 100%|██████████| 196/196 [00:41<00:00,  4.78it/s]
rank:1: 100%|██████████| 196/196 [00:41<00:00,  4.78it/s]
2023-05-15 18:12:37,815 - INFO - rank: 2, epoch: 0, loss: 562.134765625
2023-05-15 18:12:37,815 - INFO - rank: 0, epoch: 0, loss: 562.4710083007812
2023-05-15 18:12:37,822 - INFO - rank: 1, epoch: 0, loss: 564.1324462890625
2023-05-15 18:12:37,822 - INFO - rank: 3, epoch: 0, loss: 561.6441650390625
rank:2: 100%|██████████| 196/196 [00:23<00:00,  8.20it/s]
rank:0: 100%|██████████| 196/196 [00:23<00:00,  8.20it/s]
rank:1: 100%|██████████| 196/196 [00:23<00:00,  8.20it/s]
rank:3: 100%|██████████| 196/196 [00:23<00:00,  8.20it/s]
2023-05-15 18:13:01,812 - INFO - rank: 2, epoch: 1, loss: 475.28961181640625
2023-05-15 18:13:01,812 - INFO - rank: 1, epoch: 1, loss: 469.89666748046875
2023-05-15 18:13:01,812 - INFO - rank: 0, epoch: 1, loss: 471.66217041015625
2023-05-15 18:13:01,819 - INFO - rank: 3, epoch: 1, loss: 471.5682373046875
rank:2: 100%|██████████| 196/196 [00:25<00:00,  7.58it/s]
rank:3: 100%|██████████| 196/196 [00:25<00:00,  7.59it/s]
rank:1: 100%|██████████| 196/196 [00:25<00:00,  7.58it/s]
rank:0: 100%|██████████| 196/196 [00:25<00:00,  7.58it/s]
2023-05-15 18:13:27,752 - INFO - rank: 2, epoch: 2, loss: 458.9241943359375
2023-05-15 18:13:27,752 - INFO - rank: 1, epoch: 2, loss: 459.09698486328125
2023-05-15 18:13:27,752 - INFO - rank: 0, epoch: 2, loss: 458.3254699707031
2023-05-15 18:13:27,759 - INFO - rank: 3, epoch: 2, loss: 457.5740661621094
rank:0: 100%|██████████| 196/196 [00:22<00:00,  8.68it/s]
rank:2: 100%|██████████| 196/196 [00:22<00:00,  8.68it/s]
rank:1: 100%|██████████| 196/196 [00:22<00:00,  8.68it/s]
rank:3: 100%|██████████| 196/196 [00:22<00:00,  8.68it/s]
2023-05-15 18:13:50,405 - INFO - rank: 2, epoch: 3, loss: 453.7042236328125
2023-05-15 18:13:50,405 - INFO - rank: 0, epoch: 3, loss: 453.55987548828125
2023-05-15 18:13:50,412 - INFO - rank: 1, epoch: 3, loss: 454.2198486328125
2023-05-15 18:13:50,412 - INFO - rank: 3, epoch: 3, loss: 451.8122253417969
rank:0: 100%|██████████| 196/196 [00:24<00:00,  8.10it/s]
rank:1: 100%|██████████| 196/196 [00:24<00:00,  8.10it/s]
rank:3: 100%|██████████| 196/196 [00:24<00:00,  8.10it/s]
rank:2: 100%|██████████| 196/196 [00:24<00:00,  8.10it/s]
2023-05-15 18:14:14,667 - INFO - rank: 2, epoch: 4, loss: 447.8721923828125
2023-05-15 18:14:14,668 - INFO - rank: 0, epoch: 4, loss: 448.20648193359375
2023-05-15 18:14:14,674 - INFO - rank: 1, epoch: 4, loss: 448.8049011230469
2023-05-15 18:14:14,675 - INFO - rank: 3, epoch: 4, loss: 447.05377197265625
50000
781.25
195.3125

进程已结束,退出代码0

如果要进行测试，可以使用dist.reduce()函数。

torch.distributed.reduce(tensor, dst, op=<RedOpType.SUM: 0>, group=None, async_op=False)[SOURCE]

Reduces the tensor data across all machines.

Only the process with rank dst is going to receive the final result.

Parameters:

tensor (Tensor) – Input and output of the collective. The function operates in-place.
dst (int) – Destination rank
op (optional) – One of the values from torch.distributed.ReduceOp enum. Specifies an operation used for element-wise reductions.
group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used.
async_op (bool, optional) – Whether this op should be an async op

Returns:

Async work handle, if async_op is set to True. None, if not async_op or if not part of the group

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "2,3,4,1"


def sum_tensor(rank, size, tensor):
    # 初始化进程组
    dist.init_process_group(backend='nccl', init_method='tcp://202.201.252.101:12345', rank=rank, world_size=size)

    tensor = tensor.to(rank)

    # 汇总操作
    dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)

    # 打印汇总结果
    if rank == 0:
        print("Sum of tensors:", tensor)

    # 清理进程组
    dist.destroy_process_group()


def main():
    # 设置进程数量
    world_size = 4

    # 创建随机张量
    tensor = torch.randn(2, 2)
    print(tensor)

    # 在多个进程上执行汇总操作
    mp.spawn(sum_tensor, args=(world_size, tensor), nprocs=world_size)


if __name__ == '__main__':
    main()

/home/root1/miniconda3/envs/TSW/bin/python /home/root1/data/TSW/DistributedDataParallel/test1.py 
tensor([[-0.0916, -0.7846],
        [-0.1634, -0.1967]])
Sum of tensors: tensor([[-0.3665, -3.1383],
        [-0.6537, -0.7870]], device='cuda:0')

进程已结束,退出代码0

还好我不在意

关注

2
点赞
踩
13

收藏

觉得还不错? 一键收藏
6
评论
torch.nn.parallel.DistributedDataParallel使用

1. torch.nn.parallel.DistributedDataParallel简单介绍Implements distributed data parallelism that is based on package at the module level.这个container通过在每个模型副本上同步梯度来提供数据并行性。要同步的设备是由输入process_group指定的，默认情况下它是整个world。请注意，DistributedDataParallel不会在参与的GPU之间对输入进行分块或
复制链接

扫一扫