""" ============================================================================= 第 3 讲 · 代码 3：注意力 —— 缩放点积、因果掩码、多头、GQA/MQA、滑动窗口 ============================================================================= 对应讲义：lectures/lecture03_architectures_hyperparameters.md 第 6、9 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 注意力是 Transformer 的核心：每个位置根据"我在找什么"（query）去和所有位置的"我是什么"（key） 比较相似度，再按相似度加权汇总它们的"内容"（value）。 这个文件从零实现 CS336 作业 1 要求的注意力相关函数，并回答几个"为什么"： - 为什么点积要除以 √d_k？ - 语言模型为什么要"因果掩码"？怎么验证它真的看不到未来？ - 多头注意力是怎么把一个大向量拆成多个"头"的？ - GQA / MQA 为什么能大幅减小推理时的 KV 缓存？ - 滑动窗口注意力省了什么？ ----------------------------------------------------------------------------- 【整体思路】 ----------------------------------------------------------------------------- 第 1 部分：数值稳定的 softmax；缩放点积注意力（带掩码） 第 2 部分：为什么要除以 √d_k —— 实验 第 3 部分：因果多头自注意力（带 RoPE），与 PyTorch 官方实现对比 第 4 部分：GQA / MQA 与 KV 缓存大小 第 5 部分：滑动窗口注意力；【边界】整行都被掩掉时会得到 NaN ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 03_attention.py 依赖同目录下的 02_rope.py（用其中的 RoPE）。 """ import importlib import math import torch import torch.nn.functional as F from einops import einsum, rearrange, repeat # 【Python 小知识】文件名以数字开头，不能直接 import，用 importlib 按字符串导入（见第 1 讲 05 号文件） _rope_mod = importlib.import_module("02_rope") RotaryPositionalEmbedding = _rope_mod.RotaryPositionalEmbedding _blocks = importlib.import_module("01_building_blocks") Linear = _blocks.Linear def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：softmax 与缩放点积注意力 # ============================================================================= def softmax(x, dim=-1): """ 数值稳定的 softmax（CS336 作业 1 要求实现）。 【思路】 softmax(x)_i = exp(x_i) / Σ_j exp(x_j)。 先减去最大值：exp(x_i - m) / Σ_j exp(x_j - m)，结果不变，但所有指数 ≤ 0，不会上溢。 【边界】 - 如果某一行全是 -inf（整行都被掩掉），最大值也是 -inf，x - m 得到 nan。见第 5 部分。 """ m = x.max(dim=dim, keepdim=True).values e = torch.exp(x - m) return e / e.sum(dim=dim, keepdim=True) def scaled_dot_product_attention(Q, K, V, mask=None): """ 缩放点积注意力（CS336 作业 1 的接口）： Attention(Q, K, V) = softmax( Q K^T / √d_k [+ 掩码] ) V 【形状】 Q : (..., n, d_k) n 个 query K : (..., m, d_k) m 个 key V : (..., m, d_v) m 个 value mask : (n, m) 的布尔张量（可广播），True 表示"允许关注"，False 表示"禁止关注" 返回 : (..., n, d_v) 【思路】 1. 相似度：每个 query 与每个 key 做点积，得到 (n, m) 的分数矩阵，除以 √d_k； 2. 掩码：把禁止关注的位置设成 -inf（softmax 后变成 0）； 3. 对每个 query，在 m 个 key 上做 softmax，得到"注意力权重"（每行和为 1）； 4. 用权重对 value 加权求和。 【边界】 - mask 中 True/False 的含义在不同框架里可能相反（例如有的实现里 True 表示"掩掉"）， 使用别人的代码时一定要确认。 """ d_k = Q.shape[-1] scores = einsum(Q, K, "... n d, ... m d -> ... n m") / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(~mask, float("-inf")) weights = softmax(scores, dim=-1) return einsum(weights, V, "... n m, ... m d -> ... n d") def causal_mask(n, device=None): """ 因果掩码：位置 i 只能关注位置 0 … i（下三角，含对角线）。 【返回】(n, n) 的布尔张量，mask[i, j] = (j <= i) """ return torch.tril(torch.ones(n, n, dtype=torch.bool, device=device)) def demo_softmax_and_sdpa(): """演示 1：softmax 与注意力的基本检查。""" print_header("演示 1：数值稳定的 softmax 与缩放点积注意力") x = torch.tensor([[1000.0, 1001.0, 1002.0]]) naive = torch.exp(x) / torch.exp(x).sum(-1, keepdim=True) print(f" logits [1000, 1001, 1002]：朴素 softmax = {naive.tolist()}，稳定 softmax = " f"{[round(v, 4) for v in softmax(x)[0].tolist()]}") torch.manual_seed(0) Q, K, V = torch.randn(2, 5, 16), torch.randn(2, 7, 16), torch.randn(2, 7, 32) out = scaled_dot_product_attention(Q, K, V) ref = F.scaled_dot_product_attention(Q, K, V) print(f" Q(2,5,16), K(2,7,16), V(2,7,32) → 输出 {tuple(out.shape)}；" f"与 PyTorch 官方实现的最大差异 {(out - ref).abs().max().item():.2e}") # ============================================================================= # 第 2 部分：为什么要除以 √d_k # ============================================================================= def demo_why_scale(): """演示 2：不缩放时，softmax 会饱和。""" print_header("演示 2：为什么要除以 √d_k") torch.manual_seed(0) print(f" {'d_k':>6}{'点积的标准差':>14}{'不缩放：最大权重':>18}{'缩放后：最大权重':>18}") for d_k in [16, 64, 256, 1024]: q = torch.randn(1000, d_k) k = torch.randn(1000, d_k) dots = (q * k).sum(-1) # 每个 query 对 32 个 key 做 softmax，看最大的那个权重平均有多大 scores = torch.randn(200, 32, d_k) @ torch.randn(200, d_k, 1) scores = scores.squeeze(-1) p_raw = torch.softmax(scores, -1).max(-1).values.mean().item() p_scaled = torch.softmax(scores / math.sqrt(d_k), -1).max(-1).values.mean().item() print(f" {d_k:>6}{dots.std().item():>14.2f}{p_raw:>18.3f}{p_scaled:>18.3f}") print(""" q、k 的每个分量都是方差为 1 的随机数时，点积是 d_k 个乘积之和，方差为 d_k（标准差 √d_k）。 d_k 越大，分数的差距越大，softmax 越接近"只选最大的那一个"（权重接近 1，其余接近 0）： - 注意力退化成"硬选择"，丢失了加权平均的能力； - softmax 饱和区的梯度几乎为 0，训练很难进行。 除以 √d_k 把点积的方差拉回 1，不管 d_k 多大，softmax 都工作在正常区间。""") # ============================================================================= # 第 3 部分：因果多头自注意力 # ============================================================================= class MultiHeadSelfAttention(torch.nn.Module): """ 因果多头自注意力，可选 RoPE（CS336 作业 1 的 multihead_self_attention）。 【思路】 1. 用三个线性层把输入投影成 Q、K、V，每个都是 (batch, seq, d_model)； 2. 把 d_model 拆成 num_heads 个头，每个头的维度 d_head = d_model / num_heads； 3. 对每个头的 Q、K 施加 RoPE（同一个位置在每个头里的旋转方式相同）； 4. 每个头独立地做因果缩放点积注意力； 5. 把各头的输出拼回 d_model 维，再经过输出投影 W_O。 【为什么要多头】 一个头只能产生一种"注意力模式"（每个 query 一组权重）。多个头可以同时关注不同的东西： 有的头关注前一个词，有的头关注句子的主语，有的头关注匹配的括号…… 而且总计算量与单头（d_head = d_model）相同。 【参数】 d_model, num_heads : d_model 必须能被 num_heads 整除 max_seq_len, theta : 给 RoPE 用；theta=None 时不使用位置编码 【形状】 x : (batch, seq, d_model) → 同形状 token_positions : (batch, seq) 或 (seq,)，默认 0, 1, …, seq-1 """ def __init__(self, d_model, num_heads, max_seq_len=2048, theta=10000.0): super().__init__() if d_model % num_heads != 0: raise ValueError("d_model 必须能被 num_heads 整除") self.num_heads = num_heads self.d_head = d_model // num_heads self.q_proj = Linear(d_model, d_model) self.k_proj = Linear(d_model, d_model) self.v_proj = Linear(d_model, d_model) self.output_proj = Linear(d_model, d_model) self.rope = RotaryPositionalEmbedding(theta, self.d_head, max_seq_len) if theta else None def forward(self, x, token_positions=None): seq = x.shape[-2] q = rearrange(self.q_proj(x), "... s (h d) -> ... h s d", h=self.num_heads) k = rearrange(self.k_proj(x), "... s (h d) -> ... h s d", h=self.num_heads) v = rearrange(self.v_proj(x), "... s (h d) -> ... h s d", h=self.num_heads) if self.rope is not None: if token_positions is None: token_positions = torch.arange(seq, device=x.device) # 位置的形状 (..., seq) 要在"头"这一维上广播：插入一个大小为 1 的维度 pos = token_positions.unsqueeze(-2) if token_positions.dim() > 1 else token_positions q = self.rope(q, pos) k = self.rope(k, pos) out = scaled_dot_product_attention(q, k, v, causal_mask(seq, x.device)) return self.output_proj(rearrange(out, "... h s d -> ... s (h d)")) def demo_causal_mha(): """演示 3：因果多头注意力 —— 验证"看不到未来"，并与官方实现对比。""" print_header("演示 3：因果多头自注意力") torch.manual_seed(0) d_model, heads, seq = 64, 4, 10 mha = MultiHeadSelfAttention(d_model, heads, max_seq_len=64, theta=None) # 先不用 RoPE，便于和官方对比 x = torch.randn(2, seq, d_model) out = mha(x) # 与 PyTorch 官方的因果注意力对比 q = rearrange(mha.q_proj(x), "b s (h d) -> b h s d", h=heads) k = rearrange(mha.k_proj(x), "b s (h d) -> b h s d", h=heads) v = rearrange(mha.v_proj(x), "b s (h d) -> b h s d", h=heads) ref = mha.output_proj(rearrange(F.scaled_dot_product_attention(q, k, v, is_causal=True), "b h s d -> b s (h d)")) print(f" 输出形状 {tuple(out.shape)}；与 F.scaled_dot_product_attention(is_causal=True) 的最大差异 " f"{(out - ref).abs().max().item():.2e}") # 修改最后一个 token，前面位置的输出不应该改变 x2 = x.clone() x2[:, -1] += 10.0 out2 = mha(x2) changed = (out2 - out).abs().amax(dim=(0, 2)) print(f" 只修改第 {seq - 1} 个 token 后，各位置输出的变化量：") print(" " + ", ".join(f"{c:.1e}" for c in changed.tolist())) print(" → 只有最后一个位置变了：位置 i 看不到 i 之后的 token（因果性）。") print(""" 为什么语言模型需要因果掩码？训练时我们让每个位置预测"下一个 token"（第 2 讲 get_batch）， 如果位置 i 能看到位置 i+1，它直接"抄答案"就行了，什么也学不到。 生成时模型本来就只有过去的 token，训练时也必须和生成时一样。""") mha_rope = MultiHeadSelfAttention(d_model, heads, max_seq_len=64, theta=10000.0) print(f" 带 RoPE 的版本：输出形状 {tuple(mha_rope(x).shape)}（RoPE 没有可训练参数）") # ============================================================================= # 第 4 部分：GQA / MQA # ============================================================================= class GroupedQueryAttention(torch.nn.Module): """ 分组查询注意力 GQA（Ainslie et al. 2023；Llama 2 70B、Llama 3、Mistral、Qwen 等使用）。 【思路】 多头注意力（MHA）中，每个查询头都有自己的一组 K、V 头。 GQA 让 num_heads 个查询头**分组共享** num_kv_heads 组 K、V： num_kv_heads = num_heads → 就是 MHA num_kv_heads = 1 → 就是多查询注意力 MQA（Shazeer 2019） 介于两者之间 → GQA 【为什么】 训练时 K、V 的大小影响不大；但**推理生成**时，每生成一个 token 都要读取之前所有 token 的 K、V （KV 缓存，第 10 讲）。解码阶段是内存受限的：搬运 KV 缓存的时间决定了速度。 K、V 头数减少 8 倍，KV 缓存就小 8 倍，同样的显存能放下更长的上下文或更大的批次。 质量损失：MQA 有小幅损失，GQA 几乎没有（讲义第 9.2 节）。 【形状】与 MultiHeadSelfAttention 相同（这里为了简洁不加 RoPE）。 """ def __init__(self, d_model, num_heads, num_kv_heads): super().__init__() if num_heads % num_kv_heads != 0: raise ValueError("num_heads 必须是 num_kv_heads 的整数倍") self.h, self.kv = num_heads, num_kv_heads self.d_head = d_model // num_heads self.q_proj = Linear(d_model, num_heads * self.d_head) self.k_proj = Linear(d_model, num_kv_heads * self.d_head) # K、V 投影变小了 self.v_proj = Linear(d_model, num_kv_heads * self.d_head) self.output_proj = Linear(num_heads * self.d_head, d_model) def forward(self, x): seq = x.shape[-2] q = rearrange(self.q_proj(x), "... s (h d) -> ... h s d", h=self.h) k = rearrange(self.k_proj(x), "... s (g d) -> ... g s d", g=self.kv) v = rearrange(self.v_proj(x), "... s (g d) -> ... g s d", g=self.kv) # 把每组 K、V 复制给组内的 h/kv 个查询头（真实的高效实现不会真的复制，这里为了清楚） k = repeat(k, "... g s d -> ... (g r) s d", r=self.h // self.kv) v = repeat(v, "... g s d -> ... (g r) s d", r=self.h // self.kv) out = scaled_dot_product_attention(q, k, v, causal_mask(seq, x.device)) return self.output_proj(rearrange(out, "... h s d -> ... s (h d)")) def kv_cache_bytes_per_token(num_layers, num_kv_heads, d_head, bytes_per_value=2): """ 每个 token 的 KV 缓存大小（字节）= 2（K 和 V）× 层数 × KV 头数 × 头维度 × 每个数的字节数。 """ return 2 * num_layers * num_kv_heads * d_head * bytes_per_value def demo_gqa(): """演示 4：GQA / MQA 的参数量与 KV 缓存。""" print_header("演示 4：GQA / MQA —— 减小 KV 缓存") torch.manual_seed(0) d_model, heads = 512, 8 x = torch.randn(2, 16, d_model) for kv in [8, 2, 1]: attn = GroupedQueryAttention(d_model, heads, kv) n = sum(p.numel() for p in attn.parameters()) name = "MHA" if kv == heads else ("MQA" if kv == 1 else "GQA") print(f" {name}（{heads} 个查询头，{kv} 个 KV 头）：输出 {tuple(attn(x).shape)}，参数 {n:,}") print("\n 真实模型每个 token 的 KV 缓存（BF16）：") models = [ # 名字, 层数, 查询头数, KV 头数, 头维度 ("Llama 2 7B（MHA）", 32, 32, 32, 128), ("Llama 2 70B（GQA）", 80, 64, 8, 128), ("假如 Llama 2 70B 用 MHA", 80, 64, 64, 128), ("Llama 3 8B（GQA）", 32, 32, 8, 128), ] for name, L, h, kv, dh in models: b = kv_cache_bytes_per_token(L, kv, dh) print(f" {name:<24} 每 token {b / 1024:>7,.0f} KB；上下文 32K 时一条序列 {b * 32768 / 1e9:>6.1f} GB") print(""" Llama 2 70B 用 8 个 KV 头代替 64 个，KV 缓存缩小到 1/8。 DeepSeek-V2/V3 的多头潜在注意力（MLA）更进一步：只缓存一个压缩后的低维"潜在向量"（讲义第 9.3 节）。""") # ============================================================================= # 第 5 部分：滑动窗口注意力 & 边界 # ============================================================================= def sliding_window_causal_mask(n, window): """ 滑动窗口 + 因果掩码：位置 i 只能关注 max(0, i - window + 1) … i 这 window 个位置。 【思路】 完整注意力的计算量和 KV 缓存都随序列长度 n 线性增长（每个 query 看 n 个 key，总共 n² 次）。 滑动窗口把每个 query 能看的范围限制在最近 window 个 token，总计算量变成 n × window。 信息仍然能"接力"传得更远：堆叠 L 层后，理论上的感受野约为 L × window。 """ i = torch.arange(n)[:, None] j = torch.arange(n)[None, :] return (j <= i) & (j > i - window) def demo_sliding_window_and_nan(): """演示 5：滑动窗口掩码，以及整行被掩掉时的 NaN。""" print_header("演示 5：滑动窗口注意力；【边界】整行被掩掉") m = sliding_window_causal_mask(8, 3) print(" n = 8，window = 3 的掩码（1 = 可以关注）：") for row in m.int().tolist(): print(" " + " ".join(str(v) for v in row)) n, w = 32768, 4096 print(f" n = {n}、window = {w} 时，需要计算的分数个数：完整因果 ≈ {n * (n + 1) // 2:,}，" f"滑动窗口 ≈ {n * w - w * (w - 1) // 2:,}") print(" 现代模型常常交替使用：大部分层用滑动窗口（局部），每隔几层用一次完整注意力（全局）。") # 边界：某一行全部被掩掉 Q, K, V = torch.randn(1, 2, 4), torch.randn(1, 3, 4), torch.randn(1, 3, 4) bad_mask = torch.tensor([[True, True, False], [False, False, False]]) # 第 2 个 query 什么也不能看 out = scaled_dot_product_attention(Q, K, V, bad_mask) print(f"\n 【边界】第 2 个 query 的掩码全为 False → 输出 {out[0, 1].tolist()}") print(" 整行都是 -inf 时，softmax 计算 (-inf) - (-inf) = nan。因果掩码不会出现这种情况（对角线总是可见），") print(" 但如果你为批次中长短不一的序列加了填充掩码（padding mask），就要小心处理全被掩掉的行。") def main(): demo_softmax_and_sdpa() demo_why_scale() demo_causal_mha() demo_gqa() demo_sliding_window_and_nan() print("\n全部演示运行完毕。") if __name__ == "__main__": main()