""" ============================================================================= CS324 · 第 6 讲共享模块：adapt_lab.py —— 在第 1 讲的 base model 上做"适配"的工具箱 ============================================================================= 【这个模块是干什么的 What this module is for】 第 6 讲要把好几种适配方法放在**同一个任务、同一个打分函数**下比较： 提示（prompting）、线性探针（linear probe）、全量微调（full fine-tuning）、只调最后一层、 BitFit（只调偏置）、adapter、prefix tuning、prompt tuning、LoRA。 这些方法在论文里各有各的代码库；这里用几百行把它们写在同一个 11 万参数的小模型上， 让每种方法"到底改了模型的哪一部分"一目了然。 共享零件： load_base() 只读加载第 1 讲的 base model（绝不重训、绝不改写它的缓存） encode_pairs / pair_logprobs 把"提示 + 答案"编码成 batch；算 log p(答案 | 提示) choice_accuracy 受限打分：在候选答案（3 个标签词 / 8 个城市 / 8 个职业）里选概率最大的 greedy_answers 自由生成：贪心续写到句号，原样返回（用来检查"模型到底写了什么"） forward_with 手写的前向传播，可以在输入前接"软提示"、在每层注意力里接"前缀" LoRALinear / Adapter 两种"插件"模块 make_adapted 按方法名把 base model 的一份**拷贝**变成可训练的适配模型 train_on_pairs 统一的训练循环：只在答案字符上算损失 fit_softmax_probe 多类别线性探针（L-BFGS） bpc_with 适配之后的模型在第 1 讲验证集上的 bit/字符（量"遗忘"） cached_model 缓存本讲自己训练的整模型（只写本目录的 checkpoints/） 【两个任务：都是"知识在模型里，但格式打不开它" The tasks】 1. 评论 verdict（代码 01、03、04 的主任务）：同一条评论换一个从没见过的格式，答案换成三个从没见过的词： "opinion: the soup was tasty but the host was rude. verdict:" → " mixed." base model 用原格式（"rating:" + good / ok / bad）打分是满分；新格式下不训练就不会。 2. 人物问答的新格式（代码 02 演示 2）：第 1 讲 §4.2 里，用训练时从没见过的格式 "question: where does batul live? answer:" 问城市，base model 答对 **0.00**。用 k 个人的标签做适配， 在**没出现在训练例子里的人**身上测试——答案只能来自预训练时学到的知识，所以这个任务检验的是 "适配能不能把已有的知识暴露出来"，而不是"能不能背下 k 个标签"。 【运行方式】 python3 adapt_lab.py # 自检：手写前向 = 原模型；LoRA 初始时不改变输出；合并前后输出相同（约 1 秒） 依赖 PyTorch、numpy；只读地导入 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py。 【PyTorch 小知识】 - copy.deepcopy(model)：把整个模型（包括参数）复制一份。之后在拷贝上训练，原模型纹丝不动。 本讲每一次适配都从 base model 的一份新拷贝开始——这样第 1 讲缓存的模型永远不会被改写。 - p.requires_grad_(False)："冻结"一个参数：反向传播时不再为它算梯度，优化器也不会更新它。 - nn.Parameter：被 nn.Module 登记为"可训练参数"的张量；把它放进优化器才会被更新。 - 前向钩子（forward hook）：module.register_forward_hook(fn) 让 fn 在这个模块每次前向之后被调用， fn 可以读取（或替换）它的输入输出。代码 05 用它读 MLP 的中间激活。 """ import copy import math import os import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from fm_world import ALPHABET, CITIES, JOBS, build_world, make_review # noqa: E402 from tiny_gpt import CharTokenizer, TinyGPT, load_or_train_base_model # noqa: E402 TOK = CharTokenizer(ALPHABET) NL = TOK.stoi["\n"] CKPT_DIR = os.path.join(HERE, "checkpoints") # 新格式（第 1 讲 §4.2 里得 0.00 的那个）：问城市、问职业 FMT_CITY = "\nquestion: where does {name} live? answer:" FMT_JOB = "\nquestion: what does {name} do? answer:" # ============================================================================= # 1. 读取 base model（只读） # ============================================================================= def load_base(): """ 只读地加载第 1 讲的 base model。 【思路】用**默认**的世界（build_world()）和默认缓存路径调用 load_or_train_base_model， 它会发现缓存的配置一致，直接读取 code/lecture01/checkpoints/base_model.pt。 之后本讲只在 copy.deepcopy 出来的拷贝上训练，绝不改动这个对象。 【返回】(world, base)：fm_world.World 和 TinyGPT（eval 模式） 【边界】如果第 1 讲的缓存不存在，load_or_train_base_model 会先训练约 1 分钟（得到完全相同的模型）。 """ world = build_world() base, _ = load_or_train_base_model(world, TOK, verbose=False) base.eval() for p in base.parameters(): p.requires_grad_(False) # base model 本身永远不训练 return world, base def split_people(world, seed): """ 把人分成三组（代码 02 演示 2 用这个划分）： train_pool 15 人：有问答文本的 30 人里随机一半；适配用的 k 个标签只从这里抽 id_test 15 人：有问答文本的另一半——"同分布"测试（和训练的人是同一类人） ood_people 30 人：只在叙述句里出现过的人——"人群偏移"测试（知识更难取出，第 1 讲 §4.2） 【思路】三组对应三种测试：和训练的人同类、知识更难取出的人；"问职业"则是第三种（任务偏移），见代码 02。 【参数】seed：决定 30 人怎样一分为二（不同 seed = 不同划分，用来估计噪声） 【返回】(train_pool, id_test, ood_people)，每个都是 list[Person] 【边界】只有 15 个可用的训练例子，k 不能超过 15。 """ rng = np.random.default_rng(100 + seed) qa = [p for p in world.people if p.name in world.qa_names] order = rng.permutation(len(qa)) pool = [qa[i] for i in order[:15]] id_test = [qa[i] for i in order[15:]] ood = [p for p in world.people if p.name not in world.qa_names] return pool, id_test, ood def items_for(people, attr, fmt): """ 把一组人变成评测项 [(提示, 正确答案)]。 【参数】people：list[Person]；attr："city" 或 "job"；fmt：带 {name} 的提示模板（例如 FMT_CITY） 【返回】list[(str, str)]，例如 ("\nquestion: where does kari live? answer:", "varno") 【Python 小知识】getattr(p, "city") 等价于 p.city——属性名放在一个变量里时用它。 """ return [(fmt.format(name=p.name), getattr(p, attr)) for p in people] def train_pairs_for(people, attr="city", fmt=FMT_CITY): """ 训练用的 (提示, 答案) 对。 【思路】答案写成 " varno." 的样子：前面一个空格、后面一个句号，和语料里 "q: ... a: varno." 的写法一致， 这样训练教的只是"新格式"，不多教一种标点习惯。 【参数】people、attr、fmt：同 items_for 【返回】list[(str, str)] """ return [(fmt.format(name=p.name), f" {getattr(p, attr)}.") for p in people] # 评论任务：同样的评论，换一个从没见过的格式 + 三个从没见过的标签词 VERDICT = {"good": "loved", "ok": "mixed", "bad": "hated"} VERDICT_WORDS = ["loved", "mixed", "hated"] def verdict_items(n, rng, food_adjs=None): """ 评论任务的样本：[(提示, 标签词)]。 【思路】语料里的评论写成 "review: the soup was tasty but the host was rude. rating: ok."； 这里把同一条评论写成从没见过的格式，答案也换成语料里**从没出现过**的三个词： "opinion: the soup was tasty but the host was rude. verdict:" → " mixed." 评分规则（两个形容词决定好 / 一般 / 差）模型在预训练时已经学会（第 1 讲：500 步就 0.97）， 适配要学的只是"新格式 + 新标签词"——W22 笔记说的 formatting 差异。 【参数】n：条数；rng：random.Random；food_adjs：可选 (正面列表, 负面列表)，限定食物形容词 （用前一半形容词造训练集、后一半造测试集，就得到"分布外"的测试：形容词在标注数据里从没出现） 【返回】list[(str, str)] 【边界】三个标签的比例跟着评论的随机生成走：好 / 一般 / 差约 1 : 2 : 1，所以"永远答 mixed"约有一半是对的。 """ out = [] for _ in range(n): text, rating, _, _ = make_review(rng, food_adjs=food_adjs) out.append((f"\nopinion: {text[len('review: '):]} verdict:", VERDICT[rating])) return out # ============================================================================= # 2. 编码与打分：所有方法共用同一套 # ============================================================================= def encode_pairs(pairs): """ 把一组 (提示, 答案) 编码成一个右侧补齐的 batch。 【思路】整句 s = 提示 + 答案；输入 x = s[:-1]，目标 y = s[1:]（预测下一个字符）。 mask 只在"答案"那几个字符的位置为 1——训练时只在答案上算损失，打分时只把答案的概率加起来。 右侧补齐（用 "\\n"）不影响结果：因果注意力里，前面的位置看不到后面补的字符。 【参数】pairs：list[(prompt, completion)]，prompt 非空 【返回】(x, y, mask)：前两个是 (B, T) 的 LongTensor，mask 是 (B, T) 的 float 【边界】整句超过 96 个字符时 TinyGPT 会报错（位置嵌入只有 96 个）；本讲的提示都在 50 个字符以内。 """ seqs = [TOK.encode(p + c) for p, c in pairs] T = max(len(s) for s in seqs) - 1 x = torch.full((len(seqs), T), NL, dtype=torch.long) y = torch.full((len(seqs), T), NL, dtype=torch.long) mask = torch.zeros(len(seqs), T) for i, ((p, _), s) in enumerate(zip(pairs, seqs)): n = len(s) - 1 x[i, :n] = torch.tensor(s[:-1]) y[i, :n] = torch.tensor(s[1:]) mask[i, len(p) - 1:n] = 1.0 # 位置 len(p)-1 的输出预测的是答案的第一个字符 return x, y, mask def pair_logprobs(fwd, pairs, batch=512): """ 每一对的 log p(答案 | 提示)（自然对数，答案各字符的对数概率之和）。 【思路】每 batch 条一起算（encode_pairs 右侧补齐），用 mask 只把答案字符的对数概率加起来。 【参数】fwd：函数 idx(B, T) → logits(B, T, V)（不同的适配方法给出不同的 fwd）；batch：每次前向的条数 【返回】一维张量，长度 len(pairs) 【边界】在 torch.no_grad() 里算，不能用来训练（训练用 train_on_pairs）。 """ out = [] with torch.no_grad(): for i in range(0, len(pairs), batch): x, y, m = encode_pairs(pairs[i:i + batch]) logp = F.log_softmax(fwd(x), dim=-1).gather(-1, y[..., None])[..., 0] out.append((logp * m).sum(1)) return torch.cat(out) def choice_predictions(fwd, prompts, candidates): """ 受限打分：对每个提示，把每个候选（" varno."）接上去算概率，选最大的那个。 【思路】和 GPT-3 论文做多项选择题的方法相同（第 1 讲 completion_logprob）。 好处：不管模型"会不会停在句号"，只问"在这 8 个答案里它更相信哪个"； 8 个候选时瞎猜是 1/8 = 0.125。 【参数】fwd：idx → logits；prompts：list[str]；candidates：list[str]（不含前面的空格和后面的句号） 【返回】list[str]：每个提示选中的候选 【边界】候选的长度不同时，长的候选要多乘几个概率，天然吃亏；本讲的候选都是同一类词，长度相近（第 1 讲 §4.4 提到的"长度归一化"问题）。 """ pairs = [(p, f" {c}.") for p in prompts for c in candidates] lp = pair_logprobs(fwd, pairs).view(len(prompts), len(candidates)) return [candidates[j] for j in lp.argmax(1).tolist()] def choice_accuracy(fwd, items, candidates): """ 受限打分的准确率。 【参数】items：list[(提示, 正确答案)]；candidates：候选列表（正确答案必须在里面） 【返回】float，0 到 1 【FM 联系】本讲所有方法共用这一个打分函数（第 3 讲 §5.2 的教训：比较不同模型时，评测协议必须相同）。 """ preds = choice_predictions(fwd, [p for p, _ in items], candidates) return float(np.mean([pr == a for pr, (_, a) in zip(preds, items)])) def greedy_answers(fwd, prompts, max_new=10): """ 自由生成：对每个提示贪心续写，直到写出句号或换行（最多 max_new 个字符），原样返回续写。 【思路】把**长度相同**的提示放进同一个 batch 一起生成（名字有 4、5 个字母，所以大约分两组）， 比逐个生成快十几倍；长度不同的提示不能直接放在一起，否则位置会错开。 【返回】list[str]，顺序与 prompts 相同，例如 " varno." 或 " the staff" 【边界】max_new 个字符内没写到句号的，返回已经写出的部分（不加标记）；提示加续写不能超过 96 个字符。 【FM 联系】自由生成 + 完全匹配（exact match）是更苛刻的评测协议：模型必须自己写出正确的格式。 第 1 讲 §4.2 的 0.00 就是这种协议下的数字。 """ out = [None] * len(prompts) by_len = {} for i, p in enumerate(prompts): by_len.setdefault(len(p), []).append(i) with torch.no_grad(): for _, idxs in sorted(by_len.items()): seqs = torch.tensor([TOK.encode(prompts[i]) for i in idxs]) done = [""] * len(idxs) finished = [False] * len(idxs) for _ in range(max_new): nxt = fwd(seqs)[:, -1].argmax(-1) seqs = torch.cat([seqs, nxt[:, None]], 1) for j, c in enumerate(nxt.tolist()): if not finished[j]: ch = TOK.alphabet[c] done[j] += ch finished[j] = ch in ".\n" if all(finished): break for j, i in enumerate(idxs): out[i] = done[j] return out def bpc_with(fwd, data, T=88): """ 适配后的模型（适配器开着）在一段文本上的 bit/字符。 【思路】和第 1 讲的 bits_per_char 一样切成不重叠的窗口；窗口长 T = 88（而不是 96）， 给 prompt tuning 的 8 个软 token 留出位置，让所有方法用同一个窗口长度，数字才可比。 【参数】fwd：idx → logits；data：一维 LongTensor；T：窗口长度 【返回】float 【边界】和第 1 讲 bits_per_char（窗口 96）的数字略有不同（base model 是 0.412 而不是 0.409）：窗口短一点， 每个窗口开头"上下文不足"的位置占比就大一点。 【FM 联系】"适配后还剩多少原来的能力"——灾难性遗忘（catastrophic forgetting）的最简单的量法。 """ starts = range(0, len(data) - T - 1, T) xs = torch.stack([data[i:i + T] for i in starts]) ys = torch.stack([data[i + 1:i + T + 1] for i in starts]) total = 0.0 with torch.no_grad(): for i in range(0, len(xs), 128): logits = fwd(xs[i:i + 128]) total += F.cross_entropy(logits.reshape(-1, logits.size(-1)), ys[i:i + 128].reshape(-1), reduction="sum").item() return total / ys.numel() / math.log(2) # ============================================================================= # 3. 手写的前向传播：可以接软提示（prompt tuning）和每层前缀（prefix tuning） # ============================================================================= def forward_with(model, idx, soft=None, prefix=None): """ 和 TinyGPT.forward 完全相同的计算，多两个可选的"插口"。 【思路】 soft ：(m, d) 的张量。把它当作 m 个"虚拟 token 的嵌入"接在输入前面（prompt tuning，Lester et al. 2021）。 它们占用位置 0..m−1，真实文本从位置 m 开始；输出时把这 m 个位置的 logits 去掉。 prefix ：每层一对 (P_k, P_v)，各是 (m, d)。直接接在这一层注意力的 key / value 前面 （prefix tuning，Li & Liang 2021）：每个真实位置都能"看到"这 m 个额外的键值， 但它们不占位置嵌入、也不经过前面的层——比 soft 多了"每一层都能直接调"的自由度。 两者都不传时，结果与 model(idx) 逐位相同（自检里验证）。 模块（qkv、proj、mlp）通过 model 的属性调用，所以 LoRA / adapter 替换掉的模块也会自动生效。 【参数】model：TinyGPT（或被替换过模块的拷贝）；idx：(B, T) LongTensor 【返回】logits：(B, T, V) 【边界】T + m 不能超过 96（位置嵌入的个数）。 """ B, T = idx.shape x = model.tok_emb(idx) m = 0 if soft is not None: m = soft.shape[0] x = torch.cat([soft.unsqueeze(0).expand(B, -1, -1), x], dim=1) L = x.shape[1] x = x + model.pos_emb(torch.arange(L)) causal = torch.tril(torch.ones(L, L, dtype=torch.bool)) for li, blk in enumerate(model.blocks): h = blk.ln1(x) att = blk.attn D, H = h.shape[-1], att.n_head q, k, v = att.qkv(h).split(D, dim=2) mask = causal if prefix is not None: pk, pv = prefix[li] P = pk.shape[0] k = torch.cat([pk.unsqueeze(0).expand(B, -1, -1), k], dim=1) v = torch.cat([pv.unsqueeze(0).expand(B, -1, -1), v], dim=1) mask = torch.cat([torch.ones(L, P, dtype=torch.bool), causal], dim=1) # 前缀对所有位置可见 S = k.shape[1] q = q.view(B, L, H, D // H).transpose(1, 2) k = k.view(B, S, H, D // H).transpose(1, 2) v = v.view(B, S, H, D // H).transpose(1, 2) scores = (q @ k.transpose(-2, -1)) / math.sqrt(D // H) scores = scores.masked_fill(~mask, float("-inf")) y = (F.softmax(scores, dim=-1) @ v).transpose(1, 2).reshape(B, L, D) x = x + att.proj(y) x = x + blk.mlp(blk.ln2(x)) logits = model.head(model.ln_f(x)) return logits[:, m:] # ============================================================================= # 4. 两种"插件"：LoRA 与 adapter # ============================================================================= class LoRALinear(nn.Module): """ LoRA（Hu et al. 2021）：冻结原来的线性层 W，旁边加一个低秩的修正 ΔW = (α/r)·B·A。 y = W x + b + (α/r) · B (A x) A：(r, d_in)，B：(d_out, r)，r ≪ min(d_in, d_out) 【思路】 - A 用小的随机数初始化，B 初始化为 0 → 一开始 ΔW = 0，模型的输出和 base model **完全相同**； 训练从"原模型"出发，而不是从一个被随机扰动过的模型出发。 - 缩放 α/r：换一个秩 r 时，不用重新调学习率（论文 §4.1：α 设成第一次试的 r 之后就不再调）。 本讲取 α = r，即缩放为 1。 - merge()：训练完以后把 (α/r)·B·A 直接加进 W，插件就"消失"了——推理时没有任何额外计算。 【参数】base：nn.Linear；r：秩；alpha：缩放（默认等于 r）；gen：torch.Generator（初始化 A） 【返回】一个 nn.Module，可以直接替换原来的 nn.Linear（输入输出形状不变） 【边界】merge 之后再训练 A、B 没有意义（W 已经包含了它们）；需要"拔掉插件"时，在合并前保留一份 W。 【FM 联系】LoRA 的假设：适配时权重的**变化量**是低秩的（"内在维度低"，Aghajanyan et al. 2020）。 代码 03 检验这个假设，并演示合并、换插件。 """ def __init__(self, base, r, alpha=None, gen=None): super().__init__() self.base = base for p in self.base.parameters(): p.requires_grad_(False) self.r = r self.alpha = float(r if alpha is None else alpha) d_in, d_out = base.in_features, base.out_features self.A = nn.Parameter(torch.randn(r, d_in, generator=gen) / math.sqrt(d_in)) self.B = nn.Parameter(torch.zeros(d_out, r)) self.merged = False def delta_weight(self): """ΔW = (α/r)·B·A，形状和 W 相同 (d_out, d_in)。""" return (self.alpha / self.r) * (self.B @ self.A) def forward(self, x): y = self.base(x) if not self.merged: y = y + (self.alpha / self.r) * ((x @ self.A.t()) @ self.B.t()) return y @torch.no_grad() def merge(self): """把 ΔW 加进 W（就地修改）；之后 forward 只算 W x + b。""" if not self.merged: self.base.weight += self.delta_weight() self.merged = True class Adapter(nn.Module): """ Adapter（Houlsby et al. 2019）：在一个子层的输出后面接一个"瓶颈"小网络，带残差： h = sub(x); out = h + W_up · GELU(W_down · h) W_down：(m, d)，W_up：(d, m)，m ≪ d 【思路】W_up 初始化为 0 → 一开始 out = h，等于没有插件（论文强调"接近恒等的初始化"才能训得动）。 每个 Transformer 块放两个：注意力输出之后、MLP 输出之后（和论文的位置相同）。 【参数】sub：被包住的子层（nn.Module）；d：宽度；m：瓶颈宽度；gen：torch.Generator 【返回】一个 nn.Module，输入输出形状和 sub 相同 【边界】adapter 是串联在计算路径上的，推理时多两次小矩阵乘法，**不能**像 LoRA 那样合并掉。 """ def __init__(self, sub, d, m, gen=None): super().__init__() self.sub = sub self.down = nn.Parameter(torch.randn(m, d, generator=gen) * 0.01) self.up = nn.Parameter(torch.zeros(d, m)) def forward(self, x): h = self.sub(x) return h + F.gelu(h @ self.down.t()) @ self.up.t() # ============================================================================= # 5. 按方法名构造一个"适配模型" # ============================================================================= class Adapted: """ 一个适配方法的实例：模型拷贝 + 要训练的参数 + 前向函数。 【字段】 name : 方法名，例如 "lora4" model : base model 的拷贝（可能被替换了模块） params : list[nn.Parameter]，要训练的参数 fwd : 函数 idx → logits（prompt / prefix tuning 会走 forward_with） n_trainable : 要训练（也就是每个任务要额外存）的参数个数 """ def __init__(self, name, model, params, fwd): self.name, self.model, self.params, self.fwd = name, model, params, fwd self.n_trainable = sum(p.numel() for p in params) def _init_text_ids(m, gen, world): """ 从训练语料里随机切一段长 m 的真实文本，返回 (1, m) 的字符编号。 【思路】用来初始化软提示 / 前缀：Lester et al. 发现随机初始化最差、用真实词的嵌入好得多； Li & Liang 也发现用真实词的激活值初始化前缀，在数据少时明显更好。 【边界】world 必须传进来（make_adapted 的 prompt / prefix 方法需要它）。 """ text = world.train_text() start = int(torch.randint(0, len(text) - m - 1, (1,), generator=gen)) return torch.tensor([TOK.encode(text[start:start + m])]) def make_adapted(base, method, seed=0, world=None): """ 从 base model 的一份**新拷贝**出发，按方法名决定"训练哪些参数"。 【方法名】 "full" 全量微调：所有 109,056 个参数 "top" 只调最后一个 Transformer 块 + 最后的 LayerNorm（"只调最后几层"的常见做法） "bitfit" 只调所有偏置（bias）和 LayerNorm 的偏置（Zaken et al. 2021） "adapter{m}" 每个块两个瓶颈宽度为 m 的 adapter，例如 "adapter4" "lora{r}" 在每个块的 4 个线性层（qkv、注意力输出、MLP 两层）上加秩 r 的 LoRA，例如 "lora1" "loraqv{r}" 只在 qkv 上加 LoRA（论文主实验只调 W_q、W_v；本模型 q、k、v 是一个矩阵） "prompt{m}" m 个软 token 接在输入前（prompt tuning） "prefix{m}" 每层 m 个可训练的 key / value（prefix tuning） 【参数】base：TinyGPT（不会被改动）；seed：初始化插件用的随机种子；world：prompt / prefix 初始化要用语料 【返回】Adapted 【边界】未知的方法名抛出 ValueError。 """ gen = torch.Generator().manual_seed(seed) m = copy.deepcopy(base) for p in m.parameters(): p.requires_grad_(False) fwd = m # 默认：直接调用模型 if method == "full": params = list(m.parameters()) elif method == "top": params = list(m.blocks[-1].parameters()) + list(m.ln_f.parameters()) elif method == "bitfit": params = [p for n, p in m.named_parameters() if n.endswith("bias")] elif method.startswith("adapter"): width = int(method[len("adapter"):]) for blk in m.blocks: blk.attn.proj = Adapter(blk.attn.proj, m.cfg.d_model, width, gen) blk.mlp = Adapter(blk.mlp, m.cfg.d_model, width, gen) params = [p for p in m.parameters() if p.requires_grad] elif method.startswith("loraqv"): r = int(method[len("loraqv"):]) for blk in m.blocks: blk.attn.qkv = LoRALinear(blk.attn.qkv, r, gen=gen) params = [p for p in m.parameters() if p.requires_grad] elif method.startswith("lora"): r = int(method[len("lora"):]) for blk in m.blocks: blk.attn.qkv = LoRALinear(blk.attn.qkv, r, gen=gen) blk.attn.proj = LoRALinear(blk.attn.proj, r, gen=gen) blk.mlp[0] = LoRALinear(blk.mlp[0], r, gen=gen) blk.mlp[2] = LoRALinear(blk.mlp[2], r, gen=gen) params = [p for p in m.parameters() if p.requires_grad] elif method.startswith("prompt"): n = int(method[len("prompt"):]) ids = _init_text_ids(n, gen, world) soft = nn.Parameter(m.tok_emb(ids)[0].detach().clone()) # 用真实字符的嵌入初始化 params = [soft] def fwd(idx, _m=m, _s=soft): return forward_with(_m, idx, soft=_s) elif method.startswith("prefix"): n = int(method[len("prefix"):]) ids = _init_text_ids(n, gen, world) pre = [] with torch.no_grad(): # 用真实文本在各层的 key / value 初始化（Li & Liang 的做法） x = m.tok_emb(ids) + m.pos_emb(torch.arange(n)) for blk in m.blocks: _, k, v = blk.attn.qkv(blk.ln1(x)).split(m.cfg.d_model, dim=2) pre.append((nn.Parameter(k[0].clone()), nn.Parameter(v[0].clone()))) x = blk(x) params = [t for pair in pre for t in pair] def fwd(idx, _m=m, _p=pre): return forward_with(_m, idx, prefix=_p) else: raise ValueError(f"未知的方法：{method}") for p in params: p.requires_grad_(True) return Adapted(method, m, params, fwd) # ============================================================================= # 6. 统一的训练循环 # ============================================================================= def train_on_pairs(ad, pairs, steps, lr, batch_size=16, seed=0, weight_decay=0.0, callback=None, every=50): """ 在 (提示, 答案) 对上训练一个适配模型：只在答案字符上算交叉熵。 【思路】这就是监督微调（supervised fine-tuning, SFT）的最小形式。所有方法用同一个循环， 唯一的区别是 ad.params（训练哪些参数）和 ad.fwd（前向怎样走）。 k ≤ batch_size 时每一步都用全部 k 个样本；k 更大时每步随机抽 batch_size 个（有放回）。 【参数】ad：Adapted；pairs：list[(prompt, completion)]；steps、lr：步数和学习率（AdamW，恒定学习率） batch_size、seed：小批量大小与抽样种子 callback：可选，callback(step, ad) 每 every 步调用一次（包括第 0 步和最后一步），只应"读"模型 【返回】最后一步的训练损失（float，nats/字符） 【边界】k = 1、2 时一个 batch 只有一两个样本，几十步就能把它们背下来——这正是"少样本微调会过拟合"的来源。 """ x, y, mask = encode_pairs(pairs) gen = torch.Generator().manual_seed(seed) opt = torch.optim.AdamW(ad.params, lr=lr, weight_decay=weight_decay) ad.model.train() loss_val = float("nan") for step in range(steps + 1): if callback is not None and (step % every == 0 or step == steps): ad.model.eval() callback(step, ad) ad.model.train() if step == steps: break if len(pairs) <= batch_size: xb, yb, mb = x, y, mask else: bi = torch.randint(0, len(pairs), (batch_size,), generator=gen) xb, yb, mb = x[bi], y[bi], mask[bi] logits = ad.fwd(xb) per_pos = F.cross_entropy(logits.reshape(-1, logits.size(-1)), yb.reshape(-1), reduction="none") loss = (per_pos * mb.reshape(-1)).sum() / mb.sum() opt.zero_grad() loss.backward() opt.step() loss_val = loss.item() ad.model.eval() return loss_val # ============================================================================= # 7. 线性探针（多类别） # ============================================================================= def fit_softmax_probe(X, y, n_classes, l2=1e-2): """ 多类别逻辑回归（softmax 回归）：p(y = c | x) ∝ exp(w_c·x + b_c)，用 L-BFGS 解。 【思路】第 1 讲的探针是二分类；这里要在 8 个城市里选一个，所以每个类别一个权重向量。 L2 正则让样本很少时权重不会跑到无穷大。 【参数】X：(k, d) 张量；y：(k,) LongTensor（类别编号）；n_classes：类别数 【返回】函数 predict(X_test) → LongTensor 类别编号 【PyTorch 小知识】LBFGS 需要一个 closure：每次调用时重新算一遍损失和梯度；用 float64 让结果更稳定。 【边界】训练标签里没出现过的类别，探针几乎不可能预测出来——它的权重从没被推向那个类别。 这是探针和"在语言模型自己的输出上做选择"的根本区别：后者天生知道 8 个城市名。 """ Xd = X.double() W = torch.zeros(n_classes, X.shape[1], dtype=torch.float64, requires_grad=True) b = torch.zeros(n_classes, dtype=torch.float64, requires_grad=True) opt = torch.optim.LBFGS([W, b], max_iter=200, line_search_fn="strong_wolfe") def closure(): opt.zero_grad() loss = F.cross_entropy(Xd @ W.t() + b, y) + l2 * (W ** 2).sum() loss.backward() return loss opt.step(closure) Wd, bd = W.detach(), b.detach() return lambda Xt: (Xt.double() @ Wd.t() + bd).argmax(1) def last_hidden(fwd_model, prompts): """ 每个提示最后一个位置的最后一层表示（LayerNorm 之后），形状 (n, d)。 【思路】按长度分组成 batch，用 TinyGPT 的 return_hidden=True。读最后一个位置，是因为因果模型里 只有它看到了整个提示（第 3 讲 §5.4）。 【参数】fwd_model：TinyGPT（必须支持 return_hidden）；prompts：list[str] 【返回】(n, d) 的张量 【边界】只对普通的 TinyGPT 有意义；prompt / prefix tuning 的模型要走 forward_with，这里不支持。 """ out = [None] * len(prompts) by_len = {} for i, p in enumerate(prompts): by_len.setdefault(len(p), []).append(i) with torch.no_grad(): for _, idxs in sorted(by_len.items()): x = torch.tensor([TOK.encode(prompts[i]) for i in idxs]) _, h = fwd_model(x, return_hidden=True) for j, i in enumerate(idxs): out[i] = h[j, -1] return torch.stack(out) # ============================================================================= # 8. 缓存本讲自己训练的整模型（只写本目录的 checkpoints/） # ============================================================================= def cached_model(name, meta, build_and_train): """ 读取 code/lecture06/checkpoints/.pt；不存在或配置（meta）对不上，就调用 build_and_train() 训练并保存。 【思路】代码 04 的继续预训练要几十秒一个模型；第一次运行时训练并缓存，之后只读。 meta 记录了"决定这个模型的一切"（初始化方式、数据、步数、学习率、种子），配置一变旧缓存自动作废。 **只写本讲自己的 checkpoints/ 目录**，绝不碰第 1 讲的 base_model.pt。 【参数】name：文件名（不含 .pt）；meta：可 JSON 化的 dict； build_and_train：无参数函数，返回 (model, info)，info 是 dict（例如训练中途的评测记录） 【返回】(model, info, from_cache) 【边界】训练是确定性的（单线程 + 固定种子），删掉缓存重跑会得到同样的模型（同一台机器、同一个 PyTorch 版本）。 """ import json path = os.path.join(CKPT_DIR, name + ".pt") if os.path.exists(path): saved = torch.load(path, weights_only=False) if saved.get("meta") == meta: model = TinyGPT(saved["config"]) model.load_state_dict(saved["state_dict"]) model.eval() return model, saved["info"], True model, info = build_and_train() os.makedirs(CKPT_DIR, exist_ok=True) torch.save(dict(meta=meta, config=model.cfg, state_dict=model.state_dict(), info=info), path) with open(os.path.join(CKPT_DIR, name + "_meta.json"), "w") as f: json.dump(meta, f, indent=2) model.eval() return model, info, False # ============================================================================= # 9. 自检 # ============================================================================= if __name__ == "__main__": world, base = load_base() x, _, _ = encode_pairs([(FMT_CITY.format(name="kari"), " varno.")]) with torch.no_grad(): d0 = (forward_with(base, x) - base(x)).abs().max().item() print(f"1. 手写前向 forward_with 与 TinyGPT.forward 的最大差：{d0:.1e}") assert d0 < 1e-5 ad = make_adapted(base, "lora4", seed=0) with torch.no_grad(): d1 = (ad.fwd(x) - base(x)).abs().max().item() print(f"2. LoRA 刚初始化（B = 0）时输出与 base model 的最大差：{d1:.1e}；可训练参数 {ad.n_trainable:,}") assert d1 == 0.0 with torch.no_grad(): for p in ad.params: p.add_(0.05 * torch.randn(p.shape, generator=torch.Generator().manual_seed(1))) before = ad.fwd(x).clone() for mod in ad.model.modules(): if isinstance(mod, LoRALinear): mod.merge() d2 = (ad.fwd(x) - before).abs().max().item() print(f"3. 随便改一改 A、B 之后合并进 W：合并前后输出的最大差 {d2:.1e}") assert d2 < 1e-4 for name in ["full", "top", "bitfit", "adapter4", "lora1", "loraqv1", "prompt8", "prefix8"]: a = make_adapted(base, name, seed=0, world=world) with torch.no_grad(): same = (a.fwd(x) - base(x)).abs().max().item() print(f" {name:9s} 可训练参数 {a.n_trainable:7,}；训练前输出与 base 的最大差 {same:.2e}") print("自检通过。base model 未被改动：", all(not p.requires_grad for p in base.parameters()))