神经网络Transformer架构中的模型并行 神经网络Transformer架构中的模型并行是一种分布式训练技术，用于解决当模型规模过大而无法在单个处理器（如GPU）的内存中容纳时的问题。它将单个模型的参数分割到多个处理器上，每个处理器负责计算模型的一部分。 步骤1：理解模型并行的基本需求 当Transformer模型（例如拥有数百或数千亿参数的大语言模型）的参数数量超过单个处理器的内存容量时，无法进行正常的加载和训练。模型并行的核心目标就是将这个庞大的模型“拆分”开，分布到多个设备上。这与数据并行（每个设备都有完整的模型副本，但处理不同的数据批次）形成对比。 步骤2：模型并行的常见拆分维度 在Transformer架构中，模型并行主要通过两种主要维度进行参数拆分： 层内并行（Intra-layer Parallelism）：这是最常用的方式。以Transformer中的关键组件——前馈网络（FFN）和注意力层——为例进行拆分。 对于前馈网络（FFN）：一个标准的FFN包含两个线性层（例如，从隐藏维度 d_model 投影到中间维度 d_ff ，再投影回 d_model ）。我们可以将第一个线性层的权重矩阵按列分割（即分割输出神经元），将拆分后的部分放置在不同的设备上。每个设备计算得到部分结果，然后通过一个集合通信操作（如All-Reduce）对结果进行求和或拼接，再输入到下一个阶段。 对于多头注意力（MHA）：可以将不同的“头”分布到不同的设备上。每个设备计算自己分配到的注意力头的查询、键、值和输出投影。最后，所有设备的输出需要被拼接起来，形成完整的注意力输出。 层间并行（Inter-layer Parallelism）：也称为管道并行（Pipeline Parallelism）。这种方式将Transformer模型的不同层放置在不同的设备上。例如，设备1负责第1到第4层，设备2负责第5到第8层，以此类推。数据（输入序列的激活值）像一个流水线一样，从一个设备传递到下一个设备。为了减少设备间的空闲等待时间，通常需要将训练数据分成多个微批次（Micro-batches）。 步骤3：模型并行的执行流程与通信 模型并行训练严重依赖设备间的通信。 前向传播：输入数据从第一个设备开始。在每个设备上，计算本地存储的参数所对应的部分结果。如果该部分结果是整个计算的中间环节（如在层内并行中），则需要将结果发送到下一个设备，或者通过集合通信（如All-Gather）来整合所有部分结果，以进行后续计算。 反向传播：损失函数的梯度从最后一个设备开始，反向传播。在计算参数梯度时，由于参数分布在不同的设备上，经常需要通过集合通信（如Reduce-Scatter）来收集和分发梯度信息，以确保每个设备能正确更新其存储的那部分参数。 步骤4：模型并行的挑战与优化策略 通信开销：设备间的数据传输是主要瓶颈。优化策略包括使用高速互联网络（如NVLink）、重叠通信与计算（在计算的同时传输数据）以及精心设计模型拆分策略以最小化通信量。 计算负载均衡：理想情况下，每个设备上的计算量和内存占用应尽可能均衡。如果拆分不均，某个设备会成为瓶颈，拖慢整体训练速度。 流水线气泡（针对管道并行）：在管道并行中，当第一个微批次在后续设备上计算时，前面的设备会处于空闲状态，等待新的任务，这被称为“气泡”。通过增加微批次数量，可以使流水线更加饱满，减小气泡大小。 步骤5：模型并行的实际应用与组合 在实际中，尤其是训练超大规模模型时，模型并行通常会与数据并行结合使用。例如，在一个拥有成千上万个处理器的集群中，可以先用模型并行将一个大模型拆分到一组设备（如8个GPU）上，形成一个“模型并行组”。然后，再使用数据并行，将这个“模型并行组”视为一个单元，复制多份到不同的“数据并行组”中，每个组处理不同的数据批次。这种混合并行策略是当前训练万亿参数级别模型的主流方法。