pytorch 1.0 分布式

2024年1月12日 207次阅读 来源: lbin

Pytorch 1.0 分布式

美好的开始

1.0

PyTorch 1.0稳定版终于正式发布了！新版本增加了JIT编译器、全新的分布式包、C++ 前端，以及Torch Hub等新功能，支持AWS、谷歌云、微软Azure等云平台。

DISTRIBUTED NEWS

torch.distributed软件包和torch.nn.parallel.DistributedDataParallel模块由全新的、重新设计的分布式库提供支持。新的库的主要亮点有：

  • 新的 torch.distributed 是性能驱动的，并且对所有后端 (Gloo，NCCL 和 MPI) 完全异步操作
  • 显着的分布式数据并行性能改进，尤其适用于网络较慢的主机，如基于以太网的主机
  • 为torch.distributedpackage中的所有分布式集合操作添加异步支持
  • 在Gloo后端添加以下CPU操作：send，recv，reduce，all_gather，gather，scatter
  • 在NCCL后端添加barrier操作
  • 为NCCL后端添加new_group支持

WHAT

1.0的多机多卡的计算模型并没有采用主流的Parameter Server结构，而是直接用了Uber Horovod的形式，也是百度开源的RingAllReduce算法。采用PS计算模型的分布式，通常会遇到网络的问题，随着worker数量的增加，其加速比会迅速的恶化，例如resnet50这样的模型，目前的TF在10几台机器的时候，加速比已经开始恶化的不可接受了。因此，经常要上RDMA、InfiniBand等技术，并且还带来了一波网卡的升级，有些大厂直接上了100GBs的网卡，有钱任性。而Uber的Horovod，采用的RingAllReduce的计算方案，其特点是网络通信量不随着worker（GPU）的增加而增加，是一个恒定值。简单看下图理解下，GPU 集群被组织成一个逻辑环，每个GPU有一个左邻居、一个右邻居，每个GPU只从左邻居接受数据、并发送数据给右邻居。即每次梯度每个gpu只获得部分梯度更新，等一个完整的Ring完成，每个GPU都获得了完整的参数。

HOW

torch.distributed包提供了一个启动实用程序torch.distributed.launch，此帮助程序可用于为每个节点启动多个进程以进行分布式训练，它在每个训练节点上产生多个分布式训练进程。

该工具既可以用来做单节点多GPU训练，也可用于多节点多GPU训练。如果是单节点多GPU，将会在单个GPU上运行一个分布式进程，据称可以非常好地改进单节点训练性能。如果用于多节点分布式训练，则通过在每个节点上产生多个进程来获得更好的多节点分布式训练性能。如果有Infiniband接口则加速比会更高。

在单节点分布式训练或多节点分布式训练的两种情况下，该工具将为每个节点启动给定数量的进程（–nproc_per_node）。如果用于GPU培训，则此数字需要小于或等于当前系统上的GPU数量（nproc_per_node），并且每个进程将在从GPU 0到GPU（nproc_per_node – 1）的单个GPU上运行。

  1. Single-Node multi-process distributed training
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)
  1. Multi-Node multi-process distributed training: (e.g. two nodes)

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)

需要注意的地方：

  • 后端最好用“NCCL”，才能获取最好的分布式性能
  • 训练代码必须从命令行解析–local_rank=LOCAL_PROCESS_RANK
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

torch.cuda.set_device(arg.local_rank)
  • torch.distributed初始化方式
torch.distributed.init_process_group(backend='nccl',
                                     init_method='env://')
  • model
model = torch.nn.parallel.DistributedDataParallel(model,
                                                  device_ids=[arg.local_rank],
                                                  output_device=arg.local_rank)

其他地方一般就不用修改了，真的顺滑了非常多

PERFERMENCE

简单测了下MaskRCNN-R50FPN，32GPU用时3hours15mins，加速接近4倍。

    原文作者：lbin
    原文地址: https://zhuanlan.zhihu.com/p/52110617
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞