撰文 | 郑建华
更新|赵露阳
上文中讲到的类似于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来结构。
下面的末节将通过一个示例(docs.oneflow.org/master/para…),
展示从一般数据结构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 flow
p = 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(oneflow.readthedocs.io/en/master/g…) 的方式,来创立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
(github.com/Oneflow-Inc… )
包界说了split等3种类型。其C++ binding代码在sbp_symbol.cpp(github.com/Oneflow-Inc… ) 中。这些类型都是SbpParallel(github.com/Oneflow-Inc… ) 类型,是protobuf message目标。三种类型通过oneof parallel_type(github.com/Oneflow-Inc… ) 同享存储。
其间broadcast和partial_sum都是空消息,赋值时需求调用mutable办法
(github.com/Oneflow-Inc… )显式标明oneof字段详细是哪种类型。split的值标明在tensor的哪个轴上切分数据。轴的index值是一个[[0, 5]之间的整数]。一切的split SbpParallel目标被保存到一个静态vector
(github.com/Oneflow-Inc… )中。
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环境下,一般一个进程对应一块设备(docs.oneflow.org/master/para… )。WORLD_SIZE标明一切节点的设备(进程)总数。
在通过import oneflow
初始化oneflow时,会依据环境变量在各个节点间树立控制面通讯连接(github.com/Oneflow-Inc… ),以及数据面通讯连接。这样每个进程就知道有多少个节点、有多少个设备/进程、当时进程在整个集群的方位。
通过placement的结构函数绑定(github.com/Oneflow-Inc… )能够知道,其对应的C++类型是ParallelDesc (github.com/Oneflow-Inc… )。目标结构由函数CreateParallelDescSymbol(github.com/Oneflow-Inc… )完结。首要调用流程如下:
3.1 确认machine和device
ParseAndFormatRanks
(github.com/Oneflow-Inc… )会将ranks数组[0, 1]转为形如”machine_id:device_id”的字符串数组,供后续处理运用。这儿的逻辑决定了如何依据ranks中的id,确认tensor数据在节点和设备上的散布:
-
machine_id=rank / NumOfProcessPerNode (github.com/Oneflow-Inc… )
-
device_id=rank % NumOfProcessPerNode (github.com/Oneflow-Inc… )
从上述公式能够看出,各个节点的设备/进程数量需求是共同的。
3.2 结构并缓存ParallelDesc目标
CreateParallelDesc (github.com/Oneflow-Inc… )函数完结ParallelDesc的结构。其间MakeParallelConf (github.com/Oneflow-Inc… )会先依据”machine_id:device_id”等数据结构一个cfg::ParallelConf目标,这是一个类似oneflow::ParallelConf(github.com/Oneflow-Inc… )的类型,文件坐落build/oneflow/core/job/placement.cfg.h,是cmake构建过程中主动生成的文件。
cfg::ParallelConf等目标的接口类似protobuf message,但实现了hash办法,能够作为hash map的key。
之后的PhysicalRun (github.com/Oneflow-Inc… )虽然涉及虚拟机,但实践履行的op指令应该是空的,实质性的逻辑仅仅调用builder的GetParallelDescSymbol(github.com/Oneflow-Inc… ),其间的中心逻辑是FindOrCreate(github.com/Oneflow-Inc… ),从缓存中查找ParallelDesc或创立新的缓存。
4、Global Tensor结构调用流程
下面以本文开始的比如分析一下结构global tensor的调用流程。这可能不是一个典型的场景,仅仅人为指定一个简单的数据便于展示和debug。
通过之前讨论local tensor时的类联系图能够知道,EagerGlobalTensorImpl内含一个local tensor的变量(github.com/Oneflow-Inc… )。能够幻想,结构global tensor时,会先结构一个local tensor、再做一些后续处理。
Python端创立tensor目标时,如果像本文开始的比如那样指定placement、sbp和数据,对应的Functor是GlobalTensorWithDataCtorFunctor
(github.com/Oneflow-Inc… )。中心逻辑在MakeGlobalTensorFromData(github.com/Oneflow-Inc… )中,其首要调用流程如下:
上述各个部分的首要功能如下:
-
DataConsistencyCheck (github.com/Oneflow-Inc… )会在tensor的placement涉及的各个节点间复制数据、校验数据是否共同。
-
functional::Empty (github.com/Oneflow-Inc… )会依据shape和dtype结构一个local tensor,并等待随后填充数据(这儿和之前讨论local tensor的过程共同)。
-
SwitchCopyLocalTensorFromUntypedArray (github.com/Oneflow-Inc… )为empty的local tensor填充数据,数据既能够是本例中的python list,也能够是numpy的ndarray。
-
functional::Cast (github.com/Oneflow-Inc… )进行数据类型dtype的转换。
-
functional::LocalToGlobal (github.com/Oneflow-Inc… )把local tensor转为global tensor,但这个仅仅用于broadcast 至指定placement的暂时的global tensor(sbp list悉数为broadcast,用于播送)。
-
functional::ToGlobal (github.com/Oneflow-Inc… )将暂时的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=1
import oneflow as flow
p = 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 (github.com/Oneflow-Inc… )中首要dispatch了”normal” op,在Eager Global的mode下, 会交给
EagerGlobalInterpreter
进行各种推导和准备工作(Interpret[github.com/Oneflow-Inc…] ),并在Interpret
办法里通过PhysicalRun
,将normal op履行的指令交给虚拟机调度并履行。 -
EagerGlobalTensorImpl::New (github.com/Oneflow-Inc… )时会调用GetPhysicalShape(github.com/Oneflow-Inc… )获取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源码
(github.com/Oneflow-Inc… )
- OneFlow源码解析1:算子签名的主动推断
- OneFlow源码解析2:Op、Kernel与解说器
- OneFlow源码解析3:Op指令在虚拟机中的履行
- OneFlow源码解析4:tensor系统与local tensor
- Global Tensor:docs.oneflow.org/master/para…
- 集群的大局视角:docs.oneflow.org/master/para…
- Global View的概念和实现
- OneFlow的Global Tensor笔记和实习总结
欢迎下载体验 OneFlow v0.8.0 最新版别:
github.com/Oneflow-Inc…