PyTorch 分布式训练和启动脚本torch.distributed.launch torchrun slurm

1、DataParallel

如果当前有4个GPU,batch_size=16,那么模型将被复制到每一个GPU上,在前向传播时,每一个gpu将分到4个batch,每个gpu独立计算依据分到的batch计算出结果的梯度,然后将梯度返回到第一个GPU上,第一个GPU再进行梯度融合、模型更新。在下一次前向传播的时候,将更新后的模型再复制给每一个GPU。

1、DP在每个训练批次(batch)中,因为模型的权重都是在 一个进程上先算出来 然后再把他们分发到每个GPU上,所以网络通信就成为了一个瓶颈,而GPU使用率也通常很低。

2、因为它在每一次的前向传播的时候把模型也复制了(即每次更新都复制一遍模型),并且单进程多线程会造成GIL contention(全局解释器锁争用) 这里进程计算权重使通信成为瓶颈造成了大量的时间浪费,因此引入了DDP。

2、DistributedDataParallel

DDP采用多进程控制多GPU,共同训练模型,一份代码会被pytorch自动分配到n个进程并在n个GPU上运行。 DDP运用Ring-Reduce通信算法在每个GPU间对梯度进行通讯,交换彼此的梯度,从而获得所有GPU的梯度。对比DP,不需要在进行模型本体的通信,因此可以加速训练。

参考https://zhuanlan.zhihu.com/p/489011749

在所有节点上运行命令来初始化上面创建的 DDP 作业:

torchrun --nnodes**=**2 --nproc_per_node**=**8 --rdzv_id**=**100 --rdzv_backend**=**c10d --rdzv_endpoint**=**$MASTER_ADDR:29400 elastic_ddp.py

这里torchrun将启动8个进程并调用elastic_ddp.py 其启动的节点上的每个进程,但用户还需要应用slurm等集群管理工具才能在2个节点上实际运行此命令。

srun --nodes=2 ./torchrun_script.sh

启动脚本

无论 DDP 应用程序如何启动,每个进程都需要一种机制来了解其rank等,使用torch提供的分布式脚本可以通过环境变量将世界大小、全局等级、主地址和主端口以及本地等级作为命令行参数传递给每个实例,初始化的时候选择环境变量初始化就很方便 (就不应该使用启动子进程torch.multiprocessing.spawn 了)。

torch.distributed.launch

python -m torch.distributed.launch --nproc_per_node 8 test.py

https://github.com/pytorch/pytorch/blob/main/torch/distributed/launch.py

jsx 复制代码
import torch

import os
import time

print("|| MASTER_ADDR:",os.environ["MASTER_ADDR"],
     "|| MASTER_PORT:",os.environ["MASTER_PORT"],
     "|| LOCAL_RANK:",os.environ["LOCAL_RANK"],
     "|| RANK:",os.environ["RANK"], 
     "|| WORLD_SIZE:",os.environ["WORLD_SIZE"])

单机启动八个worker的运行结果

jsx 复制代码
/home  $ python -m torch.distributed.launch   --nproc_per_node 8 test.py                                            2 ↵
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
*****************************************
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 0 || RANK: 0 || WORLD_SIZE: 8
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 1 || RANK: 1 || WORLD_SIZE: 8
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 5 || RANK: 5 || WORLD_SIZE: 8
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 4 || RANK: 4 || WORLD_SIZE: 8
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 3 || RANK: 3 || WORLD_SIZE: 8
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 2 || RANK: 2 || WORLD_SIZE: 8
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 6 || RANK: 6 || WORLD_SIZE: 8
|| MASTER_ADDR: 127.0.0.1 || MASTER_PORT: 29500 || LOCAL_RANK: 7 || RANK: 7 || WORLD_SIZE: 8

多机

jsx 复制代码
/home# python -m torch.distributed.launch --nproc_per_node 8 --nnodes 2 --node_rank 0 --master_addr "替换成你的IP地址" --master_port 1234 test.py

|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 3 || RANK: 3 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 4 || RANK: 4 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 5 || RANK: 5 || WORLD_SIZE: 16
|| MASTER_ADDR:|| MASTER_ADDR: IP地址  IP地址|| MASTER_PORT:  || MASTER_PORT:1234  1234|| LOCAL_RANK:  || LOCAL_RANK:1  7|| RANK:  || RANK:1  7|| WORLD_SIZE:  || WORLD_SIZE:16
16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 6 || RANK: 6 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 2 || RANK: 2 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 0 || RANK: 0 || WORLD_SIZE: 16

第二个机器运行同样的命令(除了noderank

❯ python -m torch.distributed.launch --nproc_per_node 8 --nnodes 2 --node_rank 1 --master_addr "IP地址" --master_port 1234 test.py

*****************************************
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 0 || RANK: 8 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 1 || RANK: 9 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 3 || RANK: 11 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 4 || RANK: 12 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 2 || RANK: 10 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 5 || RANK: 13 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 6 || RANK: 14 || WORLD_SIZE: 16
|| MASTER_ADDR: IP地址 || MASTER_PORT: 1234 || LOCAL_RANK: 7 || RANK: 15 || WORLD_SIZE: 16

torchrun

torchrun 提供了 torch.distributed.launch 功能的超集,并具有一些附加功能:

从 torch.distributed.launch 过渡到 torchrun

https://pytorch.org/docs/stable/elastic/run.html

代码中需要改成从环境变量中读取

slurm多机启动

上面的torch.distributed.launch或者torchrun启动如果多机的话要手动在多个机器执行,使用slurm提交一个命令脚本slurm会分发到多个node来执行。

srun --nodes=2 ./torchrun_script.sh

https://github.com/pytorch/tutorials/blob/main/intermediate_source/ddp_tutorial.rst

相关推荐
阿里云大数据AI技术7 分钟前
AI 时代,零售商超如何用多模态数据"看见"每一排货架?
人工智能
一RTOS一18 分钟前
【无标题】
人工智能·鸿道实时操作系统·国产嵌入式操作系统选型·智算控一体
cian_26 分钟前
我把 Anthropic 官方的 Claude Cookbook 拆成了 30+ 个能直接跑的例子
人工智能
秦先生在广东29 分钟前
`claude-video /watch`:给 Claude 装上“眼睛“看视频的工程实现与边界分析
人工智能
极客猴子42 分钟前
会议记录APP怎么选 2026年实测功能对比与使用指南
人工智能·xcode
秦先生在广东1 小时前
Impeccable:给 AI 编码 Agent 装上设计判断力的工具包
人工智能
hnsoyon1 小时前
智慧城市方案强调预防供水在线监测系统预警水质异常
人工智能·智慧城市·智慧城市方案·供水在线监测系统
Nontee221 小时前
开源 AI 模型到底该不该禁?Anthropic 的一场"不情愿的澄清",撕开了硅谷最深的裂痕
人工智能
韩楚风1 小时前
【参天引擎】一次宕机后的数据恢复,让我把 Cantian 持久化与恢复的六大机制全搞明白了
服务器·网络·数据库·分布式·mysql·架构·cantian
无敌秋1 小时前
无线接入网 RAN
人工智能