Pytorc分布式训练

当模型训练过程显存占用较大无法在单个gpu上完成的时候，必须使用分布式训练（这个结论太武断，实际上还有一种思路，有兴趣的读者可以百度“BERT重计算”），将模型或者数据分配到不同的节点上，降低单个gpu的负荷。

这里重点介绍最简单的数据并行范式的分布式训练（DistributedDataParallel，ddp）,基本思想就是将一个大的batch的训练数据的loss计算拆分成多个小的batch，分配到多个不同的gpu上分别计算loss，然后收集加总，再进行反向传播。

多机多卡数据并行训练

多机多卡数据并行训练最方便的实现是利用torch.nn.parallel.DistributedDataParallel和torch.distributed.launch脚本

修改自官方example的示例代码

pytorch-example-distributed

不使用torch.distributed.launch.py

# 单个节点
# 直接使用默认设置
python example.py

# 多个节点 按实际需求设置环境变量
#  第1个节点
# master节点IP
export MASTER_ADDR=192.168.0.40
# master节点端口
export MASTER_PORT=12669
# 当前节点编号
export RANK=0
# 总节点数量
export WORLD_SIZE=3
# 分布式训练开始运行
python example.py


#  第2个节点
# export MASTER_ADDR=192.168.0.40
# export MASTER_PORT=12669
# export RANK=1
# export WORLD_SIZE=2
# python example.py

使用torch.distributed.launch.py

# 使用torch.distributed.launch.py 详情请参考torch.distributed.launch.py的注释

#单个节点
python -m torch.distributed.launch --nnode=1 --node_rank=0 --nproc_per_node=1 simple_example.py --local_rank=0 --local_world_size=1

# 多个节点
# 按实际需求修改命令中的参数
#  nnode 为总节点数
#  node_rank 当前节点编号
#  nproc_per_node 每个节点进程数
#  master_addr 主节点IP
#  master_port 主节点端口

#  第1个节点
# python -m torch.distributed.launch --nnode=2 --node_rank=0 --nproc_per_node=1 --master_addr "192.168.0.1" --master_port 12669 example.py --local_rank=0 --local_world_size=1

#  第2个节点
# python -m torch.distributed.launch --nnode=2 --node_rank=1 --nproc_per_node=1 --master_addr "192.168.0.1" --master_port 12669 example.py --local_rank=0 --local_world_size=1

来自transformers的example

transformers分布式训练example

官方参考文档

个性化的需求可以参考pytorch官方文档：

• PyTorch Distributed Overview

• PyTorch ddp example