共计 9958 个字符,预计需要花费 25 分钟才能阅读完成。
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 Signature
本节以下文字摘自 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。短少(局部)输出,或(局部)输入,不能形成签名。
- 所以 SbpSignature.bn_in_op2sbp_parallel 是一个 map 构造,key 就是各个 input 和 output 的标识。
- 输出与输入的 sbp 组合,在算子的运算法令下必须是非法的。
- 算子的作者须要列出非法 SBP Signature 的候选集。
- 如果推导出的 SBP Signature,sbp 与 inputs 不统一,会通过 GetBoxingOutput 转换为 sbp 统一的 tensor。
1.1 NdSbp 及 NdSbpSignature
以下文字次要参考官网文档 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 的第一维度执行播送,将数据别离拷贝到
[0, 1]
和[2, 3]
。 -
在 ranks 的第二维度别离执行
split(0)
。- 例如,对于
[0, 1]
这个 sub group,将上一步中调配给它的数据按行拆分给0
和1
。
- 例如,对于
如果 Tensor 的数据分布模式是多维的,如[[0, 1], [2, 3]]
,算子对应的 SBP Signature 也是多维的,所以 NdSbpSignature 中,每个 input/output 对应的 sbp_parallel 都是数组。
2 placement.hierarchy
placement 对应的 C++ 类型是 ParallelDesc。
结构 placement 的 ranks 能够是多维数组,示意设施的多维阵列散布。placement.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 中定义的几个函数,会在 SBP Signature 推导过程中用到。
4 一维 SBP 的推导过程
SBP Signature 推导相干的类关系如下:
示例代码中的 tensor add 操作(t1 + t2
),执行到 Interpret 中调用 GetOrInfer 时,会进行 SBP Signature 的推导。
在 ConsistentTensorInferCache::GetOrInfer 中,会把推导后果存起来,不须要每次都进行推导。
ConsistentTensorMetaInferArgs
的 hash 函数次要依赖输出 tensor 的如下信息:
- shape
- dtype
- nd_sbp
- placement
- consumer_nd_sbp_constraint
不同的 tensor 对象,只有这些元信息雷同,就能够复用同一个推导后果。
UserOpExpr
通过 ConsistentTensorInferCache
持有所有推导过的后果。
4.1 ConsistentTensorInferCache 中的推导筹备
理论的推导在 ConsistentTensorInferCache::Infer 中进行。
4.1.1 推导 output 的 shape 和 dtype
user_op_expr.InferLogicalTensorDesc
的作用次要是推导 output 的 shape 和 data_type,后果保留到 output_mut_metas。
这里波及到 UserOpExpr 和 Op 两个模块之间的交互关系。前面会总结一下几个模块之间的契约协定。
user_op_expr.InferLogicalTensorDesc 中用到的两个函数对象,是在结构时从 OpRegistryResult 中拷贝的。OpRegistryResult 的函数对象来自 Op 注册。示例代码中 tensor add 对应的 Op 是 AddNOp。
AddNOp 场景的理论调用程序示例如下:
-
user_op_expr.InferLogicalTensorDesc
-
tensor_desc_infer_fn_
->
AddNOp::InferLogicalTensorDesc- [out.shape = in[0].shape](https://github.com/Oneflow-In…)
-
dtype_infer_fn_
->
AddNOp::InferDataType- [out.data_type = in[0].data_type](https://github.com/Oneflow-In…)
-
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
推导结束后,ConsistentTensorInferCache
会将 inputs/outputs 的元信息,连同推导失去的 NdSbp,一起保留到ConsistentTensorInferResult
。
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 函数对象,查到 ConsistentTensorInferCache 中创立的 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 的候选集
调用 GetSbpSignaturesIf 会获取 SbpSignature 的候选集。
在这个函数中,先调用 GetSbpSignatures 获取初步的候选集,再补充 broadcast 的候选集。候选集都保留到sbp_sig_list
。
GetSbpSignatures
是一个虚函数,UserOp 实现了本人的版本。这个函数中最外围的操作就是 val_->get_sbp_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 中会补充。
候选集数据示例如下:
{"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
束缚,也就是从ConsistentTensorInferCache
一路传过来的参数 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 中,须要时再检索、转换。
下图展现了不同模块之间的局部关系:
7 断点
启动命令
source /mnt/oneflow/build/source.sh
gdb --args python3 /mnt/oneflow/test.py
# set breakpoints
# run
断点示例
set breakpoint pending on
break oneflow::AddNOp::GetSbp
break oneflow::(anonymous namespace)::CheckAndConstructOp
break oneflow::Operator::FillBlobParallelDesc
break oneflow::Operator::InferNdSbpSignature
break oneflow::one::ConsistentTensorMetaInferArgs::MakeNdSbpConstraints
break oneflow::one::ConsistentTensorMetaInferArgs::MakeNdSbpInferHints
break oneflow::one::(anonymous namespace)::Interpret
break oneflow::ParallelDesc::MaybeInit
break oneflow::one::functional::impl::AddFunctor::operator()
8 参考资料
- oneflow v0.8.0
- SBP Signature
- 2D SBP
- placement api