斯坦福CS336作业一：Transformer语言模型实现与解析

本次我们聚焦在斯坦福CS336课程作业一的第二部分，也是整个作业的重头戏——**Transformer语言模型架构的实现**。确切地说，是一个纯解码器（Decoder-Only）的Transformer，也就是GPT系列模型的核心结构。我们会从零开始，逐步搭建起多头注意力、旋转位置编码（RoPE）、RMSNorm归一化、SwiGLU激活函数等关键组件。 ### 章节阅读路线图 1. **Transformer语言模型架构** → 解码器结构、多头注意力、RoPE、RMSNorm、SwiGLU 2. **完整可运行代码** → 整合所有组件的完整实现 3. **总结** → 回顾核心要点 ### 1. Transformer语言模型架构 #### 1.1 整体架构 作业一的第二部分，核心是实现一个纯解码器的Transformer。听起来很高大上，其实就是我们常说的GPT风格模型。 它的几个核心特点非常鲜明： - **解码器架构**：只有解码器，没有编码器。 - **旋转位置编码（RoPE）**：用旋转矩阵来编码位置信息，好处是天然支持长度外推。 - **RMSNorm**：用均方根归一化替代传统的LayerNorm，训练起来更稳定。 - **SwiGLU激活函数**：前馈网络里用了这个门控线性单元，效果比传统的ReLU要好。 - **自回归生成**：这是Decoder-Only模型的天性，一个词一个词地往外蹦。 整体数据流可以这样理解：输入的文本先变成Token ID，然后经过Embedding层变成向量，再叠加上RoPE位置信息，接着穿过N个Transformer Block，最后经过一个RMSNorm和线性层，输出预测的logits。在每个Block内部，则是先做多头自注意力，再做SwiGLU前馈网络。 ``` 输入文本 → Token IDs → Embedding → + RoPE → Transformer Block × N → RMSNorm → 输出 logits ↓ Multi-Head Self-Attention ↓ SwiGLU FFN ``` #### 1.2 基础通用层 ##### 1.2.1 无偏置线性层（Linear） 首先是线性层。不过这里有个有意思的细节——移除了传统的偏置项（bias）。为什么现在的大语言模型都要这么做？简单来说，在Pre-Norm加残差连接的架构下，偏置项有点“画蛇添足”。 | 原因 | 说明 | | :--- | :--- | | RMSNorm 已覆盖偏置功能 | RMSNorm的可学习缩放参数已经承担了归一化功能，再加偏置就是多余的了。 | | 残差连接使 bias 失效 | 多层堆叠后，前一层的偏置偏移会被后续的RMSNorm给“洗掉”。 | | 训练稳定性提升 | PaLM论文里明确指出，移除bias可以提升大模型的训练稳定性。 | | 参数节省 | 每层能省下 `d_model` 个参数，几十层堆叠下来，省出来的量还是很可观的。 | 来看一下这个无偏置线性层的实现： ``` import torch # 导入 PyTorch 核心库 import torch.nn as nn # 导入神经网络模块 from einops import einsum # 导入 einops 用于清晰的张量维度映射 # 无偏置线性层：y = x @ W^T，移除偏置项以匹配现代 LLM 架构 # 参数: in_features 输入维度，out_features 输出维度 # 示例: linear = Linear(in_features=768, out_features=3072) class Linear(nn.Module): """无偏置线性层，执行 y = x @ W^T 变换。 参数: in_features: 输入特征维度 out_features: 输出特征维度 device: 计算设备 dtype: 数据类型 返回: 无 示例: >>> linear = Linear(in_features=768, out_features=3072) >>> linear.weight.shape torch.Size([3072, 768]) """ def __init__(self, in_features: int, out_features: int, device=None, dtype=None): super().__init__() self.weight = nn.Parameter(torch.empty((out_features, in_features), device=device, dtype=dtype)) # 权重矩阵：(out_features, in_features) std = (2 / (in_features + out_features)) ** 0.5 # 计算类 Xa vier 初始化标准差 nn.init.trunc_normal_(self.weight, std=std, a=-3*std, b=3*std) # 截断正态分布初始化，限制在 ±3σ """前向传播计算无偏置线性变换。 参数: x: 输入张量，形状为 [..., in_features] 返回: 输出张量，形状为 [..., out_features] 示例: >>> x = torch.randn(2, 10, 768) # [batch=2, seq_len=10, in_features=768] >>> y = linear(x) >>> y.shape torch.Size([2, 10, 3072]) """ def forward(self, x: torch.Tensor) -> torch.Tensor: return einsum(x, self.weight, "... in_features, out_features in_features -> ... out_features") # y = x @ W^T: [2,10,768] × [768,3072] → [2,10,3072] ``` 有两篇不错的文章可以作为扩展阅读：Sebastian Raschka写的一篇关于为什么有些LLM要移除线性层偏置项的文章，以及PaLM的原始论文。 ##### 1.2.2 词嵌入层（Embedding） 词嵌入层的工作很简单，就是把离散的Token ID映射成连续的稠密向量。这个向量的维度要与模型的隐藏层维度（`d_model`）保持一致。 ``` # 词嵌入层：token ID → 稠密向量 # 参数: num_embeddings 词汇表大小，embedding_dim 嵌入维度 # 示例: embedding = Embedding(num_embeddings=32000, embedding_dim=768) class Embedding(nn.Module): # 初始化嵌入权重矩阵 # 参数: num_embeddings 词汇表大小，embedding_dim 嵌入维度，示例：32000 个 token，每个 768 维 # 返回: 无 # 示例: self.weight.shape = (32000, 768) def __init__(self, num_embeddings: int, embedding_dim: int, device=None, dtype=None): super().__init__() self.vocab_size = num_embeddings self.d_model = embedding_dim # 嵌入权重矩阵：形状为 (词汇表大小, 嵌入维度) self.weight = nn.Parameter(torch.empty((self.vocab_size, self.d_model), device=device, dtype=dtype)) nn.init.trunc_normal_(self.weight, std=1, a=-3, b=3) # 前向传播获取 token 的嵌入向量 # 参数: token_ids 输入 token ID 张量 [batch_size, seq_len]，示例：[2, 10] # 返回: embeddings 嵌入向量 [batch_size, seq_len, embedding_dim]，示例：[2, 10, 768] # 示例: embeddings = embedding(token_ids) def forward(self, token_ids: torch.LongTensor) -> torch.Tensor: # 输入：(batch_size, seq_len) → 输出：(batch_size, seq_len, embedding_dim) # 数据流动：token_ids[2,10] → embeddings[2,10,768]（通过索引获取对应 token 的嵌入向量） return self.weight[token_ids] ``` ##### 1.2.3 RMS归一化层（RMSNorm） RMSNorm是LayerNorm的一个简化版。它将均值中心化这一步去掉了，只对输入的均方根进行归一化。好处很明显：计算量更小，而且训练稳定性反而更好，这也是LLaMA、GPT-4这些模型选择它的原因。 ``` # RMS归一化层：仅对均方根进行归一化 # 参数: d_model 输入维度，eps 防止除零的微小值 # 示例: rmsnorm = RMSNorm(d_model=768, eps=1e-5) class RMSNorm(nn.Module): # 初始化 RMSNorm 参数 # 参数: d_model 输入维度，eps 防止除零的微小值（默认 1e-5），示例：d_model=768 表示归一化 768 维向量 # 返回: 无 # 示例: self.weight.shape = (768,)，初始化为 1 def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None): super().__init__() self.d_model = d_model self.eps = eps # 可学习的缩放参数：形状为 (d_model,)，初始化为1（不改变归一化结果） self.weight = nn.Parameter(torch.ones(self.d_model, device=device, dtype=dtype)) # 前向传播应用 RMS 归一化 # 参数: x 输入张量 [batch_size, seq_len, d_model]，示例：[2, 10, 768] # 返回: out 归一化后的张量 [2, 10, 768] # 示例: out = rmsnorm(x) def forward(self, x: torch.Tensor) -> torch.Tensor: # 输入：(batch_size, seq_len, d_model) → 输出：同输入形状 in_dtype = x.dtype # 保存输入数据类型，避免精度损失 x = x.to(dtype=torch.float32) # 转为 float32 计算，提升数值稳定性 # 1. 计算最后一维的均方根（RMS） # 数据流动：x[2,10,768] → x.pow(2)[2,10,768] → mean[2,10,1] → +eps → sqrt → rms[2,10,1] rms = (x.pow(2).mean(-1, keepdim=True) + self.eps) ** 0.5 # 2. 归一化 + 应用缩放参数 # 数据流动：x[2,10,768] / rms[2,10,1] → out[2,10,768] × weight[768] → out[2,10,768] out = x / rms * self.weight return out.to(dtype=in_dtype) # 恢复原数据类型 ``` | 特性 | LayerNorm | RMSNorm | | :--- | :--- | :--- | | 均值中心化 | ✓ | ✗ | | 方差归一化 | ✓ | 仅 RMS | | 计算量 | 较大 | 较小 | | 训练稳定性 | 好 | 更好 | | 使用场景 | 传统 Transformer | LLaMA、GPT-4 等现代模型 | #### 1.3 激活函数与前馈网络 ##### 1.3.1 SwiGLU激活函数 SwiGLU是门控线性单元（GLU）的一个变体。它通过Sigmoid门控来筛选信息，相比ReLU，能更好地捕捉特征之间复杂的非线性关系。 ``` # SwiGLU前馈网络：门控线性单元 # 参数: d_model 输入维度，d_ff 前馈网络中间层维度 # 示例: swiglu = SwiGLU(d_model=768, d_ff=3072) class SwiGLU(nn.Module): # 初始化 SwiGLU 的三个线性层 # 参数: d_model 输入维度，d_ff 中间层维度（通常为 d_model 的 4 倍），示例：768 → 3072 # 返回: 无 # 示例: self.w1: 768→3072, self.w2: 3072→768, self.w3: 768→3072 def __init__(self, d_model: int, d_ff: int, device=None, dtype=None): super().__init__() self.d_model = d_model self.d_ff = d_ff # 定义三个线性层：W1/W3用于生成门控与候选特征，W2用于输出投影 self.w1 = Linear(d_model, d_ff, device=device, dtype=dtype) # 门控分支 self.w2 = Linear(d_ff, d_model, device=device, dtype=dtype) # 输出投影 self.w3 = Linear(d_model, d_ff, device=device, dtype=dtype) # 候选分支 # SiLU（Sigmoid Linear Unit）激活函数 # 参数: x 输入张量，示例：[2, 10, 3072] # 返回: x * sigmoid(x)，示例：[2, 10, 3072] # 示例: silu_output = self._silu(x) def _silu(self, x: torch.Tensor) -> torch.Tensor: return x * torch.sigmoid(x) # SiLU = x × σ(x) # GLU（Gated Linear Unit）门控机制 # 参数: x 输入张量 [batch, seq_len, d_model]，示例：[2, 10, 768] # 返回: SiLU(W1(x)) × W3(x)，示例：[2, 10, 3072] # 示例: glu_output = self._glu(x) def _glu(self, x: torch.Tensor) -> torch.Tensor: return self._silu(self.w1(x)) * self.w3(x) # SiLU门控 × W3线性变换结果 # 前向传播计算 SwiGLU # 参数: x 输入张量 [batch, seq_len, d_model]，示例：[2, 10, 768] # 返回: 输出张量 [batch, seq_len, d_model]，示例：[2, 10, 768] # 示例: output = swiglu(x) def forward(self, x: torch.Tensor) -> torch.Tensor: # 输入：(batch_size, seq_len, d_model) → 输出：同输入形状 # 数据流动：x[2,10,768] → _glu → [2,10,3072] → w2 → [2,10,768] return self.w2(self._glu(x)) # 门控结果通过 W2 投影回 d_model 维度 ``` SwiGLU的内部数据流可以清晰地看到三门分支和候选分支是如何融合的： ``` 输入 x [batch, seq_len, d_model=768] ↓ ┌────┴────┐ W1 W3 ↓ ↓ [2,10,3072] [2,10,3072] ↓ SiLU (x * sigmoid(x)) ↓ × (逐元素相乘) ↓ [2,10,3072] (门控结果) ↓ W2 ↓ 输出 [2,10,768] ``` #### 1.4 位置编码：旋转位置编码（RoPE） Transformer本身不“知道”词的先后顺序，所以需要位置编码来注入序列信息。RoPE的做法很巧妙：它设计了一个旋转矩阵，将位置信息“旋转”到token的嵌入向量中。这种编码方式不仅能处理相对位置，还天然支持长度外推。 ``` import math # 导入数学库，用于计算 # 旋转位置编码（RoPE） # 参数: theta 基础频率（通常 10000），d_k 注意力头维度，max_seq_len 最大序列长度 # 示例: rope = ROPE(theta=10000, d_k=64, max_seq_len=1024) class RoPE(nn.Module): # 初始化 RoPE 参数并预计算 cos/sin 缓存 # 参数: theta 基础频率，d_k 头维度，max_seq_len 最大序列长度，示例：d_k=64 表示每个头 64 维 # 返回: 无 # 示例: self.cos_cached.shape = (1024, 32) def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None): super().__init__() self.theta = theta self.d_k = d_k self.max_seq_len = max_seq_len self.device = device # 预计算cos/sin缓存（仅在首次初始化时计算，避免重复计算） # 1. 计算频率矩阵：shape (d_k//2,) # 数据流动：arange[0,2,...,62] → /64 → theta^ → 1/ → freqs_d[32] freqs_d = 1 / (theta ** (torch.arange(0, d_k, 2, device=device).float() / d_k)) # 2. 计算位置矩阵：shape (max_seq_len,) # 数据流动：arange[0,1,...,1023] → pos_i[1024] pos_i = torch.arange(max_seq_len, device=device).float() # 3. 频率-位置外积：shape (max_seq_len, d_k//2) # 数据流动：freqs_d[32] × pos_i[1024] → freqs[1024,32] freqs = torch.outer(pos_i, freqs_d) # 预计算cos和sin值（后续直接索引使用） self.register_buffer("cos_cached", torch.cos(freqs), persistent=False) # cos_cached[1024,32] self.register_buffer("sin_cached", torch.sin(freqs), persistent=False) # sin_cached[1024,32] # 前向传播应用旋转位置编码 # 参数: x 输入张量 [..., seq_len, d_k]，token_positions 位置索引 [..., seq_len] # 返回: 旋转后的张量 [..., seq_len, d_k] # 示例: x_rotated = rope(x, positions) def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor: # 1. 按最后一维的奇偶索引分组（d_k需为偶数） # 数据流动：x[...,10,64] → x_even[...,10,32], x_odd[...,10,32] x_even = x[..., ::2] # 偶数维度：索引 0,2,4... x_odd = x[..., 1::2] # 奇数维度：索引 1,3,5... # 2. 获取当前序列长度对应的cos/sin值 # 数据流动：cos_cached[1024,32] + token_positions[...,10] → cos[...,10,32] cos = self.cos_cached[token_positions] sin = self.sin_cached[token_positions] # 3. 应用旋转公式：将位置信息融入向量 # 数据流动：cos[...10,32]*x_even[...10,32] - sin[...10,32]*x_odd[...10,32] → out1[...10,32] out1 = cos * x_even - sin * x_odd # 偶数维度旋转结果 out2 = sin * x_even + cos * x_odd # 奇数维度旋转结果 # 4. 重组维度：将奇偶分组合并回原d_k维度 # 数据流动：out1[...10,32], out2[...10,32] → stack[...10,32,2] → flatten[...10,64] out = torch.stack([out1, out2], dim=-1).flatten(-2) return out ``` 旋转公式本身也很直观：对于位置 `m` 和偶数维度 `i`，旋转矩阵作用于向量对 `(x_m, i, x_m, i+1)`，其中频率 `θᵢ = 10000^(-2i/d_k)`。 \[ \begin{bmatrix} x_{m, i}^{out} \\ x_{m, i+1}^{out} \end{bmatrix} = \begin{bmatrix} cos(mθᵢ) & -sin(mθᵢ) \\ sin(mθᵢ) & cos(mθᵢ) \end{bmatrix} \times \begin{bmatrix} x_{m, i} \\ x_{m, i+1} \end{bmatrix} \] #### 1.5 多头自注意力（Multi-Head Self-Attention） 注意力机制是Transformer的绝对核心，它负责捕捉序列内部不同Token之间的依赖关系。多头则意味着从多个不同的角度去观察这种关系。 ``` # 多头自注意力机制 # 参数: d_model 模型维度，n_heads 注意力头数，max_seq_len 最大序列长度 # 示例: mha = MultiHeadAttention(d_model=768, n_heads=12, max_seq_len=1024) class MultiHeadAttention(nn.Module): # 初始化多头注意力参数 # 参数: d_model 模型维度，n_heads 头数，示例：768/12=64 表示每个头 64 维 # 返回: 无 # 示例: d_k=64, self.wq: 768→768, self.wk: 768→768, self.wv: 768→768 def __init__(self, d_model: int, n_heads: int, max_seq_len: int, device=None, dtype=None): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads # 每个头的维度 # Q、K、V 的线性变换 self.wq = Linear(d_model, d_model, device=device, dtype=dtype) # 查询投影 self.wk = Linear(d_model, d_model, device=device, dtype=dtype) # 键投影 self.wv = Linear(d_model, d_model, device=device, dtype=dtype) # 值投影 self.wo = Linear(d_model, d_model, device=device, dtype=dtype) # 输出投影 # RoPE 位置编码 self.rope = RoPE(theta=10000, d_k=self.d_k, max_seq_len=max_seq_len, device=device) # Dropout self.dropout = nn.Dropout(0.1) # 前向传播计算多头自注意力 # 参数: x 输入张量 [batch, seq_len, d_model]，示例：[2, 10, 768] # 返回: 输出张量 [batch, seq_len, d_model]，示例：[2, 10, 768] # 示例: output = mha(x) def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, seq_len, _ = x.shape # 获取 batch 和序列长度 # 1. 计算 Q、K、V # 数据流动：x[2,10,768] → wq → [2,10,768] → view → [2,10,12,64] → transpose → [2,12,10,64] Q = self.wq(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) K = self.wk(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) V = self.wv(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) # 2. 应用 RoPE 位置编码 # 数据流动：positions[10] → Q[2,12,10,64] → Q_rotated[2,12,10,64] positions = torch.arange(seq_len, device=x.device) # 位置索引 [0,1,...,9] Q = self.rope(Q, positions) K = self.rope(K, positions) # 3. 缩放点积注意力 # 数据流动：Q[2,12,10,64] × K^T[2,12,64,10] → scores[2,12,10,10] / sqrt(64) → scores[2,12,10,10] scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # 4. 因果掩码（防止看到未来信息） # 数据流动：scores[2,12,10,10] + causal_mask → scores[2,12,10,10]（上三角变为 -inf） causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device)) # 下三角矩阵 causal_mask = causal_mask.view(1, 1, seq_len, seq_len) # 广播到 [1,1,10,10] scores = scores.masked_fill(causal_mask == 0, float('-inf')) # 5. Softmax + Dropout # 数据流动：scores[2,12,10,10] → softmax → weights[2,12,10,10] → dropout → weights[2,12,10,10] attention_weights = torch.softmax(scores, dim=-1) attention_weights = self.dropout(attention_weights) # 6. 加权求和 # 数据流动：weights[2,12,10,10] × V[2,12,10,64] → output[2,12,10,64] output = torch.matmul(attention_weights, V) # 7. 拼接多头 + 输出投影 # 数据流动：output[2,12,10,64] → transpose → [2,10,12,64] → view → [2,10,768] → wo → [2,10,768] output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.wo(output) ``` 把数据流动可视化之后，过程就更清晰了： ``` 输入 x [batch=2, seq_len=10, d_model=768] ↓ ┌────┴────┬─────────┬─────────┐ wq wk wv ↓ ↓ ↓ Q[2,10,768] K[2,10,768] V[2,10,768] ↓ ↓ ↓ view+transpose (分割为 12 个头) ↓ ↓ ↓ Q[2,12,10,64] K[2,12,10,64] V[2,12,10,64] ↓ ↓ RoPE 位置编码 ↓ 缩放点积注意力：Q×K^T/√64 → softmax → ×V ↓ output[2,12,10,64] ↓ transpose + view (拼接多头) ↓ output[2,10,768] ↓ wo (输出投影) ↓ 最终输出 [2,10,768] ``` #### 1.6 Transformer Block 一个标准的Transformer解码器块，由两个子层组成：多头自注意力和前馈网络（SwiGLU）。每个子层前面都有RMSNorm，后面都有残差连接。这种“先归一化，再计算，最后加残差”的模式，就是Pre-Norm架构。 ``` # Transformer 解码器块（Pre-Norm 架构） # 参数: d_model 模型维度，n_heads 注意力头数，d_ff 前馈网络维度，max_seq_len 最大序列长度 # 示例: block = TransformerBlock(d_model=768, n_heads=12, d_ff=3072, max_seq_len=1024) class TransformerBlock(nn.Module): # 初始化 Transformer 块 # 参数: d_model 模型维度，n_heads 头数，d_ff FFN 维度，max_seq_len 最大序列长度 # 返回: 无 # 示例: self.attn_norm: RMSNorm(768), self.attn: MHA(768,12), self.ffn_norm: RMSNorm(768), self.ffn: SwiGLU(768,3072) def __init__(self, d_model: int, n_heads: int, d_ff: int, max_seq_len: int, device=None, dtype=None): super().__init__() # 注意力子层：RMSNorm → Multi-Head Attention → 残差连接 self.attn_norm = RMSNorm(d_model, device=device, dtype=dtype) # 预归一化 self.attn = MultiHeadAttention(d_model, n_heads, max_seq_len, device=device, dtype=dtype) # 前馈子层：RMSNorm → SwiGLU → 残差连接 self.ffn_norm = RMSNorm(d_model, device=device, dtype=dtype) # 预归一化 self.ffn = SwiGLU(d_model, d_ff, device=device, dtype=dtype) # 前向传播计算 Transformer 块 # 参数: x 输入张量 [batch, seq_len, d_model]，示例：[2, 10, 768] # 返回: 输出张量 [batch, seq_len, d_model]，示例：[2, 10, 768] # 示例: output = block(x) def forward(self, x: torch.Tensor) -> torch.Tensor: # 1. 注意力子层：Pre-Norm + 自注意力 + 残差 # 数据流动：x[2,10,768] → attn_norm → [2,10,768] → attn → [2,10,768] + x → [2,10,768] x = x + self.attn(self.attn_norm(x)) # 2. 前馈子层：Pre-Norm + SwiGLU + 残差 # 数据流动：x[2,10,768] → ffn_norm → [2,10,768] → ffn → [2,10,768] + x → [2,10,768] x = x + self.ffn(self.ffn_norm(x)) return x ``` | 架构 | 顺序 | 训练稳定性 | 使用场景 | | :--- | :--- | :--- | :--- | | Post-Norm | Attention → Add → Norm | 较难训练 | 原始 Transformer | | Pre-Norm | Norm → Attention → Add | 更稳定 | LLaMA、GPT-4 等现代模型 | #### 1.7 完整 Transformer 语言模型 把上面的所有组件拼起来，就是一个完整的语言模型了。它的结构很清晰：词嵌入层 → N个Transformer Block → 最终归一化层 → 输出线性层。 ``` # 完整的 Transformer 语言模型（Decoder-Only） # 参数: vocab_size 词汇表大小，d_model 模型维度，n_heads 头数，n_layers 层数，d_ff FFN 维度，max_seq_len 最大序列长度 # 示例: model = TransformerLM(vocab_size=32000, d_model=768, n_heads=12, n_layers=12, d_ff=3072, max_seq_len=1024) class TransformerLM(nn.Module): # 初始化完整的 Transformer 语言模型 # 参数: vocab_size 词汇表大小，d_model 模型维度，n_heads 头数，n_layers 层数，d_ff FFN 维度，max_seq_len 最大长度 # 返回: 无 # 示例: 768 维模型，12 个头，12 层，3072 维 FFN，支持最长 1024 序列 def __init__(self, vocab_size: int, d_model: int, n_heads: int, n_layers: int, d_ff: int, max_seq_len: int, device=None, dtype=None): super().__init__() self.d_model = d_model # 词嵌入层 self.embedding = Embedding(vocab_size, d_model, device=device, dtype=dtype) # Transformer 块堆叠 self.layers = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff, max_seq_len, device=device, dtype=dtype) for _ in range(n_layers)]) # 最终归一化层 self.final_norm = RMSNorm(d_model, device=device, dtype=dtype) # 输出投影层（与嵌入层共享权重） self.lm_head = Linear(d_model, vocab_size, device=device, dtype=dtype) # 前向传播计算语言模型 # 参数: token_ids 输入 token ID 张量 [batch, seq_len]，示例：[2, 10] # 返回: logits 输出 logits [batch, seq_len, vocab_size]，示例：[2, 10, 32000] # 示例: logits = model(token_ids) def forward(self, token_ids: torch.Tensor) -> torch.Tensor: # 1. 词嵌入 # 数据流动：token_ids[2,10] → embedding → x[2,10,768] x = self.embedding(token_ids) # 2. Transformer 块堆叠 # 数据流动：x[2,10,768] → layer1 → [2,10,768] → layer2 → ... → layer12 → [2,10,768] for layer in self.layers: x = layer(x) # 3. 最终归一化 # 数据流动：x[2,10,768] → final_norm → x[2,10,768] x = self.final_norm(x) # 4. 输出 logits（每个位置对词汇表中每个 token 的预测分数） # 数据流动：x[2,10,768] → lm_head → logits[2,10,32000] logits = self.lm_head(x) return logits ``` 整个模型的完整数据流如下： ``` 输入 token_ids [batch=2, seq_len=10] ↓ Embedding ↓ x [2,10,768] ↓ TransformerBlock × 12（每个块包含：RMSNorm → MHA → 残差 → RMSNorm → SwiGLU → 残差） ↓ x [2,10,768] ↓ RMSNorm (final_norm) ↓ x [2,10,768] ↓ Linear (lm_head) ↓ logits [2,10,32000]（每个位置对 32000 个 token 的预测分数） ``` ### 2. 完整可运行代码 把所有组件整合到一起，下面是完整的、可以直接运行的脚本。为了方便测试，我们设置了一个较小的配置。 ``` import torch # 导入 PyTorch 核心库，提供张量运算和自动求导 import torch.nn as nn # 导入神经网络模块，包含 Dropout、Parameter 等 import math # 导入数学库，用于计算平方根等 from einops import einsum # 导入 einops 用于清晰的张量维度映射 """无偏置线性层，执行 y = x @ W^T 变换 参数: in_features: 输入特征维度，示例：768 out_features: 输出特征维度，示例：3072 device: 计算设备 dtype: 数据类型 返回: 输出张量，形状为 [..., out_features] 示例: >>> linear = Linear(in_features=768, out_features=3072) >>> x = torch.randn(2, 10, 768) # [batch=2, seq_len=10, in_features=768] >>> y = linear(x) >>> y.shape torch.Size([2, 10, 3072]) """ class Linear(nn.Module): """初始化无偏置线性层权重 参数: in_features: 输入特征维度 out_features: 输出特征维度 device: 计算设备 dtype: 数据类型 返回: 无 示例: linear = Linear(in_features=768, out_features=3072) """ def __init__(self, in_features: int, out_features: int, device=None, dtype=None): super().__init__() self.weight = nn.Parameter(torch.empty((out_features, in_features), device=device, dtype=dtype)) # 权重矩阵：(out_features, in_features) std = (2 / (in_features + out_features)) ** 0.5 # 计算类 Xa vier 初始化标准差 nn.init.trunc_normal_(self.weight, std=std, a=-3*std, b=3*std) # 截断正态分布初始化，限制在 ±3σ """前向传播计算无偏置线性变换 参数: x: 输入张量，形状为 [..., in_features] 返回: 输出张量，形状为 [..., out_features] 示例: y = linear(x) """ def forward(self, x: torch.Tensor) -> torch.Tensor: return einsum(x, self.weight, "... in_features, out_features in_features -> ... out_features") # y = x @ W^T: [2,10,768] × [768,3072] → [2,10,3072] """词嵌入层，将 token ID 映射为稠密向量 参数: num_embeddings: 词汇表大小，示例：32000 embedding_dim: 嵌入维度，示例：768 device: 计算设备 dtype: 数据类型 返回: 嵌入向量，形状为 [batch_size, seq_len, embedding_dim] 示例: >>> embedding = Embedding(num_embeddings=32000, embedding_dim=768) >>> token_ids = torch.randint(0, 32000, (2, 10)) # [batch=2, seq_len=10] >>> embeddings = embedding(token_ids) >>> embeddings.shape torch.Size([2, 10, 768]) """ class Embedding(nn.Module): """初始化嵌入权重矩阵 参数: num_embeddings: 词汇表大小 embedding_dim: 嵌入维度 device: 计算设备 dtype: 数据类型 返回: 无 示例: embedding = Embedding(num_embeddings=32000, embedding_dim=768) """ def __init__(self, num_embeddings: int, embedding_dim: int, device=None, dtype=None): super().__init__() self.vocab_size = num_embeddings self.d_model = embedding_dim self.weight = nn.Parameter(torch.empty((self.vocab_size, self.d_model), device=device, dtype=dtype)) # 嵌入权重：(vocab_size, d_model) nn.init.trunc_normal_(self.weight, std=1, a=-3, b=3) # 截断正态分布初始化 """前向传播获取 token 的嵌入向量 参数: token_ids: 输入 token ID 张量 [batch_size, seq_len] 返回: 嵌入向量 [batch_size, seq_len, embedding_dim] 示例: embeddings = embedding(token_ids) """ def forward(self, token_ids: torch.LongTensor) -> torch.Tensor: return self.weight[token_ids] # 查表操作：token_ids[2,10] → embeddings[2,10,768] """RMS 归一化层，仅对均方根进行归一化 参数: d_model: 输入维度，示例：768 eps: 防止除零的微小值（默认 1e-5） device: 计算设备 dtype: 数据类型 返回: 归一化后的张量，形状与输入相同 示例: >>> rmsnorm = RMSNorm(d_model=768, eps=1e-5) >>> x = torch.randn(2, 10, 768) >>> out = rmsnorm(x) >>> out.shape torch.Size([2, 10, 768]) """ class RMSNorm(nn.Module): """初始化 RMSNorm 参数 参数: d_model: 输入维度 eps: 防止除零的微小值 device: 计算设备 dtype: 数据类型 返回: 无 示例: rmsnorm = RMSNorm(d_model=768, eps=1e-5) """ def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None): super().__init__() self.d_model = d_model self.eps = eps self.weight = nn.Parameter(torch.ones(self.d_model, device=device, dtype=dtype)) # 可学习缩放参数，初始化为 1 """前向传播应用 RMS 归一化 参数: x: 输入张量 [batch_size, seq_len, d_model] 返回: 归一化后的张量 [batch_size, seq_len, d_model] 示例: out = rmsnorm(x) """ def forward(self, x: torch.Tensor) -> torch.Tensor: in_dtype = x.dtype # 保存输入数据类型 x = x.to(dtype=torch.float32) # 转为 float32 计算 # 计算均方根（RMS） # 数据流动：x[2,10,768] → pow(2) → mean(-1) → +eps → sqrt → rms[2,10,1] rms = (x.pow(2).mean(-1, keepdim=True) + self.eps) ** 0.5 # 归一化 + 应用缩放参数 # 数据流动：x[2,10,768] / rms[2,10,1] × weight[768] → out[2,10,768] out = x / rms * self.weight return out.to(dtype=in_dtype) # 恢复原数据类型 """SwiGLU 前馈网络，门控线性单元 参数: d_model: 输入维度，示例：768 d_ff: 前馈网络中间层维度，示例：3072 device: 计算设备 dtype: 数据类型 返回: 输出张量，形状为 [batch_size, seq_len, d_model] 示例: >>> swiglu = SwiGLU(d_model=768, d_ff=3072) >>> x = torch.randn(2, 10, 768) >>> output = swiglu(x) >>> output.shape torch.Size([2, 10, 768]) """ class SwiGLU(nn.Module): """初始化 SwiGLU 的三个线性层 参数: d_model: 输入维度 d_ff: 中间层维度 device: 计算设备 dtype: 数据类型 返回: 无 示例: swiglu = SwiGLU(d_model=768, d_ff=3072) """ def __init__(self, d_model: int, d_ff: int, device=None, dtype=None): super().__init__() self.d_model = d_model self.d_ff = d_ff self.w1 = Linear(d_model, d_ff, device=device, dtype=dtype) # 门控分支 self.w2 = Linear(d_ff, d_model, device=device, dtype=dtype) # 输出投影 self.w3 = Linear(d_model, d_ff, device=device, dtype=dtype) # 候选分支 """SiLU（Sigmoid Linear Unit）激活函数 参数: x: 输入张量 返回: x * sigmoid(x) 示例: output = self._silu(x) """ def _silu(self, x: torch.Tensor) -> torch.Tensor: return x * torch.sigmoid(x) # SiLU = x × σ(x) """GLU（Gated Linear Unit）门控机制 参数: x: 输入张量 [batch, seq_len, d_model] 返回: SiLU(W1(x)) × W3(x) 示例: output = self._glu(x) """ def _glu(self, x: torch.Tensor) -> torch.Tensor: return self._silu(self.w1(x)) * self.w3(x) # SiLU 门控 × W3 线性变换 """前向传播计算 SwiGLU 参数: x: 输入张量 [batch, seq_len, d_model] 返回: 输出张量 [batch, seq_len, d_model] 示例: output = swiglu(x) """ def forward(self, x: torch.Tensor) -> torch.Tensor: # 数据流动：x[2,10,768] → _glu → [2,10,3072] → w2 → [2,10,768] return self.w2(self._glu(x)) # 门控结果通过 W2 投影回 d_model """旋转位置编码（RoPE） 参数: theta: 基础频率（通常 10000） d_k: 注意力头维度，示例：64 max_seq_len: 最大序列长度，示例：1024 device: 计算设备 返回: 旋转后的张量，形状与输入相同 示例: >>> rope = RoPE(theta=10000, d_k=64, max_seq_len=1024) >>> x = torch.randn(2, 12, 10, 64) # [batch, heads, seq_len, d_k] >>> positions = torch.arange(10) >>> x_rotated = rope(x, positions) >>> x_rotated.shape torch.Size([2, 12, 10, 64]) """ class RoPE(nn.Module): """初始化 RoPE 参数并预计算 cos/sin 缓存 参数: theta: 基础频率 d_k: 头维度 max_seq_len: 最大序列长度 device: 计算设备 返回: 无 示例: rope = RoPE(theta=10000, d_k=64, max_seq_len=1024) """ def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None): super().__init__() self.theta = theta self.d_k = d_k self.max_seq_len = max_seq_len self.device = device # 预计算频率矩阵 # 数据流动：arange[0,2,...,62] → /64 → theta^ → 1/ → freqs_d[32] freqs_d = 1 / (theta ** (torch.arange(0, d_k, 2, device=device).float() / d_k)) # 预计算位置矩阵 # 数据流动：arange[0,1,...,1023] → pos_i[1024] pos_i = torch.arange(max_seq_len, device=device).float() # 频率-位置外积 # 数据流动：freqs_d[32] × pos_i[1024] → freqs[1024,32] freqs = torch.outer(pos_i, freqs_d) # 预计算 cos 和 sin 值 self.register_buffer("cos_cached", torch.cos(freqs), persistent=False) # cos_cached[1024,32] self.register_buffer("sin_cached", torch.sin(freqs), persistent=False) # sin_cached[1024,32] """前向传播应用旋转位置编码 参数: x: 输入张量 [..., seq_len, d_k] token_positions: 位置索引 [..., seq_len] 返回: 旋转后的张量 [..., seq_len, d_k] 示例: x_rotated = rope(x, positions) """ def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor: # 按奇偶索引分组 # 数据流动：x[...,10,64] → x_even[...,10,32], x_odd[...,10,32] x_even = x[..., ::2] # 偶数维度 x_odd = x[..., 1::2] # 奇数维度 # 获取对应位置的 cos/sin 值 # 数据流动：cos_cached[1024,32] + token_positions → cos[...,10,32] cos = self.cos_cached[token_positions] sin = self.sin_cached[token_positions] # 应用旋转公式 # 数据流动：cos*x_even - sin*x_odd → out1[...10,32] out1 = cos * x_even - sin * x_odd # 偶数维度旋转 out2 = sin * x_even + cos * x_odd # 奇数维度旋转 # 重组维度 # 数据流动：out1[...10,32], out2[...10,32] → stack → flatten → [...10,64] out = torch.stack([out1, out2], dim=-1).flatten(-2) return out """多头自注意力机制 参数: d_model: 模型维度，示例：768 n_heads: 注意力头数，示例：12 max_seq_len: 最大序列长度，示例：1024 device: 计算设备 dtype: 数据类型 返回: 输出张量，形状为 [batch_size, seq_len, d_model] 示例: >>> mha = MultiHeadAttention(d_model=768, n_heads=12, max_seq_len=1024) >>> x = torch.randn(2, 10, 768) >>> output = mha(x) >>> output.shape torch.Size([2, 10, 768]) """ class MultiHeadAttention(nn.Module): """初始化多头注意力参数 参数: d_model: 模型维度 n_heads: 头数 max_seq_len: 最大序列长度 device: 计算设备 dtype: 数据类型 返回: 无 示例: mha = MultiHeadAttention(d_model=768, n_heads=12, max_seq_len=1024) """ def __init__(self, d_model: int, n_heads: int, max_seq_len: int, device=None, dtype=None): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads # 每个头的维度 # Q、K、V 的线性变换 self.wq = Linear(d_model, d_model, device=device, dtype=dtype) # 查询投影 self.wk = Linear(d_model, d_model, device=device, dtype=dtype) # 键投影 self.wv = Linear(d_model, d_model, device=device, dtype=dtype) # 值投影 self.wo = Linear(d_model, d_model, device=device, dtype=dtype) # 输出投影 # RoPE 位置编码 self.rope = RoPE(theta=10000, d_k=self.d_k, max_seq_len=max_seq_len, device=device) # Dropout self.dropout = nn.Dropout(0.1) """前向传播计算多头自注意力 参数: x: 输入张量 [batch, seq_len, d_model] 返回: 输出张量 [batch, seq_len, d_model] 示例: output = mha(x) """ def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, seq_len, _ = x.shape # 获取 batch 和序列长度 # 1. 计算 Q、K、V # 数据流动：x[2,10,768] → wq → view → transpose → Q[2,12,10,64] Q = self.wq(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) K = self.wk(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) V = self.wv(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) # 2. 应用 RoPE 位置编码 # 数据流动：positions[10] → Q[2,12,10,64] → Q_rotated[2,12,10,64] positions = torch.arange(seq_len, device=x.device) # 位置索引 [0,1,...,9] Q = self.rope(Q, positions) K = self.rope(K, positions) # 3. 缩放点积注意力 # 数据流动：Q[2,12,10,64] × K^T[2,12,64,10] → scores[2,12,10,10] / sqrt(64) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # 4. 因果掩码 # 数据流动：scores[2,12,10,10] + causal_mask → 上三角变为 -inf causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device)) # 下三角矩阵 causal_mask = causal_mask.view(1, 1, seq_len, seq_len) # 广播到 [1,1,10,10] scores = scores.masked_fill(causal_mask == 0, float('-inf')) # 5. Softmax + Dropout # 数据流动：scores → softmax → weights → dropout attention_weights = torch.softmax(scores, dim=-1) attention_weights = self.dropout(attention_weights) # 6. 加权求和 # 数据流动：weights[2,12,10,10] × V[2,12,10,64] → output[2,12,10,64] output = torch.matmul(attention_weights, V) # 7. 拼接多头 + 输出投影 # 数据流动：output[2,12,10,64] → transpose → view → wo → [2,10,768] output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.wo(output) """Transformer 解码器块（Pre-Norm 架构） 参数: d_model: 模型维度，示例：768 n_heads: 注意力头数，示例：12 d_ff: 前馈网络维度，示例：3072 max_seq_len: 最大序列长度，示例：1024 device: 计算设备 dtype: 数据类型 返回: 输出张量，形状为 [batch_size, seq_len, d_model] 示例: >>> block = TransformerBlock(d_model=768, n_heads=12, d_ff=3072, max_seq_len=1024) >>> x = torch.randn(2, 10, 768) >>> output = block(x) >>> output.shape torch.Size([2, 10, 768]) """ class TransformerBlock(nn.Module): """初始化 Transformer 块 参数: d_model: 模型维度 n_heads: 头数 d_ff: FFN 维度 max_seq_len: 最大序列长度 device: 计算设备 dtype: 数据类型 返回: 无 示例: block = TransformerBlock(d_model=768, n_heads=12, d_ff=3072, max_seq_len=1024) """ def __init__(self, d_model: int, n_heads: int, d_ff: int, max_seq_len: int, device=None, dtype=None): super().__init__() # 注意力子层：RMSNorm → Multi-Head Attention → 残差连接 self.attn_norm = RMSNorm(d_model, device=device, dtype=dtype) # 预归一化 self.attn = MultiHeadAttention(d_model, n_heads, max_seq_len, device=device, dtype=dtype) # 前馈子层：RMSNorm → SwiGLU → 残差连接 self.ffn_norm = RMSNorm(d_model, device=device, dtype=dtype) # 预归一化 self.ffn = SwiGLU(d_model, d_ff, device=device, dtype=dtype) """前向传播计算 Transformer 块 参数: x: 输入张量 [batch, seq_len, d_model] 返回: 输出张量 [batch, seq_len, d_model] 示例: output = block(x) """ def forward(self, x: torch.Tensor) -> torch.Tensor: # 1. 注意力子层：Pre-Norm + 自注意力 + 残差 # 数据流动：x[2,10,768] → attn_norm → attn → +x → [2,10,768] x = x + self.attn(self.attn_norm(x)) # 2. 前馈子层：Pre-Norm + SwiGLU + 残差 # 数据流动：x[2,10,768] → ffn_norm → ffn → +x → [2,10,768] x = x + self.ffn(self.ffn_norm(x)) return x """完整的 Transformer 语言模型（Decoder-Only） 参数: vocab_size: 词汇表大小，示例：32000 d_model: 模型维度，示例：768 n_heads: 注意力头数，示例：12 n_layers: 层数，示例：12 d_ff: 前馈网络维度，示例：3072 max_seq_len: 最大序列长度，示例：1024 device: 计算设备 dtype: 数据类型 返回: logits 输出，形状为 [batch_size, seq_len, vocab_size] 示例: >>> model = TransformerLM(vocab_size=32000, d_model=768, n_heads=12, n_layers=12, d_ff=3072, max_seq_len=1024) >>> token_ids = torch.randint(0, 32000, (2, 10)) >>> logits = model(token_ids) >>> logits.shape torch.Size([2, 10, 32000]) """ class TransformerLM(nn.Module): """初始化完整的 Transformer 语言模型 参数: vocab_size: 词汇表大小 d_model: 模型维度 n_heads: 头数 n_layers: 层数 d_ff: FFN 维度 max_seq_len: 最大序列长度 device: 计算设备 dtype: 数据类型 返回: 无 示例: model = TransformerLM(vocab_size=32000, d_model=768, n_heads=12, n_layers=12, d_ff=3072, max_seq_len=1024) """ def __init__(self, vocab_size: int, d_model: int, n_heads: int, n_layers: int, d_ff: int, max_seq_len: int, device=None, dtype=None): super().__init__() self.d_model = d_model # 词嵌入层 self.embedding = Embedding(vocab_size, d_model, device=device, dtype=dtype) # Transformer 块堆叠 self.layers = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff, max_seq_len, device=device, dtype=dtype) for _ in range(n_layers)]) # 最终归一化层 self.final_norm = RMSNorm(d_model, device=device, dtype=dtype) # 输出投影层 self.lm_head = Linear(d_model, vocab_size, device=device, dtype=dtype) """前向传播计算语言模型 参数: token_ids: 输入 token ID 张量 [batch, seq_len] 返回: logits 输出 [batch, seq_len, vocab_size] 示例: logits = model(token_ids) """ def forward(self, token_ids: torch.Tensor) -> torch.Tensor: # 1. 词嵌入 # 数据流动：token_ids[2,10] → embedding → x[2,10,768] x = self.embedding(token_ids) # 2. Transformer 块堆叠 # 数据流动：x[2,10,768] → layer1 → layer2 → ... → layer12 → [2,10,768] for layer in self.layers: x = layer(x) # 3. 最终归一化 # 数据流动：x[2,10,768] → final_norm → [2,10,768] x = self.final_norm(x) # 4. 输出 logits # 数据流动：x[2,10,768] → lm_head → logits[2,10,32000] logits = self.lm_head(x) return logits """测试 Transformer 语言模型 参数: 无 返回: logits: 模型输出 [batch_size, seq_len, vocab_size] 示例: >>> logits = test_transformer() """ def test_transformer(): # 设置随机种子 torch.manual_seed(42) # 模型配置 vocab_size = 1000 # 词汇表大小 d_model = 256 # 模型维度 n_heads = 8 # 注意力头数 n_layers = 4 # 层数 d_ff = 1024 # FFN 维度 max_seq_len = 512 # 最大序列长度 batch_size = 2 # 批次大小 seq_len = 20 # 序列长度 # 创建模型 model = TransformerLM(vocab_size=vocab_size, d_model=d_model, n_heads=n_heads, n_layers=n_layers, d_ff=d_ff, max_seq_len=max_seq_len) # 生成随机输入 token_ids = torch.randint(0, vocab_size, (batch_size, seq_len)) # [2, 20] # 前向传播 logits = model(token_ids) # [2, 20, 1000] # 打印信息 print("=" * 70) print("Transformer 语言模型测试") print("=" * 70) print(f"词汇表大小: {vocab_size}") print(f"模型维度: {d_model}") print(f"注意力头数: {n_heads}") print(f"层数: {n_layers}") print(f"FFN 维度: {d_ff}") print(f"输入形状: {token_ids.shape}") print(f"输出形状: {logits.shape}") print(f"总参数量: {sum(p.numel() for p in model.parameters()):,}") print("=" * 70) return logits if __name__ == "__main__": # 运行测试 logits = test_transformer() ``` #### 2.1 运行结果示例 执行上面的脚本后，会得到类似下面的输出： ``` Transformer 语言模型测试 ====================================================================== 词汇表大小: 1000 模型维度: 256 注意力头数: 8 层数: 4 FFN 维度: 1024 输入形状: torch.Size([2, 20]) 输出形状: torch.Size([2, 20, 1000]) 总参数量: 4,708,608 ====================================================================== ``` 可以看到，输入是 `[2, 20]` 的两个样本，每个样本20个Token；输出是 `[2, 20, 1000]`，每个位置对应词汇表里1000个词的预测分数。整个模型大约有470万参数，用于测试和验证概念绰绰有余。 ### 3. 总结 到这里，斯坦福CS336作业一关于Transformer语言模型的核心内容就完成了。关键理解在于：通过从零实现，我们能真正理解现代LLM为何采用Pre-Norm架构、为何钟情于RoPE和RMSNorm、以及为什么SwiGLU能取代ReLU。 | 组件 | 作用 | 关键技术 | | :--- | :--- | :--- | | Embedding | Token IDs → 稠密向量 | 查表操作 | | RoPE | 注入位置信息 | 旋转矩阵 | | RMSNorm | 归一化 | 均方根归一化 | | SwiGLU | 非线性变换 | 门控线性单元 | | Multi-Head Attention | 捕捉 token 依赖 | 缩放点积注意力 | | Transformer Block | 核心计算单元 | Pre-Norm + 残差 |