Skip to content

Commit

Permalink
Merge pull request #7 from AmourWaltz/parallelism-byxue
Browse files Browse the repository at this point in the history 
7.24
  • Loading branch information
@hscspring
hscspring authored Jul 28, 2024
2 parents a53013f + 634258a commit ca5e33d
Show file tree
Hide file tree
Showing 22 changed files with 172 additions and 155 deletions.
156 changes: 1 addition & 155 deletions docs/chapter8/chapter8_1.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,158 +2,4 @@

## 概述

## 数据并行

## 流水线并行


当前主流大模型的参数规模通常达到数十亿甚至成百上千亿级，单块 GPU 很难加载整个模型，因此我们需要将模型分布在不同设备上，这也称为模型并行方法。模型并行方法有两种：流水线并行和张量并行。如下图所示，流水线并行将模型按层划分到不同 GPU 上，在各层之间进行并行计算；

![Untitled](figs/Untitled.png)


## 张量并行

与流水线并行不同，张量并行是将模型中的张量进行拆分然后分配到不同的 GPU 上，每块 GPU 都可以得到所有层张量的部分参数。这样在前向计算中有效减少了流水行并行中的空置时间，提高了 GPU 的显存利用率，因此张量并行也成了当下大模型训练和推理的主流并行方法。显存效率：模型并行会根据 worker 数量成比例地减少显存使用量。至关重要的是，这是减少单个网络层的激活显存的唯一方法。DeepSpeed 通过在模型并行 worker 之间划分激活显存来进一步提高显存效率。
计算效率：由于每次前向和反向传播中都需要额外通信激活值，模型并行的计算效率很低。模型并行需要高通信带宽，并且不能很好地扩展到通信带宽受限的节点。此外，每个模型并行worker 都会减少每个通信阶段之间执行的计算量，从而影响计算效率。模型并行性通常与数据并行性结合使用，以在内存和计算效率之间进行权衡。

### 1D 张量并行

从实现上看，张量并行就是把同一层的参数矩阵分块进行相互独立的矩阵乘法计算，然后合并结果，通过不同 GPU 之间的通信，保证计算图的正确性。对一个单独的矩阵，我们可以很自然想到基于行和列进行拆分，称为行并行和列并行。

如图 (a) 所示，以一般矩阵乘法为例，假设我们有 $Y=XW$ ，其中 $X\in\mathbb R^{2\times2}$ 为输入数据，$W\in\mathbb R^{2\times2}$ 为参数矩阵，在两块 GPU 上进行并行计算，输入数据 $X$ 与权重向量 $W$ 进行矩阵相乘时，计算行列对之间的点积是相互独立的。列并行就是将权重参数 $W$ 沿列分割成 $W=[W_0\ W_1]$，每块 GPU 持有一列参数 $W_0, W_1 \in \mathbb R^{2\times1}$，如图 (b) 所示，我们将 $X$ 分别输入 rank 0 和 rank 1 的 GPU 上，然后与 GPU 上的参数进行矩阵相乘，我们将得到 $Y_0, Y_1$ ，然后通过一次拼接操作就可以得到和 (a) 等价的 $Y$。而行并行则是将 $W$ 沿行进行切分 $W=[W_0 \ W_1]^{\mathrm T}$ 放置，每块 GPU 持有一行参数 $W_0, W_1 \in \mathbb R^{1\times2}$ ，然后将输入 $X$ 也沿列切分为 $X_0, X_1 \in \mathbb R^{2\times1}$ 并输入到两块 GPU 上分别进行矩阵乘法运算得到 $Y^\prime,Y^{\prime\prime}\in\mathbb R^{2\times 2}$ ，然后按元素位置相加 $Y^\prime+Y^{\prime\prime}$ 也可以得到等价的 $Y$。

![Untitled](figs/Untitled%201.png)

![Untitled](figs/Untitled%202.png)

![Untitled](figs/Untitled%203.png)

不难看出，无论行并行还是列并行，想要得到最终结果，我们都需要收集所有 GPU 上的计算结果，而神经网络通常包含很多层，也就意味着在进行当前层的张量并行时，我们必须确保所有 GPU 都接收到上一层前向计算的完整计算图结果，而将所有设备的结果统筹收集的过程，称为全收集操作 All Reduce，这也正是张量并行对设备间通信带宽要求较高的原因。而我们可以通过灵活使用行并行和列并行，尽量减少前向计算中的全收集操作。

观察行并行和列并行的输入输出形式，不难发现在不进行全收集操作的情况下，行并行的输出形式恰好是列并行的输入形式，列并行的输出方式恰好是行并行的输入形式，因此不难想到，在前向计算中，这两种并行方式是交替使用的。以 Transformer 中的 FFN 层为例，FFN 层包含两个线性层 nn.linear 以及线性层中间的激活函数 GELU(·)，假设两个线性层参数分别为 $A$ 和 $B$。如果在第一个线性层进行行并行将 $A$ 拆分为 $A=[A_0 \ A_1]^{\mathrm T}$，由于激活单元需要完整计算结果，所以就需要进行一次全收集操作从而得到经过第一个线性层后的完整计算图结果 $Y=Y^\prime+Y^{\prime\prime}$，再经过激活函数 GELU，而在第二个线性层处无论进行行并行还是列并行，都需要在输入下个 Transformer 层的多注意力模块时进行全收集操作，这样在经过 FFN 层时就需要两次全收集操作；而如果在第一个线性层进行列并行 $A = [A_0 \ A_1]$，那么得到的计算图结果 $Y_0=[y_{00} \ y_{01}]^{\mathrm T}, Y_1=[y_{10} \ y_{11}]^{\mathrm T}$ 是可以分别独立经过激活单元的，而后在第二个线性层处进行行并行 $B=[B_0 \ B_1]^{\mathrm T}$，这样就只需要一次全收集操作 $Z=Z^\prime+Z^{\prime\prime}$，下图分别展示了 FFN 层两种张量并行计算步骤。

![Untitled](figs/Untitled%204.png)

### 2D 和 2.5D 张量并行

Nvidia Megatron-LM 使用的是 1D 张量并行，这种方法虽然将参数划分到多个处理器上，但每个处理器仍需要存储整个中间计算图结果，在每次计算中，每块 GPU 都需要与其他设备通信，在处理大模型时会浪费大量显存空间，通信成本会不断增加。对此，Colossal-AI 提供多维张量并行，这里先介绍 2D 张量并行。2D 张量并行技术将输入数据、模型权重和层输出拆分成两个维度，与 1D 张量并行相比，内存消耗更低。对于一个 2×2 的矩阵乘法，假设我们有 4 块 GPU，那就可以将矩阵乘法分块到每块 GPU 上。将输入 $X$ 和参数 $W$ 进行如下分块 $X=[X_0\ X_1]$, $W=[W_0\ W_1]^{\mathrm T}$，首先在 Step 1 进行 $X_0$ 与 $W_0$ 的矩阵乘法，将4 个算子分配到 4 块 GPU 上进行计算，同样，Step 2 进行 $X_1$ 与 $W_1$ 的运算，最后将两步的计算结果相加，便得到最终结果。

![Untitled](figs/Untitled%205.png)

```python
import colossalai
import colossalai.nn as col_nn
import torch
from colossalai.utils import print_rank_0
from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.utils import get_current_device

# 并行设置
CONFIG = dict(parallel=dict(
data=1,
pipeline=1,
tensor=dict(size=4, mode='2d'),
))

parser = colossalai.get_default_parser()
colossalai.launch(config=CONFIG,
rank=args.rank,
world_size=args.world_size,
local_rank=args.local_rank,
host=args.host,
port=args.port)

class MLP(torch.nn.Module):
def __init__(self, dim: int = 256):
super().__init__()
intermediate_dim = dim * 4
self.dense_1 = col_nn.Linear(dim, intermediate_dim)
print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.shape}')
self.activation = torch.nn.GELU()
self.dense_2 = col_nn.Linear(intermediate_dim, dim)
print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.shape}')
self.dropout = col_nn.Dropout(0.1)

def forward(self, x):
x = self.dense_1(x)
print_rank_0(f'Output of the first linear layer: {x.shape}')
x = self.activation(x)
x = self.dense_2(x)
print_rank_0(f'Output of the second linear layer: {x.shape}')
x = self.dropout(x)
return x

# 创建模型
m = MLP()

# 随机输入一些数据来运行这个模型
x = torch.randn((16, 256), device=get_current_device())

# partition input
torch.distributed.broadcast(x, src=0)
x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_COL)]
x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_ROW)]
print_rank_0(f'Input: {x.shape}')

x = m(x)
```

2.5D张量并行技术通过为矩阵添加可选的深度维度，使2D并行技术向前推进了一步。当扩展到大量设备时，这种方法可进一步减小通信量。张量被分割，使得N=S2D，其中S是正方形一侧的大小，D是立方体的深度。当深度为1时，这种方法与2D张量并行技术类似。之所以叫 2.5D 张量并行是因为在 d = 1 时，这种并行模式可以退化成 2D 张量并行；在 d = q 时，它就变成了3D 张量并行。下面我们来看看 3D 张量并行。

```python
# 并行设置
CONFIG = dict(parallel=dict(
data=1,
pipeline=1,
tensor=dict(size=8, mode='2.5d', depth=2),
))

...

# 创建模型
m = MLP()

# 随机输入一些数据来运行这个模型
x = torch.randn((16, 256), device=get_current_device())

# partition input
torch.distributed.broadcast(x, src=0)
x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_DEP)]
x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_COL)]
x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_ROW)]
print_rank_0(f'Input: {x.shape}')

x = m(x)
```

### 3D 张量并行

3D张量并行技术将张量分割成立方体形状，并对第一个和最后一个维度进行划分。这种技术可实现最佳通信成本，并均匀分配计算和内存使用量。当扩展到更多设备时，高级张量并行技术可进一步减小通信量，并与流水线并行更加兼容。与1D张量并行技术不同，1D张量并行技术是在将张量传递到下一个流水线阶段之前将其分割成若干块，并在到达目标流水线阶段后通过节点内通信进行收集，而高级张量并行技术已经为下一阶段提供了整个逻辑张量的子块，因此无须进行拆分-收集操作，从而降低通信成本。

```python
# 并行设置
CONFIG = dict(parallel=dict(
data=1,
pipeline=1,
tensor=dict(size=8, mode='3d'),
))

...

# 创建模型
m = MLP()

# 随机输入一些数据来运行这个模型
x = torch.randn((16, 256), device=get_current_device())

# partition input
torch.distributed.broadcast(x, src=0)
x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_WEIGHT)]
x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_INPUT)]
x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_OUTPUT)]
print_rank_0(f'Input: {x.shape}')

x = m(x)
```
如今的大语言模型虽然与传统预训练模型采用了相似的网络架构以及训练方法，但其通过扩展模型参数规模、训练数据量以及算力资源等方法，实现了模型性能的质变。为了提高大语言模型的推理效果，科研人员需要使用更大规模的模型并提供更多的训练数据，因此用于训练算力资源也会随之增加。当前主流的大模型如 ChatGPT 等都是在成百上千块 GPU 上训练的。为了进一步高效且经济地加速大模型训练推理，大模型的并行方法就成了人工智能领域越来越关注的重点。本文将介绍几种大模型部署中常用的并行方法。
7 changes: 7 additions & 0 deletions docs/chapter8/chapter8_2.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
## 数据并行

数据并行算法主要是用于大模型的训练过程，即使用多个计算节点并行训练同一个模型的参数，同时不同节点之间相互通信聚合更新梯度信息，从而提高大模型数据处理的并行度。而大模型推理的数据并行则实现较为简单，如下图所示，在每个计算节点上复制一份完整模型，并将输入数据分成不同 batch 送入不同节点，各计算节点独立完成推理，输出对应的结果。但这种方法由于需要在每台机器上都复制一遍完整模型，因此内存利用率较低，并且如今的大模型通常都为百亿级别的参数规模，单张 GPU 无法加载整个模型，因此数据并行无法单独适用，于是引出模型并行的方法。

![alt text](image.png)

![alt text](image-1.png)
16 changes: 16 additions & 0 deletions docs/chapter8/chapter8_3.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
## 流水线并行


由于当前主流大模型的规模超出单块 GPU 的内存，因此需要将模型分布在不同设备上，这也称为模型并行方法。模型并行方法有两种：流水线并行和张量并行。这里首先介绍流水线并行，如下图所示，流水线并行将模型按层划分到不同 GPU 上，在各层之间并行地执行各自前向计算。
流水线并行的过程以下图为例，假设一个神经网络模型共有四层神经元，我们可以将其切分成两个层，分并别放置到两个不同的计算节点中进行计算，在计算过程中，输入数据input先进入到第一个计算节点Rank 0 中，等计算完毕后，再通过通信将Rank 1计算得到的数据结果传输到第二个计算节点 Rank 1中，再进行 Rank 1 所存储网络层的前向计算。

![alt text](image-2.png)


但这种流水线并行有着明显的缺点：后续的计算节点 Rank 1 在 Rank 0 还没处理完数据的时候会一直处于闲置状态，直到Rank 0 计算节点传来数据，这使得整个计算过程中设备利用率较低，被称为 Bubble 问题。Bubble问题指的是在运行过程中有GPU闲置导致资源浪费的情况，顾名思义，其在整个流程看起来像是一个拱起来的气泡。以上图为例子，假设当前的运行的大模型被分成4个层，且将这4个层分布式的部署在4台GPU设备上，在整个推理过程中，后面阶段的层需要等待前面阶段的层执行完毕后才能继续往下执行，在这个过程中会出现设备闲置的情况。

![alt text](image-3.png)

为了使得流水线并行中的设备利用率达到最佳，流水线并行的调度方式应当让每个阶段的执行时间应尽可能平衡。当当前节点处理完数据后，应当立即有后续数据可以处理，而非让设备限制着。该方式只在相对同构的模型（如 Transformers）中实现起来较为容易，因为其网络模型大多数层的大小通常相同，但在其他模型（如 Resnets 或 UNets）中实现则很困难。

![alt text](image-4.png)
Loading

0 comments on commit ca5e33d

Please sign in to comment.