张量并行详解 - Ming's Blog

Note

多机多卡 AI 推理/训练系列博客

Tensor Parallelism 的本质是把一个”大矩阵运算”切到多个 GPU 上，让每张 GPU 只保存和计算矩阵的一部分，最后通过通信把结果拼起来。

大矩阵乘法的切分#

我们考虑如何切分矩阵乘法

Y = XW

其中

X \in \mathbb{R}^{B \times K}, \quad W \in \mathbb{R}^{K \times N}, Y \in \mathbb{R}^{B \times N}

Column Parallel Linear#

如果有 4 张 GPU，那么可以把 $W$ 沿着输出维度 $N$ 切开：

W = [W_0, W_1, W_2, W_3]

每张 GPU 保存一部分：

GPU_i: W_i \in \mathbb{R}^{K \times N / 4}

于是：

Y_i = XW_i

每张 GPU 独立计算一部分输出：

Y = [Y_0, Y_1, Y_2, Y_3]

最后把它们拼起来即可。

这就是最简单的 Column Parallel Linear。

下面是 Column Parallel Linear 的代码实现：

1
class ColumnParallelLinear(torch.nn.Module):
2
    def __init__(self, in_features: int, out_features: int, bias: bool):
3
        super().__init__()
4

5
        self.tp_world_size = get_tp_world_size()
6
        self.tp_rank = get_tp_rank()
7

8
        self.in_features = in_features
9
        self.out_features = out_features
10

11
        assert out_features % self.tp_world_size == 0
12
        self.output_size_per_partition = out_features // self.tp_world_size
13

14
        # Allocate space for the weight and bias
15
        # Note: torch.nn.functional.linear performs XW^T + b so we exchange the order of dimensions
16
        self.weight = nn.Parameter(torch.Tensor(self.output_size_per_partition, self.in_features))
17
        if bias:
18
            self.bias = nn.Parameter(torch.Tensor(self.output_size_per_partition))
19
        else:
20
            self.register_parameter("bias", None)
21

22
    def forward(self, x: torch.Tensor):
23
        return F.linear(x, self.weight, self.bias)

Row Parallel Linear#

另一种切分方式是把 $W$ 按行切：

W = \begin{bmatrix} W_0 \\ W_1 \\ W_2 \\ W_3 \end{bmatrix}

同时把输入也按照对应维度切开：

X = [X_0, X_1, X_2, X_3]

于是：

XW = X_0W_0 + X_1W_1 + X_2W_2 + X_3W_3

每张 GPU 计算：

Y_i = X_iW_i

最后：

Y = \sum_{i}Y_i

Note

总结

Column Parallel： $W$ 输出维度被切开，可以使用 AllGather 获取全部输出。
Row Parallel: $W$ 输入维度被切开，最后必须使用 AllReduce 求和。

下面是 Row Parallel Linear 的代码实现：

1
class RowParallelLinear(nn.Module):
2
    def __init__(self, in_features: int, out_features: int, bias: bool):
3
        super().__init__()
4

5
        self.tp_world_size = ...
6
        self.tp_rank = ...
7

8
        self.in_features = in_features
9
        self.out_features = out_features
10
        assert in_features % self.tp_world_size == 0
11
        self.input_size_per_partition = in_features // self.tp_world_size
12

13
        self.weight = nn.Parameter(torch.Tensor(out_features, input_size_per_partition))
14
        if bias:
15
            self.bias = nn.Parameter(torch.Tensor(out_features))
16
        else:
17
            self.register_parameter("bias", None)
18

19
    def forward(self, x: torch.Tensor):
20
        # x: [batch, input_size_per_partition]
21
        y = F.linear(x, self.weight, None)
22
        dist.all_reduce(y, op=dist.ReduceOp.SUM) # group = ...
23
        return y if self.bias is None else y + self.bias

在 Transformer Block 中使用 Tensor Parallelism#

Transformer Block 中最重要的两个组件是 FFN 和 Attention。接下来介绍如何分别在 FFN 和 Attention 中使用 Tensor Parallelism。

在 FFN 中使用 Tensor Parallelism#

假设 FFN:

\begin{aligned} H &= \sigma(XW_1) \\ Y &= HW_2 \end{aligned}

我们先把 $W_1$ 按列切分：

W_1 = [W_{1,0}, W_{1,1}]

那么：

H = [H_0, H_1]

其中：

H_i = \sigma(XW_{1,i})

然后把 $W_2$ 按行切：

W_2 = \begin{bmatrix} W_{2,0} \\ W_{2,1} \end{bmatrix}

那么：

Y = H_0 W_{2,0} + H_1 W_{2,1}

Important

先进行 Column Parallel 再进行 Row Parallel 的顺序不能颠倒，否则需要对中间计算结果进行 AllReduce。

MLP 的 Tensor Parallel 实现如下：

1
class ParallelMLP(nn.Module):
2
    def __init__(self, hidden_size: int, intermediate_size: int):
3

4
        self.fc1 = ColumnParallelLinear(in_features=hidden_size, out_features=intermediate_size, bias=False)
5
        self.fc2 = RowParallelLinear(in_features=intermediate_size, out_features=hidden_size, bias=False)
6

7
    def forward(self, x: torch.Tensor):
8
        x = self.fc1(x)
9
        x = F.gelu(x)
10
        x = self.fc2(x)
11
        return x

在 Attention 中使用 Tensor Parallelism#

假设有 num_attention_heads == 32，4 张 GPU。

最简单的做法是

GPU0 -> heads 0 ~ 7
GPU1 -> heads 8 ~ 15
GPU2 -> heads 16 ~ 23
GPU3 -> heads 24 ~ 31

这种 Head Parallel 本质上属于 Tensor Parallel。

每张 GPU 计算自己的：

Q_i, K_i, V_i

然后独立完成：

Z_i = Attention_i = \operatorname{Softmax}(\frac{Q_i K_i^{T}}{\sqrt{d}}) V_i

由于不同 head 本身是独立的，所以这里不需要 GPU 之间通信。

然后经过 output projection：

Y = Z W_O

将 $W_O$ 做 Row Parallel：

W_O = \begin{bmatrix} W_{O,0} \\ W_{O,1} \\ W_{O,2} \\ W_{O,3} \end{bmatrix}

于是：

Y = \sum_{i} Z_i W_{O,i}

最后一次 AllReduce。

Note

对于 MQA 和 GQA 也可以使用 Tensor Parallelism。

MHA 的 Tensor Parallel 实现示例如下：

1
class ParallelMultiHeadAttention(nn.Module):
2
    def __init__(self, hidden_size: int, num_heads: int):
3
        self.tp_world_size = ...
4
        self.tp_rank = ...
5

6
        self.hidden_size = hidden_size
7
        self.num_heads = num_heads
8
        self.head_dim = hidden_size // num_heads
9
        assert num_heads % self.tp_world_size == 0
10
        self.heads_per_partition = num_heads // self.tp_world_size
11

12
        self.qkv_proj = ColumnParallelLinear(in_features=hidden_size, out_features=3 * hidden_size, bias=False)
13

14
        self.o_proj = RowParallelLinear(in_features=hidden_size, out_features=hidden_size, bias=False)
15

16
    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
17
        # x: [batch, seq_len, hidden_size]
18
        batch, seq_len, _ = x.shape
19

20
        qkv = self.qkv_proj(x) # [batch, seq_len, 3 * heads_per_partition * head_dim]
21

22
        q, k, v = qkv.trunk(3, dim=-1) # [batch, seq_len, heads_per_partition * head_dim]
23

24
        q = q.view(batch, seq_len, self.heads_per_partition, head_dim)
25
        k = k.view(batch, seq_len, self.heads_per_partition, head_dim)
26
        v = v.view(batch, seq_len, self.heads_per_partition, head_dim)
27

28
        q = q.transpose(1, 2)
29
        k = k.transpose(1, 2)
30
        v = v.transpose(1, 2)
31

32
        attn_weights = torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5)
33
        if mask:
34
            attn_weights = attn_weights + mask
35
        attn_weights = F.softmax(attn_weights, dim=-1)
36

37
        out = torch.matmul(attn_weights, v) # [batch, heads_per_partition, seq_len, head_dim]
38
        out = out.transpose(1, 2).contiguous() # [batch, seq_len, heads_per_partition, head_dim]
39
        out = out.view(batch, seq_len, -1) # [batch, seq_len, heads_per_partition * head_dim]
40
        return self.o_proj(out)

整个 TP Transformer Block 就是这样的：

1
class TPTransformerBlock(nn.Module):
2
    def __init__(self, hidden_size, num_heads, intermediate_size):
3
        super().__init__()
4

5
        self.norm1 = nn.LayerNorm(hidden_size)
6

7
        self.attn = ParallelMultiHeadAttention(hidden_size, num_heads)
8

9
        self.norm2 = nn.LayerNorm(hidden_size)
10

11
        self.mlp = ParallelMLP(hidden_size, intermediate_size)
12

13
    def forward(self, x):
14

15
        # Attention
16
        residual = x
17
        x = self.norm1(x)
18
        x = self.attn(x)
19
        x = residual + x
20

21
        # MLP
22
        residual = x
23
        x = self.norm2(x)
24
        x = self.mlp(x)
25
        x = residual + x
26

27
        return x

Important

Tensor Parallelism 的缺陷

Tensor Parallelism 最大的缺陷是计算和通信没办法重叠，而且对于通信带宽的要求极高。

Sequence Parallelism#

需要注意的是 Tensor Parallelism 无法处理 Transformer Block 中的 LayerNorm 和 Residual。虽然这两个部分的计算很简单，但是它们的激活依然占据很大的空间，因此有 SP (Sequence Parallelism) 进行进一步切分。

参考#

Ultrascale Playbook

大矩阵乘法的切分#

Column Parallel Linear#

Row Parallel Linear#

在 Transformer Block 中使用 Tensor Parallelism#

在 FFN 中使用 Tensor Parallelism#

在 Attention 中使用 Tensor Parallelism#

Sequence Parallelism#

参考#

支持与分享

目录