pytorch ddp

(1) 关键概念 world_size: 集群中所有GPU的数量 rank: 范围[0, world_size-1], 表示GPU的编号 local_rank: GPU在每台机器上的编号 比如,两台机器,每台机器4块卡,那么world_size= 2*4, rank 取值范围 [0,1,2,3,4,5,6,7], local_rank 取值范围[0,1,2,3] (2) torch.distributed.launch 启动集群参数 –nnodes: 一共多少台机器 –node_rank: 当前机器编号 –nproc_per_node: 每台机器多少个进程 –master_adderss: master节点ip地址 –master_port: master节点端口 master节点的node_rank必须为0 command example: python -m torch.distributed.launch --nnodes=2 --node_rank=0 --nproc_per_node 8 \ --master_adderss $my_address --master_port $my_port main.py (3) mp.spwan 启动 PyTorch引入了torch.multiprocessing.spawn,可以使得单卡、DDP下的外部调用一致,即不用使用torch.distributed.launch。 python xxxx.py一句话搞定DDP模式。 def demo_fn(rank, world_size): dist.init_process_group("nccl", rank=rank, world_size=world_size) # lots of code. ... def run_demo(demo_fn, world_size): mp.spawn(demo_fn, args=(world_size,), nprocs=world_size, join=True) (4) 集群训练步骤 import torch....

九月 15, 2021 · 2 分钟 · Pan