关于深度学习:OneFlow源码解析Global-Tensor

撰文 | 郑建华
更新｜赵露阳

上文中讲到的相似于PyTorch中的一般Tensor，在OneFlow中称为Local Tensor。Local Tensor是单卡视角下的一般Tensor。与之绝对，OneFlow中还有一个独有的概念——Global Tensor。

Global Tensor是指被placement和SBP属性所指定的，一个全局视角下的逻辑Tensor。Global Tensor的shape是逻辑形态，其实在数据依据placement和SBP的规定散布在多个rank上。

Global Tensor既能够通过一般的Local Tensor通过tensor.to_global()转换失去，也能够间接用数据或Numpy来结构。

上面的大节将通过一个示例（https://docs.oneflow.org/mast...），
展现从一般数据结构Global Tensor的过程，以及别离形容SBP、Placement和Global Tensor结构的细节。

1、Global Tensor示例

开启2个终端，终端一、二别离设置环境变量：

# 终端一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 flowp = flow.placement("cpu", ranks=[0, 1])sbp = flow.sbp.split(0)x = flow.tensor([[1,2,3],[4,5,6]], placement=p, sbp=sbp)print(x.shape)print(x.to_local())

终端一、二的输入如下：

# 终端一oneflow.Size([2, 3])tensor([[1, 2, 3]], dtype=oneflow.int64)# 终端二oneflow.Size([2, 3])tensor([[4, 5, 6]], dtype=oneflow.int64)

这个例子中：

export xxx环境变量通知oneflow环境用于通信的IP和Port，以及全局共有2个rank（WORLD_SIZE=2），终端一所在的是rank0，终端二所在的是rank1。
p = flow.placement("cpu", ranks=[0, 1])设置了global tensor将会被搁置于rank0和rank1。
sbp = flow.sbp.split(0)设置了global tensor的sbp属性为split，即按第0维度进行切分。
x = flow.tensor([[1,2,3],[4,5,6]], placement=p, sbp=sbp)从python list数据配合sbp和placement结构了一个global tensor x。

这里，x是由[[1,2,3],[4,5,6]]结构而来，其shape为(2,3)，所以咱们print(x.shape)失去的是：oneflow.Size([2, 3])，x是一个global tensor，其shape示意全局范畴内的逻辑形态。

而后，在特定rank上执行x.to_local()示意将global tensor转为以后rank上的local tensor，因为x的sbp是split(0)，示意tensor按第0维切分，即[1,2,3]寄存于rank0；[4,5,6]寄存于rank1。

所以，print(x.to_local())失去终端一的输入为：
tensor([[1, 2, 3]], dtype=oneflow.int64)

终端二的输入为：
tensor([[4, 5, 6]], dtype=oneflow.int64)

当然，上述只是一个小例子，用于了解global tensor以及sbp和placement属性的概念，实在利用场景下，通常都会间接用local tensor通过tensor.to_global(https://oneflow.readthedocs.i...)的形式，来创立global tensor并应用。

2、SBP

SBP由split, broadcast, partial的首字母组合而成，SBP是一种规定，其形容了逻辑tensor（global tensor)在物理设施上的散布策略。

split示意global tensor在各个rank(物理设施)都存在分片，每个分片能够看作是将global tensor沿着某一维度切分失去的本rank重量（rank由placement指定）。
broadcast示意global tensor在每个rank上齐全一样，等价于从某个rank复制并播送至所有rank。
partial示意global tensor与物理设施上的tensor的形态雷同，然而物理设施上的值，只是global tensor的一部分，global tensor的值须要这些rank上的local tensor进行 sum、max、mean等相似操作。

Python端flow.sbp
（https://github.com/Oneflow-In...）
包定义了split等3种类型。其C++ binding代码在sbp_symbol.cpp（https://github.com/Oneflow-In...）中。这些类型都是SbpParallel
（https://github.com/Oneflow-In...）类型，是protobuf message对象。三种类型通过oneof parallel_type（https://github.com/Oneflow-In...）共享存储。

其中broadcast和partial_sum都是空音讯，赋值时须要调用mutable办法
（https://github.com/Oneflow-In...）显式表明oneof字段具体是哪种类型。split的值示意在tensor的哪个轴上切分数据。轴的index值是一个[[0, 5]之间的整数]。所有的split SbpParallel对象被保留到一个动态vector
（https://github.com/Oneflow-In...）中。

3、Placement的结构

placement属性指定逻辑tensor理论寄存在哪些物理设施上，更具体的，是寄存于哪些rank上。

在上述例子中：

flow.placement("cpu", ranks=[0, 1])创立了一个placement对象。第一个参数是设施类型，目前反对cpu或cuda。ranks[0, 1]示意tensor散布在rank 0和rank1上。

sbp = flow.sbp.split(0)表明tensor的数据分布是按split切分，且是沿着第0维进行切分。

ranks只列出了rank id（全局惟一），没有指定节点host。是因为rank与host关系曾经依据环境变量所确定。环境变量RANK示意全局惟一的rank id，LOCAL_RANK示意节点内的本地rank id。在GPU环境下，个别一个过程对应一块设施（https://docs.oneflow.org/mast...）。WORLD_SIZE示意所有节点的设施（过程）总数。

在通过import oneflow初始化oneflow时，会依据环境变量在各个节点间建设管制面通信连贯（https://github.com/Oneflow-In...），以及数据面通信连贯。这样每个过程就晓得有多少个节点、有多少个设施/过程、以后过程在整个集群的地位。

通过placement的构造函数绑定（https://github.com/Oneflow-In...）能够晓得，其对应的C++类型是ParallelDesc
（https://github.com/Oneflow-In...）。对象结构由函数CreateParallelDescSymbol（https://github.com/Oneflow-In...）实现。次要调用流程如下：

3.1 确定machine和device

ParseAndFormatRanks
（https://github.com/Oneflow-In...）会将ranks数组[0, 1]转为形如"machine_id:device_id"的字符串数组，供后续解决应用。这里的逻辑决定了如何依据ranks中的id，确定tensor数据在节点和设施上的散布：

machine_id=rank / NumOfProcessPerNode
（https://github.com/Oneflow-In...）
device_id=rank % NumOfProcessPerNode
（https://github.com/Oneflow-In...）

从上述公式能够看出，各个节点的设施/过程数量须要是统一的。

3.2 结构并缓存ParallelDesc对象

CreateParallelDesc
（https://github.com/Oneflow-In...）函数实现ParallelDesc的结构。其中MakeParallelConf
（https://github.com/Oneflow-In...）会先依据"machine_id:device_id"等数据结构一个cfg::ParallelConf对象，这是一个相似oneflow::ParallelConf（https://github.com/Oneflow-In...）的类型，文件位于build/oneflow/core/job/placement.cfg.h，是cmake构建过程中主动生成的文件。

cfg::ParallelConf等对象的接口相似protobuf message，但实现了hash办法，能够作为hash map的key。

之后的PhysicalRun
（https://github.com/Oneflow-In...）尽管波及虚拟机，但理论执行的op指令应该是空的，实质性的逻辑只是调用builder的GetParallelDescSymbol（https://github.com/Oneflow-In...），其中的外围逻辑是FindOrCreate（https://github.com/Oneflow-In...），从缓存中查找ParallelDesc或创立新的缓存。

4、Global Tensor结构调用流程

上面以本文开始的例子剖析一下结构global tensor的调用流程。这可能不是一个典型的场景，只是人为指定一个简略的数据便于展现和debug。

通过之前探讨local tensor时的类关系图能够晓得，EagerGlobalTensorImpl内含一个local tensor的变量（https://github.com/Oneflow-In...）。能够设想，结构global tensor时，会先结构一个local tensor、再做一些后续解决。

Python端创立tensor对象时，如果像本文开始的例子那样指定placement、sbp和数据，对应的Functor是GlobalTensorWithDataCtorFunctor
（https://github.com/Oneflow-In...）。外围逻辑在MakeGlobalTensorFromData（https://github.com/Oneflow-In...）中，其次要调用流程如下：

上述各个局部的次要职能如下：

DataConsistencyCheck（https://github.com/Oneflow-In...）会在tensor的placement波及的各个节点间拷贝数据、校验数据是否统一。
functional::Empty
（https://github.com/Oneflow-In...）会依据shape和dtype结构一个local tensor，并期待随后填充数据（这里和之前探讨local tensor的过程统一）。
SwitchCopyLocalTensorFromUntypedArray（https://github.com/Oneflow-In...）为empty的local tensor填充数据，数据既能够是本例中的python list，也能够是numpy的ndarray。
functional::Cast
（https://github.com/Oneflow-In...）进行数据类型dtype的转换。
functional::LocalToGlobal
（https://github.com/Oneflow-In...）把local tensor转为global tensor，但这个只是用于broadcast 至指定placement的长期的global tensor(sbp list全副为broadcast，用于播送)。
functional::ToGlobal
（https://github.com/Oneflow-In...）将长期的global tensor依据placement和sbp，ToGlobal转换为最终的global tensor。

5、用flow.randn结构Global Tensor

上面看一个通过op结构global tensor的例子

# 终端一# 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 flowp = flow.placement("cpu", ranks=[0, 1])sbp = flow.sbp.split(0)x = flow.randn(4, 5, placement=p, sbp=sbp)print(x.shape) # (4,5)print(x.to_local().shape) # (2,5)

randn op在local和global下别离对应着不同的functor实现：

# oneflow/core/functional/functional_api.yaml- name: "randn"  signature: [      "Tensor (Shape size, *, DataType dtype=None, Device device=None,      Generator generator=None, Bool requires_grad=False) => RandN",      "Tensor (Shape size, *, Placement placement, SbpList sbp, DataType dtype=None,      Generator generator=None, Bool requires_grad=False) => GlobalRandN",    ]  bind_python: True

一般的flow.randn对应RandNFunctor，而global版本（带placement和sbp参数）的randn则对应的是GlobalRandNFunctor。

能够看到：

GlobalRandNFunctor
（https://github.com/Oneflow-In...）中次要dispatch了"normal" op，在Eager Global的mode下，会交给EagerGlobalInterpreter进行各种推导和筹备工作（Interpret[https://github.com/Oneflow-In...]），并在Interpret办法里通过PhysicalRun，将normal op执行的指令交给虚拟机调度并执行。
EagerGlobalTensorImpl::New（https://github.com/Oneflow-In...）时会调用GetPhysicalShape（https://github.com/Oneflow-In...）获取local tensor的shape。

这里，咱们能够正当猜想，在每个rank上都会通过同样的Interpret、调用同样的normal op，生成本rank下局部的randn后果——local tensor，其shape都为(2, 5)，通过组装失去global tensor x，其shape为(4, 5)。通过debug验证了上述猜想是正确的。从这个例子中，大抵能够失去论断：

1.Global Tensor其实是基于Local Tensor以及SBP和placement的一层封装，其shape为全局逻辑形态；其数据由各个ranks所持有（ranks由placement指定)。

2.每个rank上的数据分片都是独立的Local Tensor，通过SBP规定的组装，失去下层的Global Tensor。

3.Global Tensor的计算实际上就是通过不同rank上数据分片（Local Tensor）独立通过kernel计算、boxing机制等组合实现的。

参考资料：

OneFlow源码
（https://github.com/Oneflow-In...）
OneFlow源码解析1：算子签名的主动推断
OneFlow源码解析2：Op、Kernel与解释器
OneFlow源码解析3：Op指令在虚拟机中的执行
OneFlow源码解析4：tensor体系与local tensor
Global Tensor：https://docs.oneflow.org/mast...
集群的全局视角：https://docs.oneflow.org/mast...
Global View的概念和实现
OneFlow的Global Tensor笔记和实习总结

欢送下载体验 OneFlow v0.8.0 最新版本：
https://github.com/Oneflow-In...