OneFlow 用一致性视角轻松填平分布式训练难的鸿沟

对于上面 PyTorch 遇到的分布式难题， OneFlow 如何轻松解决？

OneFlow 的两大独特设计是：

1. 运行时 Actor 机制

2. 编译期一致性视角 Consistent View，通过 Placement + SBP + Boxing 解决了分布式易用性的问题。

对于分布式集群环境（多机多卡训练场景），OneFlow 会把整个分布式集群抽象成为一个超级设备，用户只需要在这个超级设备上搭建深度学习模型即可。这个虚拟出来的超级设备我们就称之为逻辑视角（Logical View），而实际上的分布式集群的多机多卡就是物理视角（Physical View），OneFlow维护逻辑视角和物理视角之间的数学上的正确性就称之为一致性视角（Consistent View）。

一致性视角（Consistent View）抽象

理想情况下，抽象出来的超级设备（逻辑视角）的算力是所有物理视角下的设备算力之和（如果算力完全用满，就是线性加速比）；逻辑视角下的显存资源也是所有物理设备的显存资源之和。

1.流水并行 ，只需要配置 Placement 就够了

Placement 描述了逻辑上的 Op 和 物理上的 Op 的映射关系， 一个 Op 分布在哪些机器、哪些设备上，只需要在逻辑图里指定其 Placement 即可。在流水并行里，用户不需要像 PyTorch 一样关心什么时机 send/recv 、 到底是执行 forward 还是 forward + backward，用户只需要配置 Placement （现阶段我们还需要用户帮忙指定 stage id）就完成了流水并行，以下就是 OneFlow 来做 GPT 的流水并行的代码（其实是包含了所有的并行方式）：

class Transformer(object):
    def __call__(self, hidden_states):
        for i in range(self.num_layers):
            with distribute.layer_placement_scope(i):
                h = self.layers[i](h)

其中 layer_placement_scope 就是在 scope 中配置了 placement 和 stage id：

def layer_placement_scope(layer_idx, device="gpu"):
    dist_util = get_dist_util()
    with flow.scope.placement(
        device, dist_util.get_layer_placement(layer_idx), dist_util.parallel_hierarchy,
    ):
        if dist_util.is_pipeline_parallel():
            with flow.experimental.scope.config(
                pipeline_stage_id_hint=dist_util.get_layer_stage(layer_idx)
            ):

其余的所有流水并行的细节：如各个 stage 之间怎么通信， 做前向还是后向， 怎么流水起来，怎么保证正确性，这些全都不用用户操心。

在下一章节我们还会介绍 OneFlow 内部是怎么实现流水并行的，相较于 Megatron 的复杂实现， OneFlow 系统层面的工作对框架开发者而言也更加友好。

上图展示了一个可能的 Placement 例子，用于 GPU0 和 GPU1 之间的流水并行。图中负责在 CPU 和 GPU、GPU 与 GPU 之间进行数据搬运的Op（CopyH2D、CopyD2D）是 OneFlow 系统自动添加的。

2.数据 + 模型的混合并行，只需要配置 Variable 的 SBP 就够了

SBP 描述了 逻辑上的 Tensor 和 物理上的 Tensor 的映射关系。SBP 是三种基础映射的首字母组合：Split、Broadcast、Partial，（其中 Partial 是一个 reduce 操作，包括 PartialSum、PartialMin、PartialMax等），全称叫做 SbpParallel。其中：

Split：表示物理上的多个 Tensor 是由逻辑上的 Tensor 进行切分后得到的。Split 会包含一个参数 Axis，表示被切分的维度。如果把所有物理上的 Tensor 按照 Split 的维度进行拼接，就能还原出逻辑上的 Tensor。

Broadcast：表示物理上的多个 Tensor 是逻辑上 Tensor 的复制，两者数据完全相同。

PartialSum：表示物理上的多个 Tensor 跟逻辑上的 Tenso r的形状相同，但每个对应位置上元素的值 是逻辑Tensor对应位置元素的值的一部分。如果把所有物理上的 Tensor 按照对应位置相加（element-wise），即可还出原逻辑上的 Tensor。

下图展示了SBP的几种简单情形。

SBP 逻辑与物理 Tensor 的对应关系

需要注意的是，对于同一个逻辑上的 Tensor，其物理上的 Tensor 的映射关系可能会有多种，这取决于生产这个 Tensor 的 Op 和消费这个 Tensor 的 Ops 是如何看待这个逻辑上的 Tensor 的。

那么用 OneFlow 做数据并行、模型并行，需要用户做什么呢？其实只需要配置 Variable 的 SBP 即可。我们简单介绍一下数据并行和模型并行在 OneFlow 里的配置方式：

数据并行下，每个设备上都有整份的模型，所以 Variable 的 SbpParallel 是 Broadcast， 表示物理上的每个设备上的模型都是逻辑上的完整模型的复制。其余的用户就不用再做任何操作了（其实数据并行下，反向梯度更新的同步操作 AllReduce 是 OneFlow 系统内部根据 SBP 的推导规则自动插入的。）

模型并行下，每个设备都把模型切分并只保留一部分， 所以 Variable 的 SbpParallel 是 Split(0)， 表示物理上的每个设备上的模型都是逻辑上的完整模型经过第0维切分后的。其余的用户就不用再做任何操作了。前后向的数据同步操作也是 OneFlow 系统内部根据 SBP 推导规则自动插入的。

其实对于 Linear Layer 里的 Variable （假设是 row major 存储），Split(0) 和 Split(1) 分别表示了两种不同的模型并行方式。如果是 Split(0) 前后向就需要插入 AllReduce， 如果是 Split(1) ，前后向就需要插入 AllGather 操作了。至于为什么要插入 AllReduce 或者 AllGather，我会在下一章节介绍 SBP 推导的时候详细解释。另外，其实数据并行梯度更新要插入 AllReduce 做梯度同步，在 OneFlow 里也是自动推导出来的，并不是一个像 PyTorch DDP 一样的模块做特判。

而且 OneFlow 的 Consistent View 还保证了：任何配置 SBP 得到的并行结果， OneFlow 都保证了 其计算在数学上是完全一致的， 我们从机制上保证了分布式训练的正确性难题， 这一点是现在的 PyTorch 无法做到的。

3. 2D SBP

那么如何同时让一个 Op 既做数据并行，又做模型并行（分组）？这里就用到了 2-D SBP。

在 2-D SBP 下，（其实 OneFlow 还支持扩展到任意维度 N-D）集群和设备呈现为一个 2-D 的拓扑结构。比如我们一共有 2 机 8 卡（每台机器 4 张 GPU），我们可以将 8 个设备表示成一个 (2 x 4) 的矩阵， 那么 如何在机器间数据并行、机器内模型并行呢？用户只需要将 Variable 的 2-D SBP 配置成： [ Broadcast, Split(0) ] 即可，那么实际上 各个设备上的 物理 Tensor 跟 逻辑 Tensor 的映射关系如下图所示。在第一维上的 Broadcast， 表示 GPU0 和 GPU4、GPU1 和 GPu 5、 GPU2 和 GPU6、 GPU3 和 GPU7 在机器间做数据并行，在第二维上的 Split(0)， 表示 GPU0，1，2，3 、 GPU4，5，6，7 在机器内做模型并行。

2D-SBP 下 逻辑 Tensor 和 物理 Tensor 的对应关系

在 2D SBP 下， 设备拓扑被称之为 ParallelConf::hierarchy ，表示如何将一维的 world_size = 8 映射成为二维的层次结构：(2, 4)。而 2-D SBP 被称之为ParallelDistribution， 是由两个 SbpParallel 组成的 list: [Broadcast, Split(0)]。只需要简单配置一下 SBP， 用户就可以轻松设定网络是数据并行、模型并行还是混合并行。

3. OneFlow：让每一位算法工程师都有能力训练 GPT

综合比较 PyTorch 和 OneFlow 在分布式并行上的设计以及开放给用户的接口，有如下总结：

PyTorch 分布式的困境：

PyTorch 只有物理视角，没有逻辑视角，所以 PyTorch 的用户想要做 分布式并行，任何时候都需要自己推导深度学习模型中哪处需要跟其他的物理设备进行通信和数据同步操作，既要推导通信所在的位置，又要推导通信的操作类型，还要推导跟其他哪些设备通信。这个在简单的数据并行下可以使用 DDP 或者 Horovod 来实现，但是在复杂的模型并行、混合并行下，做并行的门槛会非常高。

PyTorch 没有将模型网络的算法逻辑和分布式并行训练的通信逻辑解耦出来，导致需要在 算子的 kernel 实现中、 搭网络的脚本里到处插入通信原语。这些手写通信原语的操作不仅繁琐、易错、而且没法复用，是根据特定模型、特定脚本位置、特定算子特判得到的。

PyTorch 在非对称的并行方式里（如流水并行），各个设备的调度逻辑需要用户自己手写，用户需要自己精细的控制每个设备上的启动以及执行逻辑，且执行逻辑把前后向执行和send/recv通信操作糅合在了一起，即使在最规整的 Transformer Layer 的流水并行下也很复杂，想要扩展到其他模型上的工作量也很大。

PyTorch 没有机制保证分布式并行训练中的正确性 和 数学一致性。即使用户写错了通信操作、插错了位置、 跟错误的设备进行了通信，PyTorch也检查不出来。

这些都使得用户使用 PyTorch 开发复杂分布式训练的脚本极为困难，以至于只有 NVIDIA 、 微软 等大企业的分布式训练专家才能开发出 GPT 的 PyTorch 版本。

OneFlow 分布式的易用性：

OneFlow 有一致性视角 Consistent View， 将分布式训练下的多机通信和 算法逻辑解耦，使得用户可以不用关心分布式训练的细节，降低了分布式训练的使用门槛。

OneFlow 通过 Placement + SBP 机制解决了分布式训练中任意并行场景的需求，用户只需要配置 op 的 Placement 就可以完成流水并行，只需要配置 Tensor 的 SBP 就可以实现数据并行、模型并行和混合并行。而且任何并行方式都是 Placement + SBP 的一种特例， OneFlow 从系统层面不需要做任何的特判，SBP 才是各种分布式并行的本质。

OneFlow 的通信逻辑可以复用，不需要为任何特定网络和特定算子实现相应的通信逻辑。通信逻辑由 OneFlow 的 Boxing 机制完成，与具体的算子和模型无关。

OneFlow 的 SBP 还保证了数学上的一致性， 相同的逻辑上的模型脚本，使用任意的并行方式（数据并行、模型并行、流水并行）、使用任意的集群拓扑，OneFlow 都从数学上保证了模型分布式训练的正确性。

因此，我们才说 OneFlow 用一套简洁设计解决了分布式并行的各种难题。对于用户而言， OneFlow 让算法工程师不需要成为一位分布式训练的专家也有能力做复杂的分布式训练， 只要有硬件资源，任何一位算法工程师都可以训练 GPT， 任何一位算法工程师都可以开发一个新的大规模分布式训练的模型。

*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。