了解Pytorch 分布式训练,这一篇足够了!
加入极市专业CV交流群,与6000+来自腾讯,华为,百度,北大,清华,中科院等名企名校视觉开发者互动交流!更有机会与李开复老师等大牛群内互动!
同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总,行业技术交流。关注 极市平台 公众号 ,回复 加群,立刻申请入群~
作者:
来源:
第一部分中,部分图片来自知乎提问部分,文中有链接,可以看更详细的讲解,侵删。
未经许可,严禁转载!!!
内容较多,整理的的有些乱,将就着看吧。
能力有限,欢迎指正批评。
参考
https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#
https://pytorch.org/tutorials/intermediate/dist_tuto.html#
https://pytorch.org/docs/master/nn.html#dataparallel-layers-multi-gpu-distributed
https://pytorch.org/docs/master/distributed.html
https://www.zhihu.com/question/57799212/answer/612786337
https://zhuanlan.zhihu.com/p/68717029
概述
torch.distributed 包支持
RingAllReduce VS TreeAllReduce
Pytorch 分布式使用流程
基本概念
group:
即进程组。默认情况下,只有一个组,一个 job 即为一个组,也即一个 world。
当需要进行更加精细的通信时,可以通过 new_group 接口,使用 word 的子集,创建新组,用于集体通信等。
world size :
表示全局进程个数。
rank:
表示进程序号,用于进程间通讯,表征进程优先级。rank = 0 的主机为 master 节点。
local_rank:
进程内,GPU 编号,非显式参数,由 torch.distributed.launch 内部指定。比方说, rank = 3,local_rank = 0 表示第 3 个进程内的第 1 块 GPU。
基本使用流程
在使用 distributed 包的任何其他函数之前,需要使用 init_process_group 初始化进程组,同时初始化 distributed 包。
如果需要进行小组内集体通信,用 new_group 创建子分组
创建分布式并行模型 DDP(model, device_ids=device_ids)
为数据集创建 Sampler
使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本,开始训练
使用 destory_process_group() 销毁进程组
使用模板
TCP 初始化方式
代码
import torch.distributed as dist
import torch.utils.data.distributed
# ......
parser = argparse.ArgumentParser(description='PyTorch distributed training on cifar-10')
parser.add_argument('--rank', default=0,
help='rank of current process')
parser.add_argument('--word_size', default=2,
help="word size")
parser.add_argument('--init_method', default='tcp://127.0.0.1:23456',
help="init-method")
args = parser.parse_args()
# ......
dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.word_size)
# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)
# ......
net = Net()
net = net.cuda()
net = torch.nn.parallel.DistributedDataParallel(net)
执行方式
# Node 1 : ip 192.168.1.201 port : 12345
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 0 --word_size 3
# Node 2 :
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 1 --word_size 3
# Node 3 :
python tcp_init.py --init_method tcp://192.168.1.201:12345 --rank 2 --
word_size 3
说明
在 TCP 方式中,在 init_process_group 中必须手动指定以下参数
rank 为当前进程的进程号
word_size 为当前 job 的总进程数
init_method 内指定 tcp 模式,且所有进程的 ip:port 必须一致,设定为主进程的 ip:port
必须在 rank==0 的进程内保存参数。
若程序内未根据 rank 设定当前进程使用的 GPUs,则默认使用全部 GPU,且以数据并行的方式使用。
每条命令表示一个进程。若已开启的进程未达到 word_size 的数量,则所有进程会一直等待
每台主机上可以开启多个进程。但是,若未为每个进程分配合适的 GPU,则同机不同进程可能会共用 GPU,应该坚决避免这种情况。
使用 gloo 后端进行 GPU 训练时,会报错。
若每个进程负责多块 GPU,可以利用多 GPU 进行模型并行。如下所示:
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)
......
dev0 = rank * 2
dev1 = rank * 2 + 1
mp_model = ToyMpModel(dev0, dev1)
ddp_mp_model = DDP(mp_model)
......
Env 初始化方式
代码
import torch.distributed as dist
import torch.utils.data.distributed
# ......
import argparse
parser = argparse.ArgumentParser()
# 注意这个参数,必须要以这种形式指定,即使代码中不使用。因为 launch 工具默认传递该参数
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
# ......
dist.init_process_group(backend='nccl', init_method='env://')
# ......
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=download, transform=transform)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, sampler=train_sampler)
# ......
# 根据 local_rank,配置当前进程使用的 GPU
net = Net()
device = torch.device('cuda', args.local_rank)
net = net.to(device)
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)
执行方式
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=0 --master_addr="192.168.1.201" --master_port=23456 env_init.py
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=1 --master_addr="192.168.1.201" --master_port=23456 env_init.py
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=3 --node_rank=2 --master_addr="192.168.1.201" --master_port=23456 env_init.py
说明
在 Env 方式中,在 init_process_group 中,无需指定任何参数
必须在 rank==0 的进程内保存参数。
该方式下,使用 torch.distributed.launch 在每台主机上,为其创建多进程,其中:
nproc_per_node 参数指定为当前主机创建的进程数。一般设定为当前主机的 GPU 数量
nnodes 参数指定当前 job 包含多少个节点
node_rank 指定当前节点的优先级
master_addr 和 master_port 分别指定 master 节点的 ip:port
若没有为每个进程合理分配 GPU,则默认使用当前主机上所有的 GPU。即使一台主机上有多个进程,也会共用 GPU。
使用 torch.distributed.launch 工具时,将会为当前主机创建 nproc_per_node 个进程,每个进程独立执行训练脚本。同时,它还会为每个进程分配一个 local_rank 参数,表示当前进程在当前主机上的编号。例如:rank=2, local_rank=0 表示第 3 个节点上的第 1 个进程。
需要合理利用 local_rank 参数,来合理分配本地的 GPU 资源
每条命令表示一个进程。若已开启的进程未达到 word_size 的数量,则所有进程会一直等待
进程组
初始化进程组
init_process_group
函数原型
torch.distributed.init_process_group(backend,
init_method=None,
timeout=datetime.timedelta(0, 1800),
world_size=-1,
rank=-1,
store=None)
函数作用
参数详解
backend :指定当前进程要使用的通信后端
小写字符串,支持的通信后端有 gloo,mpi,nccl 。建议用 nccl。
init_method : 指定当前进程组初始化方式
可选参数,字符串形式。如果未指定 init_method 及 store,则默认为 env://,表示使用读取环境变量的方式进行初始化。该参数与 store 互斥。
rank : 指定当前进程的优先级
int 值。表示当前进程的编号,即优先级。如果指定 store 参数,则必须指定该参数。
rank=0 的为主进程,即 master 节点。
world_size :
该 job 中的总进程数。如果指定 store 参数,则需要指定该参数。
timeout : 指定每个进程的超时时间
可选参数,datetime.timedelta 对象,默认为 30 分钟。该参数仅用于 Gloo 后端。
store
所有 worker 可访问的 key / value,用于交换连接 / 地址信息。与 init_method 互斥。
new_group
函数声明
torch.distributed.new_group(ranks=None,
timeout=datetime.timedelta(0, 1800),
backend=None)
函数作用
参数详解
ranks:指定新分组内的成员的 ranks 列表
list ,其中每个元素为 int 型
timeout:指定该分组进程组内的操作的超时时间
可选参数,datetime.timedelta 对象,默认为 30 分钟。该参数仅用于 Gloo 后端。
backend:指定要使用的通信后端
小写字符串,支持的通信后端有 gloo,nccl ,必须与 init_process_group() 中一致。
获取进程组属性
get_backend
torch.distributed.get_backend(group=<object>)
group:要获取信息的进程组。
给定进程组的后端,以小写字符串的形式给出
get_rank
torch.distributed.get_rank(group=<object>)
group
get_world_size
torch.distributed.get_world_size(group=<object>)
group
is_initialized
torch.distributed.is_initialized()
is_mpi_available
torch.distributed.is_mpi_available()
is_nccl_available
torch.distributed.is_nccl_available()
通信后端
概述
各种后端
Gloo 后端
NCCL 后端
MPI 后端
编译步骤
创建并激活 Anaconda 环境,安装 the guide 指定的依赖包,但是此时还不能运行 python setup.py install
选择并安装偏好的 MPI 实现。需要注意的是,开启 CUDA-aware MPI 可能需要一些额外的步骤。可以使用不提供 GPU 支持的 Open-MPI:conda install -c conda-forgeopenmpi
进入 Pytorch 源码,执行 python setup.py install
使用实例
# filename 'ptdist.py'
import torch
import torch.distributed as dist
def main(rank, world):
if rank == 0:
x = torch.tensor([1., -1.]) # Tensor of interest
dist.send(x, dst=1)
print('Rank-0 has sent the following tensor to Rank-1')
print(x)
else:
z = torch.tensor([0., 0.]) # A holder for recieving the tensor
dist.recv(z, src=0)
print('Rank-1 has recieved the following tensor from Rank-0')
print(z)
if __name__ == '__main__':
dist.init_process_group(backend='mpi')
main(dist.get_rank(), dist.get_world_size())
$ mpiexec -n 2 -ppn 1 -hosts miriad2a,miriad2b python ptdist.py
Rank-0 has sent the following tensor to Rank-1
tensor([ 1., -1.])
Rank-1 has recieved the following tensor from Rank-0
tensor([ 1., -1.])
如何选择
NCCL 是目前最快的后端,且对多进程分布式(Multi-Process Single-GPU)支持极好,可用于单节点以及多节点的分布式训练。
节点即主机。即使是单节点,由于底层机制不同,distributed 也比 DataParallel 方式要高效。
用 NCCL 进行分布式 GPU 训练
用 Gloo 进行分布式 CPU 训练
使用 NCCL,因为它是目前唯一支持 InfiniBand 和 GPUDirect 的后端
使用 NCCL,因为其目前提供最佳的分布式 GPU 训练性能。尤其是 multiprocess single-node 或 multi-node distributed 训练。
如果用 NCCL 训练有问题,再考虑使用 Cloo。(当前,Gloo 在 GPU 分布式上,相较于 NCCL 慢)
如果 InfiniBand 对 IB 启用 IP,请使用 Gloo,否则使使用 MPI。
在未来将添加 infiniBand 对 Gloo 的支持
使用 Gloo,除非有特别的原因使用 MPI。
初始化方式
TCP 初始化
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)
python mnsit.py --init-method tcp://192.168.54.179:22225 --rank 0 --world-size 2
python mnsit.py --init-method tcp://192.168.54.179:22225 --rank 1 --world-size 2
解析并验证参数
后端通过 name2channel.at() 函数进行解析,返回一个 channel 类,将用于执行数据传输
丢弃 GIL,并调用 THDProcessGroupInit() 函数,其实例化该 channel,并添加 master节点的地址
rank 0 对应的进程将会执行 master 过程,而其他的进程则作为 workers
master
为所有的 worker 创建 sockets
等待所有的 worker 连接
发送给他们所有其他进程的位置
每一个 worker
创建连接 master 的 sockets
发送自己的位置信息
接受其他 workers 的信息
打开一个新的 socket,并与其他 wokers 进行握手信号
初始化结束,所有的进程之间相互连接
共享文件系统初始化
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)
python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 0 --world-size 2
python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 1 --world-size 2
环境变量初始化
MASTER_PORT: 必须指定,表示 rank0上机器的一个空闲端口(必须设置)
MASTER_ADDR: 必须指定,除了 rank0 主机,表示主进程 rank0 机器的地址(必须设置)
WORLD_SIZE: 可选,总进程数,可以这里指定,在 init 函数中也可以指定
RANK: 可选,当前进程的 rank,也可以在 init 函数中指定
Node 1: (IP: 192.168.1.1, and has a free port: 1234)
>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
--nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
--master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
and all other arguments of your training script)
Node 2
>>> python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
--nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
--master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
and all other arguments of your training script)
Distributed Modules
DistributedDataParallel
原型
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)
功能
参数解析
module
要进行分布式并行的 module,一般为完整的 model
device_ids
int 列表或 torch.device 对象,用于指定要并行的设备。参考 DataParallel。
对于数据并行,即完整模型放置于一个 GPU 上(single-device module)时,需要提供该参数,表示将模型副本拷贝到哪些 GPU 上。
对于模型并行的情况,即一个模型,分散于多个 GPU 上的情况(multi-device module),以及 CPU 模型,该参数比必须为 None,或者为空列表。
与单机并行一样,输入数据及中间数据,必须放置于对应的,正确的 GPU 上。
output_device
int 或者 torch.device,参考 DataParallel。
对于 single-device 的模型,表示结果输出的位置。
对于 multi-device module 和 GPU 模型,该参数必须为 None 或空列表。
broadcast_buffers
bool 值,默认为 True
表示在 forward() 函数开始时,对模型的 buffer 进行同步 (broadcast)
process_group
对分布式数据(主要是梯度)进行 all-reduction 的进程组。
默认为 None,表示使用由 torch.distributed.init_process_group 创建的默认进程组 (process group)。
bucket_cap_mb
DistributedDataParallel will bucket parameters into multiple buckets so that gradient reduction of each bucket can potentially overlap with backward computation.
bucket_cap_mb controls the bucket size in MegaBytes (MB) (default: 25)
find_unused_parameters
bool 值。
Traverse the autograd graph of all tensors contained in the return value of the wrapped module’s forward function.
Parameters that don’t receive gradients as part of this graph are preemptively marked as being ready to be reduced.
Note that all forward outputs that are derived from module parameters must participate in calculating loss and later the gradient computation.
If they don’t, this wrapper will hang waiting for autograd to produce gradients for those parameters.
Any outputs derived from module parameters that are otherwise unused can be detached from the autograd graph using torch.Tensor.detach. (default: False)
check_reduction
when setting to True, it enables DistributedDataParallel to automatically check if the previous iteration’s backward reductions were successfully issued at the beginning of every iteration’s forward function.
You normally don’t need this option enabled unless you are observing weird behaviors such as different ranks are getting different gradients, which should not happen if DistributedDataParallel is correctly used. (default: False)
注意
要使用该 class,需要先对 torch.distributed 进行初进程组始化,可以通过 torch.distributed.init_process_group() 实现。
该 module 仅在 gloo 和 nccl 后端上可用。
根据分布式原理,Constructor 和 differentiation of the output (或 a function of the output of this module) 是一个分布式同步点。在不同的进程执行不同的代码时,需要考虑这一点。
该 module 假设,所有的参数在其创建时,在模型中已经注册,之后没有新的参数加入或者参数移除。对于 buffers 也是一样。(这也是由分布式原理决定)
该 module 假设,所有的参数在每个分布式进程模型中注册的顺序一致。该 module 自身将会按照该模型中参数注册的相反顺序执行梯度的 all-reduction。换言之,用户应该保证,每个分布式进程模型一样,且参数注册顺序一致。(这也是由分布式原理决定)
如果计划使用该 module,并用 NCCL 后端或 Gloo 后端 (使用 infiniband),需要与多 workers 的 Dataloader 一同使用,请修改多进程启动算法为 forkserver (python 3 only) 或 spawn 。不幸的是,Gloo (使用 infiniband) 和 NCCL2 fork 并不安全,并且如果不改变配置时,很可能会 deadlocks。
在 module 上定义的前向传播和反向传播 hooks 和其子 modules 将不会涉及,除非 hooks在 forward 中进行了初始化。
在使用 DistributedDataParallel 封装 model 后,不应该再修改模型的参数。也就是说,当使用 DistributedDataParallel 打包 model 时,DistributedDataParallel 的 constructor 将会在模型上注册额外的归约函数,该函数作用于模型的所有参数。
在进程之间,参数永远不会进行 broadcast。该 module 对梯度执行一个 all-reduce 步骤,并假设在所有进程中,可以被 optimizer 以相同的方式进行更改。在每一次迭代中,Buffers (BatchNorm stats 等) 是进行 broadcast 的,从 rank 0 的进程中的 module 进行广播,广播到系统的其他副本中。
使用
Single-Process Multi-GPU
Multi-Process Single-GPU
DistributedDataParallelCPU
原型
torch.nn.parallel.DistributedDataParallelCPU(module)
参数
module –
module to be parallelized
说明
在 module 级别上利用 CPU 实现分布式数据并行。
该 module 支持 mpi 和 gloo 后端。
该 container 通过在 batch 维度上,对输入进行分割,并分配到特定的设备上,实现模型的并行。将该 module 复制到每一台机器上,每一个副本处理输入的一部分。在反向传播阶段,每各节点的梯度求平均。
该 module 应该与 DistributedSampler 一起使用。DistributedSampler 将会为每个节点加载一个原始数据集的子集,每个子集的 batchsize 相同。因此,总 bs 的缩放如下所示。
n = 1, batch size = 8
n = 2, batch size = 16
n = 4, batch size = 32
n = 8, batch size = 64
警告
constructor, forward method 和 differentiation of the output (或 a function of the output of this module) 是一个分布式同步点。在不同节点可能执行不同代码的情况下,需要考虑这一点。
该 module 假设,所有的参数在其创建时,在模型中已经注册,之后没有新的参数加入或者参数移除。对于 buffers 也是一样。
该 module 假设所有的 buffers 和梯度都是密集型的。
在 module 上定义的前向传播和反向传播 hooks 和其子 modules 将不会涉及,除非 hooks在 forward 中进行了初始化。
参数在 __init__() 函数中,在不同节点之间进行 broadcast。该 module 在梯度上执行一个 all-reduce 步骤,并假设它们将会被 optimizer 在所有节点上以相同的方式进行更改。
DistributedSampler
原型
torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None)
参数
dataset
num_replicas
rank
说明
使用实例
# 分布式训练示例
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel
dataset = your_dataset()
datasampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, batch_size=batch_size_per_gpu, sampler=datasampler)
model = your_model()
model = DistributedDataPrallel(model, device_ids=[local_rank], output_device=local_rank)
通信方式
Point-to-Point Communication
基本概念
阻塞式 blocking
send
torch.distributed.send(tensor, dst, group=<object>, tag=0)
tensor
。tensor :要发送的
tensor
dst:目标
rank
,整数group:工作的进程组
tag :用于匹配当前
send
与远程recv
的标记
recv
torch.distributed.recv(tensor, src=None, group=<object>, tag=0)
tensor:要接收的 tensor
src:源 rank,如果未指定,则可以从任意进程接收数据
group:工作的进程组
tag:用于匹配当前 recv 与远程 send 的标记
Sender rank -1, if not part of the group
非阻塞式 Non-blocking
isend
torch.distributed.isend(tensor, dst, group=<object>, tag=0)
tensor :要发送的 tensor
dst:目标 rank,整数
group:工作的进程组
tag :用于匹配当前 send 与远程 recv 的标记
A distributed request object. None, if not part of the group
irecv
torch.distributed.irecv(tensor, src=None, group=<object>, tag=0)
tensor:要接收的 tensor
src:源 rank,如果未指定,则可以从任意进程接收数据
group:工作的进程组
tag:用于匹配当前 recv 与远程 send 的标记
A distributed request object. None, if not part of the group
实例
阻塞式
"""Blocking point-to-point communication."""
def run(rank, size):
tensor = torch.zeros(1)
if rank == 0:
tensor += 1
# Send the tensor to process 1
dist.send(tensor=tensor, dst=1)
else:
# Receive tensor from process 0
dist.recv(tensor=tensor, src=0)
print('Rank ', rank, ' has data ', tensor[0])
非阻塞式
is_completed()
wait()
"""Non-blocking point-to-point communication."""
def run(rank, size):
tensor = torch.zeros(1)
req = None
if rank == 0:
tensor += 1
# Send the tensor to process 1
req = dist.isend(tensor=tensor, dst=1)
print('Rank 0 started sending')
else:
# Receive tensor from process 0
req = dist.irecv(tensor=tensor, src=0)
print('Rank 1 started receiving')
req.wait()
print('Rank ', rank, ' has data ', tensor[0])
Collective Communication
集体通信的概念
同步操作(默认)
异步操作
is_completed() :判断是否执行完毕,若是则返回 True
wait():使用这个方法来阻塞这个进程,直到调用的 collective function 执行完毕
单 GPU 集体操作
broadcast
torch.distributed.broadcast(tensor, src, group=<object>, async_op=False)
tensor
src
group
async_op
scatter
torch.distributed.scatter(tensor, scatter_list, src, group=<object>, async_op=False)
tensor
到组内所有进程,注意与 broadcast
的区别。tensor
scatter_list
src
group
async_op
barrier
torch.distributed.barrier(group=<object>, async_op=False)
group
async_op
gather
torch.distributed.gather(tensor, gather_list, dst, group=<object>, async_op=False)
tensor
gather_list
dst
group
async_op
all_gather
torch.distributed.all_gather(tensor_list, tensor, group=<object>, async_op=False)
tensor_list
tensor
group
async_op
reduce
torch.distributed.reduce(tensor, dst, op=ReduceOp.SUM, group=<object>, async_op=False)
tensor
dst
op
torch.distributed.ReduceOp.SUM
torch.distributed.ReduceOp.PRODUCT
torch.distributed.ReduceOp.MIN
torch.distributed.ReduceOp.MAX
group
async_op
all_reduce
torch.distributed.torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=<object>, async_op=False)
tensor
op
torch.distributed.ReduceOp.SUM
torch.distributed.ReduceOp.PRODUCT
torch.distributed.ReduceOp.MIN
torch.distributed.ReduceOp.MAX
group
async_op
多 GPU 集体操作
说明
实例
import torch
import torch.distributed as dist
dist.init_process_group(backend="nccl",
init_method="file:///distributed_test",
world_size=2,
rank=0)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))
dist.all_reduce_multigpu(tensor_list)
import torch
import torch.distributed as dist
dist.init_process_group(backend="nccl",
init_method="file:///distributed_test",
world_size=2,
rank=1)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))
dist.all_reduce_multigpu(tensor_list)
broadcast_multigpu
torch.distributed.broadcast_multigpu(tensor_list, src, group=<object>, async_op=False, src_tensor=0)
tensor_list
src
group
async_op
src_tensor
all_gather_multigpu
torch.distributed.all_gather_multigpu(output_tensor_lists, input_tensor_list, group=<object>, async_op=False)
output_tensor_lists
input_tensor_list
group
async_op
reduce_multigpu
torch.distributed.reduce_multigpu(tensor_list, dst, op=ReduceOp.SUM, group=<object>, async_op=False, dst_tensor=0)
tensor_list
dst
op
torch.distributed.ReduceOp.SUM
torch.distributed.ReduceOp.PRODUCT
torch.distributed.ReduceOp.MIN
torch.distributed.ReduceOp.MAX
group
async_op
dst_tensor
all_reduce_multigpu
torch.distributed.all_reduce_multigpu(tensor_list, op=ReduceOp.SUM, group=<object>, async_op=False)
tensor_list
op
torch.distributed.ReduceOp.SUM
torch.distributed.ReduceOp.PRODUCT
torch.distributed.ReduceOp.MIN
torch.distributed.ReduceOp.MAX
group
async_op
启动工具 Launch utility
概述
参数
training_script
--nnodes
--node_rank
--nproc_per_node
--master_addr
--master_port
master
节点使用的端口号,必须与其他应用的端口号不冲突。使用
单节点多进程分布式训练
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
多节点多进程分布式
192.168.1.1:1234
)python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=1 --master_addr="192.168.1.1" --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
查看帮助
python -m torch.distributed.launch --help
CV细分方向交流群
添加极市小助手微信(ID : cv-mart),备注:研究方向-姓名-学校/公司-城市(如:目标检测-小极-北大-深圳),即可申请加入目标检测、目标跟踪、人脸、工业检测、医学影像、三维&SLAM、图像分割等极市技术交流群(已经添加小助手的好友直接私信),更有每月大咖直播分享、真实项目需求对接、干货资讯汇总,行业技术交流,一起来让思想之光照的更远吧~
△长按添加极市小助手
△长按关注极市平台
觉得有用麻烦给个在看啦~