OneFlow源码解析：Eager模式下的SBP Signature推导

作者｜郑建华
更新｜赵露阳

OneFlow 的 Global Tensor 有两个必要属性：

• Placement：决定了 tensor 数据分布在哪些设备上。

• SBP：决定了 tensor 数据在这些设备上的分布方式。例如：

• • split：将切分后的不同部分放到不同设备；同时指定切分的 axis。

  • broadcast：将数据复制到各个设备。

如果参与运算的 tensor 的 SBP 不一样，结果 tensor 的 SBP 是什么呢？例如下面的代码：

# export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=0 LOCAL_RANK=0
# export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=1 LOCAL_RANK=1
import oneflow as flow


P0 = flow.placement("cpu", ranks=[0, 1])


t1 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.split(0))
# t1 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.broadcast)
t2 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.split(1))
t3 = t1 + t2
# oneflow.placement(type="cpu", ranks=[0, 1])
print(t3.placement)
# (oneflow.sbp.split(dim=0),)
print(t3.sbp)

t1和t2是分布在相同设备上的两个 tensor。t1.sbp是S(0)，在行上切分；t2.sbp是S(1)，在列上切分。

计算结果t3的 SBP 不需要用户手动指定，系统可以自动推导出t3.sbp为S(0)。这个过程中的一个核心步骤，就是 SBP Signature 的推导。

1 

SBP相关概念

1.1 SBP

SBP是OneFlow中独有的概念，其描述了张量逻辑上的数据与张量在真实物理设备集群上存放的数据之间的一种映射关系。以下内容参考SBP官方文档(https://docs.oneflow.org/master/parallelism/02_sbp.html#sbp)：

详细而言：

• split表示物理设备上的 Tensor，是将全局视角的 Tensor 切分得到的。切分时，需要指定切分的维度。物理设备上的 Tensor ，经过拼接，可以还原得到全局视角的 Tensor 。

• broadcast表示全局视角下的 Tensor，会复制并广播到所有的物理设备上。

• partial 表示全局视角下的 Tensor 与物理设备上的 Tensor 的 形状相同，但是物理设备上的值，只是全局视角下 Tensor 的 一部分。以 partial sum 为例，如果我们将集群中所有设备的张量按位置相加，那么就可以还原得到全局视角的 Tensor。除了 sum 外，min、max 等操作也适用于 partial。

下图中分别展示了 SBP 的情况，分别是 split(0)、split(1)、broadcast 和 partial sum。

1.2 SBP Signature

SBP Signature即SBP签名，是OneFlow中独创且很重要的概念。本节以下文字摘自SBP Signature的官方文档：

• 对于一个孤立的 Tensor，我们可以随意设置它的 SBP 属性。但是，对于一个有输入、输出数据的算子，我们却不可以随意设置它的输入、输出的 SBP 属性。这是因为随意设置一个算子输入输出的 SBP 属性，可能不符合全局视角下算子的运算法则。

• 对于某个算子，其输入输出的一个特定的、合法的 SBP 属性组合，称为这个算子的一个 SBP Signature。

• 算子作者根据算子的运算法则，在开发算子时，就已经罗列并预设好该算子所有可能的 SBP Signature。

• 某一层算子只要有输入的 SBP 属性，OneFlow 就可以根据 SBP Signature 推导出该层算子输出的 SBP 属性。

• 所谓的 SBP Signature 自动推导，指的是：在给定所有算子的所有合法的 SBP Signature 的前提下，OneFlow 有一套算法，会基于传输代价为每种合法的 SBP Signature 进行打分，并选择传输代价最小的那个 SBP Signature。这样使得系统的吞吐效率最高。

• 如果 OneFlow 自动选择的 SBP Signature，上一层算子的输出与下一层算子的输入的 SBP 属性不匹配时，那怎么办呢？OneFlow 会检测到这种不一致，并且在上游的输出和下游的输入间插入一类算子，做相关的转换工作。这类自动加入做转换的算子，就称为 Boxing 算子。

总结一下，SBP Signature 的要点如下：

• 每个算子都需要设置相应的SBP签名，用于描述数据(Tensor)的分布方式。

• SBP签名包括算子的全部输入、输出的SBP。缺少（部分）输入，或（部分）输出࿰