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/mast...)：

详细而言：

• 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。缺少（部分）输入，或（部分）输出，不能构成签名。

  • 所以 SbpSignature.bn_in_op2sbp_parallel 是一个 map 结构，key 就是各个 input 和 output 的标识。
• 输入与输出的 SBP 签名组合，在算子的运算法则下必须是合法的，算子的作者需要列出合法 SBP 签名的候选集。
• 如果输入数据(input tensor)的 SBP 与该算子合法的 SBP 签名不一致，则为了得到该算子正确计算所需要的数据（tensor），OneFlow 会在上游的输出和下游的输入间插入 boxing 算子（可能包含 nccl 等集合通信操作），做自动转换工作，这类自动转换的过程，就称为 Boxing。例如，eager global 模式下的 interpreter 在 GetBoxingOutput 方法中完成 Boxing 过程。

1.3 NdSbp 及 NdSbpSignature

在上面 1.1 小节中，我们了解到 SBP 用于描述一个逻辑张量（Tensor），与其对应物理设备上的映射关系，那 OneFlow 中的 2D 甚至 ND SBP 又是什么意思呢？

简单理解就是，普通的 SBP（1D/1 维 SBP）只能比较粗粒度地对张量进行切分，譬如 split(0)就表示，沿着张量第 0 维进行切分，如果在此基础上，想进行更细粒度的切分，譬如继续沿着第 1 维再“切一刀”，那么普通的 1D SBP 就无法做到了，于是需要 2D 或者 ND SBP。

以下文字主要参考官方文档 2D SBP。

我们可以通过 ranks=[0, 1, 2, 3]指定 tensor 的数据分布在这 4 个设备上。这 4 个设备组成了一个一维的设备矩阵。对应的 SBP 如 split(1)，是单个值，即 1D SBP。

Tensor 数据的分布也可以指定为 ranks=[[0, 1], [2, 3]]。四个计算设备被划分为 2x2 的设备矩阵。这时，SBP 也必须与之对应，是一个长度为 2 的数组。对应的 NdSbp.sbp_parallel 的类型就是数组。

例如 sbp = (broadcast, split(0))。这个 2D SBP 的含义是：

• 在 ranks 的第一维度执行广播，将数据分别拷贝到 group 0(rank [0, 1])和 group 1(rank [2, 3])。

• 在 ranks 的第二维度分别执行 split(0)。

  • 例如，对于 group 0，将上一步中分配给它的数据按行拆分成(1，2)和(3，4)分别给 device 0 和 device 1。

示意图如下：

如果 Tensor 的数据分布形式是多维的，如[[0, 1], [2, 3]]，算子对应的 SBP Signature 也是多维的，所以 NdSbpSignature 中，每个 input/output 对应的 sbp_parallel 都是数组。

2、 placement.hierarchy

placement 对应的 C++ 类型是 ParallelDesc。构造 placement 的 ranks 可以是多维数组，表示设备的多维分布矩阵。

placement.hierarchy 表示了 placement 上 ranks 的层次信息。简单理解，hierarchy 就是用于描述 ranks 分布的形状（类似于 shape 可用于描述 tensor 数据分布的形状），hierarchy 存储了 ranks 在各个维度的 size 信息。

• hierarchy 数组的长度是 ranks 的维数。
• hierarchy 数组的元素值，是 ranks 对应维度的 size。
• 构造 hierarchy 的 C++ 代码可参考 GetRanksShape。

运行下面的代码可以观察 hierarchy 的值。

import oneflow as flow

placements = [
    flow.placement("cpu", ranks=[ 0, 1, 2,   3, 4, 5]),
    flow.placement("cpu", ranks=[[0, 1, 2], [3, 4, 5]]),
]
for p in placements:
    print(p.hierarchy)
# outputs:
# [6]
# [2, 3]

3 、tensor add 是哪个算子？

为了提高性能，从 v0.8.0 开始，Tensor 的接口基本都通过 C API 提供给 Python。

PyTensorObject_methods 中定义了很多 Tensor 方法。不过，add 方法是通过 Python C API 的 number protocol 实现的，指定 PyTensorObject_nb_add 实现加法操作，实际由 functional::add 实现。

functional::add 的定义在 functional_api.yaml.pybind.cpp 中，这是一个在构建期自动生成的文件。顺着这个找，容易发现示例代码对应的是 AddFunctor。Op 的名字是"add_n"，自动生成的文件 op_generated.cpp 中定义了 add_n 对应的 Op 是 AddNOp。add_n_op.cpp 中定义的 AddNOp 的几个方法，会在 SBP Signature 推导过程中用到。

4 、一维 SBP 的推导过程

SBP Signature 推导相关的类关系如下：

示例代码中的 tensor add 操作（t1 + t2），执行到 Interpreter 的中调用 GetOrInfer 时，会进行 SBP Signature 的推导。在 GlobalTensorInferCache::GetOrInfer 中，会以 GlobalTensorMetaInferArgs 作为 key 把推导结果存起来，不需要每次都进行推导。

GlobalTensorMetaInferArgs 的 hash 函数主要依赖输入 tensor 的如下信息：

• shape
• dtype
• nd_sbp
• placement
• consumer_nd_sbp_constraint

不同的 tensor 对象，只要这些元信息相同，就可以复用同一个推导结果。

UserOpExpr 通过 GlobalTensorInferCache 持有所有推导过的结果。

4.1 GlobalTensorInferCache 中的推导准备

实际的推导在 GlobalTensorInferCache::Infer 中进行。

4.1.1 推导 output 的 shape 和 dtype
user_op_expr.InferLogicalTensorDesc 的作用主要是推导 output 的 shape 和 data_type，结果保存到 output_mut_metas。这里涉及到 UserOpExpr 和 Op 两个模块之间的交互关系。后面会总结一下几个模块之间的部分交互接口。

user_op_expr.InferLogicalTensorDesc 中用到的两个函数对象，在 Op 中定义，并注册到 OpRegistry 中。OpRegistryResult 的函数对象来自 Op 注册。示例代码中 tensor add 对应的 Op 是 AddNOp。

AddNOp 场景的实际调用顺序示例如下：

• user_op_expr.InferLogicalTensorDesc

  • logical_tensor_desc_infer_fn_->AddNOp::InferLogicalTensorDesc

    • out.shape = in[0].shape

  • dtype_infer_fn_->AddNOp::InferDataType

    • out.data_type = in[0].data_type

4.1.2 构造 UserOp
MakeOp(user_op_expr...)返回一个 Operator，具体类型是 UserOp（参考之前静态图的讨论）。这个对象负责执行具体的推导。

CheckInputParallelDescIdentical 要求所有 inputs 的 placement 是一致的。因为这里是针对 UserOp 做的推导，例如 tensor add、matmul 等操作，操作数都在相同的设备时，这些操作才能直接计算，否则，就需要通过系统 Op 将数据搬运到一起，再进行计算。

既然所有 inputs 的 placement 都是一样的，那就用第一个作为代表，并赋值给 UserOp 保存。

op->InferParallelSignatureIf()的作用是将 placement 填充到 op.bn2parallel_desc_。

对于 AddNOp 来说，key 是 in_0, in_1, out_0，value 是 inputs[0].placement。

infer_args.MakeInputBlobDescs 操作用伪码表示如下：

# for each input index i
blob_descs[i].shape = inputs[i].shape
blob_descs[i].stride = inputs[i].stride
blob_descs[i].data_type = inputs[i].data_type

infer_args.MakeNdSbpInferHints 操作用伪码表示如下：

# for each input index i
hints[i].parallel_desc = inputs[i].parallel_desc
hints[i].blob_desc = blob_descs[i]
hints[i].nd_sbp = inputs[i].nd_sbp

blob_descs 的作用是为了构造 pd_infer_hints，pd_infer_hints 是为了构造 NdSbpInferHint4Ibn，将相关信息封装到这个函数对象中。这个函数对象被传递给 UserOp 进行推导。在 UserOp 中，通过这个函数对象，根据 input/output 的标识 bn（blob name），获取 NdSbpInferHint，从而可以得到上述元信息。

UserOp 推导完毕后，GlobalTensorInferCache 会将 inputs/outputs 的元信息，连同推导得到的 NdSbp ，一起保存到 GlobalensorInferResult。

4.2 Operator 中的推导准备

Operator::InferNdSbpSignatureIf 中，调用 InferNdSbpSignature 进行实际的推导，然后调用 FillNdSbpSignature 保存推导结果。

InferNdSbpSignature 是一个虚函数。UserOp 会先检查 Op 有没有定义自己的 SBP Signature 推导函数，AddNOp 没有这方面的函数，就调用 Operator::InferNdSbpSignature。

InferNdSbpSignature 中会根据 parallel_desc.hierarchy() 判断是 1D SBP，还是 ND SBP。

先只看 1D SBP 的情况。调用传入的 NdSbpInferHint4Ibn 函数对象，查到 GlobalTensorInferCache 中创建的 NdSbpInferHint，转为 NdSbpInferHint 并存到 map 中。因为是一维的，所以只需要取 sbp_parallel 的第一个元素。然后调用 InferSbpSignature（名字中少了 Nd），将推导结果写到 SbpSignature。

无论是一维还是多维，结果的类型都是 NdSbpSignature。所以要将 SbpSignature 转为 NdSbpSignature。

Operator::InferSbpSignature 的作用主要是构造两个函数对象，SbpInferHint4Ibn 和 CalcOrderValue4SbpSig，然后调用子类 override 的、同名重载的虚函数 InferSbpSignature。

SbpInferHint4Ibn 是将传入的 map 数据封装到函数对象中，用于查询输入输出的元信息。

CalcOrderValue4SbpSig 给每个 SbpSignature 计算一个序值，用于对签名进行排序。

InferSbpSignature 也是一个虚函数。因为 AddNOp 没有定义签名推导函数，会调用 Operator::InferSbpSignature。

4.3 SbpSignature 的推导

之前都是做各种准备，Operator::InferSbpSignature 里才进行真正的推导。简单讲就 3 步：

• 获取候选集
• 过滤不合适的签名
• 排序

4.3.1 SbpSignature 的候选集

调用 GetValidNdSbpSignatureList 会获取 SbpSignature 的候选集。在这个函数中，先调用 GetNdSbpSignatureList 获取初步的候选集，再通过 FilterNdSbpSignatureListByLogicalShape 过滤得到正确可用的候选集。候选集都保存到 sbp_sig_list。

GetNdSbpSignatureList 是一个虚函数，UserOp 实现了自己的版本。这个函数中最核心的操作就是 val_->get_nd_sbp_list_fn，实际调用 AddNOp::GetSbp。UserOpSbpContext 是 UserOp 与 AddNOp 等类之间的协议接口的一部分。

如前所述，提供 SBP Signature 的候选集，是算子的责任。AddNOp 这个算子比较简单，只给出两类签名：

• 对输入 tensor 的 shape 的每个 axis i，所有的 input/output 都创建一个 split(i)。

  • 对于 tensor add 来说，input/output 的 shape 一样才能直接计算，所以 split 的 axis 也都一样。
• 所有的 input/output 都创建一个 partialsum。
• broadcast 的情况会在 Operator 中默认设置，因为理论上所有 inputs/outputs 都应该支持以 broadcast 的方式进行运算。

候选集数据示例如下：

 {"sbp_signature":[{"bn_in_op2sbp_parallel":{"in_0":{"split_parallel":{"axis":"0"}},"in_1":{"split_parallel":{"axis":"0"}},"out_0":{"split_parallel":{"axis":"0"}}}},{"bn_in_op2sbp_parallel":{"in_0":{"split_parallel":{"axis":"1"}},"in_1":{"split_parallel":{"axis":"1"}},"out_0":{"split_parallel":{"axis":"1"}}}},{"bn_in_op2sbp_parallel":{"in_0":{"partial_sum_parallel":{}},"in_1":{"partial_sum_parallel":{}},"out_0":{"partial_sum_parallel":{}}}},{"bn_in_op2sbp_parallel":{"in_0":{"broadcast_parallel":{}},"in_1":{"broadcast_parallel":{}},"out_0":{"broadcast_parallel":{}}}}]}

4.3.2 过滤不合适的签名
分两步过滤不合适的签名

• FilterAndCheckValidSbpSignatureListByLogicalShape 中，对于每个输入 tensor ibn，签名中 ibn 的 split axis，必须小于 tensor ibn 的 shape axes 数量。换句话说，如果 tensor 是二维的，就无法接受 split(2)，只能是 split(0)或 split(1)。
• FilterSbpSignatureList 的作用是检验 sbp_sig_conf 约束，也就是从 GlobalTensorInferCache 一路传过来的参数 nd_sbp_constraints。这个过滤规则要求，符合条件的签名，其内容必须包含 sbp_sig_conf。

4.3.3 签名排序
SortSbpSignatureListByCopyCost 对候选签名进行排序。

• 优先按 OrderValue 比较
• OrderValue 相等时，按 CopyCost 比较二者都是较小的值优先。

OrderValue4SbpSig 是对 CalcOrderValue4SbpSig 的封装，预先计算所有签名的 OrderValue 存到 map 中，便于 sort 函数查找。IbnCopyCost4SbpSig 也是同理。

回过头来看 CalcOrderValue4SbpSig 的定义。因为 AddNOp 是有输入的，对于每个输入 tensor ibn 会加上一个权重，当 ibn 的 sbp 与 签名中对应的 sbp 相同时，权重值为-10，即增加了选中的机会，因为 sbp 一致通常就不需要数据搬运。而 parallel_num 的条件判断在 UserOp 下应该是都成立的。

当 sbp_sig_conf 不空时，CalcOrderValue4SbpSig 直接返回 0。因为如果签名不包含 sbp_sig_conf，即使 SBP 都一致，签名也不一定符合要求，所以直接返回 0。

签名成本由 ComputeIbnCopyCost4SbpSig 计算。主要是根据输入和签名的 sbp 计算 cost：

• 如果 sbp 一致，cost 为 0
• partial_sum 和 broadcast 的 cost 都是一个超大的数字。
• 否则 cost 等于 input tensor 的数据传输字节数量。

4.4 推导结果

推导得到的 nd_sbp_signature 如下：

{"bn_in_op2nd_sbp":{"in_0":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]},"in_1":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]},"out_0":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]}}}

示例代码中，如果一个输入是 split，另一个是 broadcast，推导的签名结果都是 broadcast。如果推断的 sbp 签名是 split，是否能减少数据搬运呢？

5 、NdSbp 的推导过程

NdSbp 的推导主要包括 3 步

• 调用 GetValidNdSbpSignatureList 获取有效的签名
• 剔除不能包含 nd_sbp_constraints 的签名
• 贪心搜索较优的签名

重点看一下有效签名的获取。主要是两步：

• GetNdSbpSignatureList: 获取全部签名FilterNdSbpSignatureListByLogicalShape: 过滤不合适的签名

5.1 NdSbp 签名的候选集

GetNdSbpSignatureList 核心是两步：

• GetSbpSignaturesIf: 得到一维的签名（和 1D SBP 的情况相同）
• DfsGetNdSbpSignature: 根据一维签名拓展到多维

这个过程，如果深入到数据细节去看，会涉及 input/output、ranks、NdSbp 等多个维度，有点抽象复杂。如果从官方文档 2D SBP 中说明的 ranks 和 NdSbp 的物理含义出发，会更容易理解。

以 ranks=[[0, 1, 2], [3, 4, 5]]为例（ranks=[r1, r2]）

这是一个二维的设备矩阵/阵列。算子的每个输入、输出也都有两个 sbp，NdSbpSignature 中的 value 是二维的，有两个槽位。假设 Op 的 1D Sbp 有 n 个签名。

从形式上看，NdSbpSignature 是先按 bn 组织数据。但是从数据分布的过程看，是先按 SbpSignature 组织数据。一个 NdSbpSignature 等价于 SbpSignature 数组。NdSbp 中的每个槽位，都表示一个 1D Sbp 的数据分布（所有的 input/output 一起分布）。

• 比如第 0 个槽位，就是在 r1 和 r2 这两个 sub group 之间分布数据，这个分布必须是一个有效的 1D SbpSignature（所有的 input/output 一起分布）。
• 第 1 个槽位，对于 r1，就是将分配给它的数据子集，再根据一个 SbpSignature 进行分布（所有的 input/output 一起分布）。

所以，只需要按 SbpSignature 整体 填满两个槽位就行。每个槽位各有 n 种可能，一共有 n*n 个候选签名。这样生成的候选集是完整的，不会漏掉候选项。这应该就是 direct product of 1D sbp signatures 的含义。

6、模块间协作关系

SbpSignature 推导的实现用了大量 functional 的代码。应该是为了不同模块间的信息屏蔽，或者父类、子类之间的逻辑复用、信息传递等目的，很多信息都封装到 function 中，需要时再检索、转换。

下图展示了不同模块之间的部分关系：

参考资料

• oneflow v0.9.1（https://github.com/Oneflow-In...）
• SBP Signature（https://docs.oneflow.org/mast...）
• 2D SBP（https://docs.oneflow.org/mast...）
• placement api（https://oneflow.readthedocs.i...）
• https://segmentfault.com/a/11...

欢迎 Star、试用 OneFlow 最新版本：https://github.com/Oneflow-In...