做者｜郑建华
更新｜赵露阳
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=1import oneflow as flowP0 = 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。贫乏（部门）输入，或（部门）输出，不克不及构成签名。所以 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.hierarchyplacement 对应的 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 flowplacements = [ 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 的如下信息：
shapedtypend_sbpplacementconsumer_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.InferLogicalTensorDesclogical_tensor_desc_infer_fn_->AddNOp::InferLogicalTensorDescout.shape = in[0].shapedtype_infer_fn_->AddNOp::InferDataTypeout.data_type = in[0].data_type4.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 iblob_descs[i].shape = inputs[i].shapeblob_descs[i].stride = inputs[i].strideblob_descs[i].data_type = inputs[i].data_typeinfer_args.MakeNdSbpInferHints 操感化伪码暗示如下：
# for each input index ihints[i].parallel_desc = inputs[i].parallel_deschints[i].blob_desc = blob_descs[i]hints[i].nd_sbp = inputs[i].nd_sbpblob_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 为 0partial_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-Inc/oneflow/tree/0ea44f45b360cd21f455c7b5fa8303269f7867f8/oneflow）SBP Signature（https://docs.oneflow.org/master/parallelism/02_sbp.html#sbp-signature）2D SBP（https://docs.oneflow.org/master/parallelism/04_2d-sbp.html）placement api（https://oneflow.readthedocs.io/en/master/tensor_attributes.html?highlight=placement#oneflow-placement）https://segmentfault.com/a/1190000042625900欢送 Star、试用 OneFlow 最新版本：https://github.com/Oneflow-Inc/oneflow/
发表评论