GPT-3模型为何难以复现？这也许是分布式AI框架的最优设计（2）

3.后向重计算

Checkpointing 是陈天奇在2016年发表的论文 Training Deep Nets with Sublinear Memory Cost 中提到的，也称之为亚线性内存优化。亚线性内存优化有两种思路，Checkpointing 和 CPU offload：

Checkpointing 的核心思想 是在前向网络中标记少量的 Tensor （被 Checkpointing 的 Tensor ），前向计算就只会保留这些被标记的 Tensor， 其余的前向的 activation，会通过在反向传播中根据 Checkpointing 的 Tensor 临时重新计算一遍前向得到。这样就使得大量的 activation 不需要一直保存到后向计算，有效减少了大量 Tensor 的生命周期，使得内存复用效率大幅提升。

CPU offload 的思路类比于计算机操作系统中的“虚拟内存”技术（将不常用的内存临时换入换出到磁盘上，从而增加内存总量），在深度学习中，GPU 显存（Device Memory）的特点是昂贵、高速且容量小，而 CPU 主存（Host Memory）的特点是便宜、相对低速和大容量；那么将前向计算中的一些暂时用不到的 activation 临时换出到 CPU 主存上，等到反向计算需要时再换入到 GPU 显存里，通过这种方式也可以节省显存。

两种亚线性内存优化通过不同的方式达到了显存优化：Checkpointing 是通过额外的计算开销换显存， CPU offload 通过额外的传输开销换显存。

Checkpointing 优化

上图展示了两层 Transformer Layer 在做 Checkpointing 之前和之后的计算图对比， 其中重要的区别是前后向之间的连边从很多条变成了两条。不同框架实现Checkpointing的思路不同，Megatron 是自己重载了 torch.nn.Module ，实现了自己的 checkpointed_forward，相当于定制化了 Transformer Layer 的前后向执行逻辑；OneFlow 的 Checkpointing 就是上图中的设计， 我们在整个计算图中插入了重计算的子图，并使得后向对前向的消费转移到了对重计算子图的消费。

重计算并不是单独为流水并行设计的，并且之前大多使用在单卡或者数据并行场景下。但这个优化在流水并行下就非常关键，因为它使得前向不需要缓存所有的 activation，而只需要缓存非常少个数的（比如一层 Transformer Layer 只会缓存一个 ）、被 checkpoint 的特定 Tensor ，从而大大节省了流水并行下的显存开销。

4. 1F1B 策略

除了重计算，上述 GPipe 的流水并行策略还有另外一个内存问题，就是需要缓存几份 activation，是等于一个 batch 里有多少个 micro-batch 的（梯度累加的次数）。通常，这个累加次数都比较大（为了尽可能流水，累加次数一般大于两倍的 stage 数），那么即使缓存少数 Tensor， 这种策略仍需要较多显存。

因此，在另一篇流水并行的论文PipeDream （2018） 里就提出了改进方法，称之为 1F1B （One Forward pass followed by One Backward pass）的策略。这种改进策略可以解决缓存 activation 的份数问题，使得 activation 的缓存数量只跟 stage 数相关，从而进一步节省显存，训练更大的模型。

1F1B 策略的出发点也比较直观：由于前向计算的 activation 需要等到对应的后向计算完成后才能释放（无论有没有使用 Checkpointing 技术），因此在流水并行下，如果想尽可能节省缓存 activation 的份数，就要尽量缩短每份 activation 保存的时间，也就是让每份 activation 都尽可能早的释放，所以要让每个 micro-batch 的数据尽可能早的完成后向计算，因此需要把后向计算的优先级提高，让 micro-batch 标号小的后向比 micro-batch 标号大的前向先做。因此，如果我们让最后一个 stage 在做完一次 micro-batch 的前向后，立马就做本 micro-batch 的后向，那么我们就能让其他的 stage 尽可能早的开始后向计算，这就是 1F1B 策略。其时间线如下图所示：

1F1B 策略下的 Pipeline 时间线

从上图 1F1B 和之前 GPipe 的流水线对比可知， GPipe 需要缓存 8 份的 activation 供后向使用，而 1F1B 策略只需要缓存 4 份。二者虽然空闲时间的占比是一样的，但节省显存就可以跑更多的 Layer 层数 和 更大的 micro-batch size，从而提升性能。

以上几个关键技术（GPipe、梯度累加、重计算和 1F1B）的介绍就是分布式训练 GPT 的流水并行的核心技术（数据&模型并行我们放在下一章节详细介绍）。无论是 NVIDIA 的Megatron（PyTorch），还是 OneFlow、PaddlePaddle、MindSpore ，都是通过不同的设计实现了上述相同的功能，而且 Megatron 在 NVIDIA 的深度优化下， 在 GPU 上的性能表现已经非常优异了。那么 OneFlow 再搞一套 GPT 的意义何在？别急，看了下一章节，你就知道 PyTorch 做到上述这些技术的痛点在哪儿了。

Megatron ：PyTorch 分布式训练的极限、痛点在哪儿？

NVIDIA 基于 PyTorch 开发了 Megatron，本质上是一个专用于 GPT 的模型库，所有的代码都是 Python 脚本，NVIDIA 为 GPT 专门定制了分布式训练所需的算子、 流水并行调度器、模型并行所需的通信原语等功能。可以说，NVIDIA 在使用 PyTorch 做分布式训练上已经做到极致了。

在本章节，我们会简单介绍一下 Megatron 是如何使用 PyTorch 的，当你也了解 Megatron 的设计以后，你就可以回答这个问题： PyTorch 做分布式训练，真的好用吗？

1.流水并行，PyTorch 需要人工排线和精细控制流水

PyTorch 是单卡视角，一个设备上的 Tensor、模型脚本跟另一个设备上的 Tensor、模型脚本并无直接关系，对于每个设备上的模型脚本都完全对称的（Mirror）最简单的数据并行来说，PyTorch 这样的设计没有什么明显的缺陷。每个设备上的脚本运行到相同 batch 的模型更新部分（Optimizer），统一做一次模型同步（AllReduce 操作）就完成了数据并行，这就是 PyTorch 的 DDP（DistributedDataParallel）模块。

而流水并行，模型网络分布在各个设备上是非对称的，各个设备“接力”执行网络的一部分，这种并行方式用 PyTorch 要如何实现呢？

流水并行 2 卡接力执行网络

上图展示了流水并行下，前两个 stage 分布在 GPU 0 和 GPU 1 上时，网络的拓扑关系。GPU 0 和 GPU 1 是接力执行的， GPU 0 上的 T2 Layer 的输出 Tensor 需要发给 GPU 1 上的 T3 Layer 作为输入。

首先，你需要根据 stage 阶段的不同，分别在各个设备上定义只属于自己那部分的模型网络，而由于第一个 stage 和最后一个 stage 在执行时序上的特殊性，这里 Megatron 还需要进行特判 megatron/training.py 。

def train_step(...):
    if mpu.is_pipeline_first_stage():
        unwrapped_model = model[0]
    elif mpu.is_pipeline_last_stage():
        unwrapped_model = model[-1]

在每个设备根据自己的那部分网络启动以后， Megatron 需要给每个设备上的每一次执行前后都调用 NCCL 的通信操作，前一个 stage 的输出需要通过 NCCL p2p的 ncclSend 操作发给 下一个 stage， 下一个 stage 必须同时调用 ncclRecv 进行接收。当这两个操作成对出现时，这次传输才会成功。（megatron/schedules.py）

def forward_backward_pipelining_without_interleaving(...):
    for i in range(num_microbatches_remaining):
        output_tensor = forward_step(...)
        if forward_only:
            p2p_communication.send_forward(output_tensor, timers)
        else:
            output_tensor_grad = p2p_communication.send_forward_recv_backward(output_tensor, timers)
        # Add input_tensor and output_tensor to end of list, then pop from the
        # start of the list for backward pass.
        input_tensors.append(input_tensor)
        output_tensors.append(output_tensor)
        if forward_only:
            if not last_iteration:
                input_tensor = p2p_communication.recv_forward(timers)
        else:
            input_tensor, output_tensor = input_tensors.pop(0), output_tensors.pop(0)
            input_tensor_grad = backward_step(...)
            if last_iteration:
                input_tensor = None
                p2p_communication.send_backward(input_tensor_grad, timers)
            else:
                input_tensor = p2p_communication.send_backward_recv_forward(input_tensor_grad, timers)

因此对于 PyTorch 用户而言，用户自己需要关心每个 stage 在什么时机需要 recv，什么时机要 send， 发给谁；同时根据 Pipeline 的执行时序，需要特判在前多少个 step，都是需要只做前向（因为后向还没来）， 但又有一些 step，我需要既做前向又做后向，因此你可以看到在 megatron/p2p_communication.py 里，你会发现 Megatron 向用户提供了这些操作：

def recv_forward(...):
    """Receive tensor from previous rank in pipeline (forward receive)."""
def recv_backward(...):
    """Receive tensor from next rank in pipeline (backward receive)."""
def send_forward(...):
    """Send tensor to next rank in pipeline (forward send)."""
def send_backward(...):
    """Send tensor to previous rank in pipeline (backward send)."""
def send_forward_recv_backward(...):
    """Batched send and recv with next rank in pipeline."""
def send_backward_recv_forward(...):
    """Batched send and recv with previous rank in pipeline."""
def send_forward_recv_forward(...):
    """Batched recv from previous rank and send to next rank in pipeline."""
def send_backward_recv_backward(...):
    """Batched recv from next rank and send to previous rank in pipeline."""
def send_forward_backward_recv_forward_backward(...):
    """Batched send and recv with previous and next ranks in pipeline."""

通过这些接口，你就会发现，算法工程师如果想用 PyTorch 做流水并行，他需要精细的控制所有的流水细节， 包括每个 stage 的每个时刻是只做前向，还是前向后向一起做， 同时还需要管理不同 stage 之间收/发数据的节奏，这个要求对于用户而言就太高了。

更让人头痛的是，PyTorch 并没有机制保证这些流水并行中的各个设备之间数据交互的正确性 ，所以用户不仅可能写的不高效， 还可能写错，即使写错了，PyTorch 也无从检查。 这些都给用户带来了极大的使用门槛。因此，也只有 NVIDIA 、 微软等大企业的分布式训练专家可以搞得定 PyTorch 做流水并行。

2.模型并行，PyTorch 需要用户在 kernel 中手写通信原语操作，需要用户推导所有的通信位置

GPT 的大规模训练需要同时用到数据并行、模型并行和流水并行， 对于一个逻辑上的 Transformer Layer，需要同时对一个层做数据并行和模型并行，这个在 Megatron 和 DeepSpeed 的语义里称之为 data-parallel-size 和 tensor-model-parallel-size 。

为什么要既做数据并行，又做模型并行？其实是为了节省显存，并充分利用 GPU 之间的高速互联（NVLink 和 NVSwitch）带宽与机器之间的 IB 网络带宽的差别，NVIDIA 设计了一种在机器间做数据并行， 在机器内做模型并行的混合并行。

在什么样的网络结构、参数规模、网络拓扑下该用数据并行、模型并行还是流水并行，是一个非常复杂的问题。不同的并行方式导致的设备之间、机器之间的通信量是不同的；同时又需要考虑设备显存的约束、 GPU 通信带宽和网络通信带宽的占比、 总的 Batch Size 大小对模型收敛速度的影响等等。目前还没有一个严格的理论来指导具体模型在具体网络拓扑下究竟该用哪种并行配置最优。对于并行策略的研究，我们会在未来专门出一篇文章来探讨这个话题。

对于大部分情况而言，数据并行的效率一般是最高的，但在 GPT-3 这样的网络参数规模下，单个 GPU 根本装不下这么大的模型，所以必须要用到模型并行和 流水并行来降低每个 GPU 上的显存需求。又基于 NVLink 和 IB 网络通信带宽的差别，NVIDIA 设计了一种折中的的方案，对整个集群拓扑做分组，分为机器间和机器内，机器间的网络传输速度较慢，往往是分布式并行的瓶颈，所以适合做流水并行和数据并行；机器内的 NVLink 延迟低、带宽高，正好符合模型并行的要求，由于 GPT-3 必须使用模型并行，因此被放在了机器内做。

于 GPT-3 必须使用模型并行，因此被放在了机器内做。

数据并行是在反向的梯度更新时需要插入 AllReduce 操作，而模型更新在 Gradient Accumulation 里是一个低频操作，多个 micro-batch 只会做一次，所以数据并行在机器间做是比较合适的。

模型并行（Tensor Model Parallelism）， NVIDIA 推导了 Transformer Layer 里的 MLP 和 Self-Attention 操作，模型并行下需要在特定位置插入 AllReduce 来实现前向、后向的数据同步工作。由于模型并行需要在每个 micro-batch 的前向、后向都需要做数据同步， 属于高频操作， 所以 模型并行 适合在机器内做。

NVIDIA 模型并行通信推导

流水并行的优势是带宽需求比其它并行方式低，仅需要在 stage 之间传输数据， 同时还不会阻塞整个网络的计算，因此在机器间做流水并行比较合适；但流水并行必须通过把一个 Batch 分割成若干 micro-batches 才能发挥优势， 同时它还需要额外的显存来缓存 activation，在 batch 间还会留下气泡。

NVIDIA 在论文中实验了相同的总模型并行度（ model-parallel-size = tensor-model-parallel-size * pipeline-model-parallel-size）下， 分配不同的模型并行和流水并行的 size，得出当 tensor-model-parallel-size = 8 时， 总的效率最高，这与每台机器内的卡数相同 。

模型并行度和流水并行度对性能的影响

用 PyTorch 做模型并行的痛点是什么？如果你去了解一下 Megatron 搭 GPT 的模型脚本就立马清楚了，我们知道模型并行需要在前后向插入一些数据同步的操作，但是在哪里插入？NVIDIA 给出了最主要的同步操作推导结果：在 RowParallelLinear 里需要将这个同步写在 kernel 的 forward 函数里：

class RowParallelLinear(torch.nn.Module):
    def forward(self, input_):
        output_parallel = F.linear(input_parallel, self.weight)
        # All-reduce across all the partitions.
        output_ = reduce_from_tensor_model_parallel_region(output_parallel)

这是最关键的一处数据同步操作。但即使在 GPT 这样全部由 Transformer Layer 组成的非常规整的网络里，模型并行需要插入的同步操作就包括但不限于：

AllReduce ：train_step 、 calc_params_l2_norm、CrossEntropy

Scatter ：RowParallelLinear

AllGather ：ColumnParallelLinear

等等...

那么问题来了，算法工程师怎么知道这么长的模型脚本里，到底：

哪处需要插入通信操作？（现在 GPT 的脚本里 NVIDIA 给推导了需要插入通信的位置，如果用户想改网络结构，想加/换一个Op，推导是不是都得重来？）

该插入什么通信操作？（除了 AllReduce，集合通信还有 ReduceScatter、AllGather、Reduce、Broadcast、All2All 等操作，除了集合通信，还有 Scatter 、Gather 等非对称的切分、拼接操作，切分/拼接还要考虑对 Tensor 的哪个维度操作...）

通信操作要跟谁通信？（数据并行和模型并行同时做时， 整个 GPU 集群会被分组，每一组组内做 AllReduce 同步数据， 组间在模型更新时 才同步模型梯度，这意味着每个 rank 的 GPU 想要通信时，是需要跟其他特定对应的 rank 做通信的，这更加增加了实现难度）

更要命的是，如果插入了通信操作，怎么保证正确性？PyTorch 没法保证。PyTorch 将所有的操作都交给了用户， 即使用户插入了一个错误的通信原语（比如将该插入 AllGather 操作的位置插入了 AllReduce），PyTorch 也没法检查出来。

所以这就是为什么只有 NVIDIA 可以用得了 PyTorch 做 Megatron，普通用户只能直接用 megatron/pretrain_gpt.py，想基于 Megatron 做其他模型/网络的迁移、二次开发和研究，是非常困难的。

其实，NVIDIA、 微软、 PyTorch 都被绕进一个大坑里去了：在没有一致性视角（ Consistent View ）的情况下做复杂的分布式并行是非常困难的，往往只能做一些具体网络具体场景具体算子的特判和分析，通过简单的通信原语来实现分布式。而 OneFlow 通过一致性视角下的 Placement + SBP 就非常简单的实现了通用的复杂并行支持。