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/master/parallelism/03_consistent_tensor.html），

展示从普通数据构造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(https://oneflow.readthedocs.io/en/master/generated/oneflow.Tensor.to_global.html?highlight=to_global)的方式，来创建global tensor并使用。

2

SBP