pytorch 分布式训练

1.初始化进程组
dist.init_process_group(backend="nccl")
backend是后台利用nccl进行通信

2.使样本之间能够进行通信
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)

3.创建ddp模型
model = DDP(model,device_ids=[local_rank],output_device=local_rank,find_unused_parameters=True)

获得local_rank(在运行launch时会传入一个local_rank参数)
local_rank = torch.distributed.local_rank()
torch.cuda.set_device(local_rank)

运行脚本
CUDA_VISIBLE_DEVICES=$device python -m torch.distributed.launch --nproc_per_node=1 train.py

保存模型(注意只需保存主进程上的模型，保存的是ddp模型的module)

if dist.get_rank() == 0: torch.save(model.module, "model.pth")