pytorch 分布式训练 详解

主要两种方式：DataParallel和DistributedDataParallel
DataParallel实现简单，但速度较慢，且存在负载不均衡的问题。
DistributedDataParallel本身是实现多机多卡的，但单机多卡也可以使用，配置稍复杂。

详情如下：

DataParallel : Parameter Server模式，一张卡为reducer, 负载不均衡的问题比较严重,reducer的那张卡会多出3-4g的显存占用;

command model = nn.DataParallel(model)

 

DistributedDataParallel: 采用all-reduce算法，本来设计主要是为了多机多卡使用，但是单机上也能用；

单机多卡训练具体流程：

---1) 初始化使用nccl后端 torch.distributed.init_process_group(backend="nccl")

---2) 配置每个进程的gpu

local_rank = torch.distributed.get_rank()

torch.cuda.set_device(local_rank)

device = torch.device("cuda", local_rank)

---3）使用DistributedSampler

Noting DDP并不会自动shard数据 1. 如果自己写数据流，得根据torch.distributed.get_rank()去shard数据，获取自己应用的一份 2. 如果用Dataset API，则需要在定义Dataloader的时候用DistributedSampler 去shard

from torch.utils.data.distributed import DistributedSampler

sampler = DistributedSampler(dataset) # 这个sampler会自动分配数据到各个gpu上

DataLoader(dataset, batch_size=batch_size, sampler=sampler)

Note: shuffle and distributesampler can't open simultaneously

---4) 封装之前要把模型移到对应的gpu model.to(device)

---5)模型并行化 model=torch.nn.parallel.DistributedDataParallel(model)

 

command（配置2个显卡来分布式并行跑代码）

CUDA_VISIBLE_DEVICES=6,7 nohup python -m torch.distributed.launch --nproc_per_node=2 CR/train.py > log_0326.out 2>&1 &

————————————————