在前面两篇文章当中，我们陆续介绍了 RMSNorm 层归一化「RMSNorm归一化原理与实现过程，与LayerNorm层归一化有何区别？」和用于提升自注意力后两个全连接层效果的 SwiGLU 模块「大模型中的 SwiGLU 激活函数：原理教程与 ReLU 对比实验」。在接下来的这篇文章当中将继续开始介绍一个新的技术——并行计算。
只有清楚地理解每一个基础模块的原理，我们才能更好的理解基于它而构建的上层建筑，当然这也是我们写作的一贯风格。
关键字：LLaMA、并行计算、模型并行、数据并行、张量并行、ColumnParallelLinear、RowParallelLinear
1. 并行计算动机#
在深度学习中，一些大型的网络模型往往需要大量的计算资源才能进行训练，因为每一层神经网络都需要对输入数据进行复杂的矩阵乘法和非线性变换操作。由于单个GPU的计算能力及显存有限可能无法满足大规模深度神经网络的训练需要，因此我们便需要使用多个GPU来加速网络的训练速度。
以 Transformer 为例，它是现在大语言模型的核心技术，其底层便是由大矩阵乘法构成的前馈神经网络和多头注意力机制。当参数量过大时，比如 GPT-3 这样有上千亿参数的模型，一个 GPU 显然是无法同时存储模型中所有参数和中间激活值，那这个时候就需要把模型拆分到多个 GPU 上来进行并行计算。
2. 常见并行技术#
随着模型参数量的不断增长，单块 GPU 无法容纳整个模型，这时我们需要 将模型拆分到多个 GPU 上并行训练，例如常见的并行方式包括数据并行（Data Parallelism）和模型并行（Model Parallelism）。
2.1 并行策略#
数据并行是指将输入网络的训练数据分成多个批次，每个批次在不同的GPU上进行并行计算。此时每个GPU上的模型权重都相同，只是处理的数据不同，每个GPU在训练完自己的批次数据后再将梯度更新汇总到主GPU上，从而实现模型参数的更新。这种方法的优点是简单易实现不容易出错，因此这也是实现多GPU训练中使用最多的一种策略。
当然，数据并行这样的策略也仅仅适用于传统的深度学习模型，当模型本身太大单个 GPU 无法加载时就必须使用到模型并行了。
模型并行是指将模型的不同层或不同计算逻辑分配到不同GPU上进行训练的技术，此时每个 GPU 只负责处理部分逻辑的计算。通常这种方法适用于模型较大且无法在单个GPU上容纳的情况，例如现阶段各个大模型的训练过程通常都会采用模型并行技术。常见的模型并行技术有：层并行（Layer Parallelism），把模型的不同层分配给不同的 GPU；张量并行（Tensor Parallelism），把一个层内部的张量（如 Linear 权重矩阵）按维度切块分配到多个 GPU 上，如 LaMMA 和 DeepSeek 等就是采用的张量并行策略。
一句话就是，模型太大，单卡装不下，就把模型按模块或按张量维度切开，分给多个 GPU 执行，这就是模型并行。
2.2 数据并行#
数据并行是传统深度学习领域中使用最为常见的一种并行测录。在使用数据并行策略实现多GPU训练时，首先会将整个小批量数据再划分成多个小批次并分配到不同的GPU上，同时整个模型也将被复制到每个GPU上，然后在每个GPU上模型均各自独立地完成损失和梯度的计算，随后将每个GPU上计算得到的损失和梯度汇聚到主GPU上得到整个小批量数据样本的平均梯度，最后再将该梯度分配到其它GPU中进行各自模型参数的更新以完成一次迭代训练过程。
如图1所示便是含有两个GPU的数据并行原理图。例如此时每个小批量数据都含有256个样本，那么图示中每个GPU将会被分配得到128个样本进行后续的计算处理。同时，每个GPU上也都有着一模一样的网络模型，并且它们在各自拿到128个样本后会分别计算损失和梯度，然后再将两部分的梯度汇聚到主GPU上得到256个样本的平均梯度，最后再用该梯度通过梯度下降算法并行对每个GPU上的模型进行参数更新。
更多相关内容可以参见深度学习栏目「第5.6节 多GPU训练：数据并行、模型并行与 PyTorch 实现」 内容。
2.3 模型并行#
虽然模型并行的策略听起来也十分容易理解，但是实现起来相比较于数据并行则要稍显复杂，那下面就来介绍一下 LLaMA 和 DeepSeek 中使用到的张量并行技术（Tensor Parallelism）。
张量并行是一种用于训练大规模神经网络模型的模型并行技术，主要应用于参数特别多、单张 GPU 无法容纳整个模型的情况。它通过将一个操作或张量按维度划分到多个设备上执行，使得每个设备只需处理部分数据，从而降低单卡的内存负担。张量并行最早由 NVIDIA 在 Megatron-LM 中提出，用于对大型线性层（如 Transformer 中的多头注意力的全连接层）进行列方向的并行拆分。
例如在 LLaMA 和 DeepSeek 中使用到的 ColumnParallelLinear 和 RowParallelLinear 就是比较常见的张量并行技术，下面依次进行介绍。
3. 张量并行原理#
在介绍完模型并行、数据并行等相关思想后，下面我们再来详细看一下模型并行中张量并行里常见的两者并行方法：ColumnParallelLinear 和 RowParallelLinear 的原理。
3.1 ColumnParallelLinear 原理#
假设有一个线性变换层：
$$ y=xW^T+b\tag{1} $$
其中，输入 $x$ 的形状为 [batch_size, input_size]，权重 $W$ 的形状为 [output_size, input_size]，输出 $y$ 的形状为 [batch_size, output_size]。
对于式(1)的计算过程，我们可以通过图2所示的过程来进行表示。
对于图2所示的过程，在 PyTorch 中可以通过 nn.Linear() 来完成整个计算过程。
那 ColumnParallelLinear 的原理是什么呢？
ColumnParallelLinear 翻译成中文应该叫列并行线性层。不过这里需要注意一点的是，ColumnParallelLinear 表示在列上进行并行，既然是在列上进行并行那就是多列同时进行计算，故需要在行上对权重矩阵进行划分，所以大家不要搞混淆。
如图3所示，对于 ColumnParallelLinear 来说，首先需要将权重矩阵按行进行拆分。至于拆分成多少块，可以是每一行作为一块，也可以是多行作为一块，但是一定要整除。在图3中，我们将权重矩阵 W 按行切分成了5块，即 world_size = 5。
这里顺便提一下，在分布式计算中，world_size 是一个标准术语，来源于多种分布式训练框架（如 PyTorch Distributed、Horovod、MPI 等）。这个命名源于 MPI（Message Passing Interface）规范中对“全体进程集合”的描述 —— 它被称为一个 communicator world，其中包含所有可通信的节点。于是：world指的是整个通信域（所有 GPU/进程）；size指这个通信域的大小，也就是进程数。
下一步，分别开始进行计算。
根据图4的计算过程可知，它完成的便是其中一块的计算过程，而对于每一块的计算我们便可以将它们放到不同的 GPU 上以此来达到并行计算的目的。当每一块上的计算过程都结束以后，我们将 $Y_i$ 在维度 1（列）上进行拼接（All Gather）便得到了最终的结果，如图5所示。
到这里呢，我们就把 ColumnParallelLinear 的计算原理介绍清楚了，下面再来看 RowParallelLinear 的原理。
3.2 RowParallelLinear 原理#
同 ColumnParallelLinear 类似，RowParallelLinear 翻译成中文应该叫行并行线性层，也就是在行上进行并行计算，故需要在列上对权重矩阵进行划分。不过大家现在可以仔细想想，此时可以像类似图4到图5的过程进行并行计算吗？
3
2
1
答案是显然是不可以的。为什么？因为没办法算，大家可以在脑海里想一想。所以还需要对输入同样在列上仅划分，如图6所示。
这里需要注意的是，$X$ 和 $W$ 所划分的块数需要保持一样。
进一步，分别开始进行计算。
如图7所示，它完成的便是其中一块的计算过程，同样的可以将其分别放到不同的 GPU 上来试下并行。这里为了让大家更清楚图7的计算过程，我们再来看一下同时计算两块时的过程，如图8所示。
最后再将每个部分计算得到的结果加到一起（All-Reduce）即可得到最终结果，如图9所示。
到这里呢，我们就把 RowParallelLinear 的计算原理介绍清楚了，是不是觉得原理也还算简单？
这里，请大家分别记住 ColumnParallelLinear 和 RowParallelLinear 合并或相加之前各自的形状，后面还会再次用到。
那介绍到这里，大家是不是就以为我们在实现模型的时候，选择上面其中一种并行方式就行了？
是吗？ 显然不是。
但你是不是觉得好像又是？明明两种方式都能够计算得到最后的结果，那为什么又不能够这样用呢？
所以本质上我们要回答的问题就是， ColumnParallelLinear 和 RowParallelLinear 到底是替代关系还是配合使用的关系？
4. GPU 上的并行过程#
4.1 为什么不能只用一个#
有人可能会说，我使用完 ColumnParallelLinear 后，可以把拆分的输出组合到一起以后，继续使用 ColumnParallelLinear 呀？难道不行吗？
虽然我们可以在每次 ColumnParallelLinear 后把各 GPU 的部分输出合并为完整输出，然后再广播到每个 GPU，供下一个 ColumnParallelLinear 使用，但是这样做通信代价太高，因为拼接是重操作，这要做违背了流水线并行（PipeParallel）的设计思想。
因此，最后真正要实现高效的张量并行计算，需要将 ColumnParallelLinear 和 RowParallelLinear 配合起来使用，而这也是张量并行的核心。
当然，如果你只有一个线性层，那就另当别论。
4.2 图解并行计算过程#
下面，我们以连续两个线性层为例，来完整地示例一下它们在5个GPU上的并行计算流程。
首先，我们假定网络的输入为 $X$ ，形状为 [batch_size, input_size]；隐藏层权重为 $W$，形状为 [dim_size, input_size]；隐藏层输出为 $O$，形状为 [batch_size, dim_size]；输出层权重为 $V$， 形状为 [output_size, dim_size]；网络最终输出为 $Y$，形状为 [batch_size, output_size]。
在这里大家可能会觉得奇怪，为什么要把 $W$ 的形状定义为 [dim_size, input_size]，交换一下 [input_size, dim_size] 不是更好吗？相关内容可以参见专栏文章 『为什么PyTorch中 nn.Linear()在初始化权重时要定义为 [out_features, in_features] ?』
那么对于整个计算过程来说，我们可以通过图9所示的过程来进行表示。
当然，图9中的过程是我们最终看到的抽象的计算过程，下我们我们再来看看它在5个GPU上的并行计算流程。此时，可以通过如图10所示的过程来进行表示。
在图10中，首先输入 $X$ 将会分别被广播到每个 GPU 上，也就是每个 GPU 上都拥有一个完整的 $X$；然后在每个 GPU 上都分别初始化得到得到一个完整的权重矩阵 $W$，并进一步根据并行的数量从 $W$ 中划分得到对应的 $W_i$；最后各个完成并行计算得到各自的结果 $O_i$，并同时准备进入下一阶段计算。具体的划分过程，我们将在后续的代码实现过程中详细介绍。
总结就是，ColumnParallelLinear 的能力是将权重在输出维度划分成 $N$ 份，每个 GPU 拥有完整输入，只计算对应的输出部分，因此，它可以并行计算线性层的输出部分，但不能继续处理后续层（因为输出是拆分的）。
在完成 ColumnParallelLinear 计算以后，开始进行 RowParallelLinear 计算。
在图12中，上一步 ColumnParallelLinear 计算完成以后，每个 GPU 上分别已经留有对应的计算结果 $O_i$。因此，对于 RowParallelLinear 来说只需要在每个 GPU 上初始化一个权重矩阵 $V$，并划分得到对应的 $V_i$，最后完成整个并行计算得到 $Y_i$ 即可。
总结就是，RowParallelLinear 的能力是将权重在输入维度上划分成 $N$ 份，每个 GPU 只接受部分输入并计算出部分输出。
此时我们可以看出，在 ColumnParallelLinear 计算结束以后的结果分布形式恰好就是 RowParallelLinear 所需要的，根本不需要进行合并操作。
进一步，在 RowParallelLinear 完成计算以后，我们再通过 All Reduce 得到最终的输出 $Y$。这里的具体过程是，每个 GPU 分别将其它 GPU 上计算得到的结果拿过来，然后累加得到 $Y$。
到此，我们便利用 ColumnParallelLinear 和 RowParallelLinear 完成了连续两个线性层计算过程。如果后续还有这样的操作，则继续按照先 ColumnParallelLinear 再 RowParallelLinear 计算即可。
此时，有人可能会问在 RowParallelLinear 计算完毕以后不还是有合并操作？
虽然 RowParallelLinear 计算完毕以后也有合并操作，但 All Reduce 是加和操作，有高度优化的实现（比如 NCCL 中的 ring-reduce、tree-reduce）等，而 ColumnParallelLinear 之后则 All Gather 是则是拼接操作，往往通信更重，数据要全部原样广播。
因此，以先 ColumnParallelLinear 再 RowParallelLinear 的顺序就能够完美的避开 All Gather操作，实现最大的并行效率化。
所以，通过这一配合，模型能够在每个 GPU 上同时并行执行大规模线性层计算，且通信和计算负载都相对平衡，而这也是是目前各个大型模型广泛采用的核心和主流技术之一。
到此，对于张量并行计算的原理就给大家交代清楚了。至于这两个操作的代码实现过程，由于篇幅有限我们就放到下一篇文章中给大家进行详细介绍。
引用#
[1] https://github.com/facebookresearch/fairscale