分块矩阵的乘法

矩阵乘法是线性代数中最重要的运算之一。在机器学习中，矩阵乘法也是经常用到的运算，最常见于 MLP 线性层。

而在实际的模型训练和推理系统中，模型参数和中间激活的张量可能非常大，而 GPU 显存空间有限。因此，我们需要将张量切分为多个块，以在 GPU 上实现并行计算。而这和分块矩阵的乘法有着紧密的联系。

分块矩阵的定义

一个 $m \times n$ 维的矩阵 $W$ 定义如下

W_{m \times n} = [\begin{matrix} w_{11} & w_{12} & \dots & w_{1 n} \\ w_{21} & w_{22} & \dots & w_{2 n} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ w_{m 1} & w_{m 2} & \dots & w_{m n} \end{matrix}]

现在将这个矩阵在行方向上划分为 $X$ 个块，在列方向上划分为 $Y$ 个块。这样，每个块的维度为 $x \times y$ ，其中

{\begin{cases} x = \frac{m}{X} \\ y = \frac{n}{Y} \end{cases}

于是，分块后的矩阵可记为

W_{X \times Y} = [\begin{matrix} W_{11} & W_{12} & \dots & W_{1 Y} \\ W_{21} & W_{22} & \dots & W_{2 Y} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ W_{X 1} & W_{X 2} & \dots & W_{X Y} \end{matrix}]

其中，每个块 $W_{i j} (1 ⩽ i, j ⩽ X, Y)$ 可以表示为

W_{i j} = [\begin{matrix} w_{(i - 1) x + 1, (j - 1) y + 1} & w_{(i - 1) x + 1, (j - 1) y + 2} & \dots & w_{(i - 1) x + 1, j y} \\ w_{(i - 1) x + 2, (j - 1) y + 1} & w_{(i - 1) x + 2, (j - 1) y + 2} & \dots & w_{(i - 1) x + 2, j y} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ w_{i x, (j - 1) y + 1} & w_{i x, (j - 1) y + 2} & \dots & w_{i x, j y} \end{matrix}]

设一个 $m$ 维的向量

λ = [\begin{matrix} k_{1} \\ k_{2} \\ ⋮ \\ k_{m} \end{matrix}]

要使向量 $λ$ 右乘分块矩阵 $W$ ，需要相应地将向量 $λ$ 分成 $X$ 个块，即

λ = [\begin{matrix} λ_{1} \\ λ_{2} \\ ⋮ \\ λ_{X} \end{matrix}]

其中，每个块 $λ_{i} (1 ⩽ i ⩽ X)$ 可以表示为

λ_{i} = [\begin{matrix} k_{(i - 1) x + 1} \\ k_{(i - 1) x + 2} \\ ⋮ \\ k_{i x} \end{matrix}]

因此，向量 $λ$ 右乘分块矩阵 $W$ 可以表示为

λ^{⊤} W = [\begin{matrix} \sum_{i = 1}^{X} λ_{i}^{⊤} W_{i 1} & \sum_{i = 1}^{X} λ_{i}^{⊤} W_{i 2} & \dots & \sum_{i = 1}^{X} λ_{i}^{⊤} W_{i y} \end{matrix}]

设一个 $d \times m$ 维的矩阵

A_{d \times m} = [\begin{matrix} a_{11} & a_{12} & \dots & a_{1 m} \\ a_{21} & a_{22} & \dots & a_{2 m} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ a_{d 1} & a_{d 2} & \dots & a_{d m} \end{matrix}]

仿照向量与分块矩阵相乘的方式，如果对矩阵 $A$ 作 1D 切分，即

A = [\begin{matrix} A_{1} & A_{2} & \dots & A_{X} \end{matrix}]

其中，每个分块矩阵的维度是 $d \times x$ 。那么，矩阵 $A$ 右乘分块矩阵 $W$ 可以表示为

A W = [\begin{matrix} \sum_{i = 1}^{X} A_{i} W_{i 1} & \sum_{i = 1}^{X} A_{i} W_{i 2} & \dots & \sum_{i = 1}^{X} A_{i} W_{i y} \end{matrix}]

如果对矩阵 $A$ 作 2D 切分，在行方向上切分成 $T$ 个块，在列方向上切分成 $X$ 个块。这样，每个块的维度就变成 $t \times x$ ，其中

t = \frac{d}{T}

那么矩阵 $A$ 可表示为

A_{T \times X} = [\begin{matrix} A_{11} & A_{12} & \dots & A_{1 X} \\ A_{21} & A_{22} & \dots & A_{2 X} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ A_{T 1} & A_{T 2} & \dots & A_{T X} \end{matrix}]

这样，矩阵 $A$ 右乘分块矩阵 $W$ 可以表示为

A W = [\begin{matrix} \sum_{i = 1}^{X} A_{1 i} W_{i 1} & \sum_{i = 1}^{X} A_{1 i} W_{i 2} & \dots & \sum_{i = 1}^{X} A_{1 i} W_{i y} \\ \sum_{i = 1}^{X} A_{2 i} W_{i 1} & \sum_{i = 1}^{X} A_{2 i} W_{i 2} & \dots & \sum_{i = 1}^{X} A_{2 i} W_{i y} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ \sum_{i = 1}^{X} A_{T i} W_{i 1} & \sum_{i = 1}^{X} A_{T i} W_{i 2} & \dots & \sum_{i = 1}^{X} A_{T i} W_{i y} \end{matrix}]

与普通的矩阵乘法相比，分块矩阵乘法多了一个步骤：分块累加。用规范化的数学语言描述，即

\begin{aligned} A_{m \times p} & = [\begin{array}{c} A_{11} & A_{12} & \dots & A_{1 P} \\ A_{21} & A_{22} & \dots & A_{2 P} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ A_{M 1} & A_{M 2} & \dots & A_{M P} \end{array}] \\ B_{p \times n} & = [\begin{array}{c} B_{11} & B_{12} & \dots & B_{1 N} \\ B_{21} & B_{22} & \dots & B_{2 N} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ B_{P 1} & B_{P 2} & \dots & B_{P N} \end{array}] \\ A B & = [\begin{array}{c} \sum_{i = 1}^{P} A_{1 i} B_{i 1} & \sum_{i = 1}^{P} A_{1 i} B_{i 2} & \dots & \sum_{i = 1}^{P} A_{1 i} B_{i N} \\ \sum_{i = 1}^{P} A_{2 i} B_{i 1} & \sum_{i = 1}^{P} A_{2 i} B_{i 2} & \dots & \sum_{i = 1}^{P} A_{2 i} B_{i N} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ \sum_{i = 1}^{P} A_{M i} B_{i 1} & \sum_{i = 1}^{P} A_{M i} B_{i 2} & \dots & \sum_{i = 1}^{P} A_{M i} B_{i N} \end{array}] \end{aligned}

其中，矩阵 $A$ 被 2D 切分为 $M \times P$ 份，矩阵 $B$ 被 2D 切分为 $P \times N$ 份。

而至于向量与分块矩阵的乘法，以及 1D 和 2D 切分的情况，都属于上述一般规律的特例。是否需要将分块累加，取决于是否对 $p$ 维进行切分。

设想有 2 个矩阵都被切分成 $2 \times 2$ 的块，每个块在 4 张不同的 GPU 上进行矩阵乘法运算。则每个 GPU 上的计算分别如下面 4 张图所示：

矩阵乘法是线性代数中最重要的运算之一。在机器学习中，矩阵乘法也是经常用到的运算，最常见于 MLP 线性层。

而在实际的模型训练和推理系统中，模型参数和中间激活的张量可能非常大，而 GPU 显存空间有限。因此，我们需要将张量切分为多个块，以在 GPU 上实现并行计算。

而分块矩阵乘法的意义在于，分块意味着一个完整的张量可以进行切分，在 GPU 上进行并行的矩阵乘法计算从理论上说是可行的。

不过，如果我们将矩阵在 $p$ 维上切分，最终的计算结果需要沿着 $p$ 维进行聚合。这样的操作会带来额外的通信开销。

此外，最终的结果从整体上看是多个小块矩阵的拼接，这意味着每个 GPU 可以存储相应的不同的结果。至于是否存在通信开销，取决于后面的算子是否需要完整的张量。