Transformer中的并行层架构
Transformer中的并行层架构
本文档整理了关于Transformer架构的一种变体——“并行层”（Parallel Layers）的详细讨论，旨在解释其工作原理、优势以及实现方式。
核心思想：从串行到并行
在一个标准的Transformer模块里，信息的处理是分步进行的，像流水线一样：首先经过注意力层（Attention Layer），然后其输出再经过前馈网络层（MLP Layer）。两者是串行关系。
并行层 (Parallel Layers) 的核心思想是：不再让MLP层等待Attention层，而是让它们同时进行计算。
- 标准（串行）结构:
y = x + MLP(LayerNorm(x + Attention(LayerNorm(x)))) - 并行结构:
y = x + MLP(LayerNorm(x)) + Attention(LayerNorm(x))
在这种新结构中，MLP和Attention的输入都是同一个值：LayerNorm(x)。它们之间不再有依赖关系，因此可以并行计算。
优势与权衡
- 显著提升训练速度：由于Attention和MLP可以并行计算，它们的输入矩阵乘法可以被“融合”（fuse）成一个单一的操作，从而更有效地利用硬件（如GPU）。这在大规模训练时，可以带来大约15%的速度提升。像PaLM、GPT-J和GPT-NeoX等模型都采用了这种设计来加速训练。
- 对模型精度的影响：
- 在较小模型规模（如80亿参数）上，这种并行结构会导致轻微的性能下降。
- 但随着模型规模增大（如620亿参数），这种性能退化会消失。
- 因此，研究人员推断，在更大规模上，并行层对模型质量没有负面影响。
注意：正如课程中所提到的，模型架构的趋势总是在不断演进，这种并行层架构在最近可能已不再是最新模型的主流选择，但其设计思想对于理解模型优化依然重要。
Brainstorm Q&A: 并行层是如何做到的？
问：怎么做到的？
答： 将串行计算改为并行计算，并实现“矩阵乘法融合”，主要通过以下三个步骤做到：
- 结构重组：消除计算依赖
- 权重合并：将多个矩阵乘法融合成一个
- 硬件并行：利用GPU的高效计算特性
1. 结构重组：消除计算依赖
首先，通过修改计算公式，让MLP和Attention的输入都变成同一个值LayerNorm(x)，消除了两者之间的先后依赖关系，这是实现并行计算的理论基础。
2. 权重合并：矩阵乘法融合的魔术
这是实现15%速度提升的核心。在Transformer模块中，计算量最大的部分就是大规模的矩阵乘法（MatMul）。
在自注意力 (Self-Attention) 机制中，第一步是将输入
H(即LayerNorm(x)) 乘以一个大的权重矩阵W_qkv，来生成查询（Q）、键（K）和值（V）三个向量。- 计算式：
QKV = H @ W_qkv
- 计算式：
在MLP（前馈网络）层中，第一步也是将同一个输入
H乘以一个权重矩阵W_mlp1。- 计算式：
MLP_hidden = H @ W_mlp1
- 计算式：
在并行结构中，由于它们的输入 H 完全相同，我们可以把它们“融合”（Fuse）起来。
融合方法：
我们将两个权重矩阵 W_qkv 和 W_mlp1 沿其输出维度进行拼接（Concatenate），形成一个更大、更宽的单一权重矩阵 W_fused。
W_fused= Concat(W_qkv,W_mlp1)
现在，我们只需要执行一次矩阵乘法，就能同时得到Attention和MLP所需要的结果：
Fused_output = H @ W_fused
这个 Fused_output 是一个很宽的矩阵，它的前半部分包含了Q、K、V的结果，后半部分则是MLP的中间层结果。计算完成后，程序只需将这个大矩阵按列切分（Split）开，分别送入后续的Attention和MLP计算流程即可。
3. 硬件并行：为什么融合后更快？
将两个小矩阵乘法融合成一个大矩阵乘法，之所以能提速，是因为它更符合GPU等现代并行处理器的“胃口”：
- 减少计算核（Kernel）启动开销：每次调用一次矩阵乘法，都需要CPU向GPU发起一次“计算核启动”指令，这本身有一定的时间开销。将两次启动合并为一次，就减少了这部分开销。
- 提升计算密度和硬件利用率：GPU拥有数千个计算核心，擅长处理大规模、规整的并行任务。一个大的矩阵乘法可以更好地“喂饱”所有核心，让它们同时工作，从而最大化硬件的利用率。相比之下，两个连续的小矩阵乘法可能会因为数据依赖和调度问题导致部分核心闲置。
因此，通过结构重组和权重合并，并行层架构将计算任务变得更加“GPU友好”，从而实现了显著的训练加速。
整理自: Transformer 并行层解析