第三章:Transformer架构详解
本章学习目标
- 理解Transformer架构诞生的背景和动机
- 掌握Transformer的整体架构(Encoder-Decoder结构)
- 深入理解自注意力机制(Self-Attention)的数学原理
- 掌握多头注意力(Multi-Head Attention)的设计哲学
- 理解位置编码的原理和不同方案
- 了解Transformer的各种变体(BERT、GPT、T5等)
- 掌握大模型训练的关键技术(分布式训练、混合精度等)
3.1 Transformer诞生背景与整体架构
3.1.1 从Seq2Seq到Transformer的演进
概念引入:Seq2Seq的瓶颈
在Transformer出现之前,序列到序列(Seq2Seq)任务主要使用**循环神经网络(RNN/LSTM/GRU)**来建模。
类比理解:想象你在阅读一本长篇小说。RNN就像是一个只能记住最近几页内容的读者------当小说很长时,前面章节的细节会逐渐被遗忘。虽然LSTM通过门控机制缓解了这个问题,但当序列长度超过几百个词时,性能仍然会显著下降。
RNN/LSTM的核心问题:
-
序列依赖导致无法并行训练
处理序列 [x₁, x₂, x₃, x₄, ...] RNN必须按顺序计算:h₁ → h₂ → h₃ → h₄ → ... 无法并行!训练速度慢 -
长期依赖问题(Long-Term Dependency)
- 梯度消失/爆炸使得模型难以捕捉长距离依赖
- 例如:在 "The cat, which was ... (100个词后) ... was black." 中,模型很难建立 "cat" 和 "was" 之间的语法关系
-
计算效率低下
- 每个时间步都需要完整的矩阵运算
- GPU的并行计算能力无法被充分利用
Attention机制的引入
Bahdanau等人在2014年将**注意力机制(Attention Mechanism)**引入Seq2Seq模型,使得解码器在生成每个词时,能够"关注"编码器隐藏状态的不同部分。
Attention的改进:
- ✅ 解决了瓶颈问题(不再依赖单一上下文向量)
- ✅ 提升了长距离依赖的建模能力
- ❌ 但仍然是RNN-based,无法完全并行化
Transformer的突破:完全基于Attention
Vaswani等人在2017年发表了里程碑论文《Attention is All You Need》,提出了Transformer架构,其革命性在于:
核心思想 :完全抛弃RNN,仅使用注意力机制 和前馈神经网络来构建序列模型。
优势:
- 完全可并行化:不同位置的注意力计算互不影响
- 长距离依赖建模能力强:任意两个位置的距离都是O(1)
- 训练速度快:充分利用GPU并行计算能力
- 可解释性强:注意力权重可视化能展示模型的"关注点"
python
# ===== 代码实战3.1:对比RNN和Transformer的计算模式 =====
import torch
import torch.nn as nn
import time
def compare_rnn_transformer_speed():
"""
对比RNN和Transformer的处理速度
展示Transformer的并行化优势
"""
print("=== RNN vs Transformer 计算模式对比 ===\n")
# 模拟输入 (batch_size=32, seq_len=128, hidden_dim=512)
batch_size = 32
seq_len = 128
hidden_dim = 512
# ===== RNN模式(序列依赖)=====
print("1. RNN/LSTM计算模式(无法并行):")
print(" 时间步计算顺序: t₁ → t₂ → t₃ → ... → tₙ")
print(" 每个时间步必须等待前一步完成")
# 模拟RNN计算(简化)
rnn = nn.LSTM(input_size=hidden_dim, hidden_size=hidden_dim, batch_first=True)
x = torch.randn(batch_size, seq_len, hidden_dim)
start_time = time.time()
output, (h_n, c_n) = rnn(x)
rnn_time = time.time() - start_time
print(f" RNN处理时间: {rnn_time:.4f}秒\n")
# ===== Transformer模式(完全并行)=====
print("2. Transformer计算模式(完全并行):")
print(" 所有位置同时计算: t₁, t₂, t₃, ..., tₙ 并行")
print(" 通过注意力机制直接建模任意位置间的关系")
# 模拟Self-Attention计算(简化)
class SimpleSelfAttention(nn.Module):
def __init__(self, hidden_dim):
super(SimpleSelfAttention, self).__init__()
self.W_q = nn.Linear(hidden_dim, hidden_dim)
self.W_k = nn.Linear(hidden_dim, hidden_dim)
self.W_v = nn.Linear(hidden_dim, hidden_dim)
def forward(self, x):
# x: (batch, seq_len, hidden_dim)
Q = self.W_q(x) # (batch, seq_len, hidden_dim)
K = self.W_k(x)
V = self.W_v(x)
# 计算注意力分数(完全并行!)
scores = torch.bmm(Q, K.transpose(1, 2)) / torch.sqrt(torch.tensor(hidden_dim, dtype=torch.float32))
attention_weights = torch.softmax(scores, dim=-1)
output = torch.bmm(attention_weights, V)
return output
attention = SimpleSelfAttention(hidden_dim)
start_time = time.time()
output = attention(x)
transformer_time = time.time() - start_time
print(f" Self-Attention处理时间: {transformer_time:.4f}秒")
print(f" 加速比: {rnn_time/transformer_time:.2f}x")
print("\n 说明: 实际加速效果取决于序列长度和硬件配置")
compare_rnn_transformer_speed()
3.1.2 Transformer整体架构概览
Encoder-Decoder结构
Transformer遵循经典的**编码器-解码器(Encoder-Decoder)**结构:
输入序列: [x₁, x₂, ..., xₙ]
↓
[Encoder Stack] (N层)
↓
中间表示: [z₁, z₂, ..., zₙ]
↓
[Decoder Stack] (N层)
↓
输出序列: [y₁, y₂, ..., yₘ]
架构图(文字版):
Encoder:
Input → Embedding → Positional Encoding
↓
┌────────────────────┐
│ Multi-Head │
│ Attention │
└────────────────────┘
↓
┌────────────────────┐
│ Add & Norm │
└────────────────────┘
↓
┌────────────────────┐
│ Feed Forward │
└────────────────────┘
↓
┌────────────────────┐
│ Add & Norm │
└────────────────────┘
↓
(重复N次)
Decoder:
Output → Embedding → Positional Encoding
↓
┌────────────────────┐
│ Masked Multi-Head │
│ Attention │
└────────────────────┘
↓
┌────────────────────┐
│ Add & Norm │
└────────────────────┘
↓
┌────────────────────┐
│ Multi-Head │
│ (Encoder-Decoder)│
└────────────────────┘
↓
┌────────────────────┐
│ Add & Norm │
└────────────────────┘
↓
┌────────────────────┐
│ Feed Forward │
└────────────────────┘
↓
┌────────────────────┐
│ Add & Norm │
└────────────────────┘
↓
(重复N次)
↓
Linear + Softmax
核心组件概览
| 组件 | 位置 | 作用 |
|---|---|---|
| Input/Output Embedding | Encoder/Decoder输入 | 将词转换为向量表示 |
| Positional Encoding | Embedding之后 | 为序列引入位置信息 |
| Multi-Head Attention | Encoder中1次,Decoder中2次 | 捕捉序列内部的依赖关系 |
| Add & Norm | 每个子层之后 | 残差连接 + 层归一化 |
| Feed Forward Network | 每个Attention之后 | 非线性变换 |
| Linear & Softmax | Decoder输出 | 生成最终词概率分布 |
关键洞察 :Transformer的成功不仅在于Attention机制,更在于残差连接 、层归一化 、位置编码等设计的精妙组合。这些组件共同使得训练极深的网络成为可能。
3.1.3 输入嵌入与位置编码
词嵌入(Word Embedding)
将离散的词映射到连续的向量空间。
数学表达:
给定词汇表大小为 VVV,嵌入维度为 dmodeld_{model}dmodel,嵌入层为一个矩阵 E∈RV×dmodelE \in \mathbb{R}^{V \times d_{model}}E∈RV×dmodel。
对于输入词 xix_ixi(表示为词汇表中的索引),其嵌入向量为:
ei=Exi∈Rdmodel \mathbf{e}i = Ex_i \in \mathbb{R}^{d{model}} ei=Exi∈Rdmodel
缩放因子:
在原始Transformer中,嵌入向量会乘以 dmodel\sqrt{d_{model}}dmodel :
ei′=ei×dmodel \mathbf{e}_i' = \mathbf{e}i \times \sqrt{d{model}} ei′=ei×dmodel
原因 :如果没有缩放,嵌入向量的量级会随着 dmodeld_{model}dmodel 增大而增大,导致点积注意力的值过大。
位置编码(Positional Encoding)
问题:Attention机制是**排列不变(Permutation Invariant)**的------即,打乱输入序列的顺序,Attention的输出不变。
类比理解:想象你在读一句话 "The cat sat on the mat"。如果Transformer不知道每个词的位置,它就无法理解 "cat" 是主语,"sat" 是谓语。位置编码就是给每个词贴上"我是第几个词"的标签。
解决方案 :为每个位置 pospospos 添加一个位置编码向量 ppos∈Rdmodel\mathbf{p}{pos} \in \mathbb{R}^{d{model}}ppos∈Rdmodel。
最终输入表示:
xi=ei+pi \mathbf{x}_i = \mathbf{e}_i + \mathbf{p}_i xi=ei+pi
绝对位置编码(Sinusoidal Positional Encoding)
原始Transformer使用正弦/余弦函数生成位置编码:
PE(pos,2i)=sin(pos100002i/dmodel) PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) PE(pos,2i)=sin(100002i/dmodelpos)
PE(pos,2i+1)=cos(pos100002i/dmodel) PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) PE(pos,2i+1)=cos(100002i/dmodelpos)
其中:
- pospospos:位置索引(0, 1, 2, ...)
- iii:维度索引(0, 1, ..., dmodel/2d_{model}/2dmodel/2)
特性:
- 确定性:不需要学习,可以直接计算
- 外推性:理论上可以处理比训练时更长的序列
- 相对位置信息 :对于任意固定偏移 kkk,PEpos+kPE_{pos+k}PEpos+k 可以表示为 PEposPE_{pos}PEpos 的线性组合
python
# ===== 代码实战3.2:实现Sinusoidal位置编码 =====
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
class SinusoidalPositionalEncoding(nn.Module):
"""
正弦位置编码(原始Transformer使用)
"""
def __init__(self, d_model, max_len=5000):
"""
初始化位置编码
Args:
d_model: 模型维度(嵌入维度)
max_len: 支持的最大序列长度
"""
super(SinusoidalPositionalEncoding, self).__init__()
# 创建位置编码矩阵 (max_len, d_model)
pe = torch.zeros(max_len, d_model)
# 位置索引 (max_len, 1)
position = torch.arange(0, max_len).unsqueeze(1)
# 分母项 1 / 10000^(2i/d_model)
div_term = torch.exp(
torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model)
)
# 偶数维度使用sin
pe[:, 0::2] = torch.sin(position * div_term)
# 奇数维度使用cos
pe[:, 1::2] = torch.cos(position * div_term)
# 添加batch维度 (1, max_len, d_model)
pe = pe.unsqueeze(0)
# 注册为buffer(不作为模型参数,但会被保存)
self.register_buffer('pe', pe)
def forward(self, x):
"""
添加位置编码到输入嵌入
Args:
x: 输入嵌入,形状为 (batch, seq_len, d_model)
Returns:
添加位置编码后的表示
"""
# 只取前seq_len个位置编码
return x + self.pe[:, :x.size(1), :]
def visualize_positional_encoding():
"""可视化位置编码"""
d_model = 64
max_len = 100
pe = SinusoidalPositionalEncoding(d_model, max_len)
# 获取位置编码矩阵
pe_matrix = pe.pe.squeeze(0).numpy() # (max_len, d_model)
# 绘制热力图
plt.figure(figsize=(12, 6))
plt.imshow(pe_matrix, aspect='auto', cmap='viridis')
plt.colorbar(label='Encoding Value')
plt.xlabel('Dimension', fontsize=12)
plt.ylabel('Position', fontsize=12)
plt.title('Sinusoidal Positional Encoding Heatmap', fontsize=14)
plt.tight_layout()
plt.savefig('positional_encoding_heatmap.png', dpi=150)
plt.show()
# 展示不同位置在某些维度上的值
plt.figure(figsize=(12, 4))
positions = [0, 10, 20, 30]
for i, pos in enumerate(positions):
plt.plot(pe_matrix[pos, :16], label=f'Position {pos}',
marker='o', markersize=4)
plt.xlabel('Dimension (first 16)', fontsize=12)
plt.ylabel('Encoding Value', fontsize=12)
plt.title('Positional Encoding Values for Different Positions', fontsize=14)
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('positional_encoding_values.png', dpi=150)
plt.show()
def demonstrate_positional_encoding_usage():
"""演示位置编码的实际使用"""
print("=== 位置编码使用示例 ===\n")
batch_size = 2
seq_len = 10
d_model = 512
# 模拟输入嵌入(通常来自词嵌入层)
embeddings = torch.randn(batch_size, seq_len, d_model)
print(f"输入嵌入形状: {embeddings.shape}")
# 添加位置编码
pe_layer = SinusoidalPositionalEncoding(d_model)
encoded = pe_layer(embeddings)
print(f"添加位置编码后形状: {encoded.shape}")
print(f"\n说明:")
print(f" - 位置编码通过与嵌入向量相加来注入位置信息")
print(f" - 不同位置的编码向量不同")
print(f" - 模型可以通过这些值感知词在序列中的位置")
if __name__ == "__main__":
demonstrate_positional_encoding_usage()
visualize_positional_encoding()
相对位置编码(Relative Positional Encoding)
绝对位置编码的问题:
- 无法直接建模词之间的相对位置关系
- 例如:"cat" 和 "sat" 之间的距离应该是相同的,无论它们在句子的哪个位置
相对位置编码的核心思想 :
不直接编码每个位置的绝对值,而是编码位置之间的差异。
主流方案:
| 方案 | 提出者 | 核心思想 |
|---|---|---|
| Relative Position Representations | Shaw et al. (2018) | 在Attention计算中显式添加相对位置偏置 |
| RoPE (Rotary Position Embedding) | Su et al. (2021) | 通过旋转矩阵编码相对位置 |
| ALiBi (Attention with Linear Biases) | Press et al. (2021) | 在Attention分数上添加线性衰减偏置 |
RoPE(当前主流方案):
将查询向量 q\mathbf{q}q 和键向量 k\mathbf{k}k 通过旋转矩阵进行变换:
qm′=qm⋅(cosmθ0−sinmθ0⋯sinmθ0cosmθ0⋯⋮⋮⋱) \mathbf{q}_m' = \mathbf{q}_m \cdot \begin{pmatrix} \cos m\theta_0 & -\sin m\theta_0 & \cdots \\ \sin m\theta_0 & \cos m\theta_0 & \cdots \\ \vdots & \vdots & \ddots \end{pmatrix} qm′=qm⋅ cosmθ0sinmθ0⋮−sinmθ0cosmθ0⋮⋯⋯⋱
kn′=kn⋅(cosnθ0−sinnθ0⋯sinnθ0cosnθ0⋯⋮⋮⋱) \mathbf{k}_n' = \mathbf{k}_n \cdot \begin{pmatrix} \cos n\theta_0 & -\sin n\theta_0 & \cdots \\ \sin n\theta_0 & \cos n\theta_0 & \cdots \\ \vdots & \vdots & \ddots \end{pmatrix} kn′=kn⋅ cosnθ0sinnθ0⋮−sinnθ0cosnθ0⋮⋯⋯⋱
关键性质:
⟨qm′,kn′⟩=⟨qm,kn⟩⋅cos((m−n)θ0) \langle \mathbf{q}_m', \mathbf{k}_n' \rangle = \langle \mathbf{q}_m, \mathbf{k}_n \rangle \cdot \cos((m-n)\theta_0) ⟨qm′,kn′⟩=⟨qm,kn⟩⋅cos((m−n)θ0)
即,旋转后的内积仅依赖于相对位置 m−nm-nm−n!
为什么RoPE在LLM时代成为主流?
- 优秀的外推性:可以在推理时处理比训练时更长的序列
- 计算高效:可以通过复数乘法快速实现
- 理论优雅:完美编码相对位置信息
当前主流大模型(LLaMA、ChatGLM、Qwen等)都使用RoPE。
3.1.4 编码器与解码器堆栈结构
编码器堆栈(Encoder Stack)
编码器由 NNN 个相同的层堆叠而成(原始Transformer中 N=6N=6N=6)。
每一层包含两个子层:
- Multi-Head Self-Attention
- Position-wise Feed-Forward Network
每个子层之后都有残差连接(Residual Connection)和层归一化(Layer Normalization)。
数学表达:
对于第 lll 层编码器:
z(l)=LayerNorm(x+MultiHeadAttention(x)) \mathbf{z}^{(l)} = \text{LayerNorm}(\mathbf{x} + \text{MultiHeadAttention}(\mathbf{x})) z(l)=LayerNorm(x+MultiHeadAttention(x))
z(l+1)=LayerNorm(z(l)+FFN(z(l))) \mathbf{z}^{(l+1)} = \text{LayerNorm}(\mathbf{z}^{(l)} + \text{FFN}(\mathbf{z}^{(l)})) z(l+1)=LayerNorm(z(l)+FFN(z(l)))
其中 x\mathbf{x}x 为输入,FFN\text{FFN}FFN 为前馈神经网络。
数据流转过程:
输入序列: [x₁, x₂, ..., xₙ]
↓ (词嵌入 + 位置编码)
Embedded: [e₁, e₂, ..., eₙ]
↓ (Encoder Layer 1)
Encoder Layer 1 Output: [h₁¹, h₂¹, ..., hₙ¹]
↓ (Encoder Layer 2)
Encoder Layer 2 Output: [h₁², h₂², ..., hₙ²]
↓
...
↓ (Encoder Layer N)
Encoder Output: [z₁, z₂, ..., zₙ]
解码器堆栈(Decoder Stack)
解码器也由 NNN 个相同的层堆叠而成。
每一层包含三个子层:
- Masked Multi-Head Self-Attention(带掩码的自注意力)
- Multi-Head Encoder-Decoder Attention(编码器-解码器注意力)
- Position-wise Feed-Forward Network
Masked Self-Attention的作用:
在训练时,解码器不应该"看到"未来的词(因为生成是自回归的)。通过**遮蔽(Mask)**未来位置来实现:
Maskij={0,i≥j(可见)−∞,i<j(遮蔽) \text{Mask}_{ij} = \begin{cases} 0, & i \geq j \quad \text{(可见)} \\ -\infty, & i < j \quad \text{(遮蔽)} \end{cases} Maskij={0,−∞,i≥j(可见)i<j(遮蔽)
Encoder-Decoder Attention的作用:
解码器的每个位置都可以"关注"编码器输出的所有位置。这是Seq2Seq模型的核心------使得解码器能够参考输入序列的信息。
python
# ===== 代码实战3.3:实现完整的Encoder Layer =====
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
"""
多头注意力机制
"""
def __init__(self, d_model, num_heads):
"""
初始化多头注意力
Args:
d_model: 模型维度
num_heads: 注意力头数
"""
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads # 每个头的维度
# 定义Q、K、V的线性变换
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def scaled_dot_product_attention(self, Q, K, V, mask=None):
"""
缩放点积注意力
Args:
Q: 查询矩阵 (batch, num_heads, seq_len, d_k)
K: 键矩阵 (batch, num_heads, seq_len, d_k)
V: 值矩阵 (batch, num_heads, seq_len, d_k)
mask: 可选的掩码
Returns:
注意力输出和注意力权重
"""
# 计算注意力分数
scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
# 应用掩码(如果提供)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
# Softmax归一化
attention_weights = F.softmax(scores, dim=-1)
# 加权求和
output = torch.matmul(attention_weights, V)
return output, attention_weights
def split_heads(self, x):
"""
将最后一维分割成多个头
Args:
x: (batch, seq_len, d_model)
Returns:
(batch, num_heads, seq_len, d_k)
"""
batch_size, seq_len, _ = x.size()
x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
return x.transpose(1, 2)
def combine_heads(self, x):
"""
将多个头合并回最后一维
Args:
x: (batch, num_heads, seq_len, d_k)
Returns:
(batch, seq_len, d_model)
"""
batch_size, num_heads, seq_len, d_k = x.size()
x = x.transpose(1, 2).contiguous()
return x.view(batch_size, seq_len, self.d_model)
def forward(self, query, key, value, mask=None):
"""
前向传播
Args:
query: (batch, seq_len, d_model)
key: (batch, seq_len, d_model)
value: (batch, seq_len, d_model)
mask: 可选的掩码
Returns:
注意力输出 (batch, seq_len, d_model)
"""
batch_size = query.size(0)
# 线性变换
Q = self.split_heads(self.W_q(query))
K = self.split_heads(self.W_k(key))
V = self.split_heads(self.W_v(value))
# 计算注意力
attention_output, attention_weights = self.scaled_dot_product_attention(Q, K, V, mask)
# 合并多头
output = self.combine_heads(attention_output)
# 输出线性变换
output = self.W_o(output)
return output, attention_weights
class PositionWiseFFN(nn.Module):
"""
位置级前馈神经网络
对每个位置独立地应用相同的FFN
"""
def __init__(self, d_model, d_ff):
"""
初始化FFN
Args:
d_model: 输入和输出维度
d_ff: 中间层维度(通常设为d_model的4倍)
"""
super(PositionWiseFFN, self).__init__()
self.fc1 = nn.Linear(d_model, d_ff)
self.fc2 = nn.Linear(d_ff, d_model)
self.relu = nn.ReLU()
def forward(self, x):
"""
前向传播
Args:
x: (batch, seq_len, d_model)
Returns:
(batch, seq_len, d_model)
"""
return self.fc2(self.relu(self.fc1(x)))
class EncoderLayer(nn.Module):
"""
Transformer编码器单层
"""
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
"""
初始化编码器层
Args:
d_model: 模型维度
num_heads: 注意力头数
d_ff: FFN中间层维度
dropout: Dropout比例
"""
super(EncoderLayer, self).__init__()
# 多头自注意力
self.self_attention = MultiHeadAttention(d_model, num_heads)
# 前馈神经网络
self.ffn = PositionWiseFFN(d_model, d_ff)
# 层归一化
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
# Dropout
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
"""
前向传播
Args:
x: 输入 (batch, seq_len, d_model)
mask: 可选的注意力掩码
Returns:
输出 (batch, seq_len, d_model)
"""
# ===== 子层1: Multi-Head Self-Attention (带残差连接和LayerNorm) =====
# Pre-Norm架构: LayerNorm在子层之前
normed_x = self.norm1(x)
attention_output, _ = self.self_attention(normed_x, normed_x, normed_x, mask)
x = x + self.dropout(attention_output) # 残差连接
# ===== 子层2: FFN (带残差连接和LayerNorm) =====
normed_x = self.norm2(x)
ffn_output = self.ffn(normed_x)
x = x + self.dropout(ffn_output) # 残差连接
return x
class Encoder(nn.Module):
"""
Transformer编码器(多层堆叠)
"""
def __init__(self, num_layers, d_model, num_heads, d_ff, dropout=0.1):
"""
初始化编码器
Args:
num_layers: 编码器层数
d_model: 模型维度
num_heads: 注意力头数
d_ff: FFN中间层维度
dropout: Dropout比例
"""
super(Encoder, self).__init__()
# 堆叠多个编码器层
self.layers = nn.ModuleList([
EncoderLayer(d_model, num_heads, d_ff, dropout)
for _ in range(num_layers)
])
# 最后的LayerNorm
self.norm = nn.LayerNorm(d_model)
def forward(self, x, mask=None):
"""
前向传播(通过所有编码器层)
Args:
x: 输入 (batch, seq_len, d_model)
mask: 可选的注意力掩码
Returns:
编码器输出 (batch, seq_len, d_model)
"""
for layer in self.layers:
x = layer(x, mask)
# 最后的归一化
return self.norm(x)
def test_encoder_implementation():
"""测试编码器实现"""
print("=== Transformer编码器测试 ===\n")
# 参数设置
batch_size = 2
seq_len = 10
d_model = 512
num_heads = 8
d_ff = 2048
num_layers = 6
# 创建编码器
encoder = Encoder(num_layers, d_model, num_heads, d_ff)
# 构造输入
x = torch.randn(batch_size, seq_len, d_model)
print(f"输入形状: {x.shape}")
# 前向传播
output = encoder(x)
print(f"输出形状: {output.shape}")
print(f"\n说明:")
print(f" - 编码器保持了序列长度和维度")
print(f" - 每个位置都能'看到'所有其他位置(通过Self-Attention)")
print(f" - 残差连接使得梯度可以直接传播到早期层")
if __name__ == "__main__":
test_encoder_implementation()
3.2 核心机制深度解析
3.2.1 自注意力机制(Self-Attention)数学原理
什么是自注意力?
**自注意力(Self-Attention)**是指序列中的每个元素都与序列中的所有其他元素(包括自己)计算注意力权重。
类比理解:想象一个团队合作项目。在会议中,每个成员都需要考虑其他所有成员的意见(包括自己的),来决定自己下一步该说什么。自注意力就是让序列中的每个词都能够"参考"所有其他词的信息。
数学原理详解
给定输入序列 X=x1,x2,...,xn\mathbf{X} = \\mathbf{x}_1, \\mathbf{x}_2, \\dots, \\mathbf{x}_nX=x1,x2,...,xn,其中 xi∈Rdmodel\mathbf{x}i \in \mathbb{R}^{d{model}}xi∈Rdmodel。
步骤1:生成Query、Key、Value
通过三个可学习的线性变换矩阵,将每个输入向量转换为三个表示:
Q=XWQ,K=XWK,V=XWV \mathbf{Q} = \mathbf{X}W^Q, \quad \mathbf{K} = \mathbf{X}W^K, \quad \mathbf{V} = \mathbf{X}W^V Q=XWQ,K=XWK,V=XWV
其中:
- WQ∈Rdmodel×dkW^Q \in \mathbb{R}^{d_{model} \times d_k}WQ∈Rdmodel×dk
- WK∈Rdmodel×dkW^K \in \mathbb{R}^{d_{model} \times d_k}WK∈Rdmodel×dk
- WV∈Rdmodel×dvW^V \in \mathbb{R}^{d_{model} \times d_v}WV∈Rdmodel×dv
通常 dk=dv=dmodel/num_headsd_k = d_v = d_{model} / \text{num\_heads}dk=dv=dmodel/num_heads。
步骤2:计算注意力分数
计算Query和Key之间的相似度(使用点积):
scoresij=qiTkjdk \text{scores}_{ij} = \frac{\mathbf{q}_i^T \mathbf{k}_j}{\sqrt{d_k}} scoresij=dk qiTkj
为什么要除以 dk\sqrt{d_k}dk ?
假设 qqq 和 kkk 的元素是独立同分布的,均值为0,方差为1。则点积 qTk=∑i=1dkqiki\mathbf{q}^T\mathbf{k} = \sum_{i=1}^{d_k} q_i k_iqTk=∑i=1dkqiki 的方差为 dkd_kdk。除以 dk\sqrt{d_k}dk 可以将方差重新缩放到1,避免点积过大导致Softmax函数进入梯度很小的饱和区。
步骤3:Softmax归一化
αij=exp(scoresij)∑k=1nexp(scoresik) \alpha_{ij} = \frac{\exp(\text{scores}{ij})}{\sum{k=1}^{n} \exp(\text{scores}_{ik})} αij=∑k=1nexp(scoresik)exp(scoresij)
步骤4:加权求和
yi=∑j=1nαijvj \mathbf{y}i = \sum{j=1}^{n} \alpha_{ij} \mathbf{v}_j yi=j=1∑nαijvj
矩阵形式(整个序列并行计算):
Attention(Q,K,V)=Softmax(QKTdk)V \text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=Softmax(dk QKT)V
自注意力的复杂度分析
| 操作 | 时间复杂度 | 空间复杂度 |
|---|---|---|
| 计算 Q,K,VQ, K, VQ,K,V | O(ndmodel2)O(nd_{model}^2)O(ndmodel2) | O(ndmodel)O(nd_{model})O(ndmodel) |
| 计算 QKTQK^TQKT | O(n2dk)O(n^2 d_k)O(n2dk) | O(n2)O(n^2)O(n2) |
| Softmax | O(n2)O(n^2)O(n2) | - |
| 加权求和 | O(n2dv)O(n^2 d_v)O(n2dv) | - |
关键问题 :Self-Attention的时间和空间复杂度都是 O(n2)O(n^2)O(n2)(nnn 为序列长度)。当 nnn 很大时(如 n=10000n=10000n=10000),计算和存储Attention矩阵会变得非常昂贵。
这也是为什么需要高效注意力机制(如稀疏注意力、Linear Attention等)的原因! 我们将在3.3.4节详细讨论。
python
# ===== 代码实战3.4:从零实现Self-Attention =====
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
def manual_self_attention(x):
"""
手动实现Self-Attention(用于理解原理)
Args:
x: 输入序列,形状为 (batch, seq_len, d_model)
Returns:
output: 注意力输出 (batch, seq_len, d_model)
attention_weights: 注意力权重 (batch, seq_len, seq_len)
"""
batch_size, seq_len, d_model = x.shape
d_k = d_model # 简化:假设d_k = d_model
# 步骤1:生成Q、K、V(使用简单的线性变换)
W_q = nn.Parameter(torch.randn(d_model, d_model))
W_k = nn.Parameter(torch.randn(d_model, d_model))
W_v = nn.Parameter(torch.randn(d_model, d_model))
Q = torch.matmul(x, W_q) # (batch, seq_len, d_model)
K = torch.matmul(x, W_k)
V = torch.matmul(x, W_v)
# 步骤2:计算注意力分数
scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
# scores形状: (batch, seq_len, seq_len)
# 步骤3:Softmax归一化
attention_weights = F.softmax(scores, dim=-1)
# 步骤4:加权求和
output = torch.matmul(attention_weights, V)
# output形状: (batch, seq_len, d_model)
return output, attention_weights
def compare_with_pytorch():
"""对比手动实现和PyTorch内置的Self-Attention"""
print("=== Self-Attention实现对比 ===\n")
batch_size = 2
seq_len = 5
d_model = 8
# 构造输入
x = torch.randn(batch_size, seq_len, d_model)
# 手动实现
output_manual, weights_manual = manual_self_attention(x)
print(f"手动实现输出形状: {output_manual.shape}")
print(f"注意力权重形状: {weights_manual.shape}")
print(f"\n注意力权重和(应为1.0): {weights_manual[0, 0, :].sum().item():.4f}")
# PyTorch内置实现(使用nn.MultiheadAttention)
# 注意:PyTorch的MultiheadAttention将d_model分割到多个头
mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=2, batch_first=True)
output_pytorch, weights_pytorch = mha(x, x, x)
print(f"\nPyTorch实现输出形状: {output_pytorch.shape}")
print(f"PyTorch注意力权重形状: {weights_pytorch.shape}")
print("\n说明:")
print(f" - 手动实现是单头的Self-Attention")
print(f" - PyTorch实现是多头的(num_heads=2)")
print(f" - 多头可以让模型同时关注不同位置的不同表示子空间")
if __name__ == "__main__":
compare_with_pytorch()
3.2.2 多头注意力(Multi-Head Attention)设计哲学
为什么需要多头注意力?
单头注意力的问题:
单个注意力头只能捕捉一种类型的依赖关系。例如,在处理句子 "The animal didn't cross the street because it was too tired" 时:
- 一个注意力头可能关注语法关系(如 "it" 指代 "animal")
- 另一个注意力头可能关注语义关系(如 "cross" 和 "street" 的搭配)
如果只有一个头,模型需要在同一种表示中平衡所有这些关系,这会降低模型的表达能力。
多头注意力的设计
核心思想 :将 dmodeld_{model}dmodel 维的Query、Key、Value并行地 通过 hhh 个不同的注意力头进行处理,然后将结果拼接起来。
数学表达:
给定 hhh 个注意力头,第 iii 个头的输出为:
headi=Attention(QWiQ,KWiK,VWiV) \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) headi=Attention(QWiQ,KWiK,VWiV)
其中 WiQ∈Rdmodel×dkW_i^Q \in \mathbb{R}^{d_{model} \times d_k}WiQ∈Rdmodel×dk, WiK∈Rdmodel×dkW_i^K \in \mathbb{R}^{d_{model} \times d_k}WiK∈Rdmodel×dk, WiV∈Rdmodel×dkW_i^V \in \mathbb{R}^{d_{model} \times d_k}WiV∈Rdmodel×dk,且 dk=dmodel/hd_k = d_{model} / hdk=dmodel/h。
将所有头的输出拼接:
MultiHead(Q,K,V)=Concat(head1,head2,...,headh)WO \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) W^O MultiHead(Q,K,V)=Concat(head1,head2,...,headh)WO
其中 WO∈Rhdk×dmodelW^O \in \mathbb{R}^{h d_k \times d_{model}}WO∈Rhdk×dmodel 为输出线性变换矩阵。
多头注意力的优势
- 表示子空间的多样性:不同的头可以学习关注不同类型的模式
- 增加模型的表达能力 :相当于将 dmodeld_{model}dmodel 维的空间分解为 hhh 个 dkd_kdk 维的子空间
- 并行计算:所有头的计算是完全独立的,可以利用GPU并行能力
可视化洞察:通过对训练好的Transformer的注意力权重进行可视化,可以发现:
- 某些头关注语法关系(如主谓一致)
- 某些头关注共指关系(如代词指代)
- 某些头关注位置信息(如相邻词)
这种"分工"是自发产生的,没有显式地告诉模型应该如何分配注意力头。
python
# ===== 代码实战3.5:可视化多头注意力的不同头 =====
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns
def visualize_attention_heads():
"""可视化不同注意力头的注意力权重"""
print("=== 可视化多头注意力的不同头 ===\n")
# 构造一个简单的句子(用于解释)
sentence = ["The", "cat", "sat", "on", "the", "mat", "."]
seq_len = len(sentence)
# 模拟多头注意力的输出
# 假设我们有4个头,每个头关注不同的模式
num_heads = 4
d_model = 64
# 创建模拟的注意力权重
attention_weights = torch.zeros(num_heads, seq_len, seq_len)
# 头1:关注相邻词
for i in range(seq_len):
for j in range(max(0, i-1), min(seq_len, i+2)):
attention_weights[0, i, j] = 1.0 / min(3, seq_len)
# 头2:关注第一个词(可能是主语)
attention_weights[1, :, 0] = 0.5
attention_weights[1, 0, 0] = 0.5
# 头3:关注最后一个词(可能是标点或从句结束)
attention_weights[2, :, -1] = 0.5
attention_weights[2, -1, -1] = 0.5
# 头4:均匀关注所有词
attention_weights[3, :, :] = 1.0 / seq_len
# 可视化
fig, axes = plt.subplots(2, 2, figsize=(12, 10))
fig.suptitle('Different Attention Heads Focus on Different Patterns', fontsize=14)
for idx in range(num_heads):
ax = axes[idx // 2, idx % 2]
# 绘制热力图
sns.heatmap(
attention_weights[idx].numpy(),
annot=True,
fmt='.2f',
cmap='YlOrRd',
xticklabels=sentence,
yticklabels=sentence,
ax=ax
)
ax.set_title(f'Head {idx+1}', fontsize=12)
ax.set_xlabel('Key (attended to)', fontsize=10)
ax.set_ylabel('Query (attending)', fontsize=10)
plt.tight_layout()
plt.savefig('multi_head_attention_visualization.png', dpi=150)
plt.show()
print("说明:")
print(" - Head 1: 关注相邻词(局部依赖)")
print(" - Head 2: 关注第一个词(可能是主语)")
print(" - Head 3: 关注最后一个词(可能是从句结束)")
print(" - Head 4: 均匀关注所有词(全局依赖)")
print("\n 实际训练中,这些模式是模型自动学习的!")
if __name__ == "__main__":
visualize_attention_heads()
3.2.3 前馈神经网络(FFN)与残差连接
位置级前馈神经网络(Position-wise FFN)
在Self-Attention之后,Transformer对每个位置独立地应用一个全连接前馈神经网络。
数学表达:
FFN(x)=max(0,xW1+b1)W2+b2 \text{FFN}(\mathbf{x}) = \max(0, \mathbf{x}W_1 + \mathbf{b}_1)W_2 + \mathbf{b}_2 FFN(x)=max(0,xW1+b1)W2+b2
或者写成两层线性变换 + ReLU激活:
FFN(x)=ReLU(xW1+b1)W2+b2 \text{FFN}(\mathbf{x}) = \text{ReLU}(\mathbf{x}W_1 + \mathbf{b}_1)W_2 + \mathbf{b}_2 FFN(x)=ReLU(xW1+b1)W2+b2
维度变化:
- 输入:x∈Rdmodel\mathbf{x} \in \mathbb{R}^{d_{model}}x∈Rdmodel
- 中间层:W1∈Rdmodel×dffW_1 \in \mathbb{R}^{d_{model} \times d_{ff}}W1∈Rdmodel×dff
- 输出:W2∈Rdff×dmodelW_2 \in \mathbb{R}^{d_{ff} \times d_{model}}W2∈Rdff×dmodel
通常 dff=4×dmodeld_{ff} = 4 \times d_{model}dff=4×dmodel(如 dmodel=512d_{model}=512dmodel=512, dff=2048d_{ff}=2048dff=2048)。
为什么需要FFN?
Self-Attention是一个线性变换 (虽然有Softmax,但本质上还是在做加权求和)。如果没有FFN,多层Transformer就等价于单层。FFN引入了非线性(通过ReLU/GELU),使得模型能够学习更复杂的函数。
位置级(Position-wise)的含义:
FFN对序列中的每个位置独立地应用相同的变换。这意味着:
FFN(x1),FFN(x2),...,FFN(xn) \text{FFN}(\mathbf{x}_1), \text{FFN}(\mathbf{x}_2), \dots, \text{FFN}(\mathbf{x}_n) FFN(x1),FFN(x2),...,FFN(xn)
这些计算是完全独立的,可以并行化。
残差连接(Residual Connection)
问题:深层网络训练困难(梯度消失/爆炸)。
解决方案 :在每一个子层(Self-Attention和FFN)之后添加残差连接:
y=LayerNorm(x+SubLayer(x)) \mathbf{y} = \text{LayerNorm}(\mathbf{x} + \text{SubLayer}(\mathbf{x})) y=LayerNorm(x+SubLayer(x))
其中 SubLayer\text{SubLayer}SubLayer 可以是Self-Attention或FFN。
作用:
- 梯度可以直接传播 :即使某一层的梯度很小,x\mathbf{x}x 也可以直接传播到前面的层
- 允许训练更深的网络:ResNet(He et al., 2016)证明了残差连接可以训练上千层的网络
- 提供"身份路径":如果某一层没有必要,模型可以学习将其输出设为0,直接传递输入
Post-Norm vs Pre-Norm:
- 原始Transformer(Post-Norm) :
x → SubLayer → Add → LayerNorm- 现代Transformer(Pre-Norm) :
x → LayerNorm → SubLayer → AddPre-Norm在训练稳定性上表现更好,是现代大模型(如GPT-3、LLaMA)的默认选择。
python
# ===== 代码实战3.6:实现带残差连接的FFN =====
import torch
import torch.nn as nn
class FeedForwardNetwork(nn.Module):
"""
前馈神经网络(带残差连接和LayerNorm)
"""
def __init__(self, d_model, d_ff, dropout=0.1, pre_norm=True):
"""
初始化FFN
Args:
d_model: 模型维度
d_ff: 中间层维度
dropout: Dropout比例
pre_norm: 是否使用Pre-Norm架构
"""
super(FeedForwardNetwork, self).__init__()
self.pre_norm = pre_norm
# FFN的两层线性变换
self.fc1 = nn.Linear(d_model, d_ff)
self.fc2 = nn.Linear(d_ff, d_model)
# 激活函数(现代Transformer常用GELU)
self.activation = nn.GELU() # 也可以用nn.ReLU()
# Dropout
self.dropout = nn.Dropout(dropout)
# LayerNorm
self.norm = nn.LayerNorm(d_model)
def forward(self, x):
"""
前向传播(带残差连接)
Args:
x: 输入 (batch, seq_len, d_model)
Returns:
输出 (batch, seq_len, d_model)
"""
if self.pre_norm:
# ===== Pre-Norm架构 =====
# 先归一化,再FFN,最后残差连接
normed_x = self.norm(x)
# FFN
ffn_output = self.fc2(self.dropout(self.activation(self.fc1(normed_x))))
# 残差连接
output = x + ffn_output
else:
# ===== Post-Norm架构(原始Transformer)=====
# 先FFN,再残差连接,最后归一化
ffn_output = self.fc2(self.dropout(self.activation(self.fc1(x))))
# 残差连接
residual = x + ffn_output
# LayerNorm
output = self.norm(residual)
return output
def compare_pre_norm_post_norm():
"""对比Pre-Norm和Post-Norm的训练稳定性"""
print("=== Pre-Norm vs Post-Norm ===\n")
d_model = 512
d_ff = 2048
batch_size = 2
seq_len = 10
# 创建Pre-Norm FFN
ffn_pre_norm = FeedForwardNetwork(d_model, d_ff, pre_norm=True)
# 创建Post-Norm FFN
ffn_post_norm = FeedForwardNetwork(d_model, d_ff, pre_norm=False)
# 构造输入
x = torch.randn(batch_size, seq_len, d_model)
# 前向传播
output_pre = ffn_pre_norm(x)
output_post = ffn_post_norm(x)
print(f"输入形状: {x.shape}")
print(f"Pre-Norm输出形状: {output_pre.shape}")
print(f"Post-Norm输出形状: {output_post.shape}")
print("\n说明:")
print(f" - Pre-Norm: 梯度更稳定,适合训练深层模型")
print(f" - Post-Norm: 原始Transformer使用,需要仔细的warmup策略")
print(f" - 现代大模型(GPT-3、LLaMA)都使用Pre-Norm")
if __name__ == "__main__":
compare_pre_norm_post_norm()
3.2.4 层归一化(Layer Normalization)的作用与变体
什么是Layer Normalization?
Batch Normalization(BN)的问题:
BN在CNN中非常有效,但在Transformer(尤其是RNN/Transformer)中效果不佳,因为:
- 依赖于批量大小:小批量时统计量不准确
- 训练和推理行为不一致:训练时使用当前batch的统计量,推理时使用滑动平均
- 不适合变长序列:不同样本的序列长度可能不同
Layer Normalization(LN)的解决方案:
LN对每个样本的特征维度进行归一化,而不依赖于batch。
数学表达:
给定输入 x=x1,x2,...,xH\mathbf{x} = x_1, x_2, \\dots, x_Hx=x1,x2,...,xH(一个样本的所有特征),LN计算:
μ=1H∑i=1Hxi \mu = \frac{1}{H}\sum_{i=1}^{H} x_i μ=H1i=1∑Hxi
σ2=1H∑i=1H(xi−μ)2 \sigma^2 = \frac{1}{H}\sum_{i=1}^{H} (x_i - \mu)^2 σ2=H1i=1∑H(xi−μ)2
x^i=xi−μσ2+ϵ \hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} x^i=σ2+ϵ xi−μ
yi=γix^i+βi y_i = \gamma_i \hat{x}_i + \beta_i yi=γix^i+βi
其中 γ\gammaγ 和 β\betaβ 是可学习的缩放和偏移参数。
与Batch Normalization的对比:
| 特性 | Batch Norm | Layer Norm |
|---|---|---|
| 归一化维度 | 跨batch的同一特征 | 每个样本的所有特征 |
| 训练/推理一致性 | 需要维护滑动平均 | 完全一致 |
| 适用于变长序列 | ❌ | ✅ |
| 适用于RNN/Transformer | ❌ | ✅ |
Layer Normalization的变体
1. RMSNorm(Root Mean Square Layer Normalization)
GPT-2之后的一些模型(如GPT-J、LLaMA)使用RMSNorm,它是LayerNorm的简化版本。
数学表达:
x^i=xi1H∑j=1Hxj2+ϵ \hat{x}i = \frac{x_i}{\sqrt{\frac{1}{H}\sum{j=1}^{H} x_j^2 + \epsilon}} x^i=H1∑j=1Hxj2+ϵ xi
yi=γix^i y_i = \gamma_i \hat{x}_i yi=γix^i
优势:
- 不需要计算均值 μ\muμ,计算更快
- 在实践中性能与LayerNorm相当
2. Adaptive Layer Norm(AdaLN)
在扩散模型(如Stable Diffusion)中使用的变体,允许条件信息(如文本嵌入)动态调整LayerNorm的参数。
为什么LayerNorm在Transformer中如此重要?
- 稳定训练:归一化使得每一层的输入分布保持稳定,加速收敛
- 允许更大的学习率:没有LN,大学习率会导致训练不稳定
- 减少梯度消失/爆炸:归一化使得梯度保持在合理范围内
⚠️ 企业级避坑1:LayerNorm的位置
在原始Transformer中,LayerNorm位于残差连接之后(Post-Norm):
x → Attention → Add → LayerNorm → FFN → Add → LayerNorm在现代Transformer中,LayerNorm位于子层之前(Pre-Norm):
x → LayerNorm → Attention → Add → LayerNorm → FFN → AddPre-Norm的训练稳定性显著更好,是现代大模型的默认选择。如果你在训练自己的Transformer时遇到不稳定问题,首先检查是否使用了Pre-Norm!
python
# ===== 代码实战3.7:实现RMSNorm =====
import torch
import torch.nn as nn
class RMSNorm(nn.Module):
"""
Root Mean Square Layer Normalization
简化版的LayerNorm,计算更快
"""
def __init__(self, d_model, eps=1e-6):
"""
初始化RMSNorm
Args:
d_model: 模型维度
eps: 数值稳定常数
"""
super(RMSNorm, self).__init__()
self.eps = eps
self.gamma = nn.Parameter(torch.ones(d_model)) # 可学习的缩放参数
def forward(self, x):
"""
前向传播
Args:
x: 输入 (batch, seq_len, d_model)
Returns:
归一化后的输出
"""
# 计算RMS
rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
# 归一化
normalized_x = x / rms
# 缩放
output = self.gamma * normalized_x
return output
def compare_ln_rmsnorm():
"""对比LayerNorm和RMSNorm"""
print("=== LayerNorm vs RMSNorm ===\n")
d_model = 512
batch_size = 2
seq_len = 10
# 创建LayerNorm和RMSNorm
ln = nn.LayerNorm(d_model)
rms_norm = RMSNorm(d_model)
# 构造输入
x = torch.randn(batch_size, seq_len, d_model)
# 前向传播
output_ln = ln(x)
output_rms = rms_norm(x)
print(f"输入形状: {x.shape}")
print(f"LayerNorm输出形状: {output_ln.shape}")
print(f"RMSNorm输出形状: {output_rms.shape}")
# 检查均值和方差
print(f"\nLayerNorm输出均值: {output_ln.mean(dim=-1).mean().item():.6f} (接近0)")
print(f"LayerNorm输出方差: {output_ln.var(dim=-1).mean().item():.6f} (接近1)")
print(f"\nRMSNorm输出RMS: {torch.sqrt(torch.mean(output_rms**2, dim=-1)).mean().item():.6f}")
print(f" (RMSNorm不强制均值为0,只归一化RMS)")
print("\n说明:")
print(f" - LayerNorm: 计算均值和方差,更准确")
print(f" - RMSNorm: 只计算RMS,计算更快")
print(f" - 在大模型中,RMSNorm的速度优势更明显")
if __name__ == "__main__":
compare_ln_rmsnorm()
3.3 Transformer变体与优化
3.3.1 Encoder-only架构:BERT系列
BERT:双向编码器表征
BERT(Bidirectional Encoder Representations from Transformers)是Encoder-only架构的代表。
核心特点:
- 双向上下文:同时利用左侧和右侧的上下文(通过Self-Attention)
- 预训练目标 :
- Masked Language Modeling(MLM):随机遮盖15%的词,预测被遮盖的词
- Next Sentence Prediction(NSP):预测两个句子是否连续
- 适合理解任务:文本分类、命名实体识别、问答系统等
架构图:
输入: [CLS] 句子1 [SEP] 句子2 [SEP]
↓
Embedding + Segment Embedding + Position Embedding
↓
[Encoder Layer 1] → [Encoder Layer 2] → ... → [Encoder Layer N]
↓
输出: [h_CLS, h_1, h_2, ..., h_n]
↓
[CLS] token的表示用于分类任务
CLS token的作用 :BERT在输入序列前添加一个特殊的
[CLS]token,其对应的输出向量被用作整个序列的聚合表示,用于分类任务。
BERT家族
| 模型 | 参数量 | 层数 | 隐藏维度 | 特点 |
|---|---|---|---|---|
| BERT-base | 110M | 12 | 768 | 原始BERT基准模型 |
| BERT-large | 340M | 24 | 1024 | 更大的BERT |
| RoBERTa | 125M | 12 | 768 | 去掉NSP,使用更多数据训练 |
| ALBERT | 12M | 12 | 768 | 参数共享,减少参数量 |
| ELECTRA | 110M | 12 | 768 | 使用判别式预训练目标 |
python
# ===== 代码实战3.8:使用Hugging Face加载BERT =====
from transformers import BertTokenizer, BertModel
import torch
def demonstrate_bert_usage():
"""演示BERT的使用"""
print("=== BERT使用示例 ===\n")
# 加载预训练的BERT模型和分词器
model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)
print(f"模型: {model_name}")
print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M\n")
# 准备输入
sentences = ["I love natural language processing.", "Transformer is amazing!"]
# 分词
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
print(f"分词结果:")
print(f" input_ids形状: {encoded_input['input_ids'].shape}")
print(f" attention_mask形状: {encoded_input['attention_mask'].shape}")
# 前向传播
with torch.no_grad():
output = model(**encoded_input)
# 获取最后一层的隐藏状态
last_hidden_states = output.last_hidden_state
print(f"\n隐藏状态形状: {last_hidden_states.shape}")
print(f" 解释: (batch_size={last_hidden_states.shape[0]}, "
f"seq_len={last_hidden_states.shape[1]}, "
f"hidden_dim={last_hidden_states.shape[2]})")
# 使用[CLS] token的输出进行分类
cls_output = last_hidden_states[:, 0, :] # [CLS]是第一个token
print(f"\n[CLS]输出形状: {cls_output.shape}")
print(f" 可用于下游分类任务")
if __name__ == "__main__":
demonstrate_bert_usage()
3.3.2 Decoder-only架构:GPT系列
GPT:生成式预训练Transformer
GPT(Generative Pre-trained Transformer)是Decoder-only架构的代表。
核心特点:
- 单向上下文:只能看到左侧的上下文(通过Masked Self-Attention)
- 预训练目标 :Causal Language Modeling(CLM)------预测下一个词
- 适合生成任务:文本生成、对话系统、代码补全等
架构图:
输入: [s_1, s_2, ..., s_n]
↓
Embedding + Positional Encoding
↓
[Decoder Layer 1] → [Decoder Layer 2] → ... → [Decoder Layer N]
↓ (只使用Masked Self-Attention和FFN,没有Encoder-Decoder Attention)
输出: [h_1, h_2, ..., h_n]
↓
Linear + Softmax → 预测下一个词
Masked Self-Attention:
通过上三角矩阵掩码,确保每个位置只能看到它之前的词:
Maskij={0,i≥j(可见)−∞,i<j(遮蔽) \text{Mask}_{ij} = \begin{cases} 0, & i \geq j \quad \text{(可见)} \\ -\infty, & i < j \quad \text{(遮蔽)} \end{cases} Maskij={0,−∞,i≥j(可见)i<j(遮蔽)
GPT家族演进
| 模型 | 参数量 | 层数 | 发布时间 | 特点 |
|---|---|---|---|---|
| GPT-1 | 117M | 12 | 2018.06 | 首个GPT模型,证明了预训练-微调范式 |
| GPT-2 | 1.5B | 48 | 2019.02 | 零样本学习能力的初步展现 |
| GPT-3 | 175B | 96 | 2020.05 | 少样本学习,涌现能力 |
| GPT-3.5 (InstructGPT) | 175B | 96 | 2022.03 | 引入RLHF,对话能力大幅提升 |
| GPT-4 | ~1.8T (MoE) | - | 2023.03 | 多模态,推理能力大幅提升 |
缩放定律(Scaling Law):GPT系列展示了模型性能随参数量、数据量、计算量的幂律增长。这为大模型的发展提供了理论指导。
python
# ===== 代码实战3.9:使用Hugging Face加载GPT-2 =====
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch
def demonstrate_gpt2_generation():
"""演示GPT-2的文本生成"""
print("=== GPT-2文本生成示例 ===\n")
# 加载预训练的GPT-2模型和分词器
model_name = 'gpt2'
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
print(f"模型: {model_name}")
print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M\n")
# 准备输入
prompt = "Once upon a time"
input_ids = tokenizer.encode(prompt, return_tensors='pt')
print(f"输入: {prompt}")
print(f"输入IDs形状: {input_ids.shape}\n")
# 生成文本
print("生成结果:")
output_ids = model.generate(
input_ids,
max_length=50,
num_return_sequences=1,
no_repeat_ngram_size=2,
early_stopping=True
)
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(f" {generated_text}")
if __name__ == "__main__":
demonstrate_gpt2_generation()
3.3.3 Encoder-Decoder架构:T5、BART等
T5(Text-to-Text Transfer Transformer)
T5将所有NLP任务统一为文本到文本的格式。
核心思想:
输入: "translate English to German: Hello, how are you?"
输出: "Hallo, wie geht es dir?"
输入: "summarize: [文章文本]"
输出: "[摘要文本]"
输入: "cola acceptability: First, the car drove down the street."
输出: "acceptable"
架构:完整的Encoder-Decoder Transformer。
预训练目标 :Masked Language Modeling(类似BERT,但是Encoder-Decoder架构)。
BART(Bidirectional and Auto-Regressive Transformer)
BART结合了BERT的双向编码器和GPT的单向解码器。
预训练目标 :噪声函数(如token masking、token deletion、text infilling等)。
适合任务:文本生成、摘要、翻译等。
Encoder-Decoder vs Encoder-only vs Decoder-only
| 架构 | 代表模型 | 适合任务 | 计算成本 |
|---|---|---|---|
| Encoder-only | BERT, RoBERTa | 理解任务(分类、NER、QA) | 低 |
| Decoder-only | GPT系列, LLaMA | 生成任务(对话、续写、代码) | 中 |
| Encoder-Decoder | T5, BART, mT5 | 序列到序列任务(翻译、摘要) | 高 |
为什么GPT成为大模型时代的主流?
虽然Encoder-Decoder架构在序列到序列任务上表现出色,但Decoder-only架构在以下方面更具优势:
- 统一的生成框架:所有任务都可以表示为"预测下一个词"
- 更好的零样本/少样本学习能力:通过Prompt可以适配各种任务
- 更适合大规模训练:架构简单,易于并行化
- 推理效率高:自回归生成可以逐步进行,易于控制
这也是为什么GPT-3、LLaMA、ChatGLM等主流大模型都采用Decoder-only架构。
3.3.4 高效注意力机制
问题:Self-Attention的O(n²)复杂度
标准Self-Attention的时间和空间复杂度都是 O(n2)O(n^2)O(n2)(nnn 为序列长度)。
当 nnn 很大时(如长文档理解、基因组序列分析等),标准Attention会变得非常昂贵。
解决方案 :设计高效注意力机制 ,将复杂度降低到 O(nlogn)O(n \log n)O(nlogn) 或 O(n)O(n)O(n)。
稀疏注意力(Sparse Attention)
核心思想 :每个token只关注一部分其他token,而非全部。
方案:
-
Fixed Pattern Sparse Attention(固定模式稀疏注意力)
- 每个token只关注固定窗口内的token(如前后各64个token)
- 复杂度:O(n×w)O(n \times w)O(n×w),其中 www 为窗口大小
-
Longformer(Beltagy et al., 2020)
- 结合局部窗口注意力和全局token
- 适合长文档理解
-
BigBird(Zaheer et al., 2020)
- 理论证明了稀疏注意力可以近似全注意力
- 使用随机图和全局token
线性注意力(Linear Attention)
核心思想 :重新排列Attention计算的顺序,避免显式计算 n×nn \times nn×n 的注意力矩阵。
数学原理:
标准Attention:
Attention(Q,K,V)=Softmax(QKTdk)V \text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=Softmax(dk QKT)V
问题:Softmax(QKTdk)\text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)Softmax(dk QKT) 的形状是 n×nn \times nn×n,计算和存储都很昂贵。
Linear Attention的解决方案:
使用一个核函数(Kernel Function) ϕ(⋅)\phi(\cdot)ϕ(⋅) 来近似Softmax:
LinearAttention(Q,K,V)=ϕ(Q)(ϕ(K)TV)ϕ(Q)(ϕ(K)T1) \text{LinearAttention}(Q, K, V) = \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q)(\phi(K)^T \mathbf{1})} LinearAttention(Q,K,V)=ϕ(Q)(ϕ(K)T1)ϕ(Q)(ϕ(K)TV)
关键技巧 :通过重新排列矩阵乘法顺序,将复杂度从 O(n2d)O(n^2 d)O(n2d) 降低到 O(nd2)O(n d^2)O(nd2):
标准Attention: (n×d) × (d×n) → (n×n) × (n×d) = O(n²d)
Linear Attention: (n×d) × ((d×n) × (n×d)) × (d×n) = O(nd²)
当 d≪nd \ll nd≪n 时,Linear Attention显著更快。
Flash Attention
核心思想:通过**IO感知(IO-Aware)**的算法设计,减少GPU高带宽内存(HBM)的读写次数。
关键技术:
- 分块计算(Tiling):将Attention矩阵分块计算,每个块适合GPU SRAM
- 重新计算(Recomputation):在反向传播时重新计算注意力矩阵,而非存储(节省显存)
效果:
- 显存占用从 O(n2)O(n^2)O(n2) 降低到 O(n)O(n)O(n)
- 实际运行速度提升2-4倍
Flash Attention已经成为现代大模型训练/推理的标配! PyTorch 2.0+已经内置了Flash Attention支持。
python
# ===== 代码实战3.10:使用Flash Attention =====
import torch
import torch.nn as nn
def demonstrate_flash_attention():
"""演示Flash Attention的使用"""
print("=== Flash Attention示例 ===\n")
# 检查PyTorch是否支持Flash Attention
if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
print("✅ PyTorch版本支持Flash Attention!")
# 构造输入
batch_size = 2
seq_len = 1024 # 长序列
num_heads = 8
d_k = 64
Q = torch.randn(batch_size, num_heads, seq_len, d_k).cuda()
K = torch.randn(batch_size, num_heads, seq_len, d_k).cuda()
V = torch.randn(batch_size, num_heads, seq_len, d_k).cuda()
# 使用PyTorch内置的scaled_dot_product_attention
# 它会自动选择最优实现(包括Flash Attention)
output = torch.nn.functional.scaled_dot_product_attention(Q, K, V)
print(f"输入形状: Q={Q.shape}")
print(f"输出形状: {output.shape}")
print("\n说明:")
print(f" - PyTorch会自动检测是否支持Flash Attention")
print(f" - 如果支持,会自动使用(无需额外代码)")
print(f" - Flash Attention可以显著减少显存占用并加速训练")
else:
print("❌ 当前PyTorch版本不支持Flash Attention")
print(" 请升级到PyTorch 2.0+")
if __name__ == "__main__":
demonstrate_flash_attention()
3.4 大模型训练技术
3.4.1 预训练目标函数
语言建模目标
1. Causal Language Modeling(CLM)
目标:给定前面的词,预测下一个词。
LCLM=−∑t=1TlogP(wt∣w1,...,wt−1) \mathcal{L}{\text{CLM}} = -\sum{t=1}^{T} \log P(w_t | w_1, \dots, w_{t-1}) LCLM=−t=1∑TlogP(wt∣w1,...,wt−1)
使用模型:GPT系列、LLaMA等Decoder-only模型。
2. Masked Language Modeling(MLM)
目标:随机遮盖部分词,预测被遮盖的词。
LMLM=−∑i∈MlogP(wi∣w∖M) \mathcal{L}{\text{MLM}} = -\sum{i \in \mathcal{M}} \log P(w_i | \mathbf{w}_{\setminus \mathcal{M}}) LMLM=−i∈M∑logP(wi∣w∖M)
其中 M\mathcal{M}M 为被遮盖的位置集合。
使用模型:BERT等Encoder-only模型。
3. Prefix Language Modeling(PLM)
目标:结合CLM和MLM,部分词可见,部分词需要预测。
使用模型:T5、BART等Encoder-Decoder模型。
去噪自编码目标
1. Denoising Autoencoder(DAE)
目标:向输入添加噪声(如随机遮盖、删除、打乱),让模型重建原始输入。
使用模型:BART、mBART。
2. Permutation Language Modeling(PLM)
目标:随机排列输入序列的顺序,然后预测被排列位置的词。
使用模型:XLNet。
对比学习目标
1. Contrastive Learning
目标:拉近正样本对的距离,推远负样本对的距离。
使用模型:SimCSE、Sentence-BERT。
3.4.2 分布式训练策略
问题:大模型无法放入单GPU显存
以GPT-3(175B参数)为例:
- FP32精度:~700GB显存
- 即使是A100(80GB显存),也无法放入单卡
解决方案 :分布式训练,将模型或数据分布到多个GPU/机器上。
数据并行(Data Parallelism, DP)
核心思想 :将数据分割成多个子集,每个GPU都有完整的模型副本,处理不同的数据子集。
流程:
输入数据: [Batch 1, Batch 2, ..., Batch N]
↓ (分割)
GPU 1: [Batch 1] → 模型副本1 → 梯度1
GPU 2: [Batch 2] → 模型副本2 → 梯度2
↓ (梯度聚合)
All-Reduce: 平均梯度 → 更新所有模型副本
优势 :实现简单,适合模型较小但数据量大的场景
局限:模型必须能放入单GPU显存
模型并行(Model Parallelism, MP)
核心思想 :将模型分割成多个部分,每个GPU负责一部分模型参数。
流程(以两层Transformer为例):
输入: [x₁, x₂, ..., xₙ]
↓
GPU 1: Encoder Layer 1 → 中间结果
↓ (通过网络传输)
GPU 2: Encoder Layer 2 → 输出
优势 :可以训练超大模型
局限:GPU之间需要频繁通信,训练速度慢
流水线并行(Pipeline Parallelism, PP)
核心思想 :将模型的不同层 分配到不同GPU,并通过流水线调度提高GPU利用率。
流程:
时间步1: GPU1处理Batch1的Layer1-4
时间步2: GPU1处理Batch2的Layer1-4, GPU2处理Batch1的Layer5-8
时间步3: GPU1处理Batch3, GPU2处理Batch2, GPU3处理Batch1
优势 :提高了GPU利用率
局限:需要仔细设计流水线调度策略
张量并行(Tensor Parallelism, TP)
核心思想 :将每一层的参数矩阵分割成多个分片,分布到不同GPU上。
示例(以线性层为例):
Y=XW Y = XW Y=XW
其中 WWW 可以按列分割:
W=W1,W2,Y=XW1,XW2 W = W_1, W_2, \quad Y = XW_1, XW_2 W=W1,W2,Y=XW1,XW2
GPU1计算 XW1XW_1XW1,GPU2计算 XW2XW_2XW2,最后拼接结果。
优势 :细粒度并行,适合超大模型
局限:GPU间通信开销大
3D并行(DP + PP + TP)
现代大模型训练通常组合使用多种并行策略:
总GPU数 = DP度数 × PP度数 × TP度数
示例(GPT-3训练配置):
- 数据并行:64路
- 流水线并行:4路
- 张量并行:8路
- 总GPU数:64 × 4 × 8 = 2048块GPU
python
# ===== 代码实战3.11:使用PyTorch实现数据并行 =====
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import os
def setup_distributed_training():
"""设置分布式训练环境(简化示例)"""
print("=== 分布式训练示例 ===\n")
# 检查是否有多GPU
if torch.cuda.device_count() > 1:
print(f"✅ 检测到 {torch.cuda.device_count()} 个GPU!")
print(" 可以使用DataParallel或DistributedDataParallel\n")
# 简单的DataParallel示例
model = nn.Transformer(
d_model=512,
nhead=8,
num_encoder_layers=6,
num_decoder_layers=6
)
if torch.cuda.is_available():
model = model.cuda()
model = nn.DataParallel(model) # 自动使用所有可用GPU
print(f"模型已包装为DataParallel")
print(f" 使用GPU数量: {torch.cuda.device_count()}")
else:
print("❌ 只有1个GPU或没有GPU")
print(" 无法演示分布式训练")
if __name__ == "__main__":
setup_distributed_training()
3.4.3 混合精度训练与梯度累积
混合精度训练(Mixed Precision Training)
核心思想 :在训练过程中混合使用不同数值精度(如FP16和FP32),以加速训练并减少显存占用。
为什么需要混合精度?
- FP32训练慢:FP32矩阵运算的吞吐量是FP16的1/2~1/8
- FP16易溢出:FP16的表示范围有限,容易出现下溢(梯度为0)或上溢(梯度为inf)
- 解决方案 :使用FP16进行计算 ,但使用FP32存储权重和梯度
技术细节:
- FP16前向传播和反向传播:加速计算
- FP32权重副本(Master Weights):存储高精度权重,用于更新
- Loss Scaling:将损失乘以一个缩放因子(如1024),避免梯度下溢
效果:
- 训练速度提升:1.5~3倍
- 显存占用减少:~50%
AMP(Automatic Mixed Precision) :PyTorch提供了
torch.cuda.amp自动混合精度训练工具,无需手动管理精度转换。
python
# ===== 代码实战3.12:使用AMP进行混合精度训练 =====
import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
def demonstrate_mixed_precision_training():
"""演示混合精度训练"""
print("=== 混合精度训练示例 ===\n")
# 创建模型
model = nn.Transformer(
d_model=512,
nhead=8,
num_encoder_layers=6,
num_decoder_layers=6
)
if torch.cuda.is_available():
model = model.cuda()
print("✅ 模型已移动到GPU\n")
# 定义优化器
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# 创建GradScaler(用于Loss Scaling)
scaler = GradScaler()
# 模拟训练循环
print("训练循环中...")
# 构造假数据
src = torch.randn(32, 10, 512).cuda()
tgt = torch.randn(32, 20, 512).cuda()
for epoch in range(3):
optimizer.zero_grad()
# ===== 使用autocast进行混合精度前向传播 =====
with autocast():
output = model(src, tgt)
loss = output.mean() # 假损失
# ===== 使用GradScaler进行混合精度反向传播 =====
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f" Epoch {epoch+1}: Loss = {loss.item():.4f}")
print("\n说明:")
print(f" - autocast(): 自动将计算转换为FP16")
print(f" - GradScaler: 自动进行Loss Scaling,避免梯度下溢")
print(f" - 混合精度训练可以显著加速训练并减少显存占用")
else:
print("❌ 没有GPU,无法演示")
if __name__ == "__main__":
demonstrate_mixed_precision_training()
梯度累积(Gradient Accumulation)
问题:当GPU显存不足以支持大Batch Size时,如何模拟大Batch训练?
解决方案 :梯度累积------在多次前向-反向传播中累积梯度,然后一次性更新参数。
数学原理:
假设目标Batch Size为 BBB,但GPU显存只能支持 B/kB/kB/k。
标准训练 (Batch Size = BBB):
g=1B∑i=1B∇L(xi,yi) \mathbf{g} = \frac{1}{B} \sum_{i=1}^{B} \nabla \mathcal{L}(x_i, y_i) g=B1i=1∑B∇L(xi,yi)
wt+1=wt−ηg \mathbf{w}_{t+1} = \mathbf{w}_t - \eta \mathbf{g} wt+1=wt−ηg
梯度累积 (累积 kkk 步):
gj=1B/k∑i∈mini-batchj∇L(xi,yi),j=1,...,k \mathbf{g}j = \frac{1}{B/k} \sum{i \in \text{mini-batch}_j} \nabla \mathcal{L}(x_i, y_i), \quad j = 1, \dots, k gj=B/k1i∈mini-batchj∑∇L(xi,yi),j=1,...,k
gaccumulated=1k∑j=1kgj \mathbf{g}{\text{accumulated}} = \frac{1}{k} \sum{j=1}^{k} \mathbf{g}_j gaccumulated=k1j=1∑kgj
wt+1=wt−ηgaccumulated \mathbf{w}_{t+1} = \mathbf{w}t - \eta \mathbf{g}{\text{accumulated}} wt+1=wt−ηgaccumulated
效果 :等效于使用Batch Size = BBB 的训练,但显存占用只有 B/kB/kB/k。
python
# ===== 代码实战3.13:实现梯度累积 =====
import torch
import torch.nn as nn
def demonstrate_gradient_accumulation():
"""演示梯度累积"""
print("=== 梯度累积示例 ===\n")
# 创建模型
model = nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# 目标Batch Size: 128
# 实际Batch Size: 32
# 累积步数: 4
target_batch_size = 128
actual_batch_size = 32
accumulation_steps = target_batch_size // actual_batch_size
print(f"目标Batch Size: {target_batch_size}")
print(f"实际Batch Size: {actual_batch_size}")
print(f"累积步数: {accumulation_steps}\n")
# 模拟训练循环
print("训练循环中...")
model.train()
for epoch in range(3):
optimizer.zero_grad()
# 累积梯度
total_loss = 0.0
for step in range(accumulation_steps):
# 构造假数据
x = torch.randn(actual_batch_size, 10)
y = torch.randint(0, 2, (actual_batch_size,))
# 前向传播
logits = model(x)
loss = nn.functional.cross_entropy(logits, y)
# 反向传播(不立即更新参数)
loss = loss / accumulation_steps # 平均损失
loss.backward()
total_loss += loss.item()
# 所有累积步完成后,更新参数
optimizer.step()
optimizer.zero_grad()
print(f" Epoch {epoch+1}: Loss = {total_loss:.4f}")
print("\n说明:")
print(f" - 梯度累积允许使用大Batch Size训练,即使显存有限")
print(f" - 注意:损失需要除以累积步数,以保持数值等价性")
if __name__ == "__main__":
demonstrate_gradient_accumulation()
3.4.4 大模型训练的稳定性挑战与解决方案
挑战1:梯度爆炸/消失
问题:在深层网络中,梯度可能指数级增长(爆炸)或指数级衰减(消失)。
解决方案:
- 梯度裁剪(Gradient Clipping):
python
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
-
使用Pre-Norm架构:如前所述,Pre-Norm比Post-Norm更稳定。
-
合适的权重初始化:如使用Xavier初始化或Kaiming初始化。
挑战2:Loss Spike(损失尖峰)
问题:在训练大模型时,Loss会突然出现很大的尖峰,然后恢复。
原因:可能是梯度爆炸,或者是某些batch的数据分布异常。
解决方案:
- 降低学习率
- 使用Warmup:在训练初期逐渐提高学习率
- 使用AdamW优化器:AdamW对超参数的选择更鲁棒
挑战3:显存不足(OOM)
问题:GPU显存不足以支持训练。
解决方案:
- 使用梯度累积(如前所述)
- 使用混合精度训练(如前所述)
- 使用梯度检查点(Gradient Checkpointing):
python
from torch.utils.checkpoint import checkpoint
# 在前向传播时不保存中间激活值,在反向传播时重新计算
output = checkpoint(model, input)
- 减少模型规模 :如使用更小的 dmodeld_{model}dmodel 或层数
⚠️ 企业级避坑2:大模型训练的调试技巧
训练大模型时,由于训练时间长(可能数周或数月),调试非常困难。以下是一些实用技巧:
- 使用小规模实验验证代码:在训练全规模模型前,先用小模型和小数据验证代码正确性
- 保存检查点(Checkpoint):定期保存模型参数,以便从故障中恢复
- 监控训练指标:使用TensorBoard或Wandb监控Loss、梯度范数、学习率等
- 使用确定性算法:设置随机种子,确保实验可复现
- 渐进式增加规模:先训练小模型,验证无误后再扩展到大规模
企业级考量与避坑指南
避坑1:位置编码选择不当导致长序列性能下降
问题描述:
不同的位置编码方案对长序列(如长度 > 2048)的处理能力不同。
解决方案:
| 位置编码方案 | 外推性 | 适合场景 |
|---|---|---|
| Sinusoidal | ✅ 好 | 训练时长度固定,推理时可能需要更长序列 |
| RoPE | ✅ 好(当前主流) | 大模型(LLaMA、ChatGLM等) |
| ALiBi | ✅ 非常好 | 需要极强外推能力的场景 |
| 可学习位置编码 | ❌ 差 | 训练和推理长度必须相同 |
经验法则:
- 如果需要在推理时处理比训练时更长的序列,使用RoPE 或ALiBi
- 如果训练和推理长度相同,可以使用可学习位置编码
避坑2:注意力头数设置不当导致计算效率低下
问题描述:
多头注意力的头数 hhh 必须能够整除 dmodeld_{model}dmodel。但即使满足这个条件,不同的 hhh 选择也会影响计算效率。
解决方案:
- 头数选择 :通常 h=8,12,16h = 8, 12, 16h=8,12,16 等
- 每个头的维度 :dk=dmodel/hd_k = d_{model} / hdk=dmodel/h,通常 dk≥32d_k \geq 32dk≥32(否则表示能力不足)
- 硬件效率 :某些GPU对特定的 hhh 有更好的支持(如Tensor Core要求 hhh 是8的倍数)
经验法则:
- 对于 dmodel=512d_{model} = 512dmodel=512,可以选择 h=8h = 8h=8(dk=64d_k = 64dk=64)或 h=16h = 16h=16(dk=32d_k = 32dk=32)
- 对于 dmodel=768d_{model} = 768dmodel=768,可以选择 h=12h = 12h=12(dk=64d_k = 64dk=64)
避坑3:未考虑推理时的内存优化
问题描述:
训练时可以通过梯度累积、混合精度等技术减少显存占用,但推理时没有梯度,这些技术不再适用。如果推理时显存不足,无法部署模型。
解决方案:
- 使用KV Cache:在自回归生成时,缓存已计算的Key和Value,避免重复计算
- 使用量化:将模型权重量化到INT8或INT4
- 使用模型并行:将模型分布到多个GPU上推理
- 使用推理加速框架:如vLLM、TensorRT-LLM等
⚠️ 企业级避坑3:推理吞吐量与延迟的权衡
在部署大模型时,需要在吞吐量 (Throughput)和延迟(Latency)之间进行权衡:
- 高吞吐量:使用更大的Batch Size,但延迟更高
- 低延迟:使用更小的Batch Size,但吞吐量更低
解决方案:
- 使用Continuous Batching:动态地将请求组合成Batch,提高GPU利用率
- 使用推测解码(Speculative Decoding):使用一个小的草稿模型快速生成多个候选词,然后用大模型验证
- 使用量化:减少计算量,降低延迟
本章小结
核心Takeaways
-
Transformer通过Attention机制完全替代了RNN,解决了序列依赖问题,实现了完全并行化训练。
-
Self-Attention是Transformer的核心,通过Query、Key、Value的计算,使得序列中的每个位置都能"关注"所有其他位置。
-
多头注意力通过并行计算多个注意力头,让模型能够同时关注不同类型的依赖关系,提高了模型的表达能力。
-
位置编码为序列引入位置信息,Sinusoidal编码是原始Transformer的选择,RoPE是现代大模型的主流方案。
-
残差连接和LayerNorm是训练深层Transformer的关键,Pre-Norm架构比Post-Norm更稳定。
-
FFN通过引入非线性,使得Transformer能够学习复杂的函数。位置级FFN对每个位置独立地应用相同的变换。
-
Transformer有三种主要变体:Encoder-only(BERT,适合理解任务)、Decoder-only(GPT,适合生成任务)、Encoder-Decoder(T5,适合序列到序列任务)。
-
**高效注意力机制(稀疏注意力、Linear Attention、Flash Attention)**解决了标准Attention的O(n²)复杂度问题,使得处理长序列成为可能。
-
**分布式训练(数据并行、模型并行、流水线并行、张量并行)**使得训练超大模型成为可能。
-
混合精度训练和梯度累积是训练大模型的关键技术,可以加速训练并减少显存占用。
-
大模型训练面临稳定性挑战(梯度爆炸/消失、Loss Spike、OOM等),需要通过梯度裁剪、Warmup、梯度检查点等技术来解决。
思考题
思考题1:为什么Self-Attention的计算复杂度是O(n²),有哪些方法可以降低这个复杂度?请详细解释其中一种方法的原理。
参考答案:
Self-Attention的O(n²)复杂度来源:
标准Self-Attention需要计算一个 n×nn \times nn×n 的注意力矩阵(其中 nnn 为序列长度),每个元素都需要计算Query和Key的点积。因此,时间和空间复杂度都是 O(n2d)O(n^2 d)O(n2d)(其中 ddd 为向量维度)。
降低复杂度的方法:
- 稀疏注意力(Sparse Attention) :每个token只关注固定窗口内的token,复杂度降至 O(nw)O(nw)O(nw)(www 为窗口大小)
- 线性注意力(Linear Attention) :通过重新排列矩阵乘法顺序,将复杂度降至 O(nd2)O(nd^2)O(nd2)
- Flash Attention :通过IO感知的算法设计,减少显存读写,虽然理论复杂度仍是 O(n2)O(n^2)O(n2),但实际运行更快且显存占用降至 O(n)O(n)O(n)
详细解释Linear Attention:
Linear Attention使用一个核函数 ϕ(⋅)\phi(\cdot)ϕ(⋅) 来近似Softmax:
Attention(Q,K,V)=ϕ(Q)(ϕ(K)TV)ϕ(Q)(ϕ(K)T1) \text{Attention}(Q, K, V) = \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q)(\phi(K)^T \mathbf{1})} Attention(Q,K,V)=ϕ(Q)(ϕ(K)T1)ϕ(Q)(ϕ(K)TV)
通过首先计算 ϕ(K)TV\phi(K)^T Vϕ(K)TV(复杂度 O(nd2)O(nd^2)O(nd2)),然后再与 ϕ(Q)\phi(Q)ϕ(Q) 相乘,避免了显式计算 n×nn \times nn×n 的注意力矩阵。
思考题2:在训练一个大模型时,你发现Loss出现了异常的尖峰(Spike),应该如何排查和解决这个问题?
参考答案:
排查步骤:
- 检查数据:查看出现Spike的batch是否包含异常数据(如极长序列、特殊字符等)
- 检查梯度:打印梯度的范数(Gradient Norm),查看是否出现梯度爆炸
- 检查学习率:查看学习率调度是否正确,是否在某些步骤学习率突然变大
- 检查数值稳定性:查看是否有NaN或Inf出现
解决方案:
- 降低学习率:如果Spike频繁出现,可能是学习率过大
- 使用Warmup:在训练初期逐渐提高学习率,避免初期的不稳定
- 使用梯度裁剪:限制梯度的最大范数,避免梯度爆炸
- 使用混合精度训练:确保Loss Scaling因子设置正确
- 数据清洗:移除或修正导致Spike的异常数据
- 使用更稳定的优化器:如AdamW(相比SGD更稳定)
预防措施:
- 使用小规模实验验证:在训练全规模模型前,先用小模型和小数据验证训练稳定性
- 监控训练指标:使用TensorBoard或Wandb实时监控Loss、梯度范数、学习率等
- 使用确定性算法:设置随机种子,确保实验可复现
本章已生成完毕。请回复【继续生成第四章】或提出您对当前章节的疑问。
参考文献
- Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention is all you need. Advances in Neural Information Processing Systems (NeurIPS), 30.
- Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2019). BERT: Pre-training of deep bidirectional transformers for language understanding. Proceedings of NAACL-HLT, 4171-4186.
- Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving language understanding by generative pre-training.
- Radford, A., Wu, J., Child, R., et al. (2019). Language models are unsupervised multitask learners. OpenAI Blog, 1(8), 9.
- Brown, T., Mann, B., Ryder, N., et al. (2020). Language models are few-shot learners. Advances in Neural Information Processing Systems (NeurIPS), 33, 1877-1901.
- Su, J., Lu, Y., Pan, S., et al. (2021). RoFormer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864.
- Press, O., Smith, N. A., & Lewis, M. (2021). Train short, test long: Attention with linear biases enables input length extrapolation. International Conference on Learning Representations (ICLR).
- Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems (NeurIPS), 35, 16344-16359.
- Shoeybi, M., Patwary, M., Puri, R., et al. (2019). Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053.
- Raffel, C., Shazeer, N., Roberts, A., et al. (2020). Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140), 1-67.