乐趣区

关于人工智能:Global-Tensor-和实习总结｜OneFlow-学习笔记

jiezi

撰文|李响

1

前言

为了简化分布式训练，OneFlow 提出了全局视角（Global View）的概念，在全局视角下，能够像单机单卡编程，进行分布式训练。在 OneFlow 的设计中，应用 Placement、SBP 和 SBP Signature 来实现这种形象。其中，Global Tensor 是为了可能满足 Global View 所需形象的一种 Tensor。本文重点探讨了 Global Tensor 的全局视角与物理视角的映射。

此外，从去年 11 月中旬开始，我在 OneFlow 五个月的实习告一段落，在文末，我对这段时间的工作和播种做了简略总结。

2

Global Tensor

2.1 OneFlow 分布式全局视角的根底保障

这部分内容对应 OneFlow 论文（https://arxiv.org/pdf/2110.15…）的 3.1 节，Placement 与 SBP、SBP Signature 是 OneFlow 分布式全局视角的重要保障，OneFlow 的全局视角使分布式训练与单机单卡一样易用。

首先，Global Tensor 的 Placement 属性能够指定该 Tensor 寄存在哪个物理设施上。咱们重点看下 SBP 的设计，它实现了全局张量和相应的部分张量之间映射的数学形象。

下图是一个形态为 2x2 的 Global Tensor 在 4 种 SBP 映射（每一种都被称为 SBP Signature）下被映射为 2 个部分张量（Local Tensor），别离为 split(0)，split(1)，broadcast 和 partial-sum，别离是按维度宰割、复制并播送和按地位相加的思路。

对于一个 Global Tensor，咱们能够任意设置它的 SBP。然而，对于有输出、输入数据的 Ops，不容许随便设置它的输出、输入的 SBP 属性。因为随便设置 Ops 的输入输出 SBP 属性，可能不合乎全局视角下算子的运算法令，所以 SBP Signature 的概念被提出。

以 MatMul 为例，给定一个数据张量X 和一个权重张量 W，Y=X*W 的 SBP Signature 能够从 X 和 W 的 Signature 中推断进去。如下表，是非法的非法 SBP 组合。当然，SBP Signature 是不须要被咱们显式指定的，OneFlow 存在主动推导机制。


以上介绍的是 1D SBP，为了兼容更简单的分布式训练场景，OneFlow 也提供了 2D SBP。咱们简略来看下区别，首先是 Placement 配置的区别：

>>> placement1 = flow.placement("cuda", ranks=[0, 1, 2, 3]) # 1D SBP 配置集群
>>> placement2 = flow.placement("cuda", ranks=[[0, 1], [2, 3]]) # 2D SBP 配置集群

当 placement 中的集群是 1/2 维的设施阵列时；Global Tensor 的 SBP 也必须与之对应。比方，对于下面的第二种状况，能够配置 (broadcast, split(0)) 的 2D SBP。这时，在设施阵列上，第 0 维做 broadcast 操作；在第 1 维做 split(0)。

再举个例子：对于 ranks=[[0, 1], [2, 3]])，假如看作是 ranks = [Dim0_DeviceGroupA, Dim0_DeviceGroupB]，Dim0_DeviceGroupA = [0, 1]，Dim0_DeviceGroupB=[2, 3]。那么 rank0 的数据就是先把 Global Tensor 残缺数据 broadcast 到第 0 维的两个设施小组 Dim0_DeviceGroupA 和 Dim0_DeviceGroupB；而后再把 Dim0_DeviceGroupA 小组失去的数据 split(0) 到第 1 维即 Dim0_DeviceGroupA 内的 rank 0 和 rank 1（ranks=[[0, 1]]），对于 Dim0_DeviceGroupB 也同理。

>>> sbp = (flow.sbp.broadcast, flow.sbp.split(0))
>>> tensor_to_global = tensor.to_global(placement=placement, sbp=sbp)

对于 SBP Signature，同样也有 2D 的版本，在 1D 的根底上，在两个维度上分别独立推导，做矩阵乘。

2.2 SBP 主动转换

同一 Global Tensor 的生产者和消费者（上游算子输入和上游算子输出）可能须要不同的 SBP Signature。如下图中所示，两个 MatMul Ops 通过 Global Tensor Y0 连贯。S(0) 是 M0 推断的 Y0 的 SBP Signature。然而，M1 冀望它的 SBP Signature 是 B。在这种状况下，在 M0 和 M1 之间须要一个转换 Y0 对应的的部分张量 SBP 类型的操作。这样就实现了 SBP 主动转换，也就是 OneFlow 的 Boxing 机制。

Boxing 机制有不必的办法，如 all2all、broadcast、reduce-scatter、all-reduce 和 all-gather 等，每种操作会产生不同的通信老本。下面的 split(0) 转换为 broadcast，相当于做了一次 all-gather 操作。其实很好了解，这部分内容更具体的解释对应 OneFlow 论文 （https://arxiv.org/pdf/2110.15…） 的 3.2 节，比方每种操作的通信老本大小计算。

2.3 to_global 办法

基于下面两节的内容看下代码，其中有典型的 to_global 办法应用，这部分内容也对应 OneFlow 论文的 3.4 节。

import oneflow as flow

P0 = flow.placement("cuda", ranks=[0, 1])
P1 = flow.placement("cuda", ranks=[2, 3])
a0_sbp = flow.sbp.split(0)
b0_sbp = flow.sbp.broadcast
y0_sbp = flow.sbp.broadcast
b1_sbp = flow.sbp.split(1)

A0 = flow.randn(4, 5, placement=P0, sbp=a0_sbp)
B0 = flow.randn(5, 8, placement=P0, sbp=b0_sbp)
Y0 = flow.matmul(A0, B0)
Y0 = Y0.to_global(placement=P1, sbp=y0_sbp)B1 = flow.randn(8, 6, placement=P1, sbp=b1_sbp)
Y2 = flow.matmul(Y0, B1)

下面代码中，咱们首先初始化了 Global Tensor 的 placement 和 SBP。上一层算子 matmul 的输入 SBP 原本是 split(0)，然而下一层算子 matmul 的输出，被转成了 broadcast。

此时，上一层的输入与下一层的输出，它们的 SBP 并不统一，这是不容许的。所以，基于 Boxing 机制，咱们应用 to_global 办法将 split(0) 转换为 broadcast，也就是代码参数中的 sbp=y0_sbp。

2.4 GlobalTensor 类代码跟踪

基于下面的内容，其实 Global Tensor 的概念就更清晰了，在 OneFlow 的设计中，它就是为了可能满足 Global View（全局一致性视角）所需形象的一种非凡 Tensor。

OneFlow 的 Tensor 设计更像桥接模式，把 Tensor 基类作为抽象化角色，TensorIf 作为 Tensor 的子类充当实现化角色接口，GlobalTensor 和 MirroredTensor 都给出了实现化角色接口的具体实现，如下图。

上图是简略的 Tensor 的设计模式，这里咱们间接重点定位到 GlobalTensor 实现类中，代码地位在 github.com/Oneflow-Inc/，类名为 ConsistentTensor，在 OneFlow v0.7.0 中曾经叫做 GlobalTensor。相干的代码也能够在这个地位开始追踪。在 ConsistentTensor 中，有指向 ConsistentTensorImpl 的指针，EagerConsistentTensorImpl 和 EagerConsistentTensorImpl 又别离继承 ConsistentTensorImpl，代码地位在 github.com/Oneflow-Inc/。值得注意的是，ConsistentTensor 的实现类 ConsistentTensorImpl 里，代码中存在 ConsistentTensorMeta 的成员指针，TensorMeta 系列类保护了除包含 Tensor 的设施、形态和数据类型等根本变量外，包含 placement 和 SBP 的信息。

后面说过，Global Tensor 是为了可能满足 Global View 分布式形象的一种非凡 Tensor，对于上图的 MirroredTensor，就是理论存储各个设施 Tensor data 的中央。EagerConsistentTensorImpl 成员中有指向 MirroredTensor 的指针，同样的，MirroredTensor 持有指向 MirroredTensorImpl 的指针，MirroredTensorImpl 的实现类中则持有指向 TensorStorage 的指针，Tensor 中的数据最终是存在于 TensorStorage 成员中，代码地位在https://github.com/Oneflow-In…

2.5 如何做 Global Ops 的执行测试

最初，简略看一下 OneFlow 如何实现 Global Ops 的测试工作。之前的文章「深度学习框架如何优雅地做算子对齐工作」，曾经介绍过 OneFlow 的 AutoTest 框架的依赖文章。

在这篇文章里，次要围绕 Local Ops 的单测工作介绍，Global Ops 的执行测试也基于该 AutoTest 框架。区别在于，咱们只须要应用文章后面介绍过的 to_global() 办法，由 local tensor 转换失去 global tensor。并且须要枚举 placement 和 SBP 信息。

还是以 matmul Op 为例，对于 matmul 这类 Binary Ops，同一 placement 下，别离须要遍历两次 SBP，代码如下所示。

@autotest(n=1, check_graph=False)
def _test_matmul(test_case, placement, x_sbp, y_sbp):
    x = random_tensor(ndim=2, dim0=8, dim1=16).to_global(placement=placement, sbp=x_sbp)
    y = random_tensor(ndim=2, dim0=16, dim1=8).to_global(placement=placement, sbp=y_sbp)
    return torch.matmul(x, y)

class TestMatMulModule(flow.unittest.TestCase):
    @globaltest
    def test_matmul(test_case):
        for placement in all_placement():
            for x_sbp in all_sbp(placement, max_dim=2):
                for y_sbp in all_sbp(placement, max_dim=2):
                    _test_matmul(test_case, placement, x_sbp, y_sbp)

if __name__ == "__main__":
    unittest.main()

3

总结

下面简略介绍了下我对 Global Tensor 的了解，在这一部分总结下我的实习工作和播种：

在 OneFlow 的这段实习工夫并不长，但这是我第一段工作经验，并且是近程工作，我还有忙学校的一些事件。至今为止我只意识啸宇哥（许啸宇）的脸，哈哈哈～ 还见过 BBuf（张晓雨）的结婚证件照，还有迟哥（姚迟）的个人头像。总的来讲，算是一段比拟乏味轻松（兴许因为我集体不作为：））的经验。

刚入职时，首先是提交 OneFlow 仓库几个简略的 PR，相熟 CI 的流程，还有性能测试、应用 OneFlow 实现多卡训练 U-Net 网络和 python 代码搬运等一些简略的工作，更快的适应工作节奏。在这期间，学习了 ONNX，认真读了 github.com/BBuf/onnx_le 下的一系列文章，并且参加了 OneFlow2ONNX 的一些工作。

实习中期，本人修复了 OneFlow 一些算子代码的 bug，并且自学了 C++ 的新规范和 CUDA 优化，这时才感觉本人大学三年如同什么都没学到一样，是一个没有感情的 Java、408 考试机器。此外，开发了一些算子，也写了很多 bug，这里不得不感激 OneFlow 的保安 zzk（郑泽康）。搭建了 libai 和 flowvision 仓库 RTD 文档的 CI，负责相干文档的 Oncall 工作。

实习前期，实现了 nn.Graph 反对 Local Ops 单测的我的项目，也就是不改代码实现同一段代码的 Graph 和 Eager 模式的自动化执行测试，这外面要解决很多子问题，然而通过一点一点的 hack，根本实现了这个工作。基于此，我对动态图（Graph）的构建、编译流程、训练和调试会更加相熟。

在 OneFlow 仓库大略提了 30+ 的 PR，参加开源有十分理论的益处，播种比每天实习做软件黑盒测试要大，比方成为 Committer 或提交了多少 PR，面试时就能够把这些放在简历上，但如果没有参加开源的经验，就得从其余方面破费很大精力去证实本人的能力。

在 OneFlow 工作，总体上还是挺轻松的。在新手村期间，始终有共事帮忙，让我不至于面对很多不懂的技术而苦楚。有工夫还想能有机会回 OneFlow 始终实习，如果反思本人，就是真的还须要好好积淀一下碎片化的常识。

接下来一小段时间开始做试验和写论文啦，I’m a fight until free～

参考链接

1. https://arxiv.org/pdf/2110.15…
2. https://github.com/Oneflow-In…

欢送下载体验 OneFlow v0.7.0 最新版本：
https://github.com/Oneflow-In…

退出移动版