分布式训练策略杂谈
分布式训练策略杂谈
5.1 常见的分布式并行策略
5.1.1 为什么分布式训练越来越流行
近年来，深度学习被广泛应用到各个领域，包括计算机视觉、语言理解、语音识别、广告推荐等。在这些不同的领域中，一个共同的特点就是模型规模越来越大，比如 GPT-3 模型的参数量达到1750亿。即使用1024张 80 GB 的 A100，那么完整训练 GPT-3 的时长都需要1个月。
模型规模的扩大，对硬件（算力、内存）的发展提出要求。然而，因为内存墙的存在，单一设备的算力及容量，受限于物理定律，持续提高芯片的集成越来越困难，难以跟上模型扩大的需求。
为了解决算力增速不足的问题，人们考虑用多节点集群进行分布式训练，以提升算力，分布式训练势在必行。
1. 常见的并行策略
简单的机器堆叠并不一定会带来算力的增长。因为神经网络的训练并不是单纯的“把原来一个设备做的事情，现在分给多个设备各自做”，它不仅需要多个设备进行计算，还涉及到设备之间的数据传输，只有协调好集群中的计算与通信，才能做高效的分布式训练。
将以矩阵乘法的例子，解释数据并行、模型并行的区别。
先了解以下逻辑上的矩阵乘法例子:
假设神经网络中某一层是做矩阵乘法，其中的输入 𝑥 的形状为 4×5，模型参数 𝑤 的形状为 5×8，那么，矩阵乘法输出形状为 4×8。如图5-1所示。
图5-1 矩阵乘法示例
单机单卡的训练中，以上矩阵乘法，先计算得到 𝑜𝑢𝑡，并将 𝑜𝑢𝑡 传递给下一层，并最终计算得到 𝑙𝑜𝑠𝑠，然后在反向传播过程中，得到 𝜕𝑙𝑜𝑠𝑠𝜕𝑤，用于更新 𝑤。
分布式训练中，依据是切分 𝑥 还是 𝑤 的不同，分为“数据并行”和“模型并行”策略。接下来，介绍常见的并行策略。
2. 数据并行
所谓的数据并行，就是将数据 𝑥 进行切分，而每个设备上的模型 𝑤 是完整的、一致的。如图5-2所示，𝑥 被按照第0维度平均切分到2个设备上，两个设备上都有完整的 𝑤。
这样，在两台设备上，分别得到的输出，都只是逻辑上输出的一半（形状为 2×8），将两个设备上的输出拼接到一起，才能得到逻辑上完整的输出。
图5-2 数据并行示例
注意，因为数据被分发到了2个设备上，因此反向传播过程，各自设备上得到的 𝜕𝑙𝑜𝑠𝑠𝜕𝑤 会不一样，如果直接使用各个设备上的梯度更新各自的模型，会造成2个设备上的 模型不一致，训练就失去了意义（到底用哪个模型好呢？）。
因此，数据并行策略下，在反向传播过程中，需要对各个设备上的梯度进行 AllReduce，以确保各个设备上的模型始终保持一致。
当数据集较大，模型较小时，由于反向过程中为同步梯度产生的通信代价较小，此时选择数据并行一般比较有优势，常见的视觉分类模型，如 ResNet50，比较适合采用数据并行。
3. 模型并行
当神经网络非常巨大，数据并行同步梯度的代价就会很大，甚至网络可能巨大到无法存放到单一计算设备中，这时候，可以采用模型并行策略解决问题。
所谓的模型并行，就是每个设备上的数据是完整的、一致的，而模型 𝑤 被切分到了各个设备上，每个设备只拥有模型的一部分，所有计算设备上的模型拼在一起，才是完整的模型。
如图5-3所示，𝑤 被按照第1维度平均切分到2个设备上，两个设备上都有完整的 𝑥。两个设备上的输出也需要通过拼接才能得到逻辑上的输出。
图5-3 模型并行示例
模型并行的好处是，省去了多个设备之间的梯度 AllReduce；但是，由于每个设备都需要完整的数据输入，因此，数据会在多个设备之间进行广播，产生通信代价。比如，图5-3中的最终得到的 𝑜𝑢𝑡 (4×8)，如果它作为下一层网络的输入，那么它就需要被广播发送到两个设备上。
语言模型，如 BERT，常采用模型并行。
4. 流水并行
当神经网络过于巨大，无法在一个设备上存放时，除了上述的模型并行的策略外，还可以选择流水并行。 流水并行指将网络切为多个阶段，并分发到不同的计算设备上，各个计算设备之间以“接力”的方式完成训练。
如图5-4所示，展示了一个逻辑上的4层网络（T1 至 T4)是如何做流水并行的。
4层网络被切分到2个计算设备上，其中 GPU0 上进行 T1 与 T2 的运算，GPU1 上进行 T3 与 T4 的计算。
GPU0 上完成前两层的计算后，它的输出被当作 GPU1 的输入，继续进行后两层的计算。
图5-4 流水线并行示例
5. 混合并行
网络的训练中，也可以将多种并行策略混用，以 GPT-3 为例，以下是它训练时的设备并行方案：
首先被分为 64 个阶段，进行流水并行。每个阶段都运行在 6 台 DGX-A100 主机上。在6台主机之间，进行的是数据并行训练；每台主机有 8 张 GPU 显卡，同一台机器上的8张 GPU 显卡之间是进行模型并行训练。
图5-5 混合并行示例
并行策略的选择影响着训练效率，框架对并行训练的接口支持程度，决定了算法工程师的开发效率。OneFlow 针对分布式训练所做的系统级设计和创新，为用户轻松上手分布式训练做足了铺垫。
5.1.2 集群的全局视角
OneFlow 提出了 全局视角（Global View） 的概念，用于简化分布式训练。简单而言，在 OneFlow 的全局视角下，集群被抽象为一台超级计算设备，如图5-6所示。
用户不用关心集群中计算、通信的细节，只需关心逻辑上的数据与计算，依然像单机单卡那样思考、编程，就能进行分布式训练。
图5-6 全局视觉框架示例
OneFlow 的全局视角，依赖几个重要概念：位置属性、SBP 与 SBP签名。
1. 位置属性
OneFlow 全局视角下的张量有完位置属性，通过位置属性可以指定该张量存放在哪个物理设备上。
OneFlow 会自动为集群中的计算设备编号。比如，如果集群中有 4 台主机，每台主机上有 8 张显卡，一共 32 张显卡，那么 OneFlow 会将这 32 张显卡自动编号为 0~31。
如果想将张量放置在第 0 台机器的前 4 张显卡上，只需要配置：placement("cuda", [0, 1, 2, 3])。
如果想将张量放置在第 0 台机器的后 4 张显卡上，只需要配置：placement("cuda", [4, 5, 6, 7])。
placement 使得 OneFlow 很容易支持流水并行，将在本专题的其它文章中看到与 placement 有关的实际例子。
2. SBP属性
SBP 是 OneFlow 发明的概念，描述了超级计算设备全局视角下的数据与集群中真实的物理设备上的数据的映射关系，它由 split, broadcast, partial 的首字母组合而成。
详细而言：
1）分支表示物理设备上的张量，是将全局视角的张量切分得到的。切分时，需要指定切分的维度。物理设备上的 Tensor，经过拼接，可以还原得到全局视角的张量。
2）广播表示全局视角下的 Tensor，会复制并广播到所有的物理设备上。
3）部分和表示全局视角下的张量与物理设备上的张量的形状相同，但是物理设备上的值，只是全局视角下张量的一部分。以部分和为例，如果将集群中所有设备的张量按位置相加，那么就可以还原得到全局视角的张量。除了 sum 外，min、max 等操作也适用于部分。
如图5-7所示，分别展示了 SBP 的情况，分别是 split(0)、split(1)、广播 和部分和。
图5-7 全局视觉SBP框架示例
在创建全局张量时，可以指定张量的SBP。
3. SBP签名
SBP 描述了全局视角下的数据与物理设备上的数据的映射关系，当进行分布式训练时，OneFlow 根据数据的 SBP 属性，将数据分发到各个物理设备，进行计算，并输出结果。
对于一个孤立的张量，可以随意设置它的 SBP 属性。但是，对于一个有输入、输出数据的算子，却不可以随意设置它的输入、输出的 SBP 属性。这是因为随意设置一个算子输入输出的 SBP 属性，可能不符合全局视角下算子的运算法则。
以矩阵乘法为例讨论这个问题。看看在有2个设备的分布式系统中，矩阵乘法的输入、输出的 SBP，要如何组合才合法，如何组合不合法。
假设全局视角下，一个形状为 (𝑚,𝑘) 的矩阵 𝐴 与形状为 (𝑘,𝑛) 的矩阵 𝐵 相乘得到 𝑌，𝑌 的形状必然为 (𝑚,𝑛)。
依据矩阵乘法的规律，可以将矩阵 𝐴 按第0维进行切分，切分为形状分别为 (𝑚0,𝑘)、(𝑚1,𝑘) 的两个矩阵：𝐴0 和 𝐴1，然后在2个设备上分别计算：
设备一：
𝐴0×𝐵=𝑌0(𝑚0,𝑘)(𝑘,𝑛)(𝑚0,𝑛)
设备二：
𝐴1×𝐵=𝑌1(𝑚1,𝑘)(𝑘,𝑛)(𝑚1,𝑛)
容易得到物理设备上的 𝐴0、𝐴1 与全局视角 𝐴 的关系，以及 𝑌0、𝑌1 与全局视角数据 𝑌 的关系：
𝐴=𝑐𝑜𝑛𝑐𝑎𝑡(𝐴0,𝐴1)(𝑚,𝑘)(𝑚0,𝑘)(𝑚1,𝑘)
𝑌=𝑐𝑜𝑛𝑐𝑎𝑡(𝑌0,𝑌1)(𝑚,𝑛)(𝑚0,𝑛)(𝑚1,𝑛)
注意：以上的 concat 表示拼接操作。
按照以上的方式，将全局视角的数据分发到各个物理设备上，是能够完成运算，并且最终得到全局视角上的正确结果的。以上较长的篇幅，若 使用 SBP 来描述，会变得异常简单：
𝐴 为 split(0)， 𝐵 为 broadcast，运算结果 𝑌 为 split(0)。
可见，对于矩阵乘法而言，其输入输出的 SBP，按以上方式组合，是合法的。对于矩阵乘法而言，合法的 SBP 组合不止一种，比如，还可以是这样的：
𝐴 为 broadcast， 𝐵 为 split(1)，运算结果 𝑌 为 split(1)。
或者：
𝐴 为 split(1)， 𝐵 为 split(0)，运算结果 𝑌 为 partial sum。
虽然展示了多个合法的 SBP 组合，但是并不是任意的 SBP 组合都是合法的，比如对于矩阵乘法，如果 𝐴、𝐵 均为 split(0)，那么：
𝐴=𝑐𝑜𝑛𝑐𝑎𝑡(𝐴0,𝐴1)(𝑚,𝑘)(𝑚0,𝑘)(𝑚1,𝑘)
𝐵=𝑐𝑜𝑛𝑐𝑎𝑡(𝐵0,𝐵1)(𝑘,𝑛)(𝑘0,𝑛)(𝑘1,𝑛)
在物理设备上，因为 𝐴0 与 𝐵0 的形状，并不满足矩阵乘法的要求，也就无法在物理设备上完成矩阵乘法。可以说， 𝐴 为 split(0)， 𝐵 为 split(0) 的 SBP 组合是不合法的。
对于某个算子，其输入输出的一个特定的、合法的 SBP 组合，称为这个算子的一个 SBP签名。
4. SBP签名自动推导
有了 SBP签名 的概念后，可能会提出几个问题：
1） 用户是否要知道算子的所有 SBP签名，才能用 OneFlow 做好分布式训练？
2）作为算法工程师，用户是否要为每层网络都设置输入的 SBP？
对于前一个问题，用户当然不需要知晓算子所有的 SBP签名。罗列某个算子所有可能的 SBP签名的工作，是算子工程师的责任。算子工程师根据算子的运算法则，在开发算子时，就已经罗列并预设好该算子所有可能的 SBP签名。
这顺便就解答了第二个问题：因为有预设好的 SBP签名，所以，某一层算子只要有输入的SBP，OneFlow 就可以根据 SBP签名，推导出该层算子输出的SBP。而上游算子的输出，又是下游算子的输入，这样，就确定了下游算子输入的 SBP，然后又可以根据 SBP签名，确定更下游输出的 SBP…… 这样不断推导、传播。因此，用户是不需要为每层网络都设置输入的 SBP。而只有最初输入层，或者需要强制指定某层的 SBP 时，才需要显式指定。
用户还可能会有新的问题：
1）一个算子的合法 SBP签名常常有多个，OneFlow 运行时到底会选择哪一个呢，它是依据什么做出选择的？
需要了解 OneFlow 的 SBP签名自动推导机制。所谓的 SBP签名自动推导，是指：在给定所有算子的所有合法的 SBP签名的前提下，OneFlow 有一套算法，会基于传输代价，为每种合法的 SBP签名进行打分，并选择传输代价最小的那个SBP签名。这样使得系统的吞吐效率最高。
5. 对抗机制
OneFlow 的对抗机制对于用户其实是透明的，用户使用 OneFlow 做分布式训练时，不用知晓感知不到它。
但是，鉴于某些深入思考的用户，可能了解 SBP签名自动推导后，会自然提出以下问题：
2）如果 OneFlow 自动选择的 SBP签名，上一层算子的输出与下一层算子的输入的 SBP 属性不匹配时，那怎么办呢？
举个具体例子，比如以下代码中，上一层算子 matmul 的输出 SBP 本来是 split(0)，但是下一层算子 matmul 的输入，被转成了广播。此时，上一层的输出与下一层的输入，它们的 SBP 其实就不一致了。
import oneflow as flow
P0 = flow.placement("cuda", ranks=[0, 1])
P1 = flow.placement("cuda", ranks=[2, 3])
a0_sbp = flow.sbp.split(0)
b0_sbp = flow.sbp.broadcast
y0_sbp = flow.sbp.broadcast
b1_sbp = flow.sbp.split(1)
A0 = flow.randn(4, 5, placement=P0, sbp=a0_sbp)
B0 = flow.randn(5, 8, placement=P0, sbp=b0_sbp)
Y0 = flow.matmul(A0, B0)
Y0 = Y0.to_global(placement=P1, sbp=y0_sbp)
B1 = flow.randn(8, 6, placement=P1, sbp=b1_sbp)
Y2 = flow.matmul(Y0, B1)
OneFlow 其实会检测到这种不一致，并且在上游的输出和下游的输入间插入一个算子，做相关的转换工作。这类自动加入做转换的算子，就称为对抗算子。以上代码的逻辑图和物理执行图的对应关系，如图5-8所示。
图5-8逻辑图和物理执行图的对应关系
总结
placement 与 SBP、SBP签名 是 OneFlow 分布式全局视角的重要保证，OneFlow 的全局视角使得 OneFlow 的分布式训练与单机单卡一样简单。
通常情况下，用户只需要在起始网络层设置 SBP，由此可以省略传统分布式训练中，手写通信操作的麻烦。除了介绍的 SBP签名自动推导机制外，OneFlow研发一种寻求全局最优解的自动并行方法，用户可以不做任何 SBP 配置，就得到很好的分布式训练效果。
5.2 全局张量
5.2.1 全局视角与物理视角的映射
1. 创建全局张量
要在有2张 GPU 显卡的主机上交互式体验全局张量，可以用以下方式在2个控制台分别启动 python。
分别使用以下 Terminal 0 或 Terminal 1 标签，查看2个控制台的命令/代码。
Terminal 0
export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=0 LOCAL_RANK=0
python3
Terminal 1
2. 直接创建全局张量
在两个控制台，分别导入 oneflow，并创建 x。
其中 flow.placement("cuda", [0,1]) 指定了全局张量在集群的范围。
1）cuda表示在 GPU 设备上。
2）placement 的第二个参数是一个列表，集群中的设备会被自动编号，通过数字指定使用哪些设备。[0,1] 表示全局张量使用第 0、1 张显卡上。
Terminal 0
import oneflow as flow
placement = flow.placement("cuda", [0,1])
sbp = flow.sbp.split(0)
x = flow.randn(4,5,placement=placement, sbp=sbp)
x.shape
Terminal 1
输出：
Terminal 0
oneflow.Size([4, 5])
Terminal 1
3. 由全局张量得到局部张量
通过 to_local 方法，可以查看物理设备上的局部张量：
Terminal 0
x.to_local()
tensor([[ 2.9186e-01, -3.9442e-01, 4.7072e-04, -3.2216e-01, 1.7788e-01],
[-4.5284e-01, 1.2361e-01, -3.5962e-01, 2.6651e-01, 1.2951e+00]],
device='cuda:0', dtype=oneflow.float32)
Terminal 1
4. 由局部张量转换得到全局张量
可以先创建局部张量，再利用 Tensor.to_global 方法，将局部张量转为全局张量。
下面的例子中，在2台设备上分别创建了 shape=(2,5) 的2个局部张量。注意经过 to_global 方法后，得到的全局张量的 shape 为 (4,5)。
这是因为选择的 sbp=flow.sbp.split(0)，2个形状为 (2,5) 的局部张量，需要在第0维拼接，得到 (4,5) 的全局张量。
Terminal 0
import oneflow as flow
x = flow.randn(2,5)
placement = flow.placement("cuda", [0,1])
sbp = flow.sbp.split(0)
x_global = x.to_global(placement=placement, sbp=sbp)
x_global.shape
Terminal 1
5.2.2 实践 SBP签名的作用
1. 数据并行
以下的代码对应了常见的分布式策略的数据并行，如图5-9所示。
图5-9常见的分布式策略的数据并行
Terminal 0
import oneflow as flow
placement = flow.placement("cuda", [0,1])
x = flow.randn(4,5,placement=placement, sbp=flow.sbp.split(0))
w = flow.randn(5,8,placement=placement, sbp=flow.sbp.broadcast)
y = flow.matmul(x,w)
y.sbp
y.shape
Terminal 1
可以观察到，flow.matmul 根据输入 x 与 w 的 SBP 分别为 split(0)、broadcast。OneFlow 自动推导出输出 y 的 SBP 应该为 split(0)，完成计算，得到 shape=(4,8) 的矩阵。输出如下：
Terminal 0
(oneflow.sbp.split(dim=0),)
oneflow.Size([4, 8])
Terminal 1
2. 模型并行
以下的代码对应了常见的分布式策略的模型并行，如图5-10所示。
图5-10对应了常见的分布式策略的模型并行
Terminal 0
import oneflow as flow
placement = flow.placement("cuda", [0,1])
x = flow.randn(4,5,placement=placement, sbp=flow.sbp.broadcast)
w = flow.randn(5,8,placement=placement, sbp=flow.sbp.split(1))
y = flow.matmul(x,w)
y.sbp
y.shape
Terminal 1
可以观察到，flow.matmul 根据输入 x 与 w 的 SBP 分别为 broadcast、split(1)。OneFlow 自动推导出输出 y 的 SBP 应该为 split(1)，完成计算，得到 shape=(4,8) 的矩阵。输出如下：
Terminal 0
(oneflow.sbp.split(dim=1),)
oneflow.Size([4, 8])
Terminal 1
5.2.3多机训练时的环境变量
通过设置环境变量配置分布式训练，仅仅是为了在交互式 Python 环境下，方便查看实验效果。如果不是学习、试验目的，而是生产需求，可以直接通过 oneflow.distributed.launch 启动分布式训练，该模块内部根据命令行参数，自动设置了必要的环境变量。
1）MASTER_ADDR：多机训练的第0号机器的 IP。
2）MASTER_PORT：多机训练的第0号机器的监听端口，不与已经占用的端口冲突即可。
3）WORLD_SIZE：整个集群中计算设备的数目，因为目前还不支持各个机器上显卡数目不一致，因此 WORLD_SIZE 的数目实际上是机器数目每台机器上的显卡数目，机器数目每台机器上的显卡数目机器数目×每台机器上的显卡数目。如这个例子中，是单机2卡的情况，因此 WORLD_SIZE=2。
RANK 和 LOCAL_RANK 都是对计算设备的编号，不同的是 RANK 是全局视角的编号，LOCAL_RANK 某个特定机器上的局部视角的编号。当是单机训练（单机单卡或单机多卡）时，两者是没有区别的。以上的例子中，有两个显卡，分别是0号和1号。
当多机训练时，每台机器上的 LOCAL_RANK 的上限，就是每台机器上的计算设备的数目；RANK 的上限，就是所有机器上所有计算设备的总和，它们的编号均从0开始。（因为编号从0开始，所以不包含上限）
以两台机器、每台机器上有两张显卡为例，可以整理出每张显卡的 LOCAL_RANK 与 RANK 对应情况，见表5-1。
表5-1 整理出每张显卡的 LOCAL_RANK 与 RANK 对应情况
|
|
RANK |
LOCAL_RANK |
|
机器0的第0张显卡 |
0 |
0 |
|
机器0的第1张显卡 |
1 |
1 |
|
机器1的第0张显卡 |
2 |
0 |
|
机器1的第1张显卡 |
3 |
1 |
1. Boxing（自动转换 SBP）
已经通过以上代码的例子，知道一个算子会根据输入张量的 SBP 属性以及算子内置的 SBP签名，自动设置输出张量的SBP。
如果上游算子输出张量的 SBP，与下游算子输入的需要不一致时，怎么办呢？
比如，假设在模型并行中，有2层矩阵乘法，在第一层和和第二层都做模型并行。如图5-11所示。
图5-11假设在模型并行中，有2层矩阵乘法，在第一层和和第二层都做模型并行
因为第一层的输出的 SBP（split(1)），并不是第二层输入所期待的（broadcast），这时候，OneFlow 会自动在上一层的输出和下一层的输出之间，插入对抗操作，利用集合通信进行必要的数据转换。
从 split(1) 转换为广播，相当于做了一次 AllGather 操作。如图5-12所示。
图5-12 从 split(1) 转换为广播，相当于做了一次 AllGather 操作
因为有对抗机制的存在，使得用户只用关心少数关键地方（如源算子）的 SBP 设置，剩下的全部都可以交给 OneFlow 框架。
2. 2D SBP
现在介绍 2D SBP，它能够更灵活地应对更复杂的分布式训练场景。
3. 2D 设备阵列
已经熟悉 1D SBP 的位置配置，在 1D SBP 的场景下，通过 oneflow.placement 接口配置集群，比如使用集群中的第 0~3 号 GPU 显卡：
>>> placement1 = flow.placement("cuda", ranks=[0, 1, 2, 3])
以上的 cuda指定了设备类型，ranks=[0, 1, 2, 3] 指定了集群中的计算设备。其实，ranks 不仅可以是一维的int list，还可以是多维的int数组：
placement2 = flow.placement("cuda", ranks=[[0, 1], [2, 3]])
当 ranks 是 ranks=[0, 1, 2, 3] 这种一维列表的形式时，集群中的所有设备组成了一个 1D 设备向量，这也是 1D SBP 名称的由来。
当 ranks 是多维数组的形式时，集群中的设备被分组为一个多维的设备阵列。ranks=[[0, 1], [2, 3]] 表示集群中的四个计算设备被划分为了 2×2 的设备阵列。
2D SBP
构造全局张量时，需要同时指定位置与 SBP。当位置中的集群是 2 维的设备阵列时；SBP 也必须与之对应，是一个长度为 2 的元祖，这个元祖中的第 0 个、第 1 个 元素，分别描述了全局张量张量在设备阵列第 0 维、第 1 维的分布。
以下代码，配置了 2×2 的设备阵列，并且设置 2D SBP 为 (broadcast, split(0))。
>>> a = flow.Tensor([[1,2],[3,4]])
>>> placement = flow.placement("cuda", ranks=[[0, 1], [2, 3]])
>>> sbp = (flow.sbp.broadcast, flow.sbp.split(0))
>>> a_to_global = a.to_global(placement=placement, sbp=sbp)
逻辑上的数据，在整个设备阵列上，在第 0 维度（竖着看）做 广播；在第 1 维度（横着看）做 split(0)。如图5-13所示。
图5-13 在第 0 维度（竖着看）做 广播；在第 1 维度（横着看）做 split(0)
此图的最左边是全局视角的数据，最右边是设备阵列上各个设备的数据。可以看到，从第 0 维的角度看，它们都是 broadcast 的关系：
1）(group0, device0) 与 (group1, device0) 中数据一致，互为 broadcast 关系
2）(group0, device1) 与 (group1, device1) 中数据一致，互为 broadcast 关系
从第 1 维的角度看，它们都是 split(0) 的关系：
1）(group0, device0) 与 (group0, device1) 互为 split(0) 关系。
2）(group1, device0) 与 (group1, device1) 互为 split(0) 关系。
直接理解逻辑数据和最终的设备阵列中的物理数据，对应关系可能有一定难度，在思考 2D SBP 时，可以假想一个中间状态（图5-13中灰色部分），以 (broadcast, split(0)) 为例：
1）原始逻辑张量，先经过广播，广播到 2 个 group 上，得到中间的状态。
2）在中间状态的基础上，继续在各自的 group 上，做 split(0)，得到最终设备阵列中各个物理张量的状态。
4. 2D SBP签名
类似 1D SBP 有 SBP签名的概念，算子也有 2D SBP签名，在掌握了 1D SBP 及其 签名的基础上，2D SBP签名非常简单，只需要遵循一条原则：在各自的维度上独立推导即可。
以矩阵乘法为例，先回顾 1D SBP 的情况，假定有 𝑥×𝑤=𝑦 有以下的 SBP签名：
𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡×𝑠𝑝𝑙𝑖𝑡(1)=𝑠𝑝𝑙𝑖𝑡(1)
以及
𝑠𝑝𝑙𝑖𝑡(0)×𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡=𝑠𝑝𝑙𝑖𝑡(0)
假定给 𝑥 设置了 2D SBP 为：(𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡,𝑠𝑝𝑙𝑖𝑡(0))，给 𝑤 设置 2D SBP 为 (𝑠𝑝𝑙𝑖𝑡(1),𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡)，那么，在 2D SBP 的背景下， 𝑥×𝑤=𝑦 运算，得到 𝑦 的 SBP 属性为 (𝑠𝑝𝑙𝑖𝑡(1),𝑠𝑝𝑙𝑖𝑡(0))。
也就是说，以下几个 2D SBP，构成矩阵乘法的 2D SBP签名：
(𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡,𝑠𝑝𝑙𝑖𝑡(0))×(𝑠𝑝𝑙𝑖𝑡(1),𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡)=(𝑠𝑝𝑙𝑖𝑡(1),𝑠𝑝𝑙𝑖𝑡(0))
5. 2D SBP 使用示例
将通过一个简单的例子，演示如何使用 2D SBP 进行分布式训练。假设有一个 2×2 的设备阵列，如果没有多个 GPU 设备，将使用 CPU 来模拟 2×2 设备阵列的情形，对输入张量采用上文图中 (broadcast, split(0)) 的并行策略。
首先，导入依赖：
import oneflow as flow
import oneflow.nn as nn
然后，定义要使用到的 placement 和 sbp：
PLACEMENT = flow.placement("cpu", [[0, 1], [2, 3]])
BROADCAST = (flow.sbp.broadcast, flow.sbp.broadcast)
BS0 = (flow.sbp.broadcast, flow.sbp.split(0))
PLACEMENT 的 ranks 参数是一个二维列表，代表将集群中的设备划分成 2×2 的设备阵列。SBP 需要与其对应，指定为长度为 2 的 tuple。其中，BROADCAST 表示在设备阵列的第 0 维和第 1 维都进行广播，BS0 的含义与前文的描述相同。
假设有以下模型：
model = nn.Sequential(nn.Linear(8, 4),
nn.ReLU(),
nn.Linear(4, 2))
将模型在集群上广播：
model = model.to_global(placement=PLACEMENT, sbp=BROADCAST)
然后构造数据并进行前向推理：
x = flow.randn(1, 2, 8)
global_x = x.to_global(placement=PLACEMENT, sbp=BS0)
pred = model(global_x)
创建了一个形状为 (1, 2, 8) 的局部张量，然后通过 Tensor.to_global 方法，获取对应的 全局张量，最后将其输入到模型中进行推理。
通过 Tensor.to_local 方法，获取当前物理设备上的局部张量后，可以通过输出其形状和值来验证数据是否被正确处理：
local_x = global_x.to_local()
print(f'{local_x.device}, {local_x.shape}, \n{local_x}')
输出结果为：
cpu:2, oneflow.Size([1, 2, 8]),
tensor([[[ 0.6068, 0.1986, -0.6363, -0.5572, -0.2388, 1.1607, -0.7186, 1.2161],
[-0.1632, -1.5293, -0.6637, -1.0219, 0.1464, 1.1574, -0.0811, -1.6568]]], dtype=oneflow.float32)
cpu:3, oneflow.Size([1, 2, 8]),
tensor([[[-0.7676, 0.4519, -0.8810, 0.5648, 1.5428, 0.5752, 0.2466, -0.7708],
[-1.2131, 1.4590, 0.2749, 0.8824, -0.8286, 0.9989, 0.5599, -0.5099]]], dtype=oneflow.float32)
cpu:1, oneflow.Size([1, 2, 8]),
tensor([[[-0.7676, 0.4519, -0.8810, 0.5648, 1.5428, 0.5752, 0.2466, -0.7708],
[-1.2131, 1.4590, 0.2749, 0.8824, -0.8286, 0.9989, 0.5599, -0.5099]]], dtype=oneflow.float32)
cpu:0, oneflow.Size([1, 2, 8]),
tensor([[[ 0.6068, 0.1986, -0.6363, -0.5572, -0.2388, 1.1607, -0.7186, 1.2161],
[-0.1632, -1.5293, -0.6637, -1.0219, 0.1464, 1.1574, -0.0811, -1.6568]]], dtype=oneflow.float32)
通过比较这些不同设备上局部张量，符合上文图中描述的状态，证明数据已被正确分布到各个设备上。
需要注意的是，不能直接通过 python xxx.py 的方式执行上述代码，而需要通过 oneflow.distributed.launch 启动。此模块可以方便地启动分布式训练，在终端中执行下列命令（假设上述代码已经保存至当前目录中的名为2d_sbp.py的文件中）：
python3 -m oneflow.distributed.launch --nproc_per_node=4 2d_sbp.py
通过将参数 nproc_per_node 指定为 4 来创建 4 个进程，模拟共有 4 个 GPU 的情形。
完整代码如下：
PLACEMENT = flow.placement("cpu", [[0, 1], [2, 3]])
BROADCAST = (flow.sbp.broadcast, flow.sbp.broadcast)
BS0 = (flow.sbp.broadcast, flow.sbp.split(0))
model = nn.Sequential(nn.Linear(8, 4),
nn.ReLU(),
nn.Linear(4, 2))
model = model.to_global(placement=PLACEMENT, sbp=BROADCAST)
x = flow.randn(1, 2, 8)
global_x = x.to_global(placement=PLACEMENT, sbp=BS0)
pred = model(global_x)
local_x = global_x.to_local()
print(f'{local_x.device}, {local_x.shape}, \n{local_x}')
6. 用 launch 模块启动分布式训练
OneFlow 提供了 oneflow.distributed.launch 模块，帮助用户更方便地启动分布式训练。
用户可以借助以下的形式，启动分布式训练：
python3 -m oneflow.distributed.launch [启动选项] train script.py
比如，启动单机两卡的训练：
python3 -m oneflow.distributed.launch --nproc_per_node 2 ./script.py
再比如，启动两台机器，每台机器有两张显卡的训练。
在0号机器上运行：
python3 -m oneflow.distributed.launch \
--nnodes=2 \
--node_rank=0 \
--nproc_per_node=2 \
--master_addr="192.168.1.1" \
--master_port=7788 \
script.py
在1号机器上运行：
python3 -m oneflow.distributed.launch \
--nnodes=2 \
--node_rank=1 \
--nproc_per_node=2 \
--master_addr="192.168.1.1" \
--master_port=7788 \
script.py
7. 常见选项说明
通过 python3 -m oneflow.distributed.launch -h， 可以查看 launch 模块的选项说明，以下是部分常见选项。
1）--nnodes：机器的数目(number of nodes)
2）--node_rank： 机器的编号，从0开始
3）--nproc_per_node：每台机器上要启动的进程数目(number of processes per node)，推荐与 GPU 数目一致
4）--logdir：子进程日志的相对存储路径
8. launch 模块与并行策略的关系
注意 oneflow.distributed.launch 的主要作用，是待用户完成分布式程序后，让用户可以更方便地启动分布式训练。它省去了配置集群中环境变量的繁琐。
但是 oneflow.distributed.launch 并不决定并行策略，并行策略是由设置数据、模型的分发方式、在物理设备上的放置位置决定的。
OneFlow 提供的 全局视角和全局张量，可以灵活地配置并行策略。并且针对数据并行，OneFlow 提供了 DistributedDataParallel 模块，可以在极少修改代码的前提下，将单机单卡的脚本改为数据并行的脚本。
5.2.4 数据并行训练
在 OneFlow中，提供了两种做数据并行的方式。
一种是使用 OneFlow 的原生的 SBP 概念，通过设置全局张量，进行数据并行训练，这也是用 OneFlow 做数据并行训练的推荐方式 。
此外，为了方便从 PyTorch 迁移到 OneFlow 的用户，OneFlow 提供了与 torch.nn.parallel.DistributedDataParallel 对齐一致的接口 oneflow.nn.parallel.DistributedDataParallel，它也能让用户方便地从单机训练脚本，扩展为数据并行训练。
1. 通过设置 SBP 做数据并行训练
以下代码，是通过配置设置全局张量，完成数据并行训练。
import oneflow as flow
import oneflow.nn as nn
import flowvision
import flowvision.transforms as transforms
BATCH_SIZE=64
EPOCH_NUM = 1
PLACEMENT = flow.placement("cuda", [0,1])
S0 = flow.sbp.split(0)
B = flow.sbp.broadcast
DEVICE = "cuda" if flow.cuda.is_available() else "cpu"
print("Using {} device".format(DEVICE))
training_data = flowvision.datasets.CIFAR10(
root="data",
train=True,
transform=transforms.ToTensor(),
download=True,
)
train_dataloader = flow.utils.data.DataLoader(
training_data, BATCH_SIZE, shuffle=True
)
model = flowvision.models.mobilenet_v2().to(DEVICE)
model.classifer = nn.Sequential(nn.Dropout(0.2), nn.Linear(model.last_channel, 10))
model = model.to_global(placement=PLACEMENT, sbp=B)
loss_fn = nn.CrossEntropyLoss().to(DEVICE)
optimizer = flow.optim.SGD(model.parameters(), lr=1e-3)
for t in range(EPOCH_NUM):
print(f"Epoch {t+1}\n-------------------------------")
size = len(train_dataloader.dataset)
for batch, (x, y) in enumerate(train_dataloader):
x = x.to_global(placement=PLACEMENT, sbp=S0)
y = y.to_global(placement=PLACEMENT, sbp=S0)
# Compute prediction error
pred = model(x)
loss = loss_fn(pred, y)
# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
current = batch * BATCH_SIZE
if batch % 5 == 0:
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
这个脚本与单机单卡的训练脚本几乎是一样的。除了几行与全局张量有关的配置代码外，少数的区别是：
1）设置 placement，让训练放置在集群第 0号、1号 GPU 上：
PLACEMENT = flow.placement("cuda", [0,1])
2）模型在集群上做广播
model = model.to_global(placement=PLACEMENT, sbp=B)
3）数据在集群上按 split(0) 做切分：
x = x.to_global(placement=PLACEMENT, sbp=S0)
y = y.to_global(placement=PLACEMENT, sbp=S0)
这样，按照 常见的分布式并行策略 中的介绍，就通过对数据进行 split(0) 切分，对模型进行广播，进行了分布式数据并行训练。
使用 DistributedDataParallel 做数据并行训练
可以用以下命令快速体验 oneflow.nn.parallel.DistributedDataParallel 做数据并行：
wget https://docs.oneflow.org/master/code/parallelism/ddp_train.py #下载脚本
python3 -m oneflow.distributed.launch --nproc_per_node 2 ./ddp_train.py #数据并行训练
输出：
50/500 loss:0.004111831542104483
50/500 loss:0.00025336415274068713
...
500/500 loss:6.184563972055912e-11
500/500 loss:4.547473508864641e-12
w:tensor([[2.0000],
[3.0000]], device='cuda:1', dtype=oneflow.float32,
grad_fn=<accumulate_grad>)
w:tensor([[2.0000],
[3.0000]], device='cuda:0', dtype=oneflow.float32,
grad_fn=<accumulate_grad>)
点击以下 “Code” 可以展开以上运行脚本的代码。
import oneflow as flow
from oneflow.nn.parallel import DistributedDataParallel as ddp
train_x = [
flow.tensor([[1, 2], [2, 3]], dtype=flow.float32),
flow.tensor([[4, 6], [3, 1]], dtype=flow.float32),
]
train_y = [
flow.tensor([[8], [13]], dtype=flow.float32),
flow.tensor([[26], [9]], dtype=flow.float32),
]
class Model(flow.nn.Module):
def __init__(self):
super().__init__()
self.lr = 0.01
self.iter_count = 500
self.w = flow.nn.Parameter(flow.tensor([[0], [0]], dtype=flow.float32))
def forward(self, x):
x = flow.matmul(x, self.w)
return x
m = Model().to("cuda")
m = ddp(m)
loss = flow.nn.MSELoss(reduction="sum")
optimizer = flow.optim.SGD(m.parameters(), m.lr)
for i in range(0, m.iter_count):
rank = flow.env.get_rank()
x = train_x[rank].to("cuda")
y = train_y[rank].to("cuda")
y_pred = m(x)
l = loss(y_pred, y)
if (i + 1) % 50 == 0:
print(f"{i+1}/{m.iter_count} loss:{l}")
optimizer.zero_grad()
l.backward()
optimizer.step()
print(f"\nw:{m.w}")
可以发现，它与单机单卡脚本，只有2个不同：
1）使用 DistributedDataParallel 处理一下 module 对象（m = ddp(m))。
2）使用 get_rank 获取当前设备编号，并针对设备分发数据。
然后使用 launcher 启动脚本，把剩下的一切都交给OneFlow，进行分布式训练，像单机单卡训练一样简单：
python3 -m oneflow.distributed.launch --nproc_per_node 2 ./ddp_train.py
DistributedSampler
为了简化问题，突出 DistributedDataParallel，因此使用的数据是手工分发的。在实际应用中，可以直接使用 DistributedSampler 配合数据并行使用。
DistributedSampler 会在每个进程中实例化 Dataloader，每个Dataloader实例会加载完整数据的一部分，自动完成数据的分发。
5.2.5 流水并行训练
1. 设置参数，实现流水并行
在 OneFlow 的 全局视角 下，通过简单的设置张量的位置属性，就可以实现流水并行。
以下代码是简单的示范，以流水并行的方式运行。前几层的 Module nn.Flatten、
nn.Linear(28*28, 512)、nn.ReLU() 在 GPU0 上运行；剩余的网络部分在 GPU1 上运行。
import oneflow as flow
BATCH_SIZE = 16
BROADCAST = [flow.sbp.broadcast]
P0 = flow.placement("cuda", ranks=[0])
P1 = flow.placement("cuda", ranks=[1])
class Stage0Module(flow.nn.Module):
def __init__(self):
super().__init__()
self.flatten = flow.nn.Flatten()
self.linear0 = flow.nn.Linear(28*28, 512)
self.relu0 = flow.nn.ReLU()
def forward(self, x):
out = self.flatten(x)
out = self.linear0(out)
out = self.relu0(out)
return out
class Stage1Module(flow.nn.Module):
def __init__(self):
super().__init__()
self.linear1 = flow.nn.Linear(512, 512)
self.relu1 = flow.nn.ReLU()
self.linear2 = flow.nn.Linear(512, 10)
self.relu2 = flow.nn.ReLU()
def forward(self, x):
out = self.linear1(x)
out = self.relu1(out)
out = self.linear2(out)
out = self.relu2(out)
return out
class PipelineModule(flow.nn.Module):
def __init__(self):
super().__init__()
self.m_stage0 = Stage0Module()
self.m_stage1 = Stage1Module()
self.m_stage0.to_global(placement=P0, sbp=BROADCAST)
self.m_stage1.to_global(placement=P1, sbp=BROADCAST)
def forward(self, x):
out_stage0 = self.m_stage0(x)
in_stage1 = out_stage0.to_global(placement=P1, sbp=BROADCAST)
out_stage1 = self.m_stage1(in_stage1)
return out_stage1
module_pipeline = PipelineModule()
sgd = flow.optim.SGD(module_pipeline.parameters(), lr=0.001)
class PipelineGraph(flow.nn.Graph):
def __init__(self):
super().__init__()
self.module_pipeline = module_pipeline
self.module_pipeline.m_stage0.to(nn.graph.GraphModule).set_stage(stage_id=0, placement=P0)
self.module_pipeline.m_stage1.to(nn.graph.GraphModule).set_stage(stage_id=1, placement=P1)
self.loss_fn = flow.nn.CrossEntropyLoss()
self.config.set_gradient_accumulation_steps(2)
self.add_optimizer(sgd)
def build(self, x, y):
out = self.module_pipeline(x)
loss = self.loss_fn(out, y)
loss.backward()
return loss
graph_pipeline = PipelineGraph()
x = flow.randn(BATCH_SIZE, 1, 28, 28)
x = x.to_global(P0, BROADCAST)
y = flow.randint(0, 10, (BATCH_SIZE,))
y = y.to_global(P1, BROADCAST)
for i in range(20):
loss = graph_pipeline(x, y)
print(loss.to_local())
以上代码，保存为脚本（如 pipeline.py）后，使用 launch 模块启动分布式训练：
python3 -m oneflow.distributed.launch --nproc_per_node 2 ./pipeline.py
2. 代码解读
1）设置 placement 与 sbp
将需要使用的 placement 与 sbp 设置提前准备好：
BROADCAST = [flow.sbp.broadcast]
P0 = flow.placement("cuda", ranks=[0])
P1 = flow.placement("cuda", ranks=[1])
P0、P1 分别代表集群的第 0 个 GPU 和第 1 个 GPU。
通过调用 nn.Module.to_global 或 Tensor.to_global 就可以将模型或张量分配到指定的计算设备上运行，将一个网络拆分为多个流水阶段（stage）。
在此定义了一个 PipelineModule 专门设置各阶段的流水。
class PipelineModule(flow.nn.Module):
def __init__(self):
#...
self.m_stage0.to_global(placement=P0, sbp=BROADCAST)
self.m_stage1.to_global(placement=P1, sbp=BROADCAST)
def forward(self, x):
out_stage0 = self.m_stage0(x)
in_stage1 = out_stage0.to_global(placement=P1, sbp=BROADCAST)
out_stage1 = self.m_stage1(in_stage1)
return out_stage1
2）局部张量与全局张量的转换
示例中使用了随机生成的数据作为输入。
x = flow.randn(BATCH_SIZE, 1, 28, 28)
x = x.to_global(P0, BROADCAST)
当使用 launch 模块启动训练时，因为命令行参数为 --nproc_per_node 2，launch 会启动 2 个进程。两个进程均为执行脚本中的代码。
其中 x = flow.randn(BATCH_SIZE, 1, 28, 28) 返回局部张量（只在本进程中有效的本地数据），当运行 x = x.to_global(P0, BROADCAST) 时，OneFlow 会自动将所有进程中的局部张量整合为全局张量。
在实际训练中，各个计算设备也可以加载属于各自的本地数据，然后通过 to_global 实现局部张量到全局张量的转化。
3）Stage ID 及梯度累积设置
当 nn.Module 的一个实例化网络层作为属性，加入继承于 nn.Graph 的新类时，内部会将该网络层用 ProxyModule 进行封装，利用方法 .to 得到一个 nn.graph.GraphModule 的实例化对象，然后使用方法 stage_id 设置流水线 Stage ID 和 Stage 对应的 Placement，Stage ID从0开始编号，依次加1。
调用 config.set_gradient_accumulation_steps 方法，设置梯度累积的步长。OneFlow 通过这两项配置，获取实现流水并行中的微批量技术所需的信息。
self.module_pipeline.m_stage0.to(nn.graph.GraphModule).set_stage(stage_id=0, placement=P0)
self.module_pipeline.m_stage1.to(nn.graph.GraphModule).set_stage(stage_id=1, placement=P1)
self.config.set_gradient_accumulation_steps(2)
人工智能芯片与自动驾驶
