关于人工智能: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]]。四个计算设施被划分为 2×2 的设施矩阵。这时，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…