列并行张量并行原理

可视化演示

展示FIFO与LRU页面置换算法在有限物理帧中的执行过程，可视化缺页中断、页面命中和替换策略的差异。

列并行概念：竖着撕配方

大排档里，大厨不再独自处理整本配方，而是把配方竖着撕开：帮厨 A 拿左边一列，帮厨 B 拿右边一列。所有帮厨面前仍摆着同一筐完整食材。也就是说，输入数据不切分，权重矩阵按列切分。每个帮厨只按自己手里的配方做一部分菜，最后把半成品拼在一起，就是一整桌菜。

在大模型训练中，一个线性层可以写成矩阵乘法 $Y = XW$。列并行把权重矩阵 $W$ 沿输出维度切成 $p$ 块，分别放到 $p$ 张卡上。第 $j$ 张卡拥有 $W_j$，并计算 $Y_j = XW_j$。因为 $X$ 对所有人都是完整的，所以每张卡都能独立完成自己的矩阵乘法，不需要在计算过程中找别人商量。

所以列并行最直观的口诀是“权重分列、输入全量、输出拼接”。就像几位帮厨各自做出半成品，拼成一整道菜；客人看到的是完整的一桌菜，不会关心是哪位帮厨做的。

分块矩阵乘法：为什么可以这样切

有人会问：把配方竖着撕，难道不会把一道菜做坏吗？不会，因为矩阵乘法天生就允许“一列一列地算”。每个输出列只由输入矩阵和权重矩阵的对应列决定，和其他列没有关系。这就像配方的“酱油用量”只影响酱油这一列的效果，帮厨 A 只需要把酱油相关的那几列做好。

从数学上看，矩阵乘法 $Y = XW$ 中，输出矩阵的第 $j$ 列，是输入矩阵每一列按权重第 $j$ 列加权求和的结果。也就是说，每一列输出是独立计算的。因此，把权重矩阵按列切成若干块后，每块对应的输出列也是相互独立的，最后拼起来不会有任何信息丢失。

这正好解释了为什么列并行是“零风险的切法”：数学上本来就可以分块。帮厨们不需要讨论各自配方之间怎么协调，只要最后把成品按顺序摆好即可。

前向计算流程：备料、翻炒、拼盘

列并行的前向过程，就像大排档的备菜流程。第一步是“备料”：把完整食材 $X$ 复制一份放到每口锅前。第二步是“翻炒”：每位帮厨只用自己的配方 $W_j$ 处理这些食材，得到自己的半成品 $Y_j$。第三步是“拼盘”：把所有半成品按顺序拼在一起，得到完整输出 $Y$。

整个过程里，设备之间不需要通信。因为每个设备都拿到了完整的 $X$，并且只负责输出维度上的一段；它既不需要等待别人的结果，也不需要把自己的结果发给别人。唯一要做的，只是最后在逻辑上把输出分片拼接起来。在工程实现里，这种拼接有时候是显式的 concat，有时候是让下一层直接消费分片，省去实际拼接。

前向不通信，是列并行最舒服的地方。帮厨们埋头炒菜，互不打扰；只有在上菜前才需要把拼盘摆好。

反向传播特点：菜品评价的互通

菜上桌后，客人给出反馈：哪里咸了、哪里淡了。这个反馈就是损失对输出的梯度。因为输出是拼接出来的，每位帮厨只需要拿到自己那块菜对应的评价，这不需要通信。但他还需要知道“整桌菜”的反馈对自己手里的食材 $X$ 有多大责任，这就必须把各人的理解汇总起来。

在数学上，反向传播要求计算 $\partial L / \partial X$。由于 $Y = [Y_1, \dots, Y_p]$，输入梯度是所有输出分片对输入贡献的总和。每个设备只能算出自己那一份局部输入梯度，必须通过 All-Reduce 把大家的结果求和，再广播给每个设备，这样每个设备手里的 $X$ 副本都能得到完整梯度。

这就是列并行“反向 All-Reduce”的由来。帮厨们不能只埋头炒菜，还得在客人反馈后凑在一起对一下账，否则每个人都不知道食材 $X$ 到底该往哪个方向调整。

行并行对比：换一个方向切

列并行是“竖着撕配方”，行并行则是“横着分食材”：把完整的一筐食材分成两份，帮厨 A 拿上半筐，帮厨 B 拿下半筐，但两人手里都有完整的配方。最后两人炒出来的菜必须倒进同一口锅里拌匀，因为每个人只做了半道菜。这个“拌匀”的过程，就是前向计算里的 All-Reduce。

从公式看，行并行把输入 $X$ 按特征维切成 $X_1, X_2$，把权重 $W$ 按行切成上下两块，然后做两个局部矩阵乘，再把结果相加：

对比两种切法，就像帮厨分工的两个极端：一个按配方列分工，最后拼盘；一个按食材量分工，最后混炒。理解了这两个方向，就能明白为什么现代框架要在不同位置选择不同切法。

Transformer组合：厨房流水线

实际的 Transformer 层里，MLP 通常有两个线性层：第一个把维度扩大，第二个把维度缩回去。如果两个线性层都用列并行，中间就需要一次拼接，第二个线性层的前向又要通信，整体通信次数多。更聪明的做法是让第一个线性层用列并行，第二个线性层用行并行，拼出一条流水线。

设第一个线性层的权重 $W_1$ 按列切成 $W_1^{(j)}$，第二个线性层的权重 $W_2$ 按行切成 $W_2^{(j)}$。设备 $j$ 上先算出隐藏分片 $H_j$，这个 $H_j$ 恰好就放在设备 $j$ 上，正好是第二个线性层行并行需要的输入分片。于是两个线性层之间不需要任何通信，只在最后输出时做一次 All-Reduce 求和。

这种“列并行 + 行并行”的组合，就像大排档的流水线：前半段各帮厨按配方列备料，后半段把各人手里的半成品直接倒进同一口锅，只在一开始和最后各通一次气。Transformer 中的注意力 QKV 投影也常用列并行，输出投影用行并行，原理完全一样。

AI 应用场景

Python 演示

import numpy as np def all_reduce(tensor, devices):    total = sum(devices)    return [total for _ in devices] class ColumnParallelLayer:    def __init__(self, in_features, out_features, world_size=2):        assert out_features % world_size == 0        self.in_features = in_features        self.out_features = out_features        self.world_size = world_size        self.per_device = out_features // world_size        rng = np.random.default_rng(0)        self.W = rng.normal(size=(in_features, out_features)).astype(np.float32)        # 本地权重分片列表 每块是完整W的一段列        self.local_W = [            self.W[:, j * self.per_device:(j + 1) * self.per_device]            for j in range(world_size)        ]     def forward(self, x):        self.x = x        # 每个设备只计算自己那段输出列        self.local_y = [x @ w for w in self.local_W]        # 前向把列方向拼回完整输出        return np.concatenate(self.local_y, axis=1)     def backward(self, grad_output):        grad_y_list = np.split(grad_output, self.world_size, axis=1)        # 权重梯度是本地计算 不需要通信        self.local_grad_W = [            self.x.T @ gy for gy in grad_y_list        ]        # 输入梯度分片 每台设备算一部分        grad_x_list = [            gy @ w.T for gy, w in zip(grad_y_list, self.local_W)        ]        # 输入梯度需要 All-Reduce 求和并广播        grad_x_total = sum(grad_x_list)        grad_x = all_reduce(grad_x_total, grad_x_list)[0]        return grad_x if __name__ == "__main__":    layer = ColumnParallelLayer(4, 6, world_size=2)    x = np.random.randn(2, 4).astype(np.float32)    y = layer.forward(x)    grad = np.random.randn(2, 6).astype(np.float32)    grad_x = layer.backward(grad)    print("output shape:", y.shape)    print("grad_x shape:", grad_x.shape)    print("local grad W shape:", layer.local_grad_W[0].shape)

1. self.local_W = [self.W[:, j * self.per_device:(j + 1) * self.per_device] for j in range(world_size)]把权重矩阵按列切成多块，每块分给一个“设备”。
2. y = np.concatenate(self.local_y, axis=1)前向把每个设备的输出分片按列拼回完整输出。
3. self.local_grad_W = [self.x.T @ gy for gy in grad_y_list]每个设备独立计算自己的权重梯度，不需要通信。
4. grad_x_list = [gy @ w.T for gy, w in zip(grad_y_list, self.local_W)]每个设备算出一份对输入 $X$ 的局部梯度。
5. grad_x = all_reduce(grad_x_total, grad_x_list)[0]用 All-Reduce 把所有局部输入梯度求和，并广播回所有设备。

常见误区

面试题