现代大语言模型架构与从输入到输出的完整流水线设计

摘要

现代大语言模型(Large Language Model, LLM)通常以 Transformer 为核心架构,通过大规模自监督预训练学习自然语言、代码、数学表达式和结构化文本中的统计规律,再经过指令微调、偏好对齐和系统级推理优化,形成可交互的文本生成系统。本文从经典论文出发,系统说明一个现代大模型从原始输入到最终输出的完整流水线,包括文本归一化、分词、嵌入、位置编码、注意力、前馈网络、归一化、残差连接、输出投影、概率解码、预训练目标、后训练对齐、参数高效微调、分布式训练、推理加速和服务化内存管理。本文重点不在经验性配置罗列,而在解释这些模块为什么被设计出来、它们在数学上如何工作、训练中优化什么目标、工程上如何降低计算和显存成本。

关键词:大语言模型;Transformer;自注意力;自回归语言建模;分词;RoPE;RMSNorm;SwiGLU;AdamW;RLHF;DPO;KV Cache;FlashAttention;PagedAttention。

1. 研究背景与经典论文脉络

1.1 从序列建模到注意力机制

早期神经序列模型常使用循环神经网络(Recurrent Neural Network, RNN)或长短期记忆网络(Long Short-Term Memory, LSTM)处理序列。Sutskever 等人的 Seq2Seq 模型将输入序列压缩为一个固定维度向量,再由解码器生成输出序列。这类方法奠定了“编码器-解码器”范式,但固定长度向量容易成为信息瓶颈。Bahdanau 等人在神经机器翻译中提出可学习的注意力机制,使解码器在每个时间步可以根据当前状态动态关注源序列不同位置,从而缓解固定向量瓶颈。

Transformer 的关键突破来自 Vaswani 等人的《Attention Is All You Need》。该论文提出完全基于注意力机制的架构,去除循环和卷积结构,使序列中所有位置可以并行计算,同时通过多头注意力捕获不同表示子空间中的依赖关系。现代 LLM 的主体架构几乎都可以看作 Transformer 的变体,尤其是 GPT 系列使用的 decoder-only 自回归 Transformer。

1.2 从任务模型到预训练基础模型

BERT 采用 encoder-only Transformer,通过 masked language modeling 学习双向表示,展示了大规模预训练加少量任务微调的有效性。GPT-3 采用 decoder-only 自回归语言模型,并通过大规模参数、数据和计算量证明了 few-shot 与 in-context learning 能力。LLaMA 进一步展示了使用公开数据、合适的数据规模、RMSNorm、SwiGLU 和 RoPE 等结构选择,也可以训练出高效基础模型。

1.3 从模型规模到计算最优

Kaplan 等人的 scaling laws 研究表明,语言模型损失与模型参数量、数据量和训练计算量之间存在近似幂律关系。Hoffmann 等人的 Chinchilla 论文指出,在固定训练计算预算下,模型参数量和训练 token 数应更平衡地增长,许多早期大模型相对参数量而言训练 token 不足。这个结论改变了大模型训练设计:不仅要增大参数,也要重视高质量数据规模和训练步数。

1.4 从可生成到可遵循人类意图

基础语言模型优化的是“下一个 token 预测”,这并不等同于“按照用户意图回答”。InstructGPT 通过监督微调、奖励模型和强化学习从人类反馈中学习,使模型输出更符合人类偏好。DPO 进一步将偏好优化写成直接优化语言模型的目标,避免单独训练奖励模型和执行复杂的强化学习采样流程。

2. 符号、张量形状与基本问题定义

令一段输入文本为字符串 ss。分词器将其映射为 token 序列:

x=(x1,x2,,xT),xt{1,2,,V}x = (x_1, x_2, \ldots, x_T), \quad x_t \in \{1,2,\ldots,V\}

其中 TT 是序列长度,$V$ 是词表大小。模型的目标是在给定上下文 x<tx_{<t} 时预测下一个 token 的条件概率:

pθ(xtx<t)p_\theta(x_t \mid x_{<t})

完整序列的自回归概率分解为:

pθ(x1,,xT)=t=1Tpθ(xtx<t)p_\theta(x_1,\ldots,x_T)=\prod_{t=1}^{T}p_\theta(x_t\mid x_{<t})

训练时通常最大化训练文本的对数似然,等价于最小化负对数似然或交叉熵损失:

L(θ)=t=1Tlogpθ(xtx<t)\mathcal{L}(\theta) = - \sum_{t=1}^{T} \log p_\theta(x_t \mid x_{<t})

在 batch 训练中,若 batch size 为 BB,输入 token 张量形状通常为:

XNB×TX \in \mathbb{N}^{B \times T}

嵌入后得到:

H(0)RB×T×dmodelH^{(0)} \in \mathbb{R}^{B \times T \times d_{\text{model}}}

其中 dmodeld_{\text{model}} 是模型隐藏维度。经过 LL 层 Transformer block 后得到:

H(L)RB×T×dmodelH^{(L)} \in \mathbb{R}^{B \times T \times d_{\text{model}}}

最后通过输出投影映射到词表维度:

Z=H(L)WoutRB×T×VZ = H^{(L)} W_{\text{out}}^\top \in \mathbb{R}^{B \times T \times V}

ZZ 在词表维度做 softmax 得到每个位置的下一个 token 概率分布。

3. 从输入到输出的完整流水线

现代大语言模型的推理流水线可以概括为:

  1. 输入接收:用户输入原始文本,系统同时接收系统提示词、开发者提示词、历史对话、工具返回结果等上下文。
  2. 文本规范化:将输入转成分词器可处理的字节或 Unicode 序列,处理特殊标记、换行、空白和对话模板。
  3. 分词:使用 BPE、WordPiece、Unigram 或 SentencePiece 将字符串变为 token id 序列。
  4. 嵌入:通过可学习嵌入矩阵把 token id 映射为连续向量。
  5. 位置处理:使用绝对位置编码、相对位置编码、RoPE 或 ALiBi 等方法注入顺序信息。
  6. Transformer 堆叠:每层执行归一化、自注意力、残差连接、前馈网络、残差连接。
  7. 输出投影:将最终隐藏状态投影到词表维度,得到 logits。
  8. 概率归一化:通过 softmax 得到词表概率分布。
  9. 解码策略:根据 greedy、temperature sampling、top-k、top-p、beam search 或约束解码选择下一个 token。
  10. KV Cache 更新:保存已生成 token 的 key 和 value,避免每一步重复计算完整历史。
  11. 迭代生成:将新 token 追加到上下文,重复前向计算直到遇到停止符、长度上限或外部中止条件。
  12. 反分词:将 token 序列还原为文本,处理特殊 token 和空白。
  13. 安全与格式后处理:根据应用需求执行内容过滤、结构化校验、工具调用解析或引用插入。

训练流水线与推理不同:训练时模型一次处理完整序列,通过 teacher forcing 同时预测每个位置的下一个 token;推理时模型逐 token 递增生成。

4. 分词与输入表示

4.1 为什么不能直接按词建模

自然语言存在开放词表问题。新词、专有名词、拼写错误、代码标识符、数字、URL 和多语言混排会导致按词分割的词表无限增长。若词表太小,未登录词数量高;若词表太大,嵌入矩阵和输出投影矩阵参数量过大,并且低频词训练不足。

子词分词解决的是“固定词表与开放文本之间的矛盾”。它允许高频词作为整体 token,也允许低频词被拆成多个子词 token。Sennrich 等人的 BPE 子词方法在神经机器翻译中证明了这种设计能改善罕见词处理。SentencePiece 则强调从原始文本直接训练分词器,不依赖预先按空格切词,更适合中文、日文、无空格语言和多语言场景。

4.2 BPE 的基本过程

BPE 的训练思想是从较小单位开始,例如字节或字符,然后反复合并语料中最常出现的相邻符号对。设当前符号序列集合为 D\mathcal{D},每一步统计所有相邻 pair (a,b)(a,b) 的频率:

freq(a,b)=sD#((a,b),s)\text{freq}(a,b)=\sum_{s\in\mathcal{D}}\#((a,b),s)

选择最高频 pair:

(a,b)=argmax(a,b)freq(a,b)(a^*,b^*)=\arg\max_{(a,b)}\text{freq}(a,b)

并引入新符号 c=abc=a^*b^*。重复合并直到达到目标词表大小。推理时,分词器按照已学习的合并规则将字符串切成 token。

BPE 的设计理念是压缩与统计建模的结合:高频片段被压缩为短 token 序列,低频片段保留为可组合的子词或字节。它不要求每个 token 都对应语言学上的词素,因此不能把 token 直接解释为自然语言语义单位。

4.3 Token embedding

分词后的 token id 是离散整数,神经网络不能直接从整数大小中获得语义,因为 token id 的数值大小只是索引。模型使用可学习嵌入矩阵:

ERV×dmodelE \in \mathbb{R}^{V \times d_{\text{model}}}

将 token xtx_t 映射为:

ht(0)=E[xt]h_t^{(0)} = E[x_t]

对整个序列:

H(0)=Embedding(X)H^{(0)} = \text{Embedding}(X)

嵌入向量在训练中通过反向传播更新。若两个 token 在相似上下文中经常出现,它们的嵌入向量可能逐渐形成相似的表示,但这种相似性不是手工指定的,而是由语言建模目标和网络结构共同学习出来的。

4.4 输入模板与特殊 token

现代对话模型通常不只接收用户文本,还接收角色信息和边界信息。例如:

<system>你是一个严谨的助手。</system>
<user>解释 Transformer。</user>
<assistant>

这些标记会被分词为特殊 token 或普通 token。它们的作用是让模型区分系统指令、用户消息、助手回复、工具调用和停止位置。对话模板不是模型架构本身,但它显著影响模型行为,因为训练和推理时的格式一致性决定了模型能否正确识别角色与生成边界。

5. 位置编码:让模型知道顺序

5.1 自注意力本身不含顺序

注意力计算主要依赖 token 表示之间的相似度。如果没有位置编码,输入序列中 token 的排列顺序无法被可靠区分。对于集合 {h1,h2,,hT}\{h_1,h_2,\ldots,h_T\},自注意力可以计算任意 token 间的关系,但它并不知道某个 token 是第 1 个还是第 10 个。

因此,Transformer 必须引入位置信息。原始 Transformer 使用正弦余弦绝对位置编码: $$ PE_{(pos,2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) $$

PE(pos,2i+1)=cos(pos100002i/dmodel)PE_{(pos,2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)

然后将位置编码加到 token embedding 上:

ht(0)=E[xt]+PEth_t^{(0)} = E[x_t] + PE_t

这种方法简单、可外推到未见过的位置,但现代 decoder-only LLM 更常使用 RoPE。

5.2 RoPE 的数学形式

RoPE(Rotary Position Embedding)由 RoFormer 提出。它不把位置向量简单相加,而是对 query 和 key 的二维子空间进行旋转。对二维向量 (q2i,q2i+1)(q_{2i}, q_{2i+1}),位置 mm 对应旋转角 θim\theta_i m,旋转矩阵为:

R(mθi)=[cos(mθi)sin(mθi)sin(mθi)cos(mθi)]R(m\theta_i)= \begin{bmatrix} \cos(m\theta_i) & -\sin(m\theta_i)\\ \sin(m\theta_i) & \cos(m\theta_i) \end{bmatrix}

旋转后的 query 和 key 为:

q~m=R(m)qm,k~n=R(n)kn\tilde{q}_m = R(m)q_m,\quad \tilde{k}_n = R(n)k_n

注意力打分中出现:

q~mk~n=qmR(m)R(n)kn=qmR(nm)kn\tilde{q}_m^\top \tilde{k}_n = q_m^\top R(m)^\top R(n)k_n = q_m^\top R(n-m)k_n

这说明注意力分数自然依赖相对位置 nmn-m。RoPE 的设计价值在于同时保留绝对位置注入和相对位置差建模,且对 Transformer 的主体结构改动较小。

5.3 长上下文中的位置问题

模型训练时通常有最大上下文长度 TtrainT_{\text{train}}。推理时若输入长度超过训练长度,位置编码外推可能退化。常见优化包括 RoPE scaling、插值、分段注意力、滑动窗口注意力和长上下文继续训练。需要注意:增大上下文长度不是只改一个配置项,还会增加注意力计算量、KV Cache 显存、训练样本构造复杂度和长距离监督需求。

6. Transformer Block 的核心结构

现代 decoder-only LLM 通常由多个相同或相似的 Transformer block 堆叠。一个常见 pre-norm block 可写为:

U()=H()+Attention(Norm(H()))U^{(\ell)} = H^{(\ell)} + \text{Attention}(\text{Norm}(H^{(\ell)})) H(+1)=U()+FFN(Norm(U()))H^{(\ell+1)} = U^{(\ell)} + \text{FFN}(\text{Norm}(U^{(\ell)}))

其中 \ell 表示层编号。核心组件包括归一化、自注意力、前馈网络和残差连接。

6.1 残差连接

残差连接将子层输出加回原输入:

y=x+f(x)y = x + f(x)

它的作用是改善深层网络优化。若某一层暂时学不到有效变换,网络可以让 f(x)f(x) 近似为 0,使信息沿残差路径传递。残差连接降低了深层网络必须每层都重新编码全部信息的压力,也有助于梯度从高层传回低层。

6.2 LayerNorm 与 RMSNorm

LayerNorm 对每个 token 的隐藏向量独立归一化。给定 xRdx\in\mathbb{R}^d

μ=1di=1dxi\mu = \frac{1}{d}\sum_{i=1}^{d}x_i σ2=1di=1d(xiμ)2\sigma^2 = \frac{1}{d}\sum_{i=1}^{d}(x_i-\mu)^2 LayerNorm(x)i=γixiμσ2+ϵ+βi\text{LayerNorm}(x)_i=\gamma_i\frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta_i

RMSNorm 去掉均值中心化,只按均方根缩放:

RMS(x)=1di=1dxi2+ϵ\text{RMS}(x)=\sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2+\epsilon} RMSNorm(x)i=γixiRMS(x)\text{RMSNorm}(x)_i=\gamma_i\frac{x_i}{\text{RMS}(x)}

RMSNorm 的研究动机是:LayerNorm 的重新缩放性质比重新中心化更关键。去掉均值计算可以降低计算量,并在许多大模型中保持稳定训练。LLaMA 等模型使用 RMSNorm,说明它已成为现代 decoder-only LLM 的常见选择。

6.3 Pre-norm 与 post-norm

原始 Transformer 使用 post-norm:

H(+1)=Norm(H()+f(H()))H^{(\ell+1)} = \text{Norm}(H^{(\ell)} + f(H^{(\ell)}))

许多现代大模型使用 pre-norm:

H(+1)=H()+f(Norm(H()))H^{(\ell+1)} = H^{(\ell)} + f(\text{Norm}(H^{(\ell)}))

pre-norm 的优势通常体现在深层训练稳定性上。因为残差路径不经过归一化直接连接,梯度传播更顺畅。缺点是最终层输出可能需要额外 final norm 来稳定 logits。

7. 自注意力机制

7.1 Query、Key、Value

给定隐藏状态:

HRB×T×dmodelH \in \mathbb{R}^{B\times T\times d_{\text{model}}}

通过线性映射得到:

Q=HWQ,K=HWK,V=HWVQ = HW_Q,\quad K=HW_K,\quad V=HW_V

其中:

WQ,WK,WVRdmodel×dheadW_Q,W_K,W_V \in \mathbb{R}^{d_{\text{model}}\times d_{\text{head}}}

对单个 attention head,缩放点积注意力为:

Attention(Q,K,V)=softmax(QKdhead+M)V\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^\top}{\sqrt{d_{\text{head}}}} + M\right)V

其中 MM 是 mask。对于自回归 decoder,mask 必须阻止当前位置看到未来 token:

Mij={0,ji,j>iM_{ij}= \begin{cases} 0, & j \le i\\ -\infty, & j > i \end{cases}

这样第 ii 个位置只能使用 11ii 的上下文,保证训练目标与推理时的因果生成一致。

7.2 为什么要除以 dhead\sqrt{d_{\text{head}}}

若 query 和 key 的各维近似独立,均值为 0,方差为 1,则点积:

qk=r=1dheadqrkrq^\top k = \sum_{r=1}^{d_{\text{head}}}q_rk_r

方差约为 dheadd_{\text{head}}。当 dheadd_{\text{head}} 较大时,点积值幅度变大,softmax 容易进入饱和区,使梯度变小。除以 dhead\sqrt{d_{\text{head}}} 后,打分方差回到近似常数范围,有助于稳定训练。

7.3 Softmax 的概率解释

对位置 ii,注意力权重为:

αij=exp(sij)r=1Texp(sir)\alpha_{ij}= \frac{\exp(s_{ij})}{\sum_{r=1}^{T}\exp(s_{ir})}

其中:

sij=qikjdhead+Mijs_{ij}=\frac{q_i^\top k_j}{\sqrt{d_{\text{head}}}}+M_{ij}

输出为 value 的加权和:

oi=j=1Tαijvjo_i=\sum_{j=1}^{T}\alpha_{ij}v_j

注意力权重不是人工指定的语法依存关系,而是模型在训练目标下学到的动态加权系数。它可以反映某些依赖关系,但不能简单等同于可解释的因果证据。

7.4 多头注意力

多头注意力将隐藏维度拆成多个 head:

headr=Attention(HWQ(r),HWK(r),HWV(r))\text{head}_r = \text{Attention}(HW_Q^{(r)},HW_K^{(r)},HW_V^{(r)}) MHA(H)=Concat(head1,,headh)WO\text{MHA}(H)=\text{Concat}(\text{head}_1,\ldots,\text{head}_h)W_O

多头设计允许不同 head 学习不同类型的关系,例如局部依赖、长程依赖、格式边界、复制行为或特定语法模式。严格地说,每个 head 的功能不是预先规定的,而是由训练数据和目标函数诱导出来的。

7.5 Multi-Query Attention 与 Grouped-Query Attention

标准多头注意力为每个 head 都维护独立的 Q,K,VQ,K,V。推理时,每生成一个 token 都需要缓存所有层的 K,VK,V。KV Cache 显存大约与层数、序列长度、batch size、head 数和 head 维度成正比。

Multi-Query Attention(MQA)让多个 query head 共享一组 key/value head。Grouped-Query Attention(GQA)则让一组 query head 共享 key/value,比 MQA 表达能力更强,比标准 MHA 更省显存。现代大模型常用 GQA 来降低长上下文推理成本。

8. 前馈网络 FFN

8.1 原始 FFN

Transformer block 中的 FFN 对每个位置独立应用同一个两层 MLP:

FFN(x)=W2ϕ(W1x+b1)+b2\text{FFN}(x)=W_2\phi(W_1x+b_1)+b_2

其中 ϕ\phi 可为 ReLU、GELU 或 SiLU。注意力层负责 token 间信息混合,FFN 负责每个 token 位置上的非线性特征变换和通道混合。FFN 通常占 Transformer 参数量和计算量的大部分。

8.2 GELU、SiLU 与门控 FFN

GELU 形式为:

GELU(x)=xΦ(x)\text{GELU}(x)=x\Phi(x)

其中 Φ(x)\Phi(x) 是标准正态分布的累积分布函数。SiLU 形式为: $$ \text{SiLU}(x)=x\sigma(x) $$

其中 σ(x)\sigma(x) 是 sigmoid。

GLU 及其变体在 Transformer FFN 中引入门控机制。SwiGLU 可写为:

SwiGLU(x)=(SiLU(xW1))(xW3)W2\text{SwiGLU}(x)=(\text{SiLU}(xW_1))\odot(xW_3)W_2

其中 \odot 表示逐元素乘法。门控结构让网络能够对不同通道进行输入相关的选择。现代 LLM 常用 SwiGLU,因为它在经验上优于普通 ReLU/GELU FFN,LLaMA 也采用了这一设计。

8.3 FFN 维度设计

经典 Transformer 常使用:

dff=4dmodeld_{\text{ff}} = 4d_{\text{model}}

使用 SwiGLU 时,为了保持参数量和计算量近似可比,常将中间维度设为约:

dff83dmodeld_{\text{ff}} \approx \frac{8}{3}d_{\text{model}}

这是因为 SwiGLU 有两条输入投影分支 W1W_1W3W_3,若仍使用 4dmodel4d_{\text{model}},参数量会显著增加。

9. 输出层与语言建模目标

9.1 Logits 与 softmax

经过最后一层 Transformer 后,隐藏状态 hth_t 被映射到词表维度:

zt=Woutht+bz_t = W_{\text{out}}h_t + b

其中:

ztRVz_t \in \mathbb{R}^{V}

softmax 给出下一个 token 的概率:

pθ(xt+1=vxt)=exp(zt,v)u=1Vexp(zt,u)p_\theta(x_{t+1}=v\mid x_{\le t})= \frac{\exp(z_{t,v})}{\sum_{u=1}^{V}\exp(z_{t,u})}

训练时真实下一个 token 为 yty_t,交叉熵损失为:

Lt=logpθ(ytxt)\mathcal{L}_t=-\log p_\theta(y_t\mid x_{\le t})

全序列平均损失为:

L=1BTb=1Bt=1TLb,t\mathcal{L}=\frac{1}{BT}\sum_{b=1}^{B}\sum_{t=1}^{T}\mathcal{L}_{b,t}

9.2 Weight tying

许多语言模型共享输入嵌入矩阵和输出投影矩阵:

Wout=EW_{\text{out}} = E

严格来说,形状上通常是输出 logits 使用 htEh_tE^\top。这种共享减少参数量,并使输入 token 表示与输出分类器处在同一词表表示空间中。

9.3 困惑度

语言模型常用困惑度(perplexity, PPL)评估预测质量。若平均负对数似然为 L\mathcal{L},则:

PPL=exp(L)\text{PPL}=\exp(\mathcal{L})

困惑度越低表示模型给真实文本分配的概率越高。但困惑度不是对话质量、真实性、安全性或推理能力的完整指标。一个模型可以有较低困惑度,但仍可能不遵循指令或生成错误事实。

10. 预训练:从大规模文本中学习

10.1 自监督学习

大模型预训练通常不需要人工逐条标注。文本本身提供监督信号:给定前文预测下一个 token。训练样本来自连续文本片段:

输入: x_1, x_2, ..., x_{T-1}
标签: x_2, x_3, ..., x_T

模型在每个位置同时预测下一个 token。这种方法称为 teacher forcing:训练时模型看到真实历史 token,而不是自己先前生成的 token。

10.2 数据混合与清洗

预训练语料可能包括网页、书籍、论文、代码、问答、百科、论坛和数学文本。数据质量直接影响模型能力。常见处理包括:

  1. 去重:减少重复网页和模板文本,避免模型过度记忆。
  2. 语言识别:控制不同语言比例。
  3. 质量过滤:移除乱码、低质量 SEO 页面、过短文本和异常 token 分布文本。
  4. 安全过滤:移除某些明显有害或隐私敏感内容。
  5. 域配比:平衡自然语言、代码、数学和对话数据。
  6. 文档打包:将多个文档拼接成固定长度训练序列,并使用边界 token 标识分隔。

数据处理不是附属工作。根据 scaling laws 和 Chinchilla 的结论,训练 token 数和质量与模型参数量同等重要。对于固定计算预算,过大模型配不足数据可能不是最优选择。

10.3 优化器:Adam 与 AdamW

Adam 使用一阶矩和二阶矩的指数滑动平均来调整每个参数的学习率。给定梯度 gt=θLtg_t=\nabla_\theta \mathcal{L}_t

mt=β1mt1+(1β1)gtm_t=\beta_1m_{t-1}+(1-\beta_1)g_t vt=β2vt1+(1β2)gt2v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2

偏差校正:

m^t=mt1β1t,v^t=vt1β2t\hat{m}_t=\frac{m_t}{1-\beta_1^t},\quad \hat{v}_t=\frac{v_t}{1-\beta_2^t}

参数更新:

θt=θt1ηm^tv^t+ϵ\theta_t=\theta_{t-1}-\eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}

AdamW 将权重衰减从梯度更新中解耦:

θt=θt1ηm^tv^t+ϵηλθt1\theta_t=\theta_{t-1} -\eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} -\eta\lambda\theta_{t-1}

这种做法避免 Adam 中 L2 正则与权重衰减不等价的问题,已成为训练 Transformer 的常见选择。

10.4 学习率调度

大模型训练常使用 warmup 加衰减。warmup 阶段学习率从较小值线性增加到峰值,避免训练初期参数尚未稳定时更新过大。之后使用 cosine decay 或线性衰减:

ηt=ηmaxf(t)\eta_t = \eta_{\max}\cdot f(t)

学习率调度的核心目标是控制优化过程的稳定性和最终收敛质量。过大学习率可能导致 loss spike 或发散;过小学习率会降低训练效率并可能欠拟合。

10.5 混合精度训练

大模型训练通常使用 FP16、BF16 或 FP8 等低精度格式以减少显存和提高吞吐。混合精度训练的基本思想是:矩阵乘法使用低精度以提升速度,关键状态如优化器统计量可能保留较高精度以保证稳定。

BF16 比 FP16 有更大的指数范围,因此在大模型训练中常更稳定。混合精度不是单纯减少位数,还需要 loss scaling、溢出检测、梯度裁剪和数值稳定 kernel。

11. 分布式训练与系统优化

11.1 数据并行

数据并行将同一模型复制到多个 GPU,每个 GPU 处理不同 batch shard。反向传播后对梯度做 all-reduce:

g=1Ni=1Ngig = \frac{1}{N}\sum_{i=1}^{N}g_i

然后每个 GPU 使用相同梯度更新参数。数据并行简单高效,但每个 GPU 都要保存完整模型参数、梯度和优化器状态。当模型过大时,单卡显存无法容纳完整状态。

11.2 张量并行

Megatron-LM 展示了将 Transformer 内部矩阵乘法切分到多个 GPU 的方法。例如线性层:

Y=XWY = XW

可以按列切分 W=[W1,W2]W=[W_1,W_2],得到:

Y=[XW1,XW2]Y=[XW_1,XW_2]

也可以按行切分并在结果上做 reduce。张量并行降低单卡参数和激活压力,但引入层内通信。适合大矩阵乘法密集的 Transformer。

11.3 流水线并行

流水线并行将不同层放在不同 GPU 上。第 1 组 GPU 计算前几层,第 2 组 GPU 计算中间层,后续 GPU 计算更高层。为了减少空泡,训练 batch 被拆成 micro-batch,使不同设备同时处理不同 micro-batch 的不同阶段。

流水线并行的挑战是调度复杂、激活传输成本高,并且容易产生 pipeline bubble。它通常与数据并行和张量并行组合使用。

11.4 ZeRO 优化

ZeRO 的核心思想是消除数据并行中的冗余状态。传统数据并行中,每个 GPU 都保存完整参数、梯度和优化器状态。ZeRO 将这些状态分片:

  1. ZeRO-1:切分优化器状态。
  2. ZeRO-2:切分优化器状态和梯度。
  3. ZeRO-3:切分优化器状态、梯度和参数。

这样单卡显存占用随数据并行规模下降,使训练更大模型成为可能。代价是需要更多通信与参数 gather/scatter 操作。

11.5 FlashAttention

标准注意力显式构造 T×TT\times T 注意力矩阵,显存和时间复杂度随序列长度平方增长。FlashAttention 的关键不是近似注意力,而是重新组织精确注意力计算,使其更符合 GPU 存储层次。它通过 tiling 将 Q,K,VQ,K,V 分块加载到片上 SRAM,避免反复读写高带宽显存 HBM 中的大型注意力矩阵。

标准注意力瓶颈常不只是 FLOPs,而是内存 IO。FlashAttention 的研究方法强调 IO-aware algorithm design:算法设计必须考虑数据在 HBM 和 SRAM 间移动的成本。它保持数学结果与标准 attention 一致,但显著降低显存占用和实际运行时间。

12. 后训练:让模型可用、可控、可对话

12.1 监督指令微调

预训练模型学习的是文本分布,不一定知道“用户问,助手答”的交互格式。监督指令微调(Supervised Fine-Tuning, SFT)使用人工或模型生成的指令-回答数据:

(pi,yi)(p_i, y_i)

其中 pip_i 是提示词,$y_i$ 是期望回答。训练目标仍是交叉熵:

LSFT=tlogpθ(yi,tpi,yi,<t)\mathcal{L}_{\text{SFT}}=-\sum_t \log p_\theta(y_{i,t}\mid p_i,y_{i,<t})

SFT 的作用是将基础模型的语言建模能力转化为遵循指令、保持格式、进行多轮对话和执行任务的能力。

12.2 RLHF

InstructGPT 的 RLHF 流程包括三步:

  1. 监督微调:用高质量示范回答训练初始策略模型。
  2. 奖励模型训练:对同一提示下多个回答收集人类偏好排序,训练奖励模型 rϕ(x,y)r_\phi(x,y)
  3. 强化学习优化:使用 PPO 等方法优化语言模型,使其生成高奖励回答,同时用 KL 惩罚限制其偏离 SFT 模型过远。

常见目标可写为:

maxθEyπθ(x)[rϕ(x,y)βDKL(πθ(x)πref(x))]\max_\theta \mathbb{E}_{y\sim \pi_\theta(\cdot|x)} \left[ r_\phi(x,y) -\beta D_{\text{KL}}(\pi_\theta(\cdot|x)\|\pi_{\text{ref}}(\cdot|x)) \right]

这里 πref\pi_{\text{ref}} 通常是 SFT 模型。KL 项的作用是约束模型不要为了奖励模型漏洞而偏离原始语言能力太远。

12.3 DPO

DPO 使用成对偏好数据 (x,yw,yl)(x,y_w,y_l),其中 ywy_w 是偏好回答,$y_l$ 是非偏好回答。DPO 将偏好优化转化为直接优化语言模型相对参考模型的概率差:

LDPO=logσ(β[logπθ(ywx)πref(ywx)logπθ(ylx)πref(ylx)])\mathcal{L}_{\text{DPO}} =-\log\sigma\left( \beta \left[ \log\frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)}- \log\frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)} \right] \right)

DPO 的优势是训练流程更简单,不需要显式奖励模型和在线强化学习采样。但它仍依赖高质量偏好数据,并且参考模型选择、$\beta$ 系数和数据分布会显著影响结果。

12.4 参数高效微调:LoRA 与 QLoRA

全参数微调需要更新所有参数,显存和存储成本高。LoRA 假设微调所需的参数更新 ΔW\Delta W 具有低秩结构:

W=W+ΔWW' = W + \Delta W ΔW=BA\Delta W = BA

其中:

BRdout×r,ARr×din,rmin(din,dout)B\in\mathbb{R}^{d_{\text{out}}\times r},\quad A\in\mathbb{R}^{r\times d_{\text{in}}},\quad r\ll \min(d_{\text{in}},d_{\text{out}})

训练时冻结原始 WW,只训练 A,BA,B。这样参数量从 doutdind_{\text{out}}d_{\text{in}} 降到 r(dout+din)r(d_{\text{out}}+d_{\text{in}})

QLoRA 进一步将基础模型权重量化到 4-bit,并通过低秩适配器反向传播梯度,在较小显存中微调大模型。它的关键点包括 4-bit NormalFloat、double quantization 和 paged optimizer。QLoRA 的目标不是让低精度权重直接大幅更新,而是在冻结量化模型上训练小规模适配参数。

13. 推理与解码

13.1 自回归生成

推理时,模型从上下文 xtx_{\le t} 得到下一个 token 分布: $$ p_\theta(x_{t+1}\mid x_{\le t}) $$

然后选择一个 token x^t+1\hat{x}_{t+1},追加到上下文:

xt+1=(xt,x^t+1)x_{\le t+1}=(x_{\le t},\hat{x}_{t+1})

重复该过程直到停止。与训练相比,推理存在误差累积:模型生成的 token 会成为后续上下文的一部分。如果早期生成偏离目标,后续输出可能继续沿错误方向发展。

13.2 Greedy decoding

贪心解码选择概率最高的 token:

x^t+1=argmaxvpθ(vxt)\hat{x}_{t+1}=\arg\max_v p_\theta(v\mid x_{\le t})

它确定性强、速度快,但容易产生重复、模板化或局部最优输出。

13.3 Temperature sampling

temperature 对 logits 进行缩放:

p(v)=exp(zv/τ)uexp(zu/τ)p(v)=\frac{\exp(z_v/\tau)}{\sum_u\exp(z_u/\tau)}

τ<1\tau<1 时,分布更尖锐,输出更确定;当 τ>1\tau>1 时,分布更平坦,输出更多样。temperature 不改变模型知识,只改变采样随机性。

13.4 Top-k 与 top-p

Top-k 只在概率最高的 kk 个 token 中采样。Top-p,又称 nucleus sampling,选择最小集合 SS,使:

vSp(v)pthreshold\sum_{v\in S}p(v)\ge p_{\text{threshold}}

然后在 SS 内重新归一化采样。Top-p 会根据分布形状动态调整候选集大小,通常比固定 top-k 更灵活。

13.5 重复惩罚与停止条件

生成系统通常会加入 repetition penalty、frequency penalty 或 presence penalty,降低已出现 token 的 logits。停止条件包括 EOS token、最大 token 数、停止字符串、工具调用结束标记或 JSON 结构闭合。需要注意,过强的惩罚可能损害代码、数学和引用场景,因为这些任务本身需要重复符号或变量名。

14. KV Cache 与高吞吐服务

14.1 为什么需要 KV Cache

在自回归推理中,第 t+1t+1 步需要对历史 token 做注意力。如果每一步都重新计算所有历史 token 的 K,VK,V,成本很高。KV Cache 保存每层历史 token 的 key 和 value:

Kt(),Vt()K_{\le t}^{(\ell)}, V_{\le t}^{(\ell)}

生成新 token 时,只需要计算新 token 的 query、key、value,并将新 K,VK,V 追加到 cache。注意力计算变为:

qt+1Kt+1q_{t+1}K_{\le t+1}^\top

这显著降低重复计算,但显存消耗随序列长度线性增长。

14.2 KV Cache 显存估算

若层数为 LL,batch 为 BB,序列长度为 TT,KV head 数为 hkvh_{kv},head 维度为 dhd_h,每个元素字节数为 ss,则 KV Cache 近似显存为:

MemoryKV2LBThkvdhs\text{Memory}_{KV} \approx 2 \cdot L \cdot B \cdot T \cdot h_{kv}\cdot d_h \cdot s

系数 2 来自 key 和 value。GQA/MQA 通过减少 hkvh_{kv} 降低这部分显存。

14.3 PagedAttention

vLLM 的 PagedAttention 关注服务阶段的 KV Cache 管理。多个请求长度不同,生成过程动态增长,若为每个请求分配连续大块显存,会产生碎片和浪费。PagedAttention 将 KV Cache 分成固定大小块,用类似页表的结构映射逻辑位置到物理块。这样可以:

  1. 减少显存碎片。
  2. 支持更灵活的 batch 调度。
  3. 允许多个请求共享公共前缀 cache。
  4. 提高同等延迟下的吞吐。

这说明推理性能不只由模型架构决定,还由内存管理和调度系统决定。

15. 现代 LLM 的典型模块组合

一个现代 decoder-only LLM 的常见配置如下:

  1. 分词器:BPE、SentencePiece 或 byte-level BPE。
  2. 输入嵌入:token embedding,常与输出层权重共享。
  3. 位置编码:RoPE 或其长上下文变体。
  4. 主体结构:多层 pre-norm decoder-only Transformer。
  5. 归一化:RMSNorm。
  6. 注意力:causal self-attention,常配 GQA 和 FlashAttention kernel。
  7. FFN:SwiGLU 或 GeGLU。
  8. 残差:attention residual 与 FFN residual。
  9. 输出层:final norm 加词表投影。
  10. 训练目标:自回归 next-token prediction。
  11. 后训练:SFT、RLHF、DPO 或其他偏好优化方法。
  12. 推理:KV Cache、continuous batching、PagedAttention、采样控制、停止条件。

可以将一次前向传播写成简化伪代码:

tokens = tokenizer.encode(prompt)
H = embedding[tokens]

for layer in layers:
    A = attention(norm1(H), causal_mask, position=RoPE)
    H = H + A
    F = feed_forward(norm2(H))
    H = H + F

H = final_norm(H)
logits = H @ embedding.T
next_token = decode(logits[-1])

训练时对所有位置的 logits 计算交叉熵;推理时通常只使用最后一个位置的 logits 选择下一个 token。

16. 设计理念总结

16.1 表示学习

分词器将无限开放的字符串空间映射到有限 token 空间;嵌入层将离散 token 映射到连续向量空间;Transformer 层不断更新每个 token 的上下文化表示。表示不是人工编写规则,而是在大规模预测任务中学习到的函数。

16.2 并行性

Transformer 替代 RNN 的关键优势是训练并行性。训练时,序列中所有位置的表示可以并行计算,因果 mask 保证不泄漏未来信息。并行性使大规模 GPU 训练成为可能。

16.3 可扩展性

Scaling laws 表明,模型性能可以在一定范围内随参数、数据和计算规律性提升。Chinchilla 强调参数和 token 的平衡。系统论文如 Megatron-LM、ZeRO、FlashAttention 和 PagedAttention 则说明,可扩展性同时依赖算法、数值精度、并行策略和内存 IO。

16.4 后训练与目标错配修正

预训练目标是预测文本,不是最大化答案真实性、无害性或用户满意度。SFT 和偏好优化用于缩小这个目标错配。它们不能从根本上保证模型总是正确,但能显著改变输出风格、指令遵循和拒答行为。

17. 常见误区与严谨表述

  1. “模型理解了语言”不是严格技术表述。更严谨的说法是:模型学习了从上下文到 token 分布的参数化函数,并在许多任务上表现出可迁移的语言和推理行为。
  2. Attention 权重不是完整解释。注意力权重可用于分析信息流,但不能直接等同于模型决策的因果解释。
  3. 参数越大不一定越优。固定计算预算下,参数量、数据量、训练步数和数据质量必须共同设计。
  4. 长上下文不等于长程推理能力。模型能接收更多 token,不代表一定能稳定利用远距离信息。
  5. 微调不能可靠注入所有知识。小规模微调更适合改变格式、风格和任务行为;事实知识更新可能需要检索增强、继续预训练或更系统的数据方案。
  6. 低困惑度不等于好助手。困惑度衡量语言预测能力,不直接衡量真实性、安全性、遵循性和可验证性。

18. 最小实现层面的结构分解

若从零实现一个教学版 decoder-only LLM,模块顺序应为:

  1. Tokenizer:提供 encode(text)->idsdecode(ids)->text
  2. Dataset:把 token 流切成长度为 TT 的训练样本。
  3. Embedding:nn.Embedding(V, d_model)
  4. Causal mask:构造上三角 mask,禁止访问未来位置。
  5. Attention:实现 Q,K,VQ,K,V 投影、RoPE、缩放点积注意力、多头拼接和输出投影。
  6. FFN:实现 SwiGLU 或 GELU MLP。
  7. Norm:实现 RMSNorm 或 LayerNorm。
  8. Block:组合 norm、attention、residual、norm、ffn、residual。
  9. Model:堆叠 LL 个 block,加 final norm 和 lm head。
  10. Loss:对 shifted labels 计算 cross entropy。
  11. Optimizer:AdamW、warmup、cosine decay、gradient clipping。
  12. Generate:实现 KV Cache、temperature、top-p、停止条件。

这个顺序对应理论流水线,也对应工程实现依赖关系。先实现分词和数据,再实现模型前向,最后实现训练与生成,可以降低调试复杂度。

19. 参考文献

[1] Sutskever, I., Vinyals, O., & Le, Q. V. Sequence to Sequence Learning with Neural Networks. arXiv:1409.3215. https://arxiv.org/abs/1409.3215

[2] Bahdanau, D., Cho, K., & Bengio, Y. Neural Machine Translation by Jointly Learning to Align and Translate. arXiv:1409.0473. https://arxiv.org/abs/1409.0473

[3] Sennrich, R., Haddow, B., & Birch, A. Neural Machine Translation of Rare Words with Subword Units. arXiv:1508.07909. https://arxiv.org/abs/1508.07909

[4] Vaswani, A., et al. Attention Is All You Need. arXiv:1706.03762. https://arxiv.org/abs/1706.03762

[5] Kudo, T., & Richardson, J. SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing. arXiv:1808.06226. https://arxiv.org/abs/1808.06226

[6] Devlin, J., Chang, M.-W., Lee, K., & Toutanova, K. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv:1810.04805. https://arxiv.org/abs/1810.04805

[7] Kingma, D. P., & Ba, J. Adam: A Method for Stochastic Optimization. arXiv:1412.6980. https://arxiv.org/abs/1412.6980

[8] Loshchilov, I., & Hutter, F. Decoupled Weight Decay Regularization. arXiv:1711.05101. https://arxiv.org/abs/1711.05101

[9] Shoeybi, M., et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053. https://arxiv.org/abs/1909.08053

[10] Rajbhandari, S., et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054. https://arxiv.org/abs/1910.02054

[11] Zhang, B., & Sennrich, R. Root Mean Square Layer Normalization. arXiv:1910.07467. https://arxiv.org/abs/1910.07467

[12] Kaplan, J., et al. Scaling Laws for Neural Language Models. arXiv:2001.08361. https://arxiv.org/abs/2001.08361

[13] Shazeer, N. GLU Variants Improve Transformer. arXiv:2002.05202. https://arxiv.org/abs/2002.05202

[14] Brown, T. B., et al. Language Models are Few-Shot Learners. arXiv:2005.14165. https://arxiv.org/abs/2005.14165

[15] Su, J., et al. RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864. https://arxiv.org/abs/2104.09864

[16] Hu, E. J., et al. LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09685. https://arxiv.org/abs/2106.09685

[17] Ouyang, L., et al. Training language models to follow instructions with human feedback. arXiv:2203.02155. https://arxiv.org/abs/2203.02155

[18] Hoffmann, J., et al. Training Compute-Optimal Large Language Models. arXiv:2203.15556. https://arxiv.org/abs/2203.15556

[19] Dao, T., et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135. https://arxiv.org/abs/2205.14135

[20] Touvron, H., et al. LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971. https://arxiv.org/abs/2302.13971

[21] Dettmers, T., et al. QLoRA: Efficient Finetuning of Quantized LLMs. arXiv:2305.14314. https://arxiv.org/abs/2305.14314

[22] Rafailov, R., et al. Direct Preference Optimization: Your Language Model is Secretly a Reward Model. arXiv:2305.18290. https://arxiv.org/abs/2305.18290

[23] Kwon, W., et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180. https://arxiv.org/abs/2309.06180