乐趣区

关于深度学习:一种分布式深度学习编程新范式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,别离配置如下环境变量:

Terminal 0

export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=0 LOCAL_RANK=0

Terminal 1

export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=1 LOCAL_RANK=1

以上环境变量的具体解释及借助工具做分布式启动,请参考
https://zhuanlan.zhihu.com/p/…

最初,在两个 Terminal 下别离启动一下test_randn_global.py,察看 Global Tensor 的创立后果:

这样在 Terminal 0 即 rank 0 能够看到:

Local data of global tensor:
  [[-0.07157125 -0.92717147  1.5102768   1.4611115   1.014263]
 [-0.1511031   1.570759    0.9416077   0.6184639   2.4420679]]
Global data of global tensor:
  [[-0.07157125 -0.92717147  1.5102768   1.4611115   1.014263]
 [-0.1511031   1.570759    0.9416077   0.6184639   2.4420679]
 [-0.38203463  0.453836    0.9136015   2.35773    -0.3279942]
 [-0.8570119  -0.91476554 -0.06646168  0.50022084 -0.4387695]]

在 Terminal 1 即 rank 1 能够看到:

Local data of global tensor:
  [[-0.38203463  0.453836    0.9136015   2.35773    -0.3279942]
 [-0.8570119  -0.91476554 -0.06646168  0.50022084 -0.4387695]]
Global data of global tensor:
  [[-0.07157125 -0.92717147  1.5102768   1.4611115   1.014263]
 [-0.1511031   1.570759    0.9416077   0.6184639   2.4420679]
 [-0.38203463  0.453836    0.9136015   2.35773    -0.3279942]
 [-0.8570119  -0.91476554 -0.06646168  0.50022084 -0.4387695]]

能够发现两个 rank 的 Local Tensor 在维度 0 拼接后,就是残缺的 Global Tensor 的值。

3

由 Local Tensor 失去 Global Tensor

能够先创立 Local Tensor,再利用 Tensor.to_global(https://oneflow.readthedocs.i…)办法,将 Local Tensor 转为 Global Tensor。

创立如下程序,采纳上文同样的形式启动:

import oneflow as flow

x = flow.randn(2, 5).cuda()
print(x.is_local) # True
print(x.is_global) # False
placement = flow.placement(type="cuda", ranks=[0, 1])
sbp = flow.sbp.split(0)
x_global = x.to_global(placement=placement, sbp=sbp)
print(x_global.shape) # (4, 5)
print(x.is_local) # True
print(x_global.is_global) # True

该程序在 2 个 GPU 设施上别离创立了 shape=(2,5) 的 Local Tensor,即 x。

而后定义 placement 为 rank 0 和 1 上的 CUDA 设施,SBP 为 tensor 第 0 维的切分,本来 Local Tensor 通过 to_global 变换后,就失去一个名为 x_global 的 Global Tensor。

能够察看到,x_global 的 shape 变为了 (4, 5),这是 Global Tensor 的 shape(global shape)。

Global Tensor 与 Local Tensor 之间为总量与重量的关系。Local Tensor 是总量在本 rank 的重量。重量和总量的具体关系由 Placement 和 SBP 确定,比方这里的关系是在 0 和 1 号 GPU 上,x_global 在第 0 维 split 而失去 x

to_global 办法依据如上关系能够从 x.shape 推理出 x_global.shape:把两个 GPU 上的 Local Tensor x 在第 0 维拼接后失去 x_global

Global Tensor 除了 shape,还有数据局部。一个 Global Tensor 的外部,在每个 rank 上都内含了一个 Local Tensor 作为其本地重量,这个 Local Tensor 就是 Global Tensor 在每个 rank 的物理数据。这合乎期待,每个 rank 只需保留一部分物理数据。

4

由 Global Tensor 失去 Local Tensor

如果想得到 Global Tensor 的本地重量,能够通过 to_local 办法失去。例如:

import oneflow as flow

placement = flow.placement(type="cuda", ranks=[0, 1])
sbp = flow.sbp.split(0)
x = flow.randn(4, 5, placement=placement, sbp=sbp)
print(x.to_local())

当执行 x.to_local() 时,两个不同的 rank 将别离失去一个 shape 为 (2, 5) 的本地重量 tensor。

在 Terminal 0 即 rank 0 能够看到:

tensor([[-0.2730,  1.8042,  0.0721, -0.5024, -1.2583],
        [-0.3379,  0.9371,  0.7981, -0.5447, -0.5629]],
       dtype=oneflow.float32)

在 Terminal 1 即 rank 1 能够看到:

tensor([[0.6829,  0.4849,  2.1611,  1.4059,  0.0934], 
        [-0.0301, -0.6942, -0.8094, -1.3050, -0.1778]], 
       dtype=oneflow.float32)

to_local() 没有任何参数,因为 Global Tensor 曾经通过 placement 和 SBP 确定好了它的本地重量,所以间接取本地重量对应的 Local Tensor 就好。

5

由 Global Tensor 转成另一个 Global Tensor

分布式计算通常都须要在失常的计算逻辑之间插入通信操作,而应用 OneFlow 时只须要做 Global Tensor 的数据分布类型转换。

相比一般的 Local Tensor,从类型上讲 ,Global Tensor 的 最大区别是带有全局数据分布类型(Global Data Distribution Type)。全局数据分布类型指定了 Global Tensor 在每个过程(Rank)的数据分布状况,由 Placement 和 SBP 组成。

全局数据分布类型中的 Placement 指定了数据分布的设施汇合:

  • 参数 type 指定了物理设施的类型,cuda 示意 GPU 设施内存, cpu 示意 CPU 设施内存;
  • 参数 ranks 指定了过程 ID 汇合,因为隐含了一个 Rank 对应一个物理设施,所以 ranks 就是设施 ID 汇合; 实际上 ranks 是一个由 rank id 组成 nd-array,反对高维设施排布。

详情请参考 oneflow.placement(https://oneflow.readthedocs.i…)。

全局数据分布类型中的 SBP 指定了全局数据和部分数据的关系:

  • S 即 split(dim),部分和全局是切分关系,示意在 dim 维度做了切分的数据分布关系;
  • B 即 broadcast,部分和全局是播送关系,示意做了播送的数据分布关系;
  • P 即 partial_sum,部分和全局是局部关系,示意做了 element-wise 累加的数据分布关系。

详情请参考 oneflow.sbp.sbp(https://oneflow.readthedocs.i…)。

数据重散布(Re-distribution)是并行计算中常常要解决的,即变换数据分布,比方把分片数据聚合到一起。在 MPI 编程范式(SPMD)下,数据重散布须要写显式的通信操作,如 AllReduce、AllGather、ReduceScatter。在 OneFlow 的 Global View 编程范式(SPSD) 下,数据重散布能够通过 Global Tensor 的全局数据分布类型转换实现。

全局数据分布类型的转换相似惯例编程语言中的(显式)类型转换。类型转换时,只需指定要变换到的类型,外面隐含的操作会被零碎主动实现。比方 double 类型到 int 类型的转换,去掉小数点局部的操作就是零碎主动实现的。

同样,只需指定 Global Tensor 要转换的新全局数据分布类型,外面隐含的通信操作会被 OneFlow 主动实现。全局数据分布类型转换的接口是 Tensor.to_global,to_globalplacementsbp 两个参数,这两个参数即冀望转换成的新全局数据分布类型。

全局数据分布类型转换中隐含的次要操作是通信的推理和执行,背地的实现机制是 OneFlow 的 Boxing,这是一种主动做数据 Re-distribution 的机制

上面看一个例子,该例子能够把一个按 split 散布的 Global Tensor 转换为一个按 broadcast 散布的 Global Tensor:

import oneflow as flow

x = flow.randn(2, 5).cuda()
placement = flow.placement(type="cuda", ranks=[0, 1])
sbp = flow.sbp.split(0)
x_global = x.to_global(placement=placement, sbp=sbp)
print(x_global.shape) # (4, 5)
print(x_global.to_local())
sbp_b = flow.sbp.broadcast
x_global_b = x_global.to_global(placement=placement, sbp=sbp_b)
print(x_global_b.shape) # (4, 5)
print(x_global_b.to_local())

能够看到,x_globalx_global_b 的全局数据分布类型变动就是 sbp 从 flow.sbp.split(0) 变成了 flow.sbp.broadcast。它们的 global shape 都是 (4, 5),然而本地重量从一个分片变成了一个残缺的数据,这个变动能够从对 to_local() 的打印后果察看到。

这里的 to_global 变换实现了对 local tensor 的归并。通常来讲,SPMD 编程模式要求用户手写一个 all-gather 汇合通信来实现。而在 OneFlow Global View 中,只需做一下类型转换。

通过 Global Tensor 的类型转换,就主动实现通信操作的推理和执行。让算法开发者能够 思考数据的散布(Thinking in data distribution)而不是思考如何通信(Thinking in data communication operation),实现了所想即所得,从而进步分布式程序的开发效率。

这里补充介绍一下 Global Tensor 的 numpy() 办法。对于任意的 Global Tensor 如 x_globalx_global.numpy() 等价于 x_global.to_global(spb=flow.sbp.broadcast).to_local().numpy(),即外部隐含了一次将原 Global Tensor 转成 SBP 为 flow.sbp.broadcast() 的 Global Tensor,而后进行一次 to_local 操作,最初对这个 Local Tensor 调用 numpy() 办法。所以 x_global.numpy() 失去的是一个残缺的数据。

6

Global Tensor 参加计算

这一节介绍 Global Tensor 如何参加理论计算。以 Global Tensor 参加矩阵乘法计算为例,结构如下程序:

import oneflow as flow

placement = flow.placement(type="cuda", ranks=[0, 1])
x = flow.randn(4, 5, placement=placement, sbp=flow.sbp.split(dim=0))
w = flow.randn(5, 8, placement=placement, sbp=flow.sbp.broadcast)
y = flow.matmul(x, w)
print(y.is_global)  # True
print(y.shape)  # (4, 8)
print(y.sbp)  # (flow.sbp.split(dim=0))
print(y.to_local().numpy())

以上程序创立了两个 Global Tensor,别离是 xw,它们参加 oneflow.matmul 计算失去 y

OneFlow 中的大部分算子都反对计算 Global Tensor。flow.matmul 执行 Global Tensor 时,在接口上并无非凡之处。能够认为 OneFlow 中的算子都是多态的。即依据输出,决定本人的行为:

  • 如果算子的输出是 Local Tensor,那么算子会依照一般的单机单设施执行模式进行计算;
  • 如果算子的输出是 Global Tensor,那么算子会采纳 Global View(多机多设施)模式进行计算;

当用户须要将单卡代码改为分布式代码时,算子反对多态执行为用户提供了极大的便当:只须要把输出的 (Local) Tensor 转换成 Global Tensor。

相似于单设施执行时要求输出数据所在设施雷同,以上程序中,flow.matmul 这一算子能够顺利执行的前置条件是:输出的 xw 的 placement 雷同。

程序中矩阵相乘的后果 y 同样是一个 Global Tensor。flow.matmul 对输出 xw 做计算时,会主动进行输入数据的 placement 和 SBP 的推理,规定如下:

  • Placement:输入和输出的 placement 雷同;
  • SBP:输入的 SBP 的推理规定,因算子类型而异,这个推理规定是 OneFlow 内置的,详情可见: SBP Signature

此处,flow.sbp.split(0)flow.sbp.broadcast 相乘的输入数据会被推理成 flow.sbp.split(0)x 在每个 rank 上是一个分片数据,w 是一个残缺的数据,二者矩阵乘法失去的 y 是一个分片的数据。看到这里,理解常见并行执行形式的用户能够发现:这里实现了一个数据并行的前向计算,x 是切片的数据,w 是残缺的参数。

7

结语

上文介绍了:

  • Global View 提供的 SPSD 编程视角;
  • Global Tensor 的跨过程可见的执行特点;
  • Global Tensor 和 Local Tensor 的互转;
  • 通过 Global Tensor 的全局数据分布类型转换来实现分布式通信;
  • OneFlow 算子的多态个性反对了 Global Tensor 的执行。

至此,本文从 Global Tensor 的创立开始,最终实现了一个基于 Global Tensor 的数据并行计算流程。更多并行形式和 SBP 的推理逻辑,将在后续内容介绍。

扩大浏览:OneFlow 多机多卡启动和依赖的环境变量

OneFlow 的 Global Tensor 执行采纳的是 多 客户端模式 (Multi-Client),每个设施对应一个过程。n 机 m 卡 的环境,就对应 n * m 个过程。每个过程都有一个过程 rank 编号,Global Tensor 中的 placement 参数中的 ranks 对应的就是这个 rank 编号。

2 机 2 卡 为例,0 号机器中两张卡别离对应编号 0 和 1,第 1 号机器中两张卡别离对应编号 2 和 3。此时 flow.placement(type="cuda", ranks=[2]) 能够惟一标识第 1 号机器中的第 0 卡。

个别地,对于 n 机 m 卡 的环境,flow.placement(type="cuda", ranks=[k]) 惟一标识第 k / n 号机器的第 k % m 张卡。

因为采纳多客户端模式,所以须要对应每个设施都启动一个过程。在 OneFlow 中,所有过程都只须要启动雷同的脚本程序。不同过程之间通过不同的环境变量来辨别过程编号和建设通信连贯。

环境变量阐明:

  • MASTER_ADDR:多机训练的第 0 号机器的 IP;
  • MASTER_PORT:多机训练的第 0 号机器的监听端口,不与曾经占用的端口抵触即可;
  • WORLD_SIZE:整个集群中计算设施的数目,因为目前还不反对各个机器上显卡数目不统一,因而 WORLD_SIZE 的数目实际上是 $ 机器数目 \times 每台机器上的显卡数目 $。创立 Global Tensor 中的示例是单机 2 卡的状况,因而 WORLD_SIZE=2;
  • RANK:集群内所有机器下的过程编号;
  • LOCAL_RANK:单个机器内的过程编号。

RANKLOCAL_RANK 的区别:

  • 在单机训练(单机单卡或单机多卡)时,两者相等;
  • 在多机训练时,每台机器上的 LOCAL_RANK 的下限,就是每台机器上的计算设施的数目;RANK 的下限,就是所有机器上所有计算设施的总和,它们的编号均从 0 开始(因为编号从 0 开始,所以不蕴含下限)。

2 机 2 卡 为例,每张显卡的 LOCAL_RANKRANK 对应状况如下:

应用环境变量启动尽管繁琐,然而适用性广,能够采纳任意的形式来启动过程。另外为了方便使用,OneFlow 也提供了一个分布式启动多过程且主动构建环境变量的工具 oneflow.distributed.launch(https://docs.oneflow.org/mast…)。

(原文:https://docs.oneflow.org/mast…

欢送下载体验 OneFlow v0.8.0 最新版本https://github.com/Oneflow-In…

退出移动版