作者：台运鹏 (正在寻找internship...)
主页：https://yunpengtai.top
鉴于网上此类教程有不少模糊不清，对原理不得其法，代码也难跑通，故而花了几天细究了一下相关原理和实现，欢迎批评指正！
关于此部分的代码，可以去
https://github.com/sherlcok314159/dl-tools查看
「在开始前，我需要特别致谢一下一位挚友，他送了我双显卡的机器来赞助我做个人研究，否则多卡的相关实验就得付费在云平台上跑了，感谢好朋友一路以来的支持，这份恩情值得一辈子铭记！这篇文章作为礼物赠与挚友。」
[上一篇] DataParallel里为什么会显存不均匀以及如何解决
实现原理
与DataParallel不同的是，Distributed Data Parallel会开设多个进程而非线程，进程数 = GPU数，每个进程都可以独立进行训练，也就是说代码的所有部分都会被每个进程同步调用，如果你某个地方print张量，你会发现device的差异
sampler会将数据按照进程数切分，「确保不同进程的数据不同」
每个进程独立进行前向训练
每个进程利用Ring All-Reduce进行通信，将梯度信息进行聚合
每个进程同步更新模型参数，进行新一轮训练
按进程切分
如何确保数据不同呢？不妨看看DistributedSampler的源码
# 判断数据集长度是否可以整除GPU数
# 如果不能，选择舍弃还是补全，进而决定总数
# If the dataset length is evenly divisible by # of replicas
# then there is no need to drop any data, since the dataset
# will be split equally.
if (self.drop_last and
len(self.dataset) % self.num_replicas != 0):
# num_replicas = num_gpus
self.num_samples = math.ceil((len(self.dataset) -
self.num_replicas) /self.num_replicas)
else:
self.num_samples = math.ceil(len(self.dataset) /
self.num_replicas)
self.total_size = self.num_samples * self.num_replicas
# 根据是否shuffle来创建indices
if self.shuffle:
# deterministically shuffle based on epoch and seed
g = torch.Generator()
g.manual_seed(self.seed + self.epoch)
indices = torch.randperm(len(self.dataset), generator=g).tolist()
else:
indices = list(range(len(self.dataset)))
if not self.drop_last:
# add extra samples to make it evenly divisible
padding_size = self.total_size - len(indices)
if padding_size <= len(indices):
# 不够就按indices顺序加
# e.g., indices为[0, 1, 2, 3 ...]，而padding_size为4
# 加好之后的indices[..., 0, 1, 2, 3]
indices += indices[:padding_size]
else:
indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
else:
# remove tail of data to make it evenly divisible.
indices = indices[:self.total_size]
assert len(indices) == self.total_size
# subsample
# rank代表进程id
indices = indices[self.rank:self.total_size:self.num_replicas]
return iter(indices)Ring All-Reduce
那么什么是「Ring All-Reduce」呢？又为啥可以降低通信成本呢？
首先将每块GPU上的梯度拆分成四个部分，比如，如下图（此部分原理致谢下王老师，讲的很清晰[1]：
所有GPU的传播都是「同步」进行的，传播的规律有两条：
只与自己
下一个位置的GPU进行通信，比如0 > 1，3 > 0四个部分，哪块GPU上占的多，就由该块GPU往它下一个传，初始从主节点传播，即
GPU0，你可以想象跟接力一样，a传b，b负责传给c
第一次传播如下：
那么结果就是：
那么，按照谁多谁往下传的原则，此时应该是GPU1往GPU2传a0和a1，GPU2往GPU3传b1和b2，以此类推
接下来再传播就会有GPU3 a的部分全有，GPU0上b的部分全有等，就再往下传
再来几遍便可以使得每块GPU上都获得了来自其他GPU的梯度啦
代码使用
基础概念
第一个是后端的选择，即数据传输协议，从下表可以看出[2]，当使用CPU时可以选择gloo而GPU则可以是nccl
| 「Backend」 | 「gloo」 | 「mpi」 | 「nccl」 | |||
|---|---|---|---|---|---|---|
| Device | CPU | GPU | CPU | GPU | CPU | GPU |
| send | ✓ | ✘ | ✓ | ? | ✘ | ✓ |
| recv | ✓ | ✘ | ✓ | ? | ✘ | ✓ |
| broadcast | ✓ | ✓ | ✓ | ? | ✘ | ✓ |
| all_reduce | ✓ | ✓ | ✓ | ? | ✘ | ✓ |
| reduce | ✓ | ✘ | ✓ | ? | ✘ | ✓ |
| all_gather | ✓ | ✘ | ✓ | ? | ✘ | ✓ |
| gather | ✓ | ✘ | ✓ | ? | ✘ | ✓ |
| scatter | ✓ | ✘ | ✓ | ? | ✘ | ✘ |
| reduce_scatter | ✘ | ✘ | ✘ | ✘ | ✘ | ✓ |
| all_to_all | ✘ | ✘ | ✓ | ? | ✘ | ✓ |
| barrier | ✓ | ✘ | ✓ | ? | ✘ | ✓ |
接下来是一些参数的解释[3]：
| Arg | Meaning |
|---|---|
| group | 一次发起的所有进程构成一个group，除非想更精细通信，创建new_group |
| world_size | 一个group中进程数目，即为GPU的数量 |
| rank | 进程id，主节点rank=0，其他的在0和world_size-1之间 |
| local_rank | 进程在本地节点/机器的id |
举个例子，假如你有两台服务器（又被称为node），每台服务器有4张GPU，那么，world_size即为8，rank=[0, 1, 2, 3, 4, 5, 6, 7], 每个服务器上的进程的local_rank为[0, 1, 2, 3]
然后是「初始化方法」的选择，有TCP和共享文件两种，一般指定rank=0为master节点
TCP显而易见是通过网络进行传输，需要指定主节点的ip（可以为主节点实际IP，或者是localhost）和空闲的端口
import torch.distributed as dist
dist.init_process_group(backend, init_method='tcp://ip:port',
rank=rank, world_size=world_size)共享文件的话需要手动删除上次启动时残留的文件，加上官方有一堆警告，还是建议使用TCP
dist.init_process_group(backend, init_method='file://Path',
rank=rank, world_size=world_size)launch方法
「初始化」
这里先讲用launch的方法，关于torch.multiprocessing留到后面讲
在启动后，rank和world_size都会自动被DDP写入环境中，可以提前准备好参数类，如argparse这种
args.rank = int(os.environ['RANK'])
args.world_size = int(os.environ['WORLD_SIZE'])
args.local_rank = int(os.environ['LOCAL_RANK'])首先，在使用distributed包的任何其他函数之前，按照tcp方法进行初始化，需要注意的是需要手动指定一共可用的设备CUDA_VISIBLE_DEVICES
def dist_setup_launch(args):
# tell DDP available devices [NECESSARY]
os.environ['CUDA_VISIBLE_DEVICES'] = args.devices
args.rank = int(os.environ['RANK'])
args.world_size = int(os.environ['WORLD_SIZE'])
args.local_rank = int(os.environ['LOCAL_RANK'])
dist.init_process_group(args.backend,
args.init_method,
rank=args.rank,
world_size=args.world_size)
# this is optional, otherwise you may need to specify the
# device when you move something e.g., model.cuda(1)
# or model.to(args.rank)
# Setting device makes things easy: model.cuda()
torch.cuda.set_device(args.rank)
print('The Current Rank is %d | The Total Ranks are %d'
%(args.rank, args.world_size))「DistributedSampler」
接下来创建DistributedSampler，是否pin_memory，根据你本机的内存决定。pin_memory的意思是提前在内存中申请一部分专门存放Tensor。假如说你内存比较小，就会跟虚拟内存，即硬盘进行交换，这样转义到GPU上会比内存直接到GPU耗时。
因而，如果你的内存比较大，可以设置为True；然而，如果开了导致卡顿的情况，建议关闭
from torch.utils.data import DataLoader, DistributedSampler
train_sampler = DistributedSampler(train_dataset, seed=args.seed)
train_dataloader = DataLoader(train_dataset,
pin_memory=True,
shuffle=(train_sampler is None),
batch_size=args.per_gpu_train_bs,
num_workers=args.num_workers,
sampler=train_sampler)
eval_sampler = DistributedSampler(eval_dataset, seed=args.seed)
eval_dataloader = DataLoader(eval_dataset,
pin_memory=True,
batch_size=args.per_gpu_eval_bs,
num_workers=args.num_workers,
sampler=eval_sampler)「加载模型」
然后加载模型，跟DataParallel不同的是需要提前放置到cuda上，还记得上面关于设置cuda_device的语句嘛，因为设置好之后每个进程只能看见一个GPU，所以直接model.cuda()，不需要指定device
同时，我们必须给DDP提示目前是哪个rank
from torch.nn.parallel import DistributedDataParallel as DDP
model = model.cuda()
# tell DDP which rank
model = DDP(model, find_unused_parameters=True, device_ids=[rank])注意，当模型带有Batch Norm时：
if args.syncBN:
nn.SyncBatchNorm.convert_sync_batchnorm(model).cuda()「训练相关」
每个epoch开始训练的时候，记得用sampler的set_epoch，这样使得每个epoch打乱顺序是不一致的
关于梯度回传和参数更新，跟正常情况无异
for epoch in range(epochs):
# record epochs
train_dataloader.sampler.set_epoch(epoch)
outputs = model(inputs)
loss = loss_fct(outputs, labels)
loss.backward()
optimizer.step()
optimizer.zero_grad()这里有一点需要小心，这个loss是各个进程的loss之和，如果想要存储每个step平均损失，可以进行all_reduce操作，进行平均，不妨看官方的小例子来理解下：
>>> # All tensors below are of torch.int64 type.
>>> # We have 2 process groups, 2 ranks.
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> tensor
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.all_reduce(tensor, op=ReduceOp.SUM)
>>> tensor
tensor([4, 6]) # Rank 0
tensor([4, 6]) # Rank 1@torch.no_grad()
def reduce_value(value, average=True):
world_size = get_world_size()
if world_size < 2: # 单GPU的情况
return value
dist.all_reduce(value)
if average:
value /= world_size
return value看到这，肯定有小伙伴要问，那这样我们是不是得先求平均损失再回传梯度啊，不用，因为，当我们回传loss后，DDP会自动对所有梯度进行平均[4]，也就是说回传后我们更新的梯度和DP或者单卡同样batch训练都是一致的
loss = loss_fct(...)
loss.backward()
# 注意在backward后面
loss = reduce_value(loss, world_size)
mean_loss = (step * mean_loss + loss.item()) / (step + 1)还有个注意点就是学习率的变化，这个是和batch size息息相关的，如果batch扩充了几倍，也就是说step比之前少了很多，还采用同一个学习率，肯定会出问题的，这里，我们进行线性增大[5]
N = world_size
lr = args.lr * N肯定有人说，诶，你线性增大肯定不能保证梯度的variance一致了，正确的应该是正比于，关于这个的讨论不妨参考[6]
「evaluate相关」
接下来，细心的同学肯定好奇了，如果验证集也切分了，metric怎么计算呢？此时就需要咱们把每个进程得到的预测情况集合起来，t就是一个我们需要gather的张量，最后将每个进程中的t按照第一维度拼接，先看官方小例子来理解all_gather
>>> # All tensors below are of torch.int64 dtype.
>>> # We have 2 process groups, 2 ranks.
>>> tensor_list = [torch.zeros(2, dtype=torch.int64) for _ in range(2)]
>>> tensor_list
[tensor([0, 0]), tensor([0, 0])] # Rank 0 and 1
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> tensor
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.all_gather(tensor_list, tensor)
>>> tensor_list
[tensor([1, 2]), tensor([3, 4])] # Rank 0
[tensor([1, 2]), tensor([3, 4])] # Rank 1def sync_across_gpus(t, world_size):
gather_t_tensor = [torch.zeros_like(t) for _ in
range(world_size)]
dist.all_gather(gather_t_tensor, t)
return torch.cat(gather_t_tensor, dim=0)可以简单参考我前面提供的源码的evaluate部分，我们首先将预测和标签比对，把结果为bool的张量存储下来，最终gather求和取平均。
这里还有个有趣的地方，tensor默认的类型可能是int，bool型的res拼接后自动转为0和1了，另外bool型的张量是不支持gather的
def eval(...)
results = torch.tensor([]).cuda()
for step, (inputs, labels) in enumerate(dataloader):
outputs = model(inputs)
res = (outputs.argmax(-1) == labels)
results = torch.cat([results, res], dim=0)
results = sync_across_gpus(results, world_size)
mean_acc = (results.sum() / len(results)).item()
return mean_acc「模型保存与加载」
模型保存，参考部分官方教程[7]，我们只需要在主进程保存模型即可，注意，这里是被DDP包裹后的，DDP并没有state_dict，这里barrier的目的是为了让其他进程等待主进程保存模型，以防不同步
def save_checkpoint(rank, model, path):
if is_main_process(rank):
# All processes should see same parameters as they all
# start from same random parameters and gradients are
# synchronized in backward passes.
# Therefore, saving it in one process is sufficient.
torch.save(model.module.state_dict(), path)
# Use a barrier() to keep process 1 waiting for process 0
dist.barrier()加载的时候别忘了map_location，我们一开始会保存模型至主进程，这样就会导致cuda:0显存被占据，我们需要将模型remap到其他设备
def load_checkpoint(rank, model, path):
# remap the model from cuda:0 to other devices
map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
model.module.load_state_dict(
torch.load(path, map_location=map_location)
)进程销毁
运行结束后记得销毁进程：
def cleanup():
dist.destroy_process_group()
cleanup()如何启动
在终端输入下列命令【单机多卡】
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS
main.py (--arg1 --arg2 --arg3 and all other
arguments of your training script)目前torch 1.10以后更推荐用run
torch.distributed.launch -> torch.distributed.run / torchrun多机多卡是这样的：
# 第一个节点启动
python -m torch.distributed.launch \
--nproc_per_node=NUM_GPUS \
--nnodes=2 \
--node_rank=0 \
--master_addr="192.168.1.1" \
--master_port=1234 main.py
# 第二个节点启动
python -m torch.distributed.launch \
--nproc_per_node=NUM_GPUS \
--nnodes=2 \
--node_rank=1 \
--master_addr="192.168.1.1" \
--master_port=1234 main.pymp方法
第二个方法就是利用torch的多线程包
import torch.multiprocessing as mp
# rank mp会自动填入
def main(rank, arg1, ...):
pass
if __name__ == '__main__':
mp.spawn(main, nprocs=TOTAL_GPUS, args=(arg1, ...))这种运行的时候就跟正常的python文件一致：
python main.py优缺点
「优点」：相比于DP而言，不需要反复创建和销毁线程；Ring-AllReduce算法提高通信效率；模型同步方便
「缺点」：操作起来可能有些复杂，一般可满足需求的可先试试看DataParallel
参考资料
[1]
并行计算与机器学习: https://www.youtube.com/watch?v=rj-hjS5L8Bw
[2]Backends: https://pytorch.org/docs/stable/distributed.html#backends
[3]In distributed computing, what are world size and rank?: https://stackoverflow.com/questions/58271635/in-distributed-computing-what-are-world-size-and-rank
[4]Average loss in DP and DDP: https://discuss.pytorch.org/t/average-loss-in-dp-and-ddp/93306/4
[5]Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour: https://arxiv.org/abs/1706.02677
[6]How to scale learning rate with batch size for DDP training?: https://github.com/Lightning-AI/lightning/discussions/3706
[7]GETTING STARTED WITH DISTRIBUTED DATA PARALLEL: https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
📝论文解读投稿，让你的文章被更多不同背景、不同方向的人看到，不被石沉大海，或许还能增加不少引用的呦~ 投稿加下面微信备注“投稿”即可。
最近文章
COLING'22 | SelfMix：针对带噪数据集的半监督学习方法
ACMMM 2022 | 首个针对跨语言跨模态检索的噪声鲁棒研究工作
ACM MM 2022 Oral | PRVR: 新的文本到视频跨模态检索子任务
统计机器学习方法 for NLP：基于CRF的词性标注
统计机器学习方法 for NLP：基于HMM的词性标注
点击这里进群—>加入NLP交流群
