一种分布式深度学习编程新范式：Global Tensor

撰文｜姚迟、许啸宇、左益豪、程国良

Global Tensor 是指多机多设备执行的 Tensor，它是实现全局视角（Global View）编程的接口。

当前的并行程序，大都采用单程序多数据（SPMD）的方式来编程。并行执行同样的程序，但是处理的是不同数据，以此实现数据的并行处理。以 PyTorch DistributedDataParallel（DDP） 为例，每个进程执行同样的神经网络计算逻辑，但是每个进程加载数据集的不同分片。

单程序多数据（SPMD）编程的缺陷是多数据的通信很繁琐。在深度学习场景下，SPMD 编程需要在原计算代码中插入通信操作，比如数据并行时对梯度汇总（AllReduce 操作），模型并行时需要 AllGather/ReduceScatter 操作。如果并行模式复杂，或者需要试验新的并行模式，插入通信操作就变得难以开发和维护。

全局视角（Global View）编程提供了单程序单数据（SPSD）的编程视角。 与 SPMD 编程不同的是，Global View的数据 是同一个逻辑数据， 从编程接口层面看是单一数据，其实更简洁自然。

当我们把一个单进程程序扩展到并行执行时，一个单进程数据被扩展成多进程数据，多个进程上的这些数据都对应原单进程程序中的同一个逻辑数据。 这个逻辑数据在 OneFlow 中叫 Global Tensor。

编程时，Global Tensor 让用户可以用 SPSD 的接口来编程，即按照单机单设备的逻辑视角来写程序。然后 OneFlow 框架内部会自动地转换成物理的 SPMD/MPMD 方式来做并行/分布式执行。

使用 Global Tensor，就可以采用比较自然的 Global View 视角，把多机多设备看作 一个设备来编程，实现 SPSD 编程。

1

Global Tensor概述

在编程语言中，Global 的含义通常是进程内的全局可见，比如全局变量（Global Variable）。

但 Global Tensor 中 “Global” 的含义是进程间全局可见，所以 Global Tensor 更为准确的的说法是 Global (on all processes) Tensor，即所有进程可见的 Tensor。

Global Tensor 在每个进程上都存在，在所有进程上被某算子执行时，就自动完成了对该 Tensor 的多机多设备执行。

当前常用的 Tensor，只在单个进程内可见，存在于一个设备上， OneFlow 中把这种 Tensor 叫做 Local Tensor 。Local 是相对 Global 而言的，所以 Local Tensor 可以认为是 Local (on one process) Tensor。

OneFlow 的算子大部分兼容 Local Tensor 和 Global Tensor 的执行。Local Tensor 可以便捷地转化为 Global Tensor。如此，单机单卡执行的代码可以平滑地转换成多机多卡执行的代码。

使用 Global Tensor，可以非常便捷地进行多机多卡的模型开发，相比使用原始通信算子，可以成倍提高并行执行模型的开发效率。

2

创建Global Tensor

现在，尝试在有 2 张 GPU 的主机上创建一个 Global Tensor。以  randn  算子为例，创建一个 Python 文件  test_randn_global.py ，加入以下内容：

import oneflow as flow


# Place a global tensor on cuda device of rank(process) 0 and 1
placement = flow.placement(type="cuda", ranks=[0, 1])
# Each rank's local data is a part data as a result of spliting global data on dim 0
sbp = flow.sbp.split(dim=0)
# Create a global tensor by randn
x = flow.randn(4, 5, placement=placement, sbp=sbp)
# Print local data
print("Local data of global tensor:\n ", x.to_local().numpy())
# Print global data
print("Global data of global tensor:\n ", x.numpy())

在上述代码中有一些新出现的概念：

placement  表示 Global Tensor 分布的物理设备，参数  type  指定了物理设备的类型，这里使用 “cuda”  表示 GPU 设备，参数  ranks  指定了设备 ID。对于没有 2 张 GPU 的用户，在这里可以将  type  指定为  "cpu" ，这样可以使用 CPU 模拟多个设备，下文的代码同样适用。

sbp  表示 Global Tensor 分布的方式，代码中的  sbp = flow.sbp.split(dim=0)  表示把 Global Tensor 在维度 0 均匀切分。

to_local()  可以从 Global Tensor 中获取它在当前 rank 的 Local Tensor，因为 Global Tensor 在每个 rank 都内含了一个 Local Tensor 作为实际存在的本地分量。

然后配置下多进程启动依赖的环境变量。这里是两卡执行，对应两个进程启动，所以需要打开两个 Terminal，分别配置如下环境变量：