Pytorch - 分布式通信原语（附源码）

前言

由于工作需要，最近在补充分布式训练方面的知识。经过一番理论学习后仍觉得意犹未尽，很多知识点无法准确get到（例如：分布式原语scatter、all reduce等代码层面应该是什么样的，ring all reduce 算法在梯度同步时是怎么使用的，parameter server参数是如何部分更新的）。

著名物理学家，诺贝尔奖得主Richard Feynman办公室的黑板上写了："What I cannot create, I do not understand."。在程序员界也经常有"show me the code"的口号。 因此，我打算写一系列的分布式训练的文章，将以往抽象的分布式训练的概念以代码的形式展现出来，并保证每个代码可执行、可验证、可复现，并贡献出来源码让大家相互交流。

经过调研发现pytorch对于分布式训练做好很好的抽象且接口完善，因此本系列文章将以pytorch为主要框架进行，文章中的例子很多都来自pytorch的文档，并在此基础上进行了调试和扩充。

最后，由于分布式训练的理论介绍网络上已经很多了，理论部分的介绍不会是本系列文章的重点，我会将重点放在代码层面的介绍上面。

1 基本介绍

近些年随着深度学习的火爆，模型的参数规模也飞速增长，OpenAI数据显示：

• 2012年以前，模型计算耗时每2年增长一倍，和摩尔定律保持一致；
• 2012年后，模型计算耗时每3.4个月翻一倍，远超硬件发展速度；

近一年来，百亿、千亿级的参数模型陆续面世，谷歌、英伟达、阿里、智源研究院更是发布了万亿参数模型。因此，大模型已经成为了未来深度学习的趋势。提到大模型，就不得不提分布式训练，由于模型参数和训练数据的不断增多，只有通过分布式训练才能完成大模型的训练任务。

分布式训练可以分为数据并行、模型并行，流水线并行和混合并行。分布式算法又有典型的parameter server和ring all-reduce。无论是哪一种分布式技术一个核心的关键就是如何进行communication，这是实现分布式训练的基础，因此要想掌握分布式训练或当前流行的大模型训练务必对worker间的通信方式有所了解。

互联网上已经有很多关于分布式训练的通信方面的文章，但是均没有代码层面的例子。我是属于比较愚钝类型的，只有通过自己手动实现一下方能对一些抽象的概念有较深的理解。

Pytorch的分布式训练的通信是依赖torch.distributed模块来实现的，torch.distributed提供了point-2-point communication 和collective communication两种通信方式。

• point-2-point communication提供了send和recv语义，用于任务间的通信
• collective communication主要提供了scatter/broadcast/gather/reduce/all_reduce/all_gather 语义，不同的backend在提供的通信语义上具有一定的差异性。

Device	CPU	GPU	CPU	GPU	CPU	GPU
send	✓	✘	✓	?	✘	✘
recv	✓	✘	✓	?	✘	✘
broadcast	✓	✓	✓	?	✘	✓
all_reduce	✓	✓	✓	?	✘	✓
reduce	✓	✘	✓	?	✘	✓
all_gather	✓	✘	✓	?	✘	✓
gather	✓	✘	✓	?	✘	✘
scatter	✓	✘	✓	?	✘	✘
reduce_scatter	✘	✘	✘	✘	✘	✓
all_to_all	✘	✘	✓	?	✘	✓
barrier	✓	✘	✓	?	✘	✓

2 P2P communication

下面通过torch.distributed的send/recv接口实现一个简易的ping-pong 程序。程序功能如下：

• tensor 初始值为0
• process 0 （或叫rank 0)：对tensor加1，然后发送给process 1(或叫rank1）；
• process 1：接收到tensor后，对tensor 加2，然后在发送给process 0;
• process 0：接收process1发送的tensor；

2.1 初始化

pytorch中在分布式通信原语使用之前，需要对分布式模块进行初始化。pytorch的分布式模块通过torch.distributed.init_process_group来完成

• 通过环境变量MASTER_ADDR和MASTER_PORT设置rank0的IP和PORT信息，rank0的作用相当于是协调节点，需要其他所有节点知道其访问地址;
• 本例中后端选择的是gloo，通过设置NCCL_DEBUG环境变量为INFO，输出NCCL的调试信息；
• init_process_group：执行网络通信模块的初始化工作
  • backend：设置后端网络通信的实现库，可选的为gloo、nccl和mpi；本例选择gloo作为backend(注：nccl不支持p2p通信，mpi需要重新编译pytorch源码才能使用）；
  • rank：为当前rank的index，用于标记当前是第几个rank，取值为0到work_size - 1之间的值；
  • world_size: 有多少个进程参与到分布式训练中;

def init_process(rank_id, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank_id, world_size=size)
    fn(rank_id, size)

2.2 通信逻辑

下面的代码展示了rank0和rank1进行ping-pong通信的实现：

• 通过rank_id来区分当前应该执行哪一个rank的业务逻辑；
• pytorch 中通过torch.distributed.send(tensor, dst, group=None, tag=0) 和torch.distributed.isend(tensor, dst, group=None, tag=0) 来实现tensor的发送，其中send是同步函数，isend是异步函数；
  • tensor：要发送的数据
  • dst：目标rank，填写目标rank id即可
• pytorch中通过torch.d