PyTorch分布式DPP启动方式(包含完整用例)

pytorch 专栏收录该内容
2 篇文章 0 订阅


本篇主要讲解单卡到分布式中DDP(DistributeDataParallel )的使用基础,包括如何使用DDP和相关的一些基础问题。
主要涵盖如下问题:

1、单卡到分布式的过程中,用户代码需要怎么调整?
单卡变成分布式中处理过程差异,以及用户要干什么。
2、DDP的启动方式有几种? 差异在哪?
DDP的启动方式与代码修改如何对应,单进程单卡、单进程多卡、单卡多进程的模式如何启动。

选择分布式前用户首先需要明白自己的场景,是需要数据并行与还是需要考虑模型并行。因为DDP指的是数据并行的分布式,一般适用于单张GPU能够加载一个完全的模型,这一点在PyTorch的-DDP官网文档有说明。明确这点后,再继续往下看:

1 基本使用

直接放一个简单的列子来看一下,单卡到多卡用户的代码如何修改。这里采用一个简单的MNIST分类例子来分析,从单卡代码到分布式变动了哪些地方。

1.1 单卡版本

import torch
import torchvision
import torch.utils.data.distributed
from torchvision import transforms


def main():
    # 数据加载部分,直接利用torchvision中的datasets
    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
    data_set = torchvision.datasets.MNIST("./", train=True, transform=trans, target_transform=None, download=True)
    data_loader_train = torch.utils.data.DataLoader(dataset=data_set, batch_size=256)
    # 网络搭建,调用torchvision中的resnet
    net = torchvision.models.resnet101(num_classes=10)
    net.conv1 = torch.nn.Conv1d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
    net = net.cuda()
    # 定义loss与opt
    criterion = torch.nn.CrossEntropyLoss()
    opt = torch.optim.Adam(net.parameters(), lr=0.001)
    # 网络训练	
    for epoch in range(10):
        for i, data in enumerate(data_loader_train):
            images, labels = data
            images = torch.tensor(images).cuda()
            labels = torch.tensor(labels).cuda()
            opt.zero_grad()
            outputs = net(images)
            loss = criterion(outputs, labels)
            loss.backward()
            opt.step()
            if i % 10 == 0:
                print("loss: {}".format(loss.item()))
    # 保存checkpoint
    torch.save(net, "my_net.pth")


if __name__ == "__main__":
    main()

1.2 多卡分布式

import torch
import torchvision
import torch.utils.data.distributed
from torchvision import transforms


def main():
    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
    data_set = torchvision.datasets.MNIST("./", train=True, transform=trans, target_transform=None, download=True)
    data_loader_train = torch.utils.data.DataLoader(dataset=data_set, batch_size=256)

    net = torchvision.models.resnet101(num_classes=10)
    net.conv1 = torch.nn.Conv1d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
    net = net.cuda()
    criterion = torch.nn.CrossEntropyLoss()
    opt = torch.optim.Adam(net.parameters(), lr=0.001)
    for epoch in range(10):
        for i, data in enumerate(data_loader_train):
            images, labels = data
            images = torch.tensor(images).cuda()
            labels = torch.tensor(labels).cuda()
            opt.zero_grad()
            outputs = net(images)
            loss = criterion(outputs, labels)
            loss.backward()
            opt.step()
            if i % 10 == 0:
                print("loss: {}".format(loss.item()))
    torch.save(net, "my_net.pth")

用Comparer对比,如下图所示:
在这里插入图片描述

可以看出,主要变动的位置包括:

  1. 启动的方式引入了一个多进程机制;
  2. 引入了几个环境变量;
  3. DataLoader多了一个sampler参数;
  4. 网络被一个DistributedDataParallel(net)又包裹了一层;
  5. ckpt的保存方式发生了变化。

这些修改都是为了适配数据流动过程的变化,来对比一下单卡与分布式的过程中的数据流的变化:
在这里插入图片描述

从图片的对比可知,分布式使得数据的切分(batch_data)以及前向计算后的数据传递发生了变化(all_reduce),对应到ddp里面:
1、由于分布式采用的是多进程模式,要保证不同进程拿到的是不同的数据,需要在正向传播时对数据的分配进行调整,所以dataloader里面多了一个sampler参数。
2、反向传播计算后需要对参数进行共享通信,所以多了一个allreduce的处理。
还有一些需要处理的地方,比如大部分情况下,分布式的进程之间执行的都是相同的代码,但是有些只需处理一次的工作,比如ckpt的保存,就需要指定单个进程来完成。 于是需要对进程进行编号,指定特定编号的进程去完成特定的工作,这个编号是rank。
单机多卡、多机模式如何启动? 继续看:

2 分布式用例

DDP的启动方式形式上有多种,内容上是统一的:都是启动多进程来完成运算。
先来看一般单机多卡的形式:

2.1 单机多卡

单机多卡,可以直接用Process启动,也可以用torch.multiprocessing.spawn,还可以用launch启动(多机启动中会介绍)。这里主要介绍前两种,在前面提到MNIST用例中用的是Process格式,摘取启动的位置:

# Process格式:
if __name__ == "__main__":
    world_size= 3
    processes = []
    # 创建进程组
    for rank in range(world_size):
        p = Process(target=main, args=(rank, world_size))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

spawn本质上就是简化了Process的抒写,其spawn格式如下:

# spawn格式:
def main():
    world_size= 3
    mp.spawn(main,
        args=(world_size,),
        nprocs=world_size,
        join=True)

来一份采用spawn的完整用例

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'

def example(rank, world_size):
    # 初始化
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # 创建模型
    model = nn.Linear(10, 10).to(rank)
    # 放入DDP
    ddp_model = DDP(model, device_ids=[rank]) 
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    # 进行前向后向计算
    for i in range(1000):
        outputs = ddp_model(torch.randn(20, 10).to(rank))
        labels = torch.randn(20, 10).to(rank)
        loss_fn(outputs, labels).backward()
        optimizer.step()

def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    main()

2.2 多机分布式

多机的启动方式可以是直接传递参数并在代码内部解析环境变量,或者通过torch.distributed.launch来启动,两者在格式上面有一定的区别,总之要保证代码与启动方式对应。

2.2.1 方式一:每个进程占用一张卡

首先,看一下如何实现多机的一个进程占用一张卡的使用,需要注意的位置:

  1. dist.init_process_group里面的rank需要根据node以及GPU的数量计算;
  2. world_size的大小=节点数 x GPU 数量。
  3. ddp 里面的device_ids需要指定对应显卡。
    示例代码: demo.py
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--world_size", type=int)
parser.add_argument("--node_rank", type=int)
parser.add_argument("--master_addr", default="127.0.0.1", type=str)
parser.add_argument("--master_port", default="12355", type=str)
args = parser.parse_args()


def example(local_rank, node_rank, local_size, world_size):
    # 初始化
    rank = local_rank + node_rank * local_size
    dist.init_process_group("nccl",
                            init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
                            rank=rank,
                            world_size=world_size)
    # 创建模型
    model = nn.Linear(10, 10).to(rank)
    # 放入DDP
    ddp_model = DDP(model, device_ids=[local_rank],output_device=local_rank) 
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    # 进行前向后向计算
    for i in range(1000):
        outputs = ddp_model(torch.randn(20, 10).to(rank))
        labels = torch.randn(20, 10).to(rank)
        loss_fn(outputs, labels).backward()
        optimizer.step()


def main():
    local_size = torch.cuda.device_count()
    print("local_size: %s" % local_size)
    mp.spawn(example,
        args=(args.node_rank, local_size, args.world_size,),
        nprocs=local_size,
        join=True)


if __name__=="__main__":
    main()

启动方式:
假设一共有两台机器(节点1和节点2),每个节点上有8张卡,节点1的IP地址为127.0.0.1 占用的端口22335(端口可以更换),启动的方式如下:

>>> #节点1
>>>python python demo.py --world_size=16 --node_rank=0 --master_addr="127.0.0.1" --master_port=22335
>>> #节点2
>>>python python demo.py --world_size=16 --node_rank=2 --master_addr="127.0.0.1" --master_port=22335

2.2.2 方式二:单个进程占用多张卡

进一步,看单进程多卡的代码和启动方式,代码中需要注意的位置:

  1. dist.init_process_group里面的rank等于节点编号;
  2. world_size等于节点的总数量;
  3. DDP不需要指定device。
    示例代码
import torchvision
from torchvision import transforms
import torch.distributed as dist
import torch.utils.data.distributed
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--rank", default=0, type=int)
parser.add_argument("--world_size", default=1, type=int)
parser.add_argument("--master_addr", default="127.0.0.1", type=str)
parser.add_argument("--master_port", default="12355", type=str)
args = parser.parse_args()


def main(rank, world_size):
    # 一个节点就一个rank,节点的数量等于world_size
    dist.init_process_group("gloo",
                            init_method="tcp://{}:{}".format(args.master_addr, args.master_port),
                            rank=rank,
                            world_size=world_size)
    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
    data_set = torchvision.datasets.MNIST('~/DATA/', train=True,
                                          transform=trans, target_transform=None, download=True)
    train_sampler = torch.utils.data.distributed.DistributedSampler(data_set)
    data_loader_train = torch.utils.data.DataLoader(dataset=data_set,
                                                    batch_size=256,
                                                    sampler=train_sampler,
                                                    num_workers=16,
                                                    pin_memory=True)
    net = torchvision.models.resnet101(num_classes=10)
    net.conv1 = torch.nn.Conv1d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
    net = net.cuda()
    # net中不需要指定设备!
    net = torch.nn.parallel.DistributedDataParallel(net)
    criterion = torch.nn.CrossEntropyLoss()
    opt = torch.optim.Adam(net.parameters(), lr=0.001)
    for epoch in range(1):
        for i, data in enumerate(data_loader_train):
            images, labels = data
            images, labels = images.cuda(), labels.cuda()
            opt.zero_grad()
            outputs = net(images)
            loss = criterion(outputs, labels)
            loss.backward()
            opt.step()
            if i % 10 == 0:
                print("loss: {}".format(loss.item()))


if __name__ == '__main__':
    main(args.rank, args.world_size)
    

启动方式

>>> #节点1
>>>python python demo.py --world_size=16 --node_rank=0 --master_addr="127.0.0.1" --master_port=22335
>>> #节点2
>>>python python demo.py --world_size=16 --node_rank=2 --master_addr="127.0.0.1" --master_port=22335

2.2.3 方式三:利用launch

从torch.distributed.launch的源码,可以看出launch实际上主要完成的工作:
1、参数定义与传递。解析环境变量,并将变量传递到子进程中。
2、起多进程。调用subprocess.Popen启动多进程。
用launch方式需要注意的位置:
需要添加一个解析 local_rank的参数:
parser.add_argument("–local_rank", type=int)
dist初始化的方式 int_method取env:
dist.init_process_group(“gloo”, init_method=‘env://’)
DDP的设备都需要指定local_rank
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)
示例代码: MNIST.py

import torch
import torchvision
import torch.utils.data.distributed
import argparse
import torch.distributed as dist
from torchvision import transforms

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)  # 增加local_rank
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)

def main():
    dist.init_process_group("nccl", init_method='env://')    # init_method方式修改
    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
    data_set = torchvision.datasets.MNIST('~/DATA/', train=True,
                                          transform=trans, target_transform=None, download=True)
    data_loader_train = torch.utils.data.DataLoader(dataset=data_set,
                                                    batch_size=256,
                                                    sampler=train_sampler,
                                                    num_workers=16,
                                                    pin_memory=True)
    net = torchvision.models.resnet101(num_classes=10)
    net.conv1 = torch.nn.Conv1d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
    net = net.cuda()
    # DDP 输出方式修改:
    net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank],
                                                    output_device=args.local_rank)
    criterion = torch.nn.CrossEntropyLoss()
    opt = torch.optim.Adam(net.parameters(), lr=0.001)
    for epoch in range(1):
        for i, data in enumerate(data_loader_train):
            images, labels = data 
            # 要将数据送入指定的对应的gpu中
            images.to(args.local_rank, non_blocking=True)
            labels.to(args.local_rank, non_blocking=True)
            opt.zero_grad()
            outputs = net(images)
            loss = criterion(outputs, labels)
            loss.backward()
            opt.step()
            if i % 10 == 0:
                print("loss: {}".format(loss.item()))


if __name__ == "__main__":
    main()
    

启动方式: 假设一共有两台机器(节点1和节点2),每个节点上有8张卡,节点1的IP地址为192.168.1.1 占用的端口12355(端口可以更换),启动的方式如下:

>>> #节点1
>>>python -m torch.distributed.launch --nproc_per_node=8
           --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
           --master_port=12355 MNIST.py
>>> #节点2
>>>python -m torch.distributed.launch --nproc_per_node=8
           --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
           --master_port=12355 MNIST.py

上例子中,如果只要启动一台机器,只需要将nnodes数量改为1,就是单机多卡的另一种方式。设置如下:

>>> #节点1
>>>python -m torch.distributed.launch --nproc_per_node=8
           --nnodes=1 --node_rank=0 --master_addr="192.168.1.1"
           --master_port=12355 MNIST.py

原文首发于:PyTorch分布式训练基础–DDP使用
参考:
https://developer.nvidia.com/nccl
https://pytorch.apachecn.org/docs
https://pytorch.org/docs/stable/distributed.html#launch-utility
https://pytorch.org/docs/master/notes/ddp.html
https://discuss.pytorch.org/t/cuda-error-out-of-memory-when-load-models/38011

  • 3
    点赞
  • 3
    评论
  • 6
    收藏
  • 打赏
    打赏
  • 扫一扫,分享海报

©️2021 CSDN 皮肤主题: 数字20 设计师:CSDN官方博客 返回首页

打赏作者

9eKY

你的鼓励将是我创作的最大动力

¥2 ¥4 ¥6 ¥10 ¥20
输入1-500的整数
余额支付 (余额:-- )
扫码支付
扫码支付:¥2
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值