""" transformer.py —— 第 2 讲：从零实现 Transformer（注意力、掩码、位置编码、三种结构） ===================================================================== 【这个文件要回答的问题】 第 1 讲的语言模型里，"神经网络"是一个只看最近 5 个词的固定窗口 MLP；前置讲 P2 换成了 RNN，能看完整的历史，但要一个词一个词地串行计算。Transformer 换了一个思路： **每个位置直接去"查"前文里任意一个位置，查的方式是可学习的。** 这个"查"就是注意力（attention）。本文件把它从最小的零件开始拼起来： scaled_dot_product_attention 单头注意力：打分 → 缩放 → 掩码 → softmax → 加权求和 MultiHeadAttention 多头：把 d 维切成 h 份，各查各的，再拼回来 FeedForward 逐位置的两层 MLP（注意力负责"搬运信息"，它负责"加工"） TransformerBlock 注意力 + MLP，各自带残差连接和 LayerNorm sinusoidal_positions 原版 Transformer 的正弦位置编码 causal_mask / prefix_mask 掩码：决定"谁能看见谁"——三种结构的区别几乎全在这里 TransformerLM decoder-only（因果）/ encoder-only（双向）/ 前缀 LM，共用一套代码 EncoderDecoder encoder-decoder（T5 / 原版 Transformer），多了交叉注意力 【怎样使用】 from transformer import TransformerLM model = TransformerLM(vocab_size=100, d_model=64, n_heads=4, n_layers=2, d_ff=256, attn="causal") logits = model(ids) # ids: (B, T) 整数张量 → logits: (B, T, V) logits, attn = model(ids, return_attn=True) # attn: 每层一个 (B, h, T, T) 的注意力权重 第 2 讲的四个脚本都从这里导入；第 4 讲会在此基础上换零件（激活函数、归一化、位置编码、GQA）。 【约定】 - 掩码一律用布尔张量，True 表示"允许看"，False 表示"禁止看"。 （PyTorch 的 nn.MultiheadAttention 用的是相反的约定：True 表示"屏蔽"。这是一个常见的 bug 来源， 01 号脚本里和它核对时会专门处理。） - 形状记号：B = 批大小，T = 序列长度，d = d_model，h = 头数，hd = 每个头的维度 = d / h，V = 词表大小。 【LLM 联系】 GPT 系列、Llama、Qwen、DeepSeek 都是 decoder-only；BERT 是 encoder-only；T5、原版 Transformer 是 encoder-decoder。它们的"积木"完全相同，差别只在 (1) 掩码，(2) 有没有交叉注意力，(3) 训练目标（04 号脚本）。 """ import math import torch import torch.nn as nn import torch.nn.functional as F # ===================================================================== # 1. 注意力的核心：一个函数 # ===================================================================== def scaled_dot_product_attention(q, k, v, mask=None): """ 缩放点积注意力（Vaswani et al. 2017 的公式 1）：softmax(Q Kᵀ / √d_k) V。 【思路】 把它想成"软的字典查询"： - 每个位置拿着一个"问题"向量 q（query）； - 每个位置挂着一个"标签"向量 k（key）和一个"内容"向量 v（value）； - q 和每个 k 做点积得到"匹配分数"；分数过 softmax 变成和为 1 的权重； - 输出 = 各个 v 按权重的加权平均。 普通字典是"硬"查询（只取完全匹配的那一项）；注意力是"软"查询（按相似度混合所有项）， 因为处处可导，所以"该查什么、该挂什么标签"都能用梯度下降学出来。 【为什么除以 √d_k】 若 q、k 的每个分量独立、均值 0、方差 1，则 q·k 是 d_k 项之和，方差为 d_k，标准差 √d_k。 d_k 越大，分数的数值越大，softmax 越接近 one-hot（"饱和"），梯度几乎为 0、很难训练。 除以 √d_k 把分数的标准差拉回 1，与维度无关。01 号脚本会实测这一点。 【参数】 q：(..., Tq, d_k)；k：(..., Tk, d_k)；v：(..., Tk, d_v)。"..." 可以是 (B,) 或 (B, h) 等任意前缀维度。 mask：可选的布尔张量，能广播到 (..., Tq, Tk)；True = 允许 query i 看 key j。 【返回】(out, weights)：out 形状 (..., Tq, d_v)；weights 形状 (..., Tq, Tk)，每行和为 1。 【边界】 如果某一行的 mask 全是 False（这个 query 一个都不许看），softmax 会得到 0/0 = NaN。 这种行在实践中只出现在"补齐（padding）的位置"上，它们的输出反正会被损失掩码丢掉； 这里把 NaN 换成 0，避免 NaN 通过后续层传染到正常位置。 【Python 小知识】 k.transpose(-1, -2) 交换最后两个维度（矩阵转置）；@ 是矩阵乘法，对前缀维度自动"批量"进行。 masked_fill(条件, 值)：把条件为 True 的位置填成给定的值，这里是把禁止的位置填成 -inf， softmax 之后 e^{-inf} = 0，于是这些位置的权重恰好为 0。 """ d_k = q.shape[-1] scores = q @ k.transpose(-1, -2) / math.sqrt(d_k) # (..., Tq, Tk) if mask is not None: scores = scores.masked_fill(~mask, float("-inf")) weights = torch.softmax(scores, dim=-1) if mask is not None: weights = torch.nan_to_num(weights, nan=0.0) # 全被屏蔽的行：0/0 → 0 return weights @ v, weights # ===================================================================== # 2. 掩码：决定"谁能看见谁" # ===================================================================== def causal_mask(T, device=None): """ 因果掩码（下三角）：位置 i 只能看位置 0..i。 【返回】(T, T) 布尔张量，mask[i, j] = (j <= i)。 【例子】T = 4： 1 0 0 0 1 1 0 0 1 1 1 0 1 1 1 1 【LLM 联系】decoder-only 模型（GPT、Llama）的每一层都用这个掩码。正是它保证了 "预测第 i+1 个 token 时看不到答案"，于是一条长度为 T 的序列能**一次前向** 同时提供 T 个训练样本（teacher forcing），这是 Transformer 训练快的关键之一。 """ return torch.ones(T, T, dtype=torch.bool, device=device).tril() def prefix_mask(T, prefix_len, device=None): """ 前缀语言模型（prefix LM / non-causal decoder）的掩码：前缀内部互相全看见，前缀之后保持因果。 【参数】T：序列长度；prefix_len：整数，或形状 (B,) 的整数张量（每条序列的前缀长度可以不同）。 【返回】prefix_len 是整数时返回 (T, T)；是张量时返回 (B, 1, T, T)（中间的 1 用来广播到各个头）。 【例子】T = 5，prefix_len = 2： 1 1 0 0 0 ← 前缀里的位置也能看见前缀里"后面"的位置 1 1 0 0 0 1 1 1 0 0 1 1 1 1 0 1 1 1 1 1 【思路】mask = 因果掩码 OR (key 位置 j < prefix_len)。 【边界】prefix_len = 0 时退化成因果掩码；prefix_len >= T 时变成全 1（双向）。 """ causal = causal_mask(T, device) j = torch.arange(T, device=device) if isinstance(prefix_len, int): return causal | (j[None, :] < prefix_len) in_prefix = j[None, None, :] < prefix_len[:, None, None] # (B, 1, T)：key j 在前缀里 return (causal[None] | in_prefix)[:, None] # (B, 1, T, T) def padding_mask(pad): """ 由"哪些位置是补齐"得到 key 方向的掩码：任何 query 都不许看补齐的 key。 【参数】pad：(B, T) 布尔张量，True 表示该位置是补齐（padding）。 【返回】(B, 1, 1, T)，可以和 (T, T) 或 (B, 1, T, T) 的掩码做 & 运算。 【为什么只屏蔽 key、不屏蔽 query】补齐位置作为 query 算出来的输出是"垃圾"，但我们在损失里会把它们丢掉； 要紧的是不能让真实位置去"看"补齐位置——否则补多少会影响真实位置的结果。 """ return (~pad)[:, None, None, :] # ===================================================================== # 3. 多头注意力 # ===================================================================== class MultiHeadAttention(nn.Module): """ 多头注意力（Multi-Head Attention）。 【思路】 单头注意力的一行权重只能表达"一种"关注方式。一个词常常要同时关心好几件事： 比如物主代词 their 要找句首的主语（决定 my / her / their），也要看前一个词（from / with）。 多头 = 把 d 维的 Q、K、V 各切成 h 份，每份 hd = d / h 维，h 个头各自独立做注意力， 最后把 h 个输出拼回 d 维，再过一个输出矩阵 W_O 混合。 **参数量和计算量与单头 d 维注意力几乎一样**（切分不增加参数），换来的是 h 种不同的关注模式。 【参数】d_model：模型维度 d；n_heads：头数 h（必须整除 d）；bias：线性层是否带偏置。 【前向的参数】 x：(B, Tq, d)，query 的来源。 kv：(B, Tk, d) 或 None。None 表示自注意力（K、V 也来自 x）； 不是 None 时是交叉注意力（encoder-decoder 里 decoder 去查 encoder 的输出）。 mask：能广播到 (B, h, Tq, Tk) 的布尔张量，True = 允许看。 【返回】(out, weights)：out (B, Tq, d)；weights (B, h, Tq, Tk)。 【Python 小知识】view(B, T, h, hd) 把最后一维 d 重新解释成 h × hd（不复制数据）； transpose(1, 2) 把头的维度换到前面，变成 (B, h, T, hd)，这样 h 个头就能作为"批"一起算。 拼回去时 transpose 之后内存不连续，所以要用 reshape（必要时会复制），不能用 view。 """ def __init__(self, d_model, n_heads, bias=True): super().__init__() if d_model % n_heads != 0: raise ValueError(f"d_model={d_model} 不能被 n_heads={n_heads} 整除") self.h = n_heads self.hd = d_model // n_heads self.q_proj = nn.Linear(d_model, d_model, bias=bias) self.k_proj = nn.Linear(d_model, d_model, bias=bias) self.v_proj = nn.Linear(d_model, d_model, bias=bias) self.o_proj = nn.Linear(d_model, d_model, bias=bias) def _split(self, t): B, T, _ = t.shape return t.view(B, T, self.h, self.hd).transpose(1, 2) # (B, h, T, hd) def forward(self, x, kv=None, mask=None): src = x if kv is None else kv q = self._split(self.q_proj(x)) k = self._split(self.k_proj(src)) v = self._split(self.v_proj(src)) out, weights = scaled_dot_product_attention(q, k, v, mask) # (B, h, Tq, hd) B, _, Tq, _ = out.shape out = out.transpose(1, 2).reshape(B, Tq, self.h * self.hd) # 拼回 (B, Tq, d) return self.o_proj(out), weights # ===================================================================== # 4. 前馈网络与 Transformer 块 # ===================================================================== class FeedForward(nn.Module): """ 逐位置前馈网络（position-wise FFN）：d → d_ff → d，中间是非线性激活。 【思路】注意力只做"加权平均"——它把信息从别的位置**搬**过来，但加权平均本身是线性的； FFN 对每个位置**单独**做一次非线性加工（不同位置共享同一组参数）。 一个常见的直觉：注意力负责"交流"，FFN 负责"思考 / 查记忆"。 原版 Transformer 用 ReLU、d_ff = 4d；GPT-2 用 GELU；Llama 用 SwiGLU（第 4 讲）。 【参数】d_model；d_ff：中间层宽度；activation："relu" 或 "gelu"。 """ def __init__(self, d_model, d_ff, activation="gelu", bias=True): super().__init__() self.up = nn.Linear(d_model, d_ff, bias=bias) self.down = nn.Linear(d_ff, d_model, bias=bias) self.act = {"relu": F.relu, "gelu": F.gelu}[activation] def forward(self, x): return self.down(self.act(self.up(x))) class TransformerBlock(nn.Module): """ 一个 Transformer 块：自注意力 → (可选) 交叉注意力 → FFN，每个子层都带残差连接和 LayerNorm。 【思路】 残差连接 x = x + 子层(x)：每个子层只需要学"在原来基础上改一点什么"， 而且梯度可以沿着"+ x"这条捷径直接流回前面的层（深网络能训练的关键，ResNet 的思想）。 LayerNorm：把每个位置的向量归一化到均值 0、方差 1（再乘可学习的缩放、加偏置）， 让各层输入的数值范围稳定。 归一化放在哪里有两种做法： post-norm（原版 2017）：x = LN(x + 子层(x)) pre-norm （GPT-2 起） ：x = x + 子层(LN(x)) 本讲默认 pre-norm（norm_first=True），它更容易训练；**第 4 讲**专门比较两者为什么不同。 【参数】 d_model, n_heads, d_ff：见上。 cross：True 时多一个交叉注意力子层（encoder-decoder 的 decoder 块）。 norm_first：True = pre-norm，False = post-norm。 【前向的参数】x：(B, T, d)；mask：自注意力掩码；memory：encoder 输出 (B, S, d)（cross=True 时必须给）； memory_mask：交叉注意力掩码（能广播到 (B, h, T, S)）。 【返回】(x, weights)：weights 是字典 {"self": (B,h,T,T), "cross": (B,h,T,S) 或不存在}。 """ def __init__(self, d_model, n_heads, d_ff, cross=False, norm_first=True, activation="gelu"): super().__init__() self.norm_first = norm_first self.self_attn = MultiHeadAttention(d_model, n_heads) self.n1 = nn.LayerNorm(d_model) self.cross_attn = MultiHeadAttention(d_model, n_heads) if cross else None self.n2 = nn.LayerNorm(d_model) if cross else None self.ffn = FeedForward(d_model, d_ff, activation) self.n3 = nn.LayerNorm(d_model) def _sublayer(self, x, norm, fn): """把"残差 + 归一化"的两种摆法写在一个地方。fn 接收一个张量，返回 (输出, 权重)。""" if self.norm_first: y, w = fn(norm(x)) return x + y, w y, w = fn(x) return norm(x + y), w def forward(self, x, mask=None, memory=None, memory_mask=None): weights = {} x, weights["self"] = self._sublayer(x, self.n1, lambda t: self.self_attn(t, mask=mask)) if self.cross_attn is not None: if memory is None: raise ValueError("交叉注意力块需要 encoder 的输出 memory") x, weights["cross"] = self._sublayer( x, self.n2, lambda t: self.cross_attn(t, kv=memory, mask=memory_mask)) x, _ = self._sublayer(x, self.n3, lambda t: (self.ffn(t), None)) return x, weights # ===================================================================== # 5. 位置编码 # ===================================================================== def sinusoidal_positions(T, d_model): """ 原版 Transformer 的正弦位置编码： PE[pos, 2i] = sin(pos / 10000^{2i/d}) PE[pos, 2i+1] = cos(pos / 10000^{2i/d}) 【为什么需要位置编码】注意力对输入的顺序"视而不见"：把输入的词打乱顺序， 每个词的输出只是跟着换了位置，数值不变（排列等变性，01 号脚本验证）。 可是 "she feeds the dog" 和 "the dog feeds she" 意思完全不同，所以必须把位置信息加进去。 【为什么用正弦】 (1) 不需要学习，任意长度都能算； (2) 每对维度 (2i, 2i+1) 是一个以不同频率旋转的"时钟指针"：低维转得快（分辨相邻位置）， 高维转得慢（分辨远距离）——很像二进制计数里低位变得快、高位变得慢； (3) PE[pos + k] 可以由 PE[pos] 经过一个只依赖 k 的线性变换（旋转）得到， 原论文认为这能让模型容易学会按"相对距离"关注。第 4 讲的 RoPE 正是把这个想法做彻底。 【参数】T：位置个数；d_model：必须是偶数。 【返回】(T, d_model) 的 float32 张量。 """ if d_model % 2 != 0: raise ValueError("正弦位置编码要求 d_model 是偶数") pos = torch.arange(T, dtype=torch.float32)[:, None] # (T, 1) freq = torch.exp(-math.log(10000.0) * torch.arange(0, d_model, 2) / d_model) # (d/2,) pe = torch.zeros(T, d_model) pe[:, 0::2] = torch.sin(pos * freq) pe[:, 1::2] = torch.cos(pos * freq) return pe # ===================================================================== # 6. 三种结构 # ===================================================================== class TransformerLM(nn.Module): """ "一叠 Transformer 块"的语言模型。用 attn 参数切换三种用法： attn="causal" decoder-only（GPT）：因果掩码，用于下一个 token 预测（CLM） attn="bidirectional" encoder-only（BERT）：不加掩码，每个位置看见整条序列，用于掩码语言建模（MLM） attn="prefix" 前缀 LM（non-causal decoder）：前缀双向、之后因果，前向时要给 prefix_len 【参数】 vocab_size, d_model, n_heads, n_layers, d_ff：结构超参数。 max_len：最长序列（仅 pos="learned" 时用来建位置嵌入表；正弦编码会按需计算）。 pos："sinusoidal"、"learned" 或 "none"（不加位置信息，01 / 03 号脚本用来做对照）。 tie：输出层是否与输入嵌入共享权重（第 1 讲 §4.4）。 norm_first, activation：传给每个块。 【前向的参数】 ids：(B, T) 整数张量。 pad：可选，(B, T) 布尔张量，True = 补齐位置。 prefix_len：attn="prefix" 时必须给（整数或 (B,) 张量）。 return_attn：True 时额外返回每层的注意力权重列表 [(B, h, T, T), ...]。 【返回】logits (B, T, V)；或 (logits, attn_list)。 【边界】 - 嵌入乘以 √d：原版 Transformer 的做法，让嵌入和位置编码的数值大小相当 （nn.Embedding 默认按 N(0,1) 初始化，这里把嵌入初始化成 N(0, 1/d)，乘 √d 后约为 1）。 - pos="learned" 时 T 不能超过 max_len（没学过的位置没有嵌入可查）。 """ def __init__(self, vocab_size, d_model=64, n_heads=4, n_layers=2, d_ff=256, max_len=512, attn="causal", pos="sinusoidal", tie=True, norm_first=True, activation="gelu"): super().__init__() if attn not in ("causal", "bidirectional", "prefix"): raise ValueError(f"未知的 attn={attn}") self.attn_type, self.pos_type, self.d = attn, pos, d_model self.emb = nn.Embedding(vocab_size, d_model) nn.init.normal_(self.emb.weight, std=d_model ** -0.5) self.pos_emb = nn.Embedding(max_len, d_model) if pos == "learned" else None if self.pos_emb is not None: nn.init.normal_(self.pos_emb.weight, std=0.02 * d_model ** 0.5) self.blocks = nn.ModuleList( [TransformerBlock(d_model, n_heads, d_ff, norm_first=norm_first, activation=activation) for _ in range(n_layers)]) # pre-norm 的块输出没有归一化过，所以最后要补一个 LayerNorm（GPT-2 的 ln_f） self.norm_f = nn.LayerNorm(d_model) if norm_first else nn.Identity() self.head = None if tie else nn.Linear(d_model, vocab_size, bias=False) def embed(self, ids): """token 嵌入 × √d + 位置编码。【返回】(B, T, d)。""" T = ids.shape[1] x = self.emb(ids) * math.sqrt(self.d) if self.pos_type == "sinusoidal": x = x + sinusoidal_positions(T, self.d).to(x.device) elif self.pos_type == "learned": x = x + self.pos_emb(torch.arange(T, device=ids.device)) return x def make_mask(self, T, pad=None, prefix_len=None, device=None): """按 attn 类型（和补齐情况）组装注意力掩码；返回 None 表示"全都可以看"。""" if self.attn_type == "causal": mask = causal_mask(T, device) elif self.attn_type == "prefix": if prefix_len is None: raise ValueError("attn='prefix' 需要 prefix_len") mask = prefix_mask(T, prefix_len, device) else: mask = None if pad is not None: pm = padding_mask(pad) mask = pm if mask is None else (mask & pm) return mask def hidden_states(self, ids, pad=None, prefix_len=None, return_attn=False): """最后一层（归一化后）的隐藏状态 (B, T, d)；第 5 讲会拿它当"上下文嵌入"用。""" x = self.embed(ids) mask = self.make_mask(ids.shape[1], pad, prefix_len, ids.device) attn = [] for blk in self.blocks: x, w = blk(x, mask=mask) attn.append(w["self"]) x = self.norm_f(x) return (x, attn) if return_attn else x def forward(self, ids, pad=None, prefix_len=None, return_attn=False): x, attn = self.hidden_states(ids, pad, prefix_len, return_attn=True) logits = x @ self.emb.weight.T if self.head is None else self.head(x) return (logits, attn) if return_attn else logits class EncoderDecoder(nn.Module): """ encoder-decoder Transformer（原版 Transformer、T5）。 【思路】 encoder：一叠**双向**自注意力块，把输入（源序列）编码成一串向量 memory，每个位置都看见整条输入。 decoder：一叠块，每块有三个子层： (1) **因果**自注意力：看已经生成的目标 token； (2) **交叉注意力**：query 来自 decoder，key / value 来自 memory——"生成每个词时回头查输入"； (3) FFN。 这正是第 2 讲 slides 里 Bahdanau 注意力（2014）的直接后代：RNN encoder-decoder 只用最后一个 隐藏状态传递整句输入（信息瓶颈），注意力让 decoder 每一步都能直接访问 encoder 的每个位置。 【参数】vocab_size（源和目标共用一个词表、共用嵌入，T5 就是这样）；其余同 TransformerLM。 【前向的参数】src (B, S)、tgt_in (B, T)；src_pad / tgt_pad：可选的补齐标记 (B, S) / (B, T)。 【返回】logits (B, T, V)；return_attn=True 时再返回 {"enc": [...], "self": [...], "cross": [...]}。 """ def __init__(self, vocab_size, d_model=64, n_heads=4, n_enc=2, n_dec=2, d_ff=256, activation="gelu"): super().__init__() self.d = d_model self.emb = nn.Embedding(vocab_size, d_model) nn.init.normal_(self.emb.weight, std=d_model ** -0.5) self.enc = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff, activation=activation) for _ in range(n_enc)]) self.dec = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff, cross=True, activation=activation) for _ in range(n_dec)]) self.enc_norm = nn.LayerNorm(d_model) self.dec_norm = nn.LayerNorm(d_model) def _embed(self, ids): return self.emb(ids) * math.sqrt(self.d) + sinusoidal_positions(ids.shape[1], self.d).to(ids.device) def encode(self, src, src_pad=None, attn_log=None): """【返回】memory (B, S, d)。""" x = self._embed(src) mask = padding_mask(src_pad) if src_pad is not None else None for blk in self.enc: x, w = blk(x, mask=mask) if attn_log is not None: attn_log["enc"].append(w["self"]) return self.enc_norm(x) def decode(self, tgt_in, memory, src_pad=None, tgt_pad=None, attn_log=None): """【返回】logits (B, T, V)。""" x = self._embed(tgt_in) T = tgt_in.shape[1] mask = causal_mask(T, tgt_in.device) if tgt_pad is not None: mask = mask & padding_mask(tgt_pad) mem_mask = padding_mask(src_pad) if src_pad is not None else None for blk in self.dec: x, w = blk(x, mask=mask, memory=memory, memory_mask=mem_mask) if attn_log is not None: attn_log["self"].append(w["self"]) attn_log["cross"].append(w["cross"]) return self.dec_norm(x) @ self.emb.weight.T def forward(self, src, tgt_in, src_pad=None, tgt_pad=None, return_attn=False): log = {"enc": [], "self": [], "cross": []} if return_attn else None memory = self.encode(src, src_pad, log) logits = self.decode(tgt_in, memory, src_pad, tgt_pad, log) return (logits, log) if return_attn else logits def count_params(model): """可训练参数的总数。""" return sum(p.numel() for p in model.parameters() if p.requires_grad)