关于深度学习: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 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.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=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
  （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…