斯坦福 CS336 课程入门笔记

斯坦福 CS336 课程入门笔记

由于作者水平有限，在听课时很多部分都没有听懂，并且为自学，大部分代码实现都是作者手写完之后存在问题，扔给 AI （包括但不限于智谱、千问）进行完善和修改，代码正确性仍然存疑，并且跳过了很多部分。以及参考了一些前辈的文章。

笔记的整理思路为 CS336 所要求的 Assignment 五项代码实现进行，省略了许多对代码实现无直接帮助的概念与介绍。

一、 课程启点

为什么不用现成 API

调用 OpenAI 或 Anthropic 的 API 只是“使用”大模型，存在一些问题：

• 黑盒限制：当模型输出错误或有“幻觉”时，如果不了解内部结构，无从下手修复。
• 定制化与成本：在很多垂直领域（比如医疗、法律），或者边缘设备（如手机、机器人）上，我们需要训练特定的小模型（1B-7B），而不是依赖昂贵的云端大模型。
• 数据隐私：核心数据不能发往第三方 API，必须自己掌控从数据到权重的全流程。

---

二、 语言模型的数学本质

高中数学课本中课外学习拓展中提到，贝叶斯公式可以用于 AI 相关。也就是语言模型与概率论息息相关。在高三时模考的语文应用文卷子中经常看到一个介绍，“当前的语言模型生成言语本质是针对你当前字节对下一字节进行预测产生的”。所以很显然，我们要回归到概率学上面。

语言模型（LM）的本质是对自然语言序列的概率分布进行建模。给定一个由 \(T\) 个 token 组成的序列 \(x_1, x_2, ..., x_T\)，模型的目标是计算该序列的联合概率 \(P(x_1, ..., x_T)\)。

1. 自回归分解

利用概率论的链式法则，联合概率可以分解为条件概率的连乘：

\[P(x_1, x_2, ..., x_T) = \prod_{t=1}^{T} P(x_t | x_1, x_2, ..., x_{t-1}) \]

这就是“文字接龙”的数学表达。每次根据前文 \(x_{<t}\) 预测下一个词 \(x_t\)。

2. 训练目标：交叉熵损失

为了让模型学会这个概率分布，我们最大化真实数据的对数似然，这等价于最小化交叉熵损失：

\[\mathcal{L} = -\frac{1}{T} \sum_{t=1}^{T} \log P(x_t | x_{<t}; \theta) \]

其中 \(\theta\) 是模型的参数。由于直接计算离散概率不可导，模型输出的是 logits（未归一化的分数），通过 Softmax 转换为概率：

\[P(x_t = w | x_{<t}) = \frac{\exp(z_w / \tau)}{\sum_{k=1}^{V} \exp(z_k / \tau)} \]

• \(V\)：词表大小
• \(z_w\)：词 \(w\) 对应的 logit
• \(\tau\)：温度系数（控制分布的平滑度，训练时设为 1）

---

三、 语言模型系统的五大核心组件

我们将一个完整的 LLM 系统拆解为五个组件。

组件 1：数据

模型的能力上限由数据决定。互联网原始数据（如 Common Crawl）充斥着 SEO 垃圾、乱码和重复内容。

• 数据清洗流程：
  1. 抽取：从 HTML 中提取纯文本（如使用 resiliparse 或 trafilatura）。
  2. 过滤：使用启发式规则（如词语密度、长度）和轻量级分类器（如 fastText）剔除垃圾。这里有一个经典的代码逻辑——过滤掉那些没有实际内容的目录式网页。
  3. 去重：这是最关键的一步。使用 MinHash 和 LSH（局部敏感哈希）算法找出高度相似的文档并剔除。研究表明，去重可以显著提升模型训练的收敛速度并降低过拟合。
  4. 脱敏：使用正则表达式移除 PII（个人身份信息，如身份证号、邮箱）。
• 数据配比：代码、数学、多语言文本的比例需要精心调配。增加代码数据的比例被证明能显著提升模型的逻辑推理能力。

接下来是 MinHash 的代码实现

import random
import hashlib

class MinHash:
    def __init__(self, num_hashes=100):
        """
        初始化 MinHash 生成器
        :param num_hashes: 哈希函数的数量（决定了签名的长度和精度）
        """
        self.num_hashes = num_hashes
        # 生成 num_hashes 组随机的哈希种子 (a, b)，用于模拟不同的哈希函数
        # 公式: h(x) = (a * x + b) % MOD
        self.seeds = [(random.randint(0, 2**32 - 1), random.randint(0, 2**32 - 1)) 
                      for _ in range(num_hashes)]
        self.MOD = 2**32 - 1

    def _hash(self, token, seed):
        """将字符串 token 映射为一个整数哈希值"""
        a, b = seed
        # 使用内置 hash 获取基础值，然后进行线性变换
        base_hash = int(hashlib.md5(token.encode('utf-8')).hexdigest(), 16)
        return (a * base_hash + b) % self.MOD

    def generate_signature(self, text):
        """
        为一段文本生成 MinHash 签名（指纹）
        :param text: 输入文本
        :return: 长度为 num_hashes 的整数列表（即签名）
        """
        # 1. 简单的文本预处理：转小写并按空格分词
        tokens = set(text.lower().split())
        
        # 2. 对每一个哈希函数，寻找当前词表中的最小哈希值
        signature = []
        for seed in self.seeds:
            min_hash = float('inf')
            for token in tokens:
                h = self._hash(token, seed)
                if h < min_hash:
                    min_hash = h
            signature.append(min_hash)
            
        return signature

    @staticmethod
    def calculate_similarity(sig1, sig2):
        """
        计算两个 MinHash 签名的相似度（Jaccard 近似值）
        原理：两个签名在相同位置上值相等的比例，近似等于原文本的 Jaccard 相似度
        """
        if len(sig1) != len(sig2):
            raise ValueError("签名长度必须相同")
        matches = sum(1 for i in range(len(sig1)) if sig1[i] == sig2[i])
        return matches / len(sig1)

# ================= 测试代码 =================
if __name__ == "__main__":
    # 1. 准备测试文本
    doc1 = "the quick brown fox jumps over the lazy dog"
    doc2 = "the quick brown fox leaps over the lazy dog"  # 和 doc1 高度相似（只改了一个词）
    doc3 = "apples bananas oranges grapes watermelon"      # 和 doc1 完全不相关

    # 2. 初始化 MinHash 并生成指纹
    minhash = MinHash(num_hashes=100)
    
    sig1 = minhash.generate_signature(doc1)
    sig2 = minhash.generate_signature(doc2)
    sig3 = minhash.generate_signature(doc3)

    # 3. 计算相似度
    sim_1_2 = MinHash.calculate_similarity(sig1, sig2)
    sim_1_3 = MinHash.calculate_similarity(sig1, sig3)

    print(f"文档1与文档2的近似相似度: {sim_1_2:.2f}")  # 预期: 0.80 左右 (高度相似)
    print(f"文档1与文档3的近似相似度: {sim_1_3:.2f}")  # 预期: 0.00 左右 (毫无关联)

测试代码的输出结果为 0.82 和 0.00

组件 2：分词

计算机不认识字符，需要将文本映射为整数 ID。现代模型普遍采用 Byte-level BPE (Byte Pair Encoding)。

• 为什么用 Byte-level？：保证任何输入（包括 Emoji、罕见汉字、甚至乱码）都能被编码，彻底消除 <UNK> (Unknown Token) 问题。
• BPE 训练算法逻辑（伪代码）：

  # 初始状态：将所有文本拆分为单个 UTF-8 字节
  vocab = {bytes([i]) for i in range(256)}
  # 统计相邻字节对频率，不断合并最高频的对
  while len(vocab) < target_vocab_size:
      pair_freq = count_adjacent_pairs(corpus)
      best_pair = max(pair_freq, key=pair_freq.get)
      new_token = best_pair[0] + best_pair[1]
      vocab.add(new_token)
      corpus = merge_pair_in_corpus(corpus, best_pair, new_token)
  

我们来解释一下这个代码在干什么。

1. 初始化：准备 256 个基础“字母”

vocab = {bytes([i]) for i in range(256)}

• 在计算机里，所有的文字（无论是英文、中文还是表情符号）最终都可以被编码为 UTF-8 字节。这一步相当于把计算机底层的 256 个基础字节（0~255）全部收进我们的“基础词表”里。

2. 核心循环：不断“组词”，直到词表够大

while len(vocab) < target_vocab_size:

• 只要我们的词表大小还没达到设定的目标（比如 50000 个词），就一直重复下面的“组词”动作。

3. 统计频率：看看哪两个“字母”最爱黏在一起

pair_freq = count_adjacent_pairs(corpus)

• 去扫描整个训练语料库（corpus），统计所有相邻的两个字节（或已合并的词块）在一起出现了多少次。比如看看语料库里，t 和 h 挨在一起出现了 1000 次，而 t 和 x 只出现了 1 次。

4. 选出最频繁的组合

best_pair = max(pair_freq, key=pair_freq.get)

• 在刚才统计的结果里，找出出现频率最高的那一对。比如 t 和 h 出现次数最多，所以选择这俩。

5. 把组合变成一个新的“词”

new_token = best_pair[0] + best_pair[1]
vocab.add(new_token)

• 把 t 和 h 绑定成一个新单词 th，放进字典里。

6. 更新语料库：全局替换

corpus = merge_pair_in_corpus(corpus, best_pair, new_token)

• 回到语料库，把所有出现 t 和 h 挨着的地方，全部替换成新的词块 th。

---

总结它的运行过程

假设我们要处理文本 "low lower lowest"：

1. 初始：词表是 256 个单字节。文本被拆成 l o w l o w e r l o w e s t。
2. 第1轮：发现 l 和 o 挨着最多，合并成 lo。文本变成 lo w lo w e r lo w e s t。
3. 第2轮：发现 lo 和 w 挨着最多，合并成 low。文本变成 low low low e s t。
4. 第3轮：发现 e 和 s 挨着最多，合并成 es。文本变成 low low low es t。
5. 第4轮：发现 es 和 t 挨着最多，合并成 est。文本变成 low low low est。

最终，模型就学会了 low 和 est 这两个词

• 特殊 Token：如 <|endoftext|>（文档分隔符）。在指令微调时，还会引入 <|im_start|> 和 <|im_end|> 来区分 System、User 和 Assistant。

完整的代码实现如下（加入了测试）

from collections import defaultdict
"""" 
defaultdict 可以让代码实现更加快捷
如果字典是空的，它会将默认第一次出现的值赋为 0
"""
def count_adjacent_pairs(corpus):
    """统计语料库中所有相邻 token 对的出现频率"""
    pair_freq = defaultdict(int)
    for tokens in corpus:
        for i in range(len(tokens) - 1):
            pair_freq[(tokens[i], tokens[i + 1])] += 1
    return pair_freq


def merge_pair_in_corpus(corpus, best_pair, new_token):
    """将语料库中所有的 best_pair 替换为合并后的 new_token"""
    new_corpus = [] # 创建一个空列表，用来存放处理后的整个新语料库
    for tokens in corpus:
        new_tokens = [] # 为当前句子创建一个空列表，用来存放合并后的新句子
        i = 0
        while i < len(tokens):
            # 当前位置不是最后一个，并且当前 token 和下一个 token 正好是 best_pair
            if i < len(tokens) - 1 and (tokens[i], tokens[i + 1]) == best_pair:
                new_tokens.append(new_token)
                i += 2
            # 条件不成立 把当前 token 原封不动地加入新句子
            else:
                new_tokens.append(tokens[i])
                i += 1
        new_corpus.append(new_tokens)
    return new_corpus


def train_bpe(text, target_vocab_size):
    """完整的 BPE 训练流程"""
    # 初始状态：将所有文本拆分为单个 UTF-8 字节 (用整数 0-255 表示)
    # vocab 记录 token id 到 bytes 的映射，便于最终还原和打印

    # 初始化基础词表，包含所有单字节 (0-255)
    vocab = {i: bytes([i]) for i in range(256)}

    # 将输入文本编码为UTF-8字节列表
    corpus = [list(text.encode('utf-8'))]

    merges = {}  # 记录合并规则，格式为: {new_token_id: (id_a, id_b)}
    next_token_id = 256 # 新 Token 的 ID 从 256 开始
    while len(vocab) < target_vocab_size:
        pair_freq = count_adjacent_pairs(corpus)
        if not pair_freq:
            break # 如果没有相邻对了，提前结束

        # 找到频率最高的 pair
        best_pair = max(pair_freq, key=pair_freq.get)
        # 分配新的 token id，并记录其对应的字节表示
        new_token_id = next_token_id
        vocab[new_token_id] = vocab[best_pair[0]] + vocab[best_pair[1]]
        merges[new_token_id] = best_pair
        # 在语料库中执行合并
        corpus = merge_pair_in_corpus(corpus, best_pair, new_token_id)
        print(
            f"合并: {best_pair} -> 新Token ID: {new_token_id}, 频率: {pair_freq[best_pair]}, 当前词表大小: {len(vocab)}")
        next_token_id += 1
    return vocab, merges


def encode(text, merges):
    """使用训练好的 merges 规则对新文本进行编码 (基于贪心优先级)"""
    tokens = list(text.encode('utf-8'))
    if not merges: return tokens
    # 构建优先级表：pair 对应的 rank 越小，优先级越高（越早被学习）

    # 构建两个辅助字典，用于快速查询
    pair_to_rank = {}
    pair_to_new_token = {}
    for rank, (new_token, pair) in enumerate(merges.items()):
    # merges.items() 会返回一个包含所有键值对的可迭代对象
    # enumerate() 函数的作用是给一个可迭代对象的每个元素自动添加一个从 0 开始的序号（索引）
        pair_to_rank[pair] = rank
        pair_to_new_token[pair] = new_token
    # 持续合并，直到无法合并为止
    while len(tokens) > 1:
        min_rank = float('inf')
        min_rank_idx = -1
        # 扫描寻找当前序列中优先级最高（rank 最小）的可合并 pair
        for i in range(len(tokens) - 1):
            pair = (tokens[i], tokens[i + 1])
            if pair in pair_to_rank:
                if pair_to_rank[pair] < min_rank:
                    min_rank = pair_to_rank[pair]
                    min_rank_idx = i
        # 如果没有可以合并的 pair，结束循环
        if min_rank_idx == -1: break
        # 执行合并，替换序列中的这两个 token
        best_pair = (tokens[min_rank_idx], tokens[min_rank_idx + 1])
        new_token = pair_to_new_token[best_pair]
        tokens = tokens[:min_rank_idx] + [new_token] + tokens[min_rank_idx + 2:]
    return tokens


# ================= 测试代码 =================
if __name__ == "__main__":
    sample_text = "low lower lowest"
    # 目标词表大小设为 262，会生成 6 个新 token
    final_vocab, final_merges = train_bpe(sample_text, target_vocab_size=262)
    print("\n学习到的合并规则:")
    for token_id, pair in final_merges.items():
        # 将 ID 还原为字符打印，直观展示合并过程
        char_a = final_vocab[pair[0]].decode('utf-8', errors='replace')
        char_b = final_vocab[pair[1]].decode('utf-8', errors='replace')
        new_char = final_vocab[token_id].decode('utf-8', errors='replace')
        print(f"Token {token_id} ('{new_char}') 是由 {pair} ('{char_a}' + '{char_b}') 合并而来的")
    # 测试编码
    test_text = "lowest"
    encoded_tokens = encode(test_text, final_merges)
    print(f"\n对文本 '{test_text}' 编码后的 Token 列表: {encoded_tokens}")
    # 解码展示
    decoded_bytes = b"".join([final_vocab[t] for t in encoded_tokens])
    print(f"解码还原后的文本: '{decoded_bytes.decode('utf-8')}'")

其中测试代码的运行结果为

合并: (108, 111) -> 新Token ID: 256, 频率: 3, 当前词表大小: 257
合并: (256, 119) -> 新Token ID: 257, 频率: 3, 当前词表大小: 258
合并: (32, 257) -> 新Token ID: 258, 频率: 2, 当前词表大小: 259
合并: (258, 101) -> 新Token ID: 259, 频率: 2, 当前词表大小: 260
合并: (257, 259) -> 新Token ID: 260, 频率: 1, 当前词表大小: 261
合并: (260, 114) -> 新Token ID: 261, 频率: 1, 当前词表大小: 262

学习到的合并规则:
Token 256 ('lo') 是由 (108, 111) ('l' + 'o') 合并而来的
Token 257 ('low') 是由 (256, 119) ('lo' + 'w') 合并而来的
Token 258 (' low') 是由 (32, 257) (' ' + 'low') 合并而来的
Token 259 (' lowe') 是由 (258, 101) (' low' + 'e') 合并而来的
Token 260 ('low lowe') 是由 (257, 259) ('low' + ' lowe') 合并而来的
Token 261 ('low lower') 是由 (260, 114) ('low lowe' + 'r') 合并而来的

对文本 'lowest' 编码后的 Token 列表: [257, 101, 115, 116]
解码还原后的文本: 'lowest'

组件 3：模型架构

摒弃高级 API，从张量级别手写 Transformer。

• 输入嵌入：
  Token ID 通过查表操作映射为向量 \(X \in \mathbb{R}^{T \times d}\)。

  # PyTorch 底层实现
  token_embeddings = nn.Embedding(vocab_size, d_model)(input_ids)
  

• 位置编码：传统绝对位置编码无法外推。现代模型（如 Llama）使用 RoPE (旋转位置编码)。其核心思想是通过旋转矩阵将绝对位置信息编码为相对位置信息：
  \[q_m^T k_n = (R_m q)^T (R_n k) = q^T R_{n-m} k \]
  其中 \(R\) 是依赖于位置差 \((n-m)\) 的旋转矩阵。
• 自注意力机制：

  # 带因果掩码的点积注意力底层代码结构
  Q = X @ W_Q  # Shape: (T, d_k)
  K = X @ W_K  # Shape: (T, d_k)
  V = X @ W_V  # Shape: (T, d_v)
  
  # 计算注意力分数 （注意：平方级时空复杂度）
  scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
  
  # 应用因果掩码：将上三角（未来的 token）设为负无穷
  mask = torch.triu(torch.ones(T, T), diagonal=1).bool()
  scores = scores.masked_fill(mask, float('-inf'))
  
  # Softmax 转换为概率分布
  attn_weights = F.softmax(scores, dim=-1)
  
  # 加权求和
  output = attn_weights @ V
  

• FFN 与 SwiGLU：现代模型采用 SwiGLU 激活函数替代传统的 ReLU，公式为 \(\text{SwiGLU}(x, W, V, W_2) = (\text{Swish}(xW) \otimes xV) W_2\)，能提升模型表达能力。
• 归一化：使用 RMSNorm 替代 LayerNorm，省去均值计算和偏移，速度更快且训练更稳定。通常采用 Pre-Norm 结构（在注意力层和 FFN 层之前进行归一化），防止深层网络梯度爆炸。

现在分析上述代码在干什么

1. 生成 Q、K、V 矩阵

Q = X @ W_Q  # Shape: (T, d_k)
K = X @ W_K  # Shape: (T, d_k)
V = W_V @ X  # Shape: (T, d_v)

• 输入文本 \(X\) 经过三个不同的线性变换（乘以权重矩阵 \(W\)），变成了三个具有不同使命的矩阵：
  • Q (Query / 查询)：相当于“提问”。当前这个词（Token）想从上下文中寻找什么信息？
  • K (Key / 键)：相当于“标签”。上下文中的其他词，身上贴着什么标签，能被搜索到？
  • V (Value / 值)：相当于“内容”。如果匹配上了，具体能提取出什么实质性的信息？

可以理解为，就像在图书馆找书。\(Q\) 是你脑子里的搜索词，\(K\) 是书脊上的标签，\(V\) 是书里的具体内容。（C++里的 map）

2. 计算注意力分数（匹配过程）

scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)

• 让每一个 \(Q\) 和所有的 \(K\) 做点积运算，算出它们之间的“匹配度”（相似度）。
• 除以 math.sqrt(d_k)：这叫缩放（Scaling）。如果维度 \(d_k\) 很大，点积的结果数值会非常大，导致下一步 Softmax 时梯度消失（所有概率都挤在 0 和 1 的极端）。除以标准差可以把数值拉回一个合理的区间，让模型训练更稳定。

3. 应用因果掩码

mask = torch.triu(torch.ones(T, T), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))

• 在预测第 \(t\) 个词时，模型绝对不能“偷看”第 \(t+1\) 及以后的词。这一步生成了一个上三角矩阵（triu），把代表“未来”的分数全部强行替换成了 -inf

4. Softmax 转换为概率分布（分配注意力）

attn_weights = F.softmax(scores, dim=-1)

首先。Softmax 是什么？

Softmax 是一个数学函数，它的主要作用是将一个包含任意实数的向量（比如模型输出的原始分数），“压缩”成一个概率分布。

它的核心特点是：

1. 输出归一化：所有输出值的总和为 1，因此每个输出值都在 0 到 1 之间，可以被直接解释为概率。
2. 保持相对大小关系：输入值越大越重要，对应的输出概率也越大，并且会指数级地放大这种差异。

假设你有一个输入向量 z \(= [z₁, z₂, ..., z_K]\)，Softmax 计算第 \(i\) 个元素的公式是：

\[\text{Softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} \]

因为刚才把未来的词设为了负无穷，而 \(e^{-\infty} = 0\)。所以，经过 Softmax 后，未来词的注意力权重变成了绝对的 0，从而保证了“因果性”。

5. 加权求和

output = attn_weights @ V

用刚才算出来的概率（权重），去对 \(V\)（内容）进行加权求和。如果某个词和当前词高度相关，它的权重就大，提取出的 \(V\) 就多。最终得到的 output，就是融合了上下文信息的、更丰富的当前词表示。

import torch
import torch.nn.functional as F
import math


def causal_self_attention(X, W_Q, W_K, W_V):
    """
    带因果掩码的缩放点积注意力机制

    Args:
        X: 输入张量, Shape: (Batch_Size, Seq_Len, Embed_Dim)
        W_Q, W_K, W_V: 线性投影权重矩阵, Shape: (Embed_Dim, Head_Dim)

    Returns:
        output: 注意力输出张量, Shape: (Batch_Size, Seq_Len, Head_Dim)
    """
    # 1. 线性投影，生成 Q, K, V
    Q = X @ W_Q  # Shape: (Batch_Size, Seq_Len, Head_Dim)
    K = X @ W_K  # Shape: (Batch_Size, Seq_Len, Head_Dim)
    V = X @ W_V  # Shape: (Batch_Size, Seq_Len, Head_Dim)

    # 获取序列长度 (T) 和 键的维度 (d_k)
    T = Q.size(1)
    d_k = Q.size(-1)

    # 2. 计算注意力分数 (Scaled Dot-Product)
    # K.transpose(-2, -1) 将最后两个维度互换，Shape 变为: (Batch_Size, Head_Dim, Seq_Len)
    # 点积后 Shape 变为: (Batch_Size, Seq_Len, Seq_Len)
    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)

    # 3. 应用因果掩码 (Causal Mask)
    # 生成一个 (Seq_Len, Seq_Len) 的上三角布尔矩阵
    # diagonal=1 表示从对角线往右上方偏移一格开始为 True
    mask = torch.triu(torch.ones(T, T, device=X.device), diagonal=1).bool()

    # 将上三角（代表未来的 token）的分数设为负无穷
    # 注意：这里需要利用广播机制，mask 会自动应用到 Batch 维度
    scores = scores.masked_fill(mask, float('-inf'))

    # 4. Softmax 转换为概率分布
    # 沿着最后一个维度（即 Key 的维度）进行归一化，确保每一行的概率和为 1
    attn_weights = F.softmax(scores, dim=-1)

    # 5. 加权求和
    # (Batch_Size, Seq_Len, Seq_Len) @ (Batch_Size, Seq_Len, Head_Dim)
    # 输出 Shape: (Batch_Size, Seq_Len, Head_Dim)
    output = attn_weights @ V

    return output


# ================= 测试代码 =================
if __name__ == "__main__":
    # 假设超参数
    batch_size = 2  # 批次大小
    seq_len = 5  # 序列长度 (T)
    embed_dim = 16  # 输入嵌入维度
    head_dim = 8  # 注意力头的维度 (d_k)

    # 1. 随机初始化输入张量 X 和 权重矩阵 W
    X = torch.randn(batch_size, seq_len, embed_dim)
    W_Q = torch.randn(embed_dim, head_dim)
    W_K = torch.randn(embed_dim, head_dim)
    W_V = torch.randn(embed_dim, head_dim)

    # 2. 运行注意力机制
    result = causal_self_attention(X, W_Q, W_K, W_V)

    # 3. 打印结果验证
    print("输入 X 的形状:", X.shape)
    print("注意力输出的形状:", result.shape)

    # 4. 验证因果掩码是否生效：检查第一个 token 的注意力权重
    # 重新计算一次以获取 attn_weights（为了演示，这里直接再算一遍）
    Q = X @ W_Q
    K = X @ W_K
    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(head_dim)
    T = Q.size(1)
    mask = torch.triu(torch.ones(T, T), diagonal=1).bool()
    scores = scores.masked_fill(mask, float('-inf'))
    attn_weights = F.softmax(scores, dim=-1)

    print("\n第一个样本中，第 1 个 token 对所有 token 的注意力权重:")
    print(attn_weights[0, 0, :].detach().numpy())
    # 预期结果：第一个数字接近 1.0，后面的数字全是 0.0（因为被掩码挡住了）

组件 4：训练系统

当模型达到数十亿参数时，单张 GPU 无法装下。

• 显存瓶颈分析：以 Adam 优化器为例，对于一个参数 \(w\)，需要存储：参数本身（FP16: 2 bytes）、梯度（FP16: 2 bytes）、优化器的一阶动量和二阶动量（FP32: 4+4=8 bytes）。总计每参数需要 12-16 bytes。一个 7B 模型仅状态就需要 80GB+ 显存，单卡（如 80GB A100）根本放不下。
• 混合精度训练：前向和反向传播使用 BF16 (Bfloat16)，它具有与 FP32 相同的动态范围，但精度较低，有效防止了 FP16 容易出现的梯度下溢问题。
• 分布式策略：
  • ZeRO (Zero Redundancy Optimizer)：将优化器状态、梯度和参数切分到不同的 GPU 上。例如 ZeRO-3 阶段，每张卡只保留 1/N 的参数，需要时通过 All-Gather 通信聚合。
  • 算子优化：手写底层算子（如使用 Triton 语言编写 FlashAttention），将 Attention 的计算复杂度从内存密集型的 \(O(N^2)\) 优化为计算密集型，大幅减少 HBM（显存）的读写带宽瓶颈。

FlashAttention 的运行机制，核心在于解决传统注意力机制的一个致命瓶颈：显存读写速度远慢于计算速度。

来，我们做个比喻。

• 显存 ：是空间巨大但取放食材很慢的冰箱。
• 缓存 ：是空间很小但伸手就能拿到的灶台。

传统注意力机制的问题是，它会把所有食材（Q, K, V 矩阵）从大冰箱里搬出来，在厨房里铺开一个巨大的 N x N 的案板（注意力分数矩阵），算完后再把结果（输出矩阵）放回冰箱。这个“搬运-铺开-计算-搬运”的过程，绝大部分时间都浪费在了缓慢的搬运上。所以它的复杂度很高，是平方级别的。

FlashAttention 采用了 “分块计算 (Tiling)” 和 “重计算 (Recomputation)” 两大核心策略，实现了“在灶台上边切边炒”。

FlashAttention 不再一次性计算整个 N x N 的注意力矩阵，而是将 Q, K, V 矩阵切分成一个个能轻松放进 GPU 缓存（SRAM）的小块（Block）。

现在，我们来看看这个过程，

1. 加载小块：从显存（HBM）中加载一小块 Q、一小块 K 和一小块 V 到高速缓存（SRAM）。
2. 局部计算：在缓存中直接计算这几个小块之间的注意力分数，并立即进行 Softmax 和加权求和。
3. 累加结果：将计算出的部分结果累加到最终的输出块上。
4. 重复循环：重复以上步骤，遍历所有 K 和 V 的块，直到计算出当前 Q 块对应的完整输出。然后处理下一个 Q 块。

通过这种方式，FlashAttention 不需要在显存中生成巨大的 N x N 中间矩阵。这极大地减少了从显存读写数据的次数，将内存访问量从 O(N²) 降低到了 O(N)，从根本上解决了显存带宽瓶颈。

但是有个问题，在反向传播时，不是需要前向传播的注意力分数矩阵来计算梯度吗？如果没保存，是否需要重新算一遍？

这正是 FlashAttention 的另一个关键设计。它主动放弃了保存中间结果，选择在反向传播时，利用已经保存在显存中的 Q, K, V 矩阵，重新计算一遍前向传播的过程来获取所需的中间值。

这看似增加了计算量，但实际上是“以计算换内存”，而我们损失的复杂度原比赚来的复杂度要少得多。因为 GPU 的计算速度极快，重新计算一遍的耗时，远小于从慢速显存中读取一个巨大矩阵的耗时，因此我们就实现了优化的目的与效果。

接下来是 FlashAttention 的代码实现

import torch
import triton
import triton.language as tl
import math


# --- Triton 内核实现 (增加边界保护与序列长度掩码) ---
@triton.jit
def _fwd_kernel(
        Q, K, V, sm_scale,
        Out,
        stride_qz, stride_qh, stride_qm, stride_qk,
        stride_kz, stride_kh, stride_kn, stride_kk,
        stride_vz, stride_vh, stride_vn, stride_vk,
        stride_oz, stride_oh, stride_om, stride_ok,
        Z, H, N_CTX,
        BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
        BLOCK_N: tl.constexpr,
):
    # 1. 确定当前程序负责计算输出矩阵的哪一块
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    # 2. 计算 Q, K, V, Out 在显存中的起始位置指针
    off_z = off_hz // H
    off_h = off_hz % H
    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
    v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
    o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
    Q_block_ptr = tl.make_block_ptr(
        base=Q + qvk_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_qm, stride_qk),
        offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0),
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + qvk_offset, shape=(BLOCK_DMODEL, N_CTX), strides=(stride_kk, stride_kn),
        offsets=(0, 0), block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1),
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + v_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_vn, stride_vk),
        offsets=(0, 0), block_shape=(BLOCK_N, BLOCK_DMODEL), order=(1, 0),
    )
    O_block_ptr = tl.make_block_ptr(
        base=Out + o_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_om, stride_ok),
        offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0),
    )
    # 3. 初始化在线 Softmax 的累加器
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    # 4. 加载当前的 Q 块到 SRAM (增加 boundary_check 避免 N_CTX 不整除 BLOCK_SIZE 时越界)
    q = tl.load(Q_block_ptr, boundary_check=(0, 1))
    # 5. 循环遍历所有的 K 和 V 块
    for start_n in range(0, N_CTX, BLOCK_N):
        k = tl.load(K_block_ptr, boundary_check=(0, 1))
        v = tl.load(V_block_ptr, boundary_check=(0, 1))
        # 计算 QK^T
        qk = tl.dot(q, k)
        qk *= sm_scale
        # --- 应用因果掩码与序列长度边界掩码 ---
        col_offsets = start_n + tl.arange(0, BLOCK_N)
        row_offsets = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
        causal_mask = col_offsets[None, :] > row_offsets[:, None]
        # 防止读取超出 N_CTX 实际长度的列产生脏数据
        ctx_mask = col_offsets[None, :] < N_CTX
        qk = tl.where(causal_mask | ~ctx_mask, float("-inf"), qk)
        # --- 在线 Softmax 计算 ---
        # 1. 找到新的最大值
        m_ij = tl.maximum(m_i, tl.max(qk, 1))
        # 2. 修正旧的累加值
        alpha = tl.exp(m_i - m_ij)
        # 3. 计算当前块的指数和
        l_ij = tl.exp(qk - m_ij[:, None])
        l_ij_sum = tl.sum(l_ij, 1)
        # 4. 更新全局指数和 l_i
        l_i = l_i * alpha + l_ij_sum
        # 5. 更新输出累加器 acc
        acc = acc * alpha[:, None]
        acc += tl.dot(l_ij.to(v.dtype), v)
        # 6. 更新最大值 m_i
        m_i = m_ij
        # 7. 移动指针
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
    # 6. 最终归一化
    acc = acc / l_i[:, None]
    # 7. 将结果写回显存 (增加 boundary_check)
    tl.store(O_block_ptr, acc.to(Out.type.element_ty), boundary_check=(0, 1))


# --- Python 包装函数 ---
def triton_flash_attention(Q, K, V, causal=True):
    q_shape = Q.shape
    assert q_shape == K.shape and q_shape == V.shape, "Q, K, V must have the same shape"
    assert Q.dtype == torch.float16, "Input must be float16"
    BATCH, N_HEADS, N_CTX, D_HEAD = q_shape
    sm_scale = 1.0 / math.sqrt(D_HEAD)
    O = torch.empty_like(Q)
    BLOCK_M = 64
    BLOCK_N = 64
    grid = (triton.cdiv(N_CTX, BLOCK_M), BATCH * N_HEADS)
    _fwd_kernel[grid](
        Q, K, V, sm_scale,
        O,
        Q.stride(0), Q.stride(1), Q.stride(2), Q.stride(3),
        K.stride(0), K.stride(1), K.stride(2), K.stride(3),
        V.stride(0), V.stride(1), V.stride(2), V.stride(3),
        O.stride(0), O.stride(1), O.stride(2), O.stride(3),
        BATCH, N_HEADS, N_CTX,
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=D_HEAD,
    )
    return O


# --- 补充缺失的标准注意力实现 (用于结果对比) ---
def standard_attention(Q, K, V, causal=True):
    D_HEAD = Q.shape[-1]
    sm_scale = 1.0 / math.sqrt(D_HEAD)
    attn_weight = torch.matmul(Q, K.transpose(2, 3)) * sm_scale
    if causal:
        N_CTX = Q.shape[2]
        mask = torch.tril(torch.ones(N_CTX, N_CTX, device=Q.device, dtype=torch.bool))
        attn_weight = attn_weight.masked_fill(~mask, float("-inf"))
    attn_weight = torch.softmax(attn_weight, dim=-1)
    out = torch.matmul(attn_weight, V)
    return out


# === 测试代码 ===
if __name__ == "__main__":
    torch.manual_seed(0)
    Q = torch.randn(2, 4, 128, 64, dtype=torch.float16, device='cuda')
    K = torch.randn(2, 4, 128, 64, dtype=torch.float16, device='cuda')
    V = torch.randn(2, 4, 128, 64, dtype=torch.float16, device='cuda')
    torch_out = standard_attention(Q, K, V, causal=True)
    triton_out = triton_flash_attention(Q, K, V, causal=True)
    print(f"Triton FlashAttention 输出形状: {triton_out.shape}")
    max_error = (triton_out - torch_out).abs().max().item()
    print(f"与 PyTorch 标准实现的最大绝对误差: {max_error}")
    if max_error < 1e-2:
        print("结果正确")
    else:
        print("结果有误")

组件 5：对齐与评估

预训练后的基础模型只会“续写”，不会“回答”。需要通过以下步骤对齐人类意图：

• 监督微调 (SFT)：构造高质量的 (Instruction, Response) 数据对，用交叉熵损失微调模型，使其学会遵循指令。
• 直接偏好优化 (DPO)：相比于复杂的 RLHF（需要训练奖励模型和进行 PPO 强化学习），DPO 直接利用人类偏好的数据对 \((y_{win}, y_{lose})\)，通过一个简单的损失函数优化策略模型：
  \[\mathcal{L}_{DPO} = -\log \sigma \left( \beta \log \frac{\pi_\theta(y_{win}|x)}{\pi_{ref}(y_{win}|x)} - \beta \log \frac{\pi_\theta(y_{lose}|x)}{\pi_{ref}(y_{lose}|x)} \right) \]
  其中 \(\pi_{ref}\) 是初始的 SFT 模型，\(\pi_\theta\) 是正在训练的模型。这个公式直接拉大模型对好回答和坏回答的概率差。

---

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

# ==========================================
# 0. 基础配置与加载
# ==========================================
model_id = "Qwen/Qwen2-0.5B" # 使用轻量级模型
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 加载策略模型 (待训练) 和参考模型 (冻结，用于DPO)
policy_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32).to(device)
ref_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32).to(device)
ref_model.eval() # 参考模型不参与梯度计算

# ==========================================
# 1. 手写 SFT (监督微调) 算法
# ==========================================
def sft_train_step(model, tokenizer, prompt_text, response_text, optimizer):
    """
    SFT 核心逻辑：拼接 prompt 和 response，但只在 response 部分计算 loss
    """
    model.train()
    
    # 1. 编码文本
    prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
    response_ids = tokenizer.encode(response_text, add_special_tokens=False)
    input_ids = prompt_ids + response_ids + [tokenizer.eos_token_id]
    
    # 转为 Tensor
    input_ids = torch.tensor([input_ids], dtype=torch.long).to(device)
    
    # 2. 构建 Labels：将 prompt 部分设为 -100，这样 CrossEntropy Loss 会忽略它们
    labels = input_ids.clone()
    labels[:, :len(prompt_ids)] = -100
    
    # 3. 前向传播计算 Loss
    outputs = model(input_ids=input_ids, labels=labels)
    loss = outputs.loss
    
    # 4. 反向传播与更新
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    return loss.item()

# ==========================================
# 2. 手写 DPO (直接偏好优化) 算法
# ==========================================
def get_sequence_log_probs(model, input_ids, response_start_idx):
    """
    计算模型生成指定 Response 的对数概率 (Log Probability)
    """
    outputs = model(input_ids=input_ids)
    logits = outputs.logits
    
    # Shift 操作：因为位置 i 的 logits 预测的是位置 i+1 的 token
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()
    
    # 计算所有 token 的 log_softmax
    log_probs = F.log_softmax(shift_logits, dim=-1)
    
    # 提取 Response 部分的 log_prob
    # shift_labels 的索引比原 input_ids 少 1，所以 response_start_idx 对应的是 shift_labels 中的 response_start_idx-1
    response_log_probs = log_probs[:, response_start_idx-1:, :]
    response_labels = shift_labels[:, response_start_idx-1:]
    
    # 提取实际 token 的 log_prob 并求和
    chosen_log_probs = response_log_probs.gather(2, response_labels.unsqueeze(-1)).squeeze(-1)
    return chosen_log_probs.sum(dim=-1)

def dpo_loss_fn(policy_chosen_logps, policy_rejected_logps, 
                ref_chosen_logps, ref_rejected_logps, beta=0.1):
    """
    DPO 损失函数公式实现
    """
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps
    
    logits = pi_logratios - ref_logratios
    # -log(sigmoid(beta * logits))
    loss = -F.logsigmoid(beta * logits).mean()
    return loss

def dpo_train_step(policy_model, ref_model, tokenizer, prompt, chosen, rejected, optimizer, beta=0.1):
    """
    DPO 核心逻辑：让策略模型偏向 Chosen，远离 Rejected
    """
    policy_model.train()
    
    # 1. 数据编码
    prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
    chosen_ids = tokenizer.encode(chosen, add_special_tokens=False) + [tokenizer.eos_token_id]
    rejected_ids = tokenizer.encode(rejected, add_special_tokens=False) + [tokenizer.eos_token_id]
    
    # 拼接
    chosen_input_ids = torch.tensor([prompt_ids + chosen_ids], dtype=torch.long).to(device)
    rejected_input_ids = torch.tensor([prompt_ids + rejected_ids], dtype=torch.long).to(device)
    
    response_start_idx = len(prompt_ids)
    
    # 2. 策略模型计算 log_probs
    policy_chosen_logps = get_sequence_log_probs(policy_model, chosen_input_ids, response_start_idx)
    policy_rejected_logps = get_sequence_log_probs(policy_model, rejected_input_ids, response_start_idx)
    
    # 3. 参考模型计算 log_probs (不计算梯度)
    with torch.no_grad():
        ref_chosen_logps = get_sequence_log_probs(ref_model, chosen_input_ids, response_start_idx)
        ref_rejected_logps = get_sequence_log_probs(ref_model, rejected_input_ids, response_start_idx)
        
    # 4. 计算 DPO Loss
    loss = dpo_loss_fn(policy_chosen_logps, policy_rejected_logps, 
                       ref_chosen_logps, ref_rejected_logps, beta)
    
    # 5. 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    return loss.item()

# ==========================================
# 3. 运行测试：将基础模型转变为 Chat 模型
# ==========================================
if __name__ == "__main__":
    optimizer = torch.optim.AdamW(policy_model.parameters(), lr=5e-5)
    
    # --- 第一阶段：SFT 训练 (3 个 Epoch) ---
    print("=== 阶段 1: SFT (监督微调) ===")
    sft_prompt = "User: 你好，你是谁？\nAssistant:"
    sft_response = "你好！我是一个人工智能助手，很高兴为您服务。"
    
    for epoch in range(3):
        loss = sft_train_step(policy_model, tokenizer, sft_prompt, sft_response, optimizer)
        print(f"SFT Epoch {epoch+1} | Loss: {loss:.4f}")
    
    # 同步参考模型的权重 (DPO 中的参考模型应当是 SFT 后的模型)
    print("\n同步 Reference 模型权重...")
    ref_model.load_state_dict(policy_model.state_dict())
    
    # --- 第二阶段：DPO 训练 (2 个 Epoch) ---
    print("\n=== 阶段 2: DPO (直接偏好优化) ===")
    # 假设 SFT 模型有时会输出简短或不礼貌的回答
    dpo_prompt = "User: 帮我写首诗。\nAssistant:"
    chosen_response = "好的，这是一首关于春天的诗：春风拂柳绿，细雨润花红。燕舞晴空阔，人归画梦中。"
    rejected_response = "我不会写诗。" # 劣质回答：拒绝服务
    
    # DPO 使用更小的学习率
    dpo_optimizer = torch.optim.AdamW(policy_model.parameters(), lr=1e-6)
    
    for epoch in range(2):
        loss = dpo_train_step(
            policy_model, ref_model, tokenizer, 
            dpo_prompt, chosen_response, rejected_response, 
            dpo_optimizer, beta=0.1
        )
        print(f"DPO Epoch {epoch+1} | Loss: {loss:.4f}")
        
    print("\n训练完成！基础模型已转变为经过偏好对齐的 Chat 模型。")

底层原理解析

1. SFT 的掩码机制 (labels[:, :len(prompt_ids)] = -100)
   在 PyTorch 的 CrossEntropyLoss 中，如果目标标签是 -100，该位置的损失会被自动忽略（不参与梯度计算）。通过这种方式，模型只会在“回答”部分计算预测误差，而不会在“问题”部分浪费学习容量。

2. DPO 的 Shift 操作 (shift_logits = logits[:, :-1, :])
   自回归语言模型在时间步 \(t\) 输出的 logits，是用来预测时间步 \(t+1\) 的 token 的。因此在计算序列概率时，必须把 logits 和 labels 错开一位对齐（即用第 0 位预测第 1 位，第 1 位预测第 2 位）。

3. DPO 的 Loss 计算 (-F.logsigmoid(beta * logits))

   • logits = (Policy_Chosen - Policy_Rejected) - (Ref_Chosen - Ref_Rejected)
   • 这表示策略模型相比参考模型，更偏向 Chosen 的程度。
   • logsigmoid 将这个程度映射为概率。如果策略模型确实更偏向 Chosen，logsigmoid 的值趋近于 0，Loss 趋近于 0；反之 Loss 会变大，迫使模型调整。

四、 核心理论：缩放定律

这是指导大模型训练的第一性原理，由 Kaplan (2020) 和 Hoffmann (Chinchilla, 2022) 提出。

1. Kaplan 缩放定律

模型测试集的交叉熵损失 \(L\) 与计算量 \(C\)、模型参数量 \(N\) 和数据量 \(D\) 呈幂律关系：

\[L(N, D, C) = \text{irreducible loss} + \text{power law terms} \]

关键结论：只要增加算力和参数量，Loss 就会可预测地下降，没有天花板。

2. Chinchilla 最优缩放定律

Kaplan 定律导致业界盲目追求大模型（如 GPT-3 的 175B），但数据量给得不够。
Hoffmann 等人发现，在给定的总算力预算 \(C\) 下，最优的参数量 \(N^*\) 和数据量 \(D^*\) 应当按比例增长。具体的比例关系推导如下：

\[N^* \propto C^{0.5}, \quad D^* \propto C^{0.5} \]

这意味着 每个参数大约应该训练 20 个 Token。

• 实操意义：如果你有一个 70B 参数的模型，你需要准备 1.4 Trillion (1.4万亿) 的高质量 Token 进行训练，才能达到算力效能的最大化。在 CS336 的作业中，我们会通过训练 10M~100M 的小模型，拟合出我们自己的 Scaling Law 曲线，然后外推预测大模型的表现。

---

参考文案与学习来源

课程为 CS336
部分代码与部分公式介绍的形成优化于 智谱、千问、DeepSeek、Chat GPT
参考文献 CS336完整学习记录