大语言模型的预训练

发表于 2025-05-09 Waline：阅读次数：本文字数： 7.8k 阅读时长 ≈ 28 分钟

大语言模型的基本结构

大语言模型（Large Language Models, LLMs）的预训练是当前自然语言处理领域的核心。理解其基本结构是后续学习的基础。

传统自回归类型 —— 以 Llama 为例

自回归（Autoregressive）模型是指在生成序列时，当前时间步的输出依赖于先前所有时间步的输出。Decoder-only 架构是目前主流大语言模型（如 GPT 系列、Llama 系列）采用的结构，它仅使用 Transformer 的解码器部分进行单向的序列建模和生成。

Llama 系列模型由 Meta 于 2023 年 3 月首次发布，其开源迅速点燃了社区的热情，并催生了众多衍生模型，如 Alpaca、Vicuna 等。

Llama 等模型在 Transformer 结构的基础上进行了一些关键改进：

前置归一化（Pre-Normalization）
- 归一化层在深度学习模型中至关重要，它有助于稳定训练过程，加速收敛，并可能提升模型性能。
- Post-Norm：原始 Transformer 结构在残差连接之后应用层归一化，即 $x_{t+1} = \operatorname{Norm}(x_t + F_t(x_t))$ 。这种结构在训练初期可能不太稳定，但最终性能可能较好。
- Pre-Norm：将层归一化置于残差连接的主分支上，在自注意力或前馈网络之前，即 $x_{t+1} = x_t + F_t(\operatorname{Norm}(x_t))$ 。Pre-Norm 结构通常更容易训练，梯度更稳定，允许更大的学习率。
- Llama 采用了 RMSNorm(Root Mean Square Normalization)，它是 LayerNorm 的一种简化形式。RMSNorm 仅通过均方根对输入进行缩放，而不进行均值中心化，计算效率更高。
  - 对于输入向量 $\bm{a} = (a_1, \dots, a_n)$ $a = (a_{1}, \dots, a_{n})$ 和可学习的缩放参数 $g_i$ $g_{i}$ ：
    - RMSNorm: $\bar{a}_i = \frac{a_i}{\operatorname{RMS}(\bm{a})} g_i$ $\overset{a}{ˉ}_{i} = \frac{a _{i}}{RMS ( a )} g_{i}$
      - 其中 $\operatorname{RMS}(\bm{a}) = \sqrt{\frac{1}{n} \sum_{j=1}^n a_j^2}$
    - LayerNorm: $\bar{a}_i = \frac{a_i - \mu}{\sigma} g_i + b_i$ $\overset{a}{ˉ}_{i} = \frac{a _{i} - μ}{σ} g_{i} + b_{i}$
      - 其中 $\mu = \frac{1}{n} \sum_{j=1}^n a_j, \sigma = \sqrt{\frac{1}{n} \sum_{j=1}^n (a_j - \mu)^2}$
  - Llama 使用 RMSNorm 能够使模型的训练过程更加稳定。
SwiGLU 激活函数
- SwiGLU 是 GLU(Gated Linear Unit) 的一种变体，目前是许多大模型中常用的激活函数。它引入了一个门控机制来调节信息流。
- 其计算公式为： $\operatorname{SwiGLU}_{\beta}(\bm{x}, \bm{W}, \bm{V}, \bm{b}, \bm{c}) = \operatorname{Swish}_{\beta}(\bm{x}\bm{W} + \bm{b}) \otimes (\bm{x}\bm{V} + \bm{c})$ $SwiGLU_{β} (x, W, V, b, c) = Swish_{β} (x W + b) \otimes (x V + c)$
  - 其中 $\operatorname{Swish}_{\beta}(\bm{x}) = \bm{x} \cdot \sigma(\beta\bm{x})$ ， $\sigma$ 是 Sigmoid 函数， $\otimes$ 表示逐元素相乘。
- SwiGLU 通常比标准的 ReLU 或 GeLU 表现更好，因为它允许数据驱动的门控，可以更灵活地控制哪些信息通过。
RoPE
- RoPE(Rotary Positional Embeddings, 旋转位置编码) 是一种新颖的位置编码方式，它将绝对位置信息以旋转的方式融入到 Query 和 Key 向量中，从而在注意力计算中自然地引入相对位置信息。这部分将在「位置编码策略」中详细介绍。

混合专家模型类型 —— 以 Mixtral 为例

混合专家模型（Mixture-of-Experts, MoE）是一种通过组合多个「专家」子模型来处理任务的架构。其核心思想是让不同的专家专注于数据的不同方面或模式，从而提高模型的容量和效率，同时在推理时只激活部分专家，降低计算成本。

Mixtral 是基于 MoE 架构的代表性模型。

基本思想：在 Transformer 的每个或部分前馈网络（FFN）层替换为 MoE 层。MoE 层包含多个独立的 FFN（称为专家）和一个门控网络或称为路由器。
工作流程：
1. 对于每个输入的 token，门控网络计算一组权重，决定将该 token 发送给哪些专家（通常是 Top-K 个，例如 Mixtral 选择 Top-2）。
2. 选定的专家处理该 token。
3. 各个专家的输出根据门控网络计算的权重进行加权组合。
核心差异：相较于标准 Transformer，MoE 模型在 FFN 部分引入了门控层和多个专家全连接层。
Token-level 动态专家选择：每个词元（token）都会根据其上下文动态地选择激活的专家。这意味着专家不是全局共享的，而是根据词元粒度进行动态混合。

优势：
- 参数高效：可以在总参数量远大于单个专家的情况下进行训练，但推理时只激活一小部分参数，计算成本较低。
- 多样性理解：每个专家可以专注于不同的语言模式、结构或风格，提高模型对多样化输入的理解能力。
- 性能优越：尤其在大规模多语言、多任务处理中表现突出。

缩放法则

缩放法则描述了模型大小（参数数量 N）、训练数据量（词元数量 D）和计算资源（C）三者之间的相互关系，以及它们如何共同影响模型性能的提升趋势。

核心观点：
- 模型越大，通常需要更多的数据和更多的计算资源来充分训练其潜力。
- 当这三者按比例增加时，模型性能的提升是可预测的。OpenAI 的研究表明，损失与 N, D, C 之间存在幂律关系。
理论指导：缩放法则是大语言模型设计、资源分配和性能预估的理论指导依据。
Chinchilla Scaling Laws：DeepMind 的研究指出，对于给定的计算预算，最佳性能并非通过尽可能大的模型实现，而是通过在模型大小和训练数据量之间取得平衡。他们提出的一个关键关系是，模型参数量 N 和训练词元数 D 应该大致等比例增长。
近似关系：一个常被引用的关系是 $C \approx 6ND$ $C \approx 6 N D$ ，表示训练一个具有 N 个参数的模型，在 D 个词元上进行一次完整的前向和后向传播（不包括梯度累积等复杂情况），大约需要 $6ND$ $6 N D$ 次浮点运算（FLOPs）。
- $C$ : 总计算量（FLOPs）
- $N$ : 模型参数量
- $D$ : 训练词元数

计算量估算

$6ND$ 这个估算中， $2ND$ 来自前向传播（矩阵乘法中参数量 $N$ 与输入序列长度和批大小相关的部分，简化为 $N \times D$ 级别）， $4ND$ 来自后向传播（梯度计算通常是前向传播的两倍）。这只是一个粗略的估计，实际计算量会因具体实现、激活函数、注意力计算等因素有所不同。

注意力机制的优化

标准自注意力机制的计算复杂度和内存占用均为 $O(N^2)$ （ $N$ 为序列长度），这使得它在处理长序列时效率低下。因此，研究者提出了多种优化方法。

稀疏注意力 —— Longformer 为例

稀疏注意力旨在通过限制每个 token 只关注序列中的一部分其他 token，来降低计算复杂度。Longformer 是一个采用稀疏注意力机制的代表模型，它能够将最大可处理长度扩展至如 4096 甚至更长。

Longformer 结合了以下几种稀疏注意力模式：

滑动窗口注意力：每个 token 只关注其邻近的 $w$ 个 token（一个固定大小的窗口）。这类似于卷积神经网络中的局部感受野。
扩张滑动窗口注意力：类似于空洞卷积，滑动窗口中的 token 不是连续的，而是带有一定的间隔。这允许在不增加计算量的情况下扩大感受野。
全局 + 滑动窗口注意力：除了滑动窗口内的局部注意力，还允许一部分预选的 token（例如任务相关的特殊 token，如 [CLS]）关注整个序列中的所有 token，并且整个序列中的所有 token 也都关注这些全局 token。

多查询注意力与分组查询注意力

在 Transformer 推理（尤其是生成任务）时，每生成一个 token，都需要将新生成的 token 与所有历史 token 的键（K）和值（V）进行注意力计算。这些 K 和 V 通常被缓存起来（称为 KV 缓存），以避免重复计算。随着序列变长，KV 缓存会变得非常大，频繁加载这些键值张量会带来巨大的内存带宽压力。

多头注意力（Multi-Head Attention, MHA）：标准做法，每个注意力头都有自己独立的 Q, K, V 投影权重。
多查询注意力（Multi-Query Attention, MQA）：所有注意力头共享同一组键（K）和值（V）的投影权重和张量。Query (Q) 仍然是每个头独立的。
- 优点：显著减少 KV 缓存的大小和内存带宽需求，加快推理速度。
- 缺点：可能导致一定的性能下降，因为所有头共享 K/V 限制了它们捕捉不同关系的能力。
分组查询注意力（Grouped-Query Attention, GQA）：一种折中方案。将 Query 头分成多组，每组内的 Query 头共享同一组键（K）和值（V）。
- 例如，如果有 32 个 Query 头，可以设 4 组，则每 8 个 Query 头共享一组 K/V。
- 优点：在 MHA 的性能和 MQA 的效率之间取得了良好的平衡。Llama 2/3 等模型采用了 GQA。

FlashAttention

FlashAttention 是一种针对标准注意力机制的 I/O 感知优化算法，旨在解决其在 GPU 上受内存带宽限制的问题。

背景：
- 标准注意力机制需要计算并存储 $N \times N$ 的注意力矩阵 $\bm{S} = \bm{Q}\bm{K}^\mathsf{T}$ 和 $\bm{P} = \operatorname{Softmax}(\bm{S})$ 。当 $N$ 很大时，这两个矩阵无法完全放入 GPU 的高速 SRAM 中，必须在较慢的 HBM（高带宽内存）中读写。
- Softmax、Dropout 等操作也涉及对 HBM 的大量读写，这些操作通常是内存带宽受限的，而非计算受限。
方法：FlashAttention 的核心思想是分块计算 和 内核融合，以避免将巨大的中间注意力矩阵 $\bm{S}$ 和 $\bm{P}$ 完整写入 HBM。
1. 分块处理：将输入矩阵 $\bm{Q}, \bm{K}, \bm{V}$ 沿序列长度维度切分成块。
2. 外层循环：遍历 $\bm{K}$ 和 $\bm{V}$ 的块。将当前 $\bm{K}$ 块和 $\bm{V}$ 块从 HBM 加载到 SRAM。
3. 内层循环：对于加载到 SRAM 的 $\bm{K}, \bm{V}$ 块，再遍历 $\bm{Q}$ 的块。将当前 $\bm{Q}$ 块加载到 SRAM。
4. SRAM 内计算：在 SRAM 中计算当前 $\bm{Q}$ 块与 $\bm{K}$ 块之间的注意力得分，应用 Softmax（使用一种在线的、数值稳定的算法），然后与 $\bm{V}$ 块相乘得到该块的输出。
5. 结果写回：将计算得到的输出块累加到 HBM 中的最终输出矩阵 $\bm{O}$ 的相应位置。
- Softmax 归一化因子缓存：为了在分块处理 Softmax 时保持数值稳定性并正确归一化，FlashAttention 会缓存并更新 Softmax 的归一化因子（即分母 $\ell(\bm{x})$ 和用于稳定计算的最大值 $m(\bm{x})$ ）。这使得反向传播时也无需读取完整的注意力矩阵。

数值稳定 Softmax

标准 Softmax 计算为 $\operatorname{softmax}(\bm{x})_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}$ 。为避免数值溢出，通常使用技巧：

$m(\bm{x}) = \max_i x_i$
$\operatorname{softmax}(\bm{x})_i = \frac{\exp(x_i - m(\bm{x}))}{\sum_j \exp(x_j - m(\bm{x}))}$

FlashAttention 在分块计算时，会迭代地更新 $m(\bm{x})$ 和分母 $\ell(\bm{x}) = \sum_j \exp(x_j - m(\bm{x}))$ .

优势：
- 更快的训练和推理速度：通过显著减少 HBM 的读写次数，FlashAttention 将注意力计算从内存带宽受限转变为更接近计算受限，从而大幅提速（例如，在 GPT-2 上可达数倍加速）。
- 内存高效：由于不需要实例化完整的 $N \times N$ 注意力矩阵，内存占用从 $O(N^2)$ 降低到 $O(N \sqrt{d})$ 或 $O(N d)$ （取决于实现细节和块大小），使得能够处理更长的序列。
- 更好的模型效果：处理更长序列的能力本身就能带来模型质量的提升和新功能的实现。

位置编码策略

Transformer 模型本身不包含任何关于序列中 token 顺序或位置的信息。位置编码（Positional Encoding, PE）的目的是向模型注入这种位置信息。

RoPE (Rotary Positional Embeddings)

旋转位置编码（RoPE）是一种新颖的相对位置编码方法，它通过对 Query 和 Key 向量进行旋转操作来融入位置信息。

背景问题：
- 标准的绝对位置编码（如学习的嵌入或固定的正弦编码）在处理超过预训练长度的序列时（外推性）表现不佳，且难以适配流式输入。
- 对于长文本，传统位置编码难以精确捕捉远距离的相对位置信息。
RoPE 的核心思想：将绝对位置信息 $m$ （对于 Query）和 $n$ （对于 Key）通过一种特殊的方式编码进它们的表示 $\bm{q}_m$ 和 $\bm{k}_n$ ，使得它们的内积（注意力得分的核心部分）仅依赖于相对位置 $m-n$ 和它们的内容。
- 具体做法是：在隐含层维度上，将每两个维度视为一个二维平面的坐标 $(x_i, x_{i+1})$ 。然后，根据 token 的绝对位置 $m$ ，将这个二维向量旋转 $m\theta_i$ 角度，其中 $\theta_i$ 是一个预设的、与维度相关的非零常数（通常 $\theta_i = 10000^{-2i/d}$ ，与原始 Transformer 的正弦 PE 类似）。
- 数学上，对于一个 $d$ $d$ 维向量 $\bm{x}_m$ $x_{m}$ 在位置 $m$ $m$ ，其 RoPE 变换后的向量 $f(\bm{x}_m, m)$ $f (x_{m}, m)$ 可以通过一个旋转矩阵 $\bm{R}_{\Theta,m}^d$ $R_{Θ, m}^{d}$ 作用于 $\bm{x}_m$ $x_{m}$ 得到：
  - $f(\bm{q}, m) = \bm{R}_{\Theta,m}^d \bm{W}_q \bm{x}_m$
  - $f(\bm{k}, n) = \bm{R}_{\Theta,n}^d \bm{W}_k \bm{x}_n$
  - 其中 $\bm{R}_{\Theta,m}^d$ 是一个块对角矩阵，每个 $2 \times 2$ 的块是一个旋转矩阵： $\begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \\ \sin m\theta_i & \cos m\theta_i \end{pmatrix}$
- 关键在于，经过 RoPE 变换后的 $\bm{q}'_m = f(\bm{q}, m)$ 和 $\bm{k}'_n = f(\bm{k}, n)$ 的内积 $(\bm{q}'_m)^\mathsf{T} \bm{k}'_n$ 可以被证明等价于原始向量 $\bm{W}_q \bm{x}_m$ 和 $\bm{W}_k \bm{x}_n$ 的内积，但其相位被旋转了 $(m-n)\theta_i$ 。这使得注意力得分天然地包含了相对位置信息。
优势：
- 捕捉相对位置信息：能够有效地建模长距离依赖。
- 支持任意长度输入：理论上对序列长度没有限制，因为旋转操作可以应用于任何位置。
- 平移等变性：注意力得分对于位置的平移具有等变性，即相对位置不变时，注意力模式相似。
- 实现简单：直接修改 Q 和 K 向量，无需额外的嵌入层。

ALiBi (Attention with Linear Biases)

ALiBi 是一种更简单的处理位置信息的方法，它直接在注意力打分阶段引入与位置差相关的线性偏置项，而无需修改词嵌入或 Q/K 向量。

出发点：
- 传统位置编码（如绝对编码）不适用于超长序列。
- RoPE 虽然强大，但仍需在输入中显式地通过旋转操作加入位置结构。
- ALiBi 旨在通过一种更简洁高效的方式融入位置信息。
核心思想：在计算注意力分数 $\bm{S} = \frac{\bm{Q}\bm{K}^\mathsf{T}}{\sqrt{d_k}}$ 之后，Softmax 之前，直接加上一个偏置矩阵 $\bm{M}$ ：
$\operatorname{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \operatorname{Softmax}\left(\frac{\bm{Q}\bm{K}^\mathsf{T}}{\sqrt{d_k}} + m \cdot \text{BiasMatrix} \right) \bm{V}$
其中，BiasMatrix 是一个预定义的矩阵，其 $(i, j)$ 位置的值通常是 $-(i-j)$ （对于因果掩码，只考虑 $j \le i$ ），表示 query $i$ 和 key $j$ 之间的距离。 $m$ 是一个可学习的、与注意力头相关的标量斜率。
- 这个偏置项会惩罚距离较远的 token，使得模型更关注近期的上下文。
特点与优势：
- 无需额外位置嵌入参数：不增加模型参数量。
- 支持动态长度与超长序列：由于偏置是根据相对距离动态计算的，因此能很好地泛化到训练时未见过的长度。
- 良好的归纳偏好：自然地偏向近期上下文，有利于结构理解和外推。
- 兼容性好：易于集成到现有的大模型架构中。

长上下文处理策略

随着应用场景对更长文本理解能力的需求增加，如何让大语言模型有效处理并利用长上下文信息成为了一个重要的研究方向。

位置插值法

当需要模型处理比预训练时更长的序列时，直接外推位置编码往往效果不佳。位置插值法是一种简单而有效的方法，用于扩展模型的上下文窗口。

核心思想：将新输入序列（长度 $L' > L_{\text{train}}$ $L^{'} > L_{train}$ ）的位置索引按比例缩放到模型预训练时支持的原始窗口长度 $L_{\text{train}}$ $L_{train}$ 内。
- 例如，如果模型预训练时最大长度为 2048，现在要处理长度为 4096 的序列，则可以将新序列中位置 $m$ 的位置编码计算为原始位置编码中 $m/2$ 位置的值。
- 对于 RoPE，这意味着旋转角度 $m\theta_i$ 变为 $(m \cdot L_{\text{train}}/L')\theta_i$ 。
效果：通过将长序列的位置「压缩」到预训练的范围内，使得模型可以利用其在预训练阶段学到的位置依赖模式。这通常比直接外推效果好得多，但可能会损失一些高频（精细）的位置信息。

基于 NTK 的方法

直接的位置插值虽然能扩展上下文长度，但可能干扰模型对短距离词元间精确关系的计算（因为高频信息被平滑了）。基于神经切线核（Neural Tangent Kernel, NTK）理论的方法试图更好地平衡远距离和近距离信息。

背景与动机：
- 外推法失效：直接使用预训练的位置编码处理超出范围的位置，会导致注意力失准。
- 位置插值的不足：简单线性插值会不成比例地压缩高频信息，影响模型对局部细节的感知。
解决思路：NTK 理论分析表明，Transformer 中不同维度的 RoPE 对应着不同的旋转频率。高频部分对近距离信息敏感，低频部分对远距离信息敏感。NTK 感知的位置插值或 RoPE 缩放旨在调整 RoPE 的基频，以一种更「感知内容」的方式扩展上下文。
RoPE 中 NTK 改进公式示例：
$b' = b \times \left(s_f \times \frac{L}{L_{\text{max}}} - (s_f - 1)\right)^{\frac{d}{d-2}}$
其中 $b$ 是原始 RoPE 的旋转基数（例如 10000）， $s_f$ 是缩放因子（例如 $L_{\text{new}}/L_{\text{old}}$ ）， $L$ 是当前上下文长度， $L_{\text{max}}$ 是最大支持长度， $d$ 是维度。这个公式调整了旋转的「速度」，试图在扩展上下文的同时保留更多细节。

LongLoRA (S²-Attn)

LongLoRA 提出了一种高效的微调方法，用于在有限的计算资源下扩展预训练模型的上下文窗口，其核心是Shift Short Attention(S²-Attn)。

方法动机：
- 标准 Transformer 的全局自注意力计算复杂度为 $O(N^2)$ ，难以处理极长序列。
- 在很多情况下，大量远距离词元间的精细交互并非都重要。
S²-Attn 核心机制：
1. 注意力头分组与序列分块：将长序列划分为多个固定大小的块（例如 2048 tokens）。注意力计算主要在块内进行。
2. 分组位移：为了促进块间的信息交互，将一部分注意力头的输入（Q, K, V）在序列维度上进行位移。例如，可以将一半的注意力头向前（或向后）位移半个块的大小。
3. 块内自注意力计算：每个 attention head（经过可能的位移后）只在其所在的（可能是重叠的）块内计算注意力。
- 通过这种方式，S²-Attn 以较低的计算成本近似了全局注意力，位移机制增强了跨组（块）之间的信息交互。

YaRN (Yet another RoPE extension methods)

YaRN 是一种进一步改进 RoPE 上下文扩展能力的方法，它声称仅用少量训练数据和训练步数就能取得优于标准位置插值法和 NTK 方法的效果。

YaRN 包含两大核心设计：

NTK-by-parts 插值法：
- 动机：传统插值法假设 RoPE 的所有维度（频率）受上下文扩展的影响相同，忽略了不同频率对应不同波长 $\lambda$ 的事实。
- 核心策略：根据 RoPE 维度对应的波长 $\lambda$ $λ$ 与上下文长度 $L$ $L$ 的关系，采取不同的插值策略：
  - 波长较短（ $\lambda \ll L$ ）：对应高频信息，不进行插值，以保留局部精度。
  - 波长较长（ $\lambda \ge L$ ）：对应低频信息，进行插值，以保留全局位置感。
  - 中间波长：采用混合插值策略（渐变）。
注意力温度调节机制：
- 动机：扩展上下文后，注意力分布的熵（或困惑度）可能会发生变化。温度调节可以帮助稳定注意力分布。
- 核心策略：在计算注意力 Softmax 之前，对 logits $\frac{\bm{q}_m^\mathsf{T} \bm{k}_n}{\sqrt{|D|}}$ $\frac{q _{m}^{T} k _{n}}{∣ D ∣}$ 除以一个温度系数 $t$ $t$ （或等效地，在 Q 和 K 对应的 RoPE 上乘以缩放系数 $1/\sqrt{t}$ $1/ t$ ）。
  - $\operatorname{AttentionScores} = \operatorname{softmax}\left(\frac{\bm{q}_m^\mathsf{T} \bm{k}_n}{t\sqrt{|D|}}\right)$
- 温度系数 $t$ 通常与上下文长度的缩放比例 $s = L_{\text{new}}/L_{\text{old}}$ 相关，例如对于 Llama 模型，一个经验取值是 $t = 0.1 \ln(s) + 1$ 。
- 这种调节具有良好的推广性，可适配不同数据样本和扩展上下文长度后的词元位置。

并行训练策略

训练超大规模语言模型需要巨大的计算资源和显存。并行训练策略是将训练任务分解到多个计算设备（如 GPU）上协同完成的技术。

数据并行

数据并行是最常用的一种分布式训练策略。

核心思想：每个计算单元（GPU）都拥有模型的完整副本。训练数据集被划分成多个子集，每个 GPU 处理其中一个子集的数据。
关键步骤：
1. 数据划分：将训练数据集平均分配到多个计算单元。
2. 并行计算：各 GPU 独立地使用自己的数据子集进行前向传播和反向传播，计算梯度。
3. 梯度聚合：将所有 GPU 计算得到的梯度进行聚合（通常是求平均）。
4. 权重同步：使用聚合后的梯度更新模型参数，并将更新后的参数广播回所有 GPU，以保持模型副本的一致性。
优势：
- 显著加速训练过程。
- 易于实现，主流深度学习框架（如 PyTorch, TensorFlow）原生支持。
- 适用广泛，支持多种模型和任务类型。
挑战：
- 通信开销大：尤其在大模型中，梯度同步（AllReduce 操作）可能成为瓶颈。
- 显存压力大：每个 GPU 都需保存完整模型副本、激活值以及优化器状态，对显存要求高。
- 扩展性受限：随着设备数量增加，通信开销占比可能越来越大，导致效率提升饱和。

模型并行

模型并行是将模型的不同部分（例如，层的不同部分或不同的层）分布到多个设备上进行训练的策略。它适用于单个设备无法容纳完整模型的场景。

核心思想：将模型的参数和计算切分到不同设备上。
- 层内并行：将单个大计算层（如一个大的 MLP 层或注意力层的权重矩阵）切分到多个 GPU 上。例如，一个矩阵乘法 $\bm{Y} = \bm{X}\bm{A}$ 可以将 $\bm{A}$ 按列切分 $(\bm{A}_1, \bm{A}_2)$ ，则 $\bm{Y} = \bm{X}(\bm{A}_1, \bm{A}_2) = (\bm{X}\bm{A}_1, \bm{X}\bm{A}_2)$ 。 $\bm{X}\bm{A}_1$ 和 $\bm{X}\bm{A}_2$ 可以在不同 GPU 上计算。
- 层间并行：将模型的不同层放置在不同 GPU 上。这通常与流水线并行结合。
工作流程：
- 模型划分：将模型按层或模块（或层内张量）切分至不同设备。
- 分布式计算：
  - 前向传播：输入数据或中间激活值在设备间传递，每个设备执行其负责部分的计算。
  - 反向传播：梯度也按相反的路径在设备间传递和计算。
- 权重更新：各设备独立更新自己负责部分的参数。
- 通信协同：跨设备的数据传递（如激活值、梯度）需要高效调度，确保计算与同步一致。
优势：能够训练单个 GPU 无法容纳的超大模型。
挑战：实现复杂，通信模式多样，需要仔细设计以最小化通信开销。

流水线并行

流水线并行是将超大语言模型划分为多个计算阶段，并在多个设备上依次串行处理这些阶段的并行计算策略。其灵感源自工业流水线。

核心思想：将模型按层顺序划分为若干阶段，每个阶段分配给一个或一组 GPU。数据（以微批次 micro-batch 的形式）像在流水线中一样依次通过各个阶段。
工作机制：
1. 模型分割：将模型（例如 $L$ 层）划分为 $K$ 个阶段，每个阶段包含若干连续的层。
2. 设备配置：根据各阶段的计算量将其分配至不同 GPU。
3. 流水执行：一个批次被进一步划分为多个微批次。GPU $i$ 完成其阶段对微批次 $j$ 的处理后，立即将其输出传递给 GPU $i+1$ 处理微批次 $j$ ，同时 GPU $i$ 开始处理微批次 $j+1$ 。这样可以实现阶段间的并行处理。
4. 同步更新：通常在处理完一个完整批次的所有微批次后，聚合所有阶段的梯度，统一更新模型参数。
挑战：
- 流水线气泡：在流水线的启动和排空阶段，部分 GPU 会处于空闲状态，降低了硬件利用率。需要通过调度（如 GPipe, PipeDream）来最小化气泡。
- 梯度同步和参数更新：需要仔细管理。

混合并行

混合并行是将数据并行、模型并行（张量并行）、流水线并行等多种并行策略灵活结合的分布式训练方法，专为超大规模语言模型设计。

关键特点：
- 高性能：通过多策略协同，大幅提升训练效率与吞吐量。
- 高灵活性：可按模型结构与硬件特性自定义并行组合。
- 强扩展性：支持跨 GPU、跨节点的大规模训练任务。
- 资源优化：内存与计算资源可按阶段或维度精准分配。
常见混合结构：
1. 数据并行（DP）：在最高层级，将整个模型（可能已应用了流水线和张量并行）复制多份，每份处理不同数据。
2. 流水线并行（PP）：在每个数据并行副本内部，将模型层切分为多个阶段，分配到不同 GPU。
3. 张量并行（TP）：在每个流水线阶段内部，如果单层过大，则使用张量并行将其参数和计算切分到该阶段内的多个 GPU。
实现挑战：并行策略调度复杂；通信与同步代价高；需依赖高效框架如 DeepSpeed、Megatron-LM、Colossal-AI。

零冗余优化

ZeRO (Zero Redundancy Optimizer) 是由微软提出的一系列内存优化技术，专为超大规模模型训练设计，通过去除数据并行中的冗余副本显著降低 GPU 内存占用。

核心策略：在数据并行的基础上，对模型状态的不同部分进行分割，每个 GPU 只保存其中一部分，在需要时通过通信动态获取。
1. ZeRO-DP Stage 1 (Optimizer States Partitioning)：优化器状态（如 Adam 中的动量和方差）被分割到各个数据并行的 GPU 上。每个 GPU 只存储和更新其负责的那部分优化器状态。
2. ZeRO-DP Stage 2 (Gradient Partitioning)：在 Stage 1 的基础上，梯度也被分割。每个 GPU 只存储和更新其负责参数对应的梯度。
3. ZeRO-DP Stage 3 (Parameter Partitioning)：在 Stage 2 的基础上，模型参数本身也被分割。每个 GPU 只存储其负责的那部分参数。在前向和后向传播过程中，所需的参数会通过 AllGather 操作动态地从其他 GPU 收集。
- ZeRO-Offload: 进一步将优化器状态、梯度甚至参数卸载到 CPU 内存或 NVMe 存储，以支持更大模型。
关键优势：
- 显著节省显存：使得在有限硬件上训练更大的模型成为可能。
- 提升计算效率（尤其 Stage 1 & 2）：优化器状态与计算负担更均衡。
- 良好的可扩展性：支持跨多个 GPU/节点高效并行。
应用场景：GPT-3、OPT、LLaMA 等超大语言模型的训练；DeepSpeed、Colossal-AI 等训练框架的核心优化组件。

DeepSpeed

DeepSpeed 是由微软研究院开发的高性能分布式训练库，专为超大规模语言模型的训练和部署设计，支持数十亿乃至上千亿参数的模型。

核心能力：
- 全面的并行训练支持：内置数据并行、张量模型并行（层内）、流水线模型并行（层间），适配多种规模和拓扑结构。
- 高效内存优化：深度集成了 ZeRO 系列优化策略（ZeRO-1/2/3/Infinity/Offload），显著降低显存需求。
- 高性能通信：通过混合精度训练、压缩通信（如 1-bit Adam）、优化的集合通信操作等手段，减少跨设备数据传输开销。
- 强兼容性：无缝集成 PyTorch 与 Hugging Face Transformers，可快速接入现有模型。
- 出色可扩展性：支持千卡级别的分布式训练，具备自动调优与弹性调度能力。

下表总结了 ZeRO 不同阶段的特点：

项目	ZeRO-1 (Optimizer States)	ZeRO-2 (Gradients)	ZeRO-3 (Parameters)	ZeRO-Infinity (Offload)
梯度分割	-	✓	✓	✓
优化器状态分割	✓	✓	✓	✓
模型参数分割	-	-	✓	✓
内存优化效率	低	中	高	极高（利用 NVMe/CPU）
通信优化（开销）	低	中	高	高

DeepSpeed 通过这些技术的组合，极大地推动了超大模型训练的可行性和效率。