""" ============================================================================= 第 20 讲 · 代码 1：精简版 GPT-2 —— 和公开权重完全对应的实现 ============================================================================= 对应讲义：lectures/lecture20_default_project.md 第 2、3 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- CS224N 的默认期末项目要求自己实现 GPT-2 的关键部分，再加载 OpenAI 公开的预训练权重，用到下游任务上。 这里给出一个和第 7 讲 transformer.py 一脉相承、但**参数的名字和形状与 Hugging Face 的 GPT-2 完全对应**的实现， 这样就能把公开的权重直接装进来。 GPT2Config 超参数（默认是 GPT-2 small：12 层、768 维、12 个头、词表 50257、最长 1024 个词元） CausalSelfAttention / MLP / Block / GPT2 模型本身（Pre-LN、GELU、权重共享） GPT2.from_pretrained("gpt2") （可选）从 Hugging Face 下载权重并装进来，再和官方实现逐数比对 GPT2.generate(...) 自回归生成（温度、top-k） GPT2Classifier 在最后一个词元的隐藏状态上接分类头（情感分类、复述检测等） 直接运行本文件会做一组**不需要联网**的自我检查（约 10 秒）： · 完整 GPT-2 small 的参数个数 = 124,439,808（和第 7 讲用公式算出的一样） · 小配置下的形状检查、因果性检查、初始损失 ≈ ln(词表大小)、生成、分类头、配合第 11 讲的 LoRA python gpt2_minimal.py # 离线自我检查 python gpt2_minimal.py --hf # 另外：下载 GPT-2 权重，和 Hugging Face 的实现比对输出（需要 transformers 和网络） 【关于 --hf】编写讲义时的环境没有 transformers、也不联网，这部分**没有运行过**；只测试了"没有安装"时的提示。 """ import math import sys from dataclasses import dataclass, replace from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F @dataclass class GPT2Config: """ GPT-2 的超参数。默认值是 GPT-2 small（124M）。其他尺寸： gpt2-medium n_layer=24, n_embd=1024, n_head=16（350M） gpt2-large n_layer=36, n_embd=1280, n_head=20（774M） gpt2-xl n_layer=48, n_embd=1600, n_head=25（1558M） """ vocab_size: int = 50257 n_positions: int = 1024 n_embd: int = 768 n_layer: int = 12 n_head: int = 12 eps: float = 1e-5 dropout: float = 0.1 class CausalSelfAttention(nn.Module): """ 多头因果自注意力（和第 7 讲的 MultiHeadSelfAttention 一样，只是名字照搬 GPT-2：c_attn、c_proj）。 【参数（可学习）】 c_attn : D → 3D，一次算出 Q、K、V c_proj : D → D，把各个头的输出混合起来 【边界】OpenAI 的原始代码（和 Hugging Face）用的是一种叫 Conv1D 的层，权重矩阵是 (输入, 输出)， 而 nn.Linear 是 (输出, 输入) —— 加载公开权重时要**转置**（见 from_pretrained）。 """ def __init__(self, cfg): super().__init__() self.n_head, self.d = cfg.n_head, cfg.n_embd self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd) self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd) self.attn_drop, self.resid_drop = nn.Dropout(cfg.dropout), nn.Dropout(cfg.dropout) mask = torch.tril(torch.ones(cfg.n_positions, cfg.n_positions, dtype=torch.bool)) self.register_buffer("mask", mask.view(1, 1, cfg.n_positions, cfg.n_positions), persistent=False) def forward(self, x, pad_mask=None): """ 【参数】x：(B, T, D)；pad_mask：None 或 (B, T) 布尔张量（True = 真实的词元） 【返回】(B, T, D) """ B, T, D = x.shape H, dh = self.n_head, D // self.n_head q, k, v = self.c_attn(x).split(D, dim=2) q, k, v = (t.view(B, T, H, dh).transpose(1, 2) for t in (q, k, v)) att = q @ k.transpose(-2, -1) / math.sqrt(dh) att = att.masked_fill(~self.mask[:, :, :T, :T], float("-inf")) if pad_mask is not None: att = att.masked_fill(~pad_mask[:, None, None, :], float("-inf")) att = self.attn_drop(torch.softmax(att, dim=-1)) y = (att @ v).transpose(1, 2).contiguous().view(B, T, D) return self.resid_drop(self.c_proj(y)) class MLP(nn.Module): """ 前馈网络：D → 4D → GELU → D。 【边界】GPT-2 用的是 GELU 的 tanh 近似（Hugging Face 里叫 "gelu_new"）： 0.5 x (1 + tanh(√(2/π) (x + 0.044715 x³))) 和精确的 GELU 差别很小，但加载公开权重、要求输出逐数一致时，必须用同一个。 """ def __init__(self, cfg): super().__init__() self.c_fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd) self.c_proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd) self.drop = nn.Dropout(cfg.dropout) def forward(self, x): return self.drop(self.c_proj(F.gelu(self.c_fc(x), approximate="tanh"))) class Block(nn.Module): """Pre-LN 的 Transformer 块：x = x + Attn(LN(x))；x = x + MLP(LN(x))（第 7 讲第 7 节）。""" def __init__(self, cfg): super().__init__() self.ln_1 = nn.LayerNorm(cfg.n_embd, eps=cfg.eps) self.attn = CausalSelfAttention(cfg) self.ln_2 = nn.LayerNorm(cfg.n_embd, eps=cfg.eps) self.mlp = MLP(cfg) def forward(self, x, pad_mask=None): x = x + self.attn(self.ln_1(x), pad_mask) return x + self.mlp(self.ln_2(x)) class GPT2(nn.Module): """ GPT-2：词嵌入 wte + 位置嵌入 wpe → n_layer 个块 → 最后的层归一化 ln_f → 输出层（和 wte 共享权重）。 【参数的名字】和 Hugging Face 的 GPT2LMHeadModel 完全一样（transformer.wte.weight、transformer.h.0.attn.c_attn.weight……）， 所以加载公开权重时只需要一一对应、个别矩阵转置。 【初始化】权重 ~ N(0, 0.02²)；每个残差分支最后的投影层（c_proj）再乘以 1/√(2·n_layer)： 每一层都往残差流上加东西，层数越多加得越多，缩小这些投影层的初始值，让残差流的方差不随层数增长（GPT-2 论文的做法）。 """ def __init__(self, cfg): super().__init__() self.cfg = cfg self.transformer = nn.ModuleDict(dict( wte=nn.Embedding(cfg.vocab_size, cfg.n_embd), wpe=nn.Embedding(cfg.n_positions, cfg.n_embd), drop=nn.Dropout(cfg.dropout), h=nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]), ln_f=nn.LayerNorm(cfg.n_embd, eps=cfg.eps), )) self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) self.apply(self._init) for name, p in self.named_parameters(): if name.endswith("c_proj.weight"): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) self.lm_head.weight = self.transformer.wte.weight # 权重共享（第 7 讲第 8.2 节） @staticmethod def _init(m): if isinstance(m, (nn.Linear, nn.Embedding)): nn.init.normal_(m.weight, mean=0.0, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.zeros_(m.bias) def hidden_states(self, idx, pad_mask=None): """词元编号 (B, T) → 最后一层（经过 ln_f）的隐藏状态 (B, T, D)。""" T = idx.shape[1] if T > self.cfg.n_positions: raise ValueError(f"序列长度 {T} 超过了 n_positions={self.cfg.n_positions}") x = self.transformer.drop(self.transformer.wte(idx) + self.transformer.wpe(torch.arange(T, device=idx.device))) for block in self.transformer.h: x = block(x, pad_mask) return self.transformer.ln_f(x) def forward(self, idx, targets=None): """ 【参数】idx：(B, T) 词元编号；targets：None 或 (B, T)（-100 的位置不计损失，和 PyTorch 的约定一致） 【返回】(logits (B, T, V), 损失或 None) """ logits = self.lm_head(self.hidden_states(idx)) loss = None if targets is not None: # 【Python 小知识】用 reshape 而不是 view：targets 常常是切片（比如 X[:, 1:]），内存不连续，view 会报错 loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-100) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, generator=None): """ 自回归生成（第 5 讲第 7 节的采样方法）：每一步取最后一个位置的分布，按温度缩放、只保留前 top_k 个，采样一个词元。 【边界】温度为 0 时取概率最大的词元（贪心）；序列超过 n_positions 时只保留最后 n_positions 个词元作为上下文。 """ for _ in range(max_new_tokens): ctx = idx[:, -self.cfg.n_positions:] logits, _ = self(ctx) logits = logits[:, -1, :] if temperature == 0: nxt = logits.argmax(-1, keepdim=True) else: logits = logits / temperature if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("inf") nxt = torch.multinomial(torch.softmax(logits, -1), 1, generator=generator) idx = torch.cat([idx, nxt], dim=1) return idx @classmethod def from_pretrained(cls, name="gpt2"): """ 从 Hugging Face 加载 OpenAI 公开的 GPT-2 权重（需要 `pip install transformers` 和网络）。 【思路】两边的参数名一一对应。四种矩阵（c_attn、c_proj、c_fc 的权重）在原始实现里是 Conv1D， 形状是 (输入, 输出)，要转置成 nn.Linear 的 (输出, 输入)。注意力里的 "bias"/"masked_bias" 是因果掩码，不是参数，跳过。 【返回】(我们的模型, Hugging Face 的模型) —— 后者用来比对输出 """ from transformers import GPT2LMHeadModel sizes = {"gpt2": dict(n_layer=12, n_head=12, n_embd=768), "gpt2-medium": dict(n_layer=24, n_head=16, n_embd=1024), "gpt2-large": dict(n_layer=36, n_head=20, n_embd=1280), "gpt2-xl": dict(n_layer=48, n_head=25, n_embd=1600)} model = cls(GPT2Config(**sizes[name], dropout=0.0)) hf = GPT2LMHeadModel.from_pretrained(name).eval() sd_hf = hf.state_dict() sd = model.state_dict() transposed = ("attn.c_attn.weight", "attn.c_proj.weight", "mlp.c_fc.weight", "mlp.c_proj.weight") with torch.no_grad(): for k in sd: if k not in sd_hf: continue src = sd_hf[k].t() if k.endswith(transposed) else sd_hf[k] assert src.shape == sd[k].shape, (k, src.shape, sd[k].shape) sd[k].copy_(src) return model.eval(), hf class GPT2Classifier(nn.Module): """ 分类器：GPT-2 主体 + 线性分类头，取**最后一个真实词元**的隐藏状态（和第 9 讲 03 号文件、第 11 讲一样）。 【参数】gpt：GPT2；n_classes：类别数 【边界】批里的句子长短不一时，右侧填充，并传入 pad_mask；最后一个真实词元的位置 = 真实长度 − 1。 """ def __init__(self, gpt, n_classes): super().__init__() self.gpt = gpt self.head = nn.Linear(gpt.cfg.n_embd, n_classes) def forward(self, idx, pad_mask): h = self.gpt.hidden_states(idx, pad_mask) last = pad_mask.sum(1) - 1 return self.head(h[torch.arange(idx.shape[0]), last]) def count_parameters(model): return sum(p.numel() for p in model.parameters()) if __name__ == "__main__": torch.manual_seed(0) # 1. 完整的 GPT-2 small：参数个数 full = GPT2(GPT2Config()) n = count_parameters(full) print(f" GPT-2 small：{n:,} 个参数（词嵌入和输出层共享）") assert n == 124_439_808, n print(" ✓ 和第 7 讲用公式算出的 124,439,808 一致") del full # 2. 小配置：形状、因果性、初始损失 cfg = replace(GPT2Config(), vocab_size=100, n_positions=32, n_embd=64, n_layer=2, n_head=4, dropout=0.0) m = GPT2(cfg).eval() X = torch.randint(0, 100, (3, 16)) logits, loss = m(X[:, :-1], X[:, 1:]) # 目标 = 往后错一位（预测下一个词元） assert logits.shape == (3, 15, 100) X2 = X.clone() X2[:, -1] = (X2[:, -1] + 1) % 100 assert torch.allclose(m(X)[0][:, :-1], m(X2)[0][:, :-1], atol=1e-5) print(f" 小配置：形状正确；改变最后一个词元不影响前面的输出（因果性）；初始损失 {loss.item():.3f}（ln 100 = {math.log(100):.3f}）") out = m.generate(X[:1, :4], max_new_tokens=6, temperature=0.8, top_k=10, generator=torch.Generator().manual_seed(0)) print(f" 生成：{X[0, :4].tolist()} → {out[0].tolist()}") # 3. 分类头 + 第 11 讲的 LoRA clf = GPT2Classifier(m, 2) pad = torch.ones(3, 16, dtype=torch.bool) pad[0, 10:] = False assert clf(X, pad).shape == (3, 2) sys.path.append(str(Path(__file__).resolve().parent.parent / "lecture11")) import peft peft.freeze_all(clf.gpt) names = peft.add_lora(clf.gpt, r=4, alpha=8, targets=("attn.c_attn", "attn.c_proj")) n_tr, n_all = peft.count_trainable(clf) print(f" 分类头 + LoRA：替换了 {len(names)} 个线性层，可训练参数 {n_tr:,} / {n_all:,}") # 4. （可选）加载公开权重并和 Hugging Face 比对 if "--hf" in sys.argv: try: from transformers import GPT2Tokenizer except ImportError: print(" （没有安装 transformers：pip install transformers 之后再运行 --hf）") sys.exit(0) model, hf = GPT2.from_pretrained("gpt2") tok = GPT2Tokenizer.from_pretrained("gpt2") ids = torch.tensor([tok.encode("Hello, my name is")]) with torch.no_grad(): diff = (model(ids)[0] - hf(ids).logits).abs().max().item() print(f" 和 Hugging Face 的 GPT-2 比对：logits 最大差别 {diff:.2e}") out = model.generate(ids, max_new_tokens=20, temperature=0) print(" 贪心生成：", tok.decode(out[0])) print(" 全部自我检查通过 ✓")