pytorch 分布式训练
如何利用多 GPU、多服务器分布式训练深度学习模型,本篇介绍 pytorch 分布式训练相关。
一些概念
- node 物理节点,即服务器、机器。注意每一个节点可以有多个 GPU,即一机多卡
- rank 进程序号,常用于进程通信。每个进程对应一个 rank
- local_rank 相对于 rank (整个分布式集群而言),local_rank 表示一台机器内部的相对进程序号。rank 和 locak_rank 独立,当是单机多卡时,rank 等同于 local_rank
- nnodes 物理节点个数
- node_rank 物理节点序号
- nproc_per_node 每个物理节点上面运行的进程数,常对应想要运行的 GPU 个数,因此取值小于物理节点上 GPU 个数。
- group 进程组,默认只有一个组
- world size 全局并行数,即一个分布式任务中,全局 rank 的数量。
例子:
每个 node 包含 8 个 GPU,且 nproc_per_node=4, nnodes=5,机器的 node_rank=10,请问 world_size 是多少?
全局并行数 world size 等于真实的进程数,即 nproc_per_node x nnodes = 4 x 5 = 20,其他没有使用的 GPU 是没有参与计算,不算并行数。
pytorch 分布式特点
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支持
torch.distributed 包提供了一个启动实用程序 torch.distributed.launch,此帮助程序可用于为每个节点启动多个进程以进行分布式训练,它在每个训练节点上产生多个分布式训练进程。
单节点多 GPU
1 | # 这里 NUM_GPUS_YOU_HAVE 为你服务器上 GPU 的个数,建议使用同型号 GPU 进行并行计算 |
多节点多 GPU
以两个节点为例,更多节点可以类推。假设节点1(IP: 192.168.1.100)为主节点,节点2(IP: 192.168.1.200)为从节点。
在节点1上:
1 | # 因为是2个节点,所以 nnodes = 2 |
在节点2上:
1 | # 在每个从节点上,指定主节点IP地址和通信端口号 |
与单节点单GPU的区别:
- 后端最好用 NCCL,才能获取最好的分布式性能
- 训练代码必须从命令行解析
--local_rank=LOCAL_PROCESS_RANK
1
2
3
4
5
6import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
torch.cuda.set_device(arg.local_rank) - 分布式初始化
1
torch.distributed.init_process_group(backend='nccl', init_method='env://')
- 模型分布式
1
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[arg.local_rank], output_device=arg.local_rank)
并行
深度模型训练加速建议:
- 当模型不大时(可以加载到一台服务器上时),使用数据并行;
- 当模型较大时,可将模型分解放在多个服务器上运行,即模型并行;
参考文献
本博客所有文章除特别声明外,均采用 CC BY-NC-SA 4.0 许可协议。转载请注明来自 J. Xu!
评论