By Xu Xiaoyu, YaoChi, Zuo Yihao, Guoliang Cheng

Global Tensor 是多机多设备执行的 Tensor,是实现全局视角(Global View)编程的接口。

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

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

全局视角(Global View)编程提供了单程序单数据(SPSD)的编程视角。和 SPMD 编程不同的是,数据从编程接口层面看也是单一的了。

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

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

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

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,可以非常便捷地进行多机多卡的模型开发,相比使用原始通信算子,可以成倍提高并行执行模型的开发效率。

创建 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())

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