Global Tensor和实习总结｜OneFlow学习笔记

撰文｜李响

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.15032.pdf） 的 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 配置的区别：