"""
peft_lab.py —— 参数高效微调（PEFT）的工具箱：LoRA、Adapter、前缀微调、只调顶层、随机子空间（第 10 讲）
=====================================================================
【这个文件要回答的问题】
一个预训练好的模型有很多参数。微调时能不能**只改很少一部分**（或者只加很少一点新参数），
却得到和"全部参数一起改"差不多的效果？slides 列了四种做法：
1. 只调一部分参数（subset）：比如只调最上面 K 层；
2. Adapter：在每层里插入一个"窄"的小网络，只训练它；
3. 前缀微调（prefix tuning）：假装序列前面还有几个"虚拟 token"，只训练它们在每层的 key / value；
4. LoRA：给每个权重矩阵加一个低秩的增量 ΔW = BA，只训练 A、B。
这个文件把四种做法都实现成"套在已有模型外面"的改装函数，作用在第 2 讲的 GPT（农场语言）和
第 5 讲的 ViT（形状世界）上——两者用的是同一个 Block / MultiHeadAttention（W_q、W_k、W_v、W_o 四个线性层）。
【本文件提供】
LoRALinear / apply_lora LoRA 线性层（可合并、可拆回）；把模型里指定的线性层换成 LoRA 版
Adapter / WithAdapter / apply_adapters 瓶颈结构的 adapter，插在注意力和前馈子层之后
PrefixAttention / apply_prefix 前缀微调：每层 n 个可训练的 key / value（可选 MLP 重参数化）
train_top_layers 只调最上面 K 个块（以及最后的 LayerNorm 和输出层）
freeze / count_trainable / trainable_state 冻结、数可训练参数、取出可训练参数的副本
Subspace 随机子空间训练（Li et al. 2018 的"内在维度"）：θ = θ₀ + P·θ_d
WINTER / front_time / make_task / TASKS 农场语言的两个"下游任务"：冬天的农场、时间前置的方言
load_base / encode_task / excess_nll / finetune 读取预训练的农场 GPT、编码、评估、带早停的微调循环
【约定】
- 所有"改装函数"都**原地**修改模型并冻结原有参数；想保留原模型请先 copy.deepcopy。
- 新加的参数都初始化成"一开始什么也不改"：LoRA 的 B = 0，adapter 的上投影 = 0。
所以改装后、训练前的模型输出和原模型逐位相同（前缀微调除外，见 PrefixAttention 的【边界】）。
"""
import contextlib
import copy
import math
import os
import sys
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
_HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture03", "lecture02", "lecture01"):
_p = os.path.join(_HERE, "..", _d)
if _p not in sys.path:
sys.path.insert(0, _p)
import toy_corpus # noqa: E402
from kv_cache import load_farm_gpt # noqa: E402
from lm_models import MultiHeadAttention, attention, causal_mask # noqa: E402
from rnn_lm import IGNORE, make_batches # noqa: E402
from toy_corpus import PROFILES, TIMES, generate_sentence # noqa: E402
# =====================================================================
# 1. LoRA
# =====================================================================
class LoRALinear(nn.Module):
"""
LoRA 线性层（Hu et al., 2021）：z = W₀x + b + (α/r)·B A x
【参数】
base : 原来的 nn.Linear（权重 W₀ 形状 (d, k)：输入 k 维、输出 d 维）；它会被冻结
r : 秩（rank），r << min(d, k)
alpha : 缩放常数 α，增量乘以 α/r（默认 α = r，也就是不缩放）
gen : torch.Generator，用来初始化 A（可复现）
init : "zero_B"（slides 的做法：A ~ N(0, σ²)、B = 0）或 "both_random"（A、B 都随机，做对照实验用）
【思路】
冻结 W₀，只学一个"增量" ΔW = BA，其中 A ∈ R^{r×k}、B ∈ R^{d×r}。ΔW 的秩最多是 r：
d×k 个数被压缩成 r(d + k) 个参数。前向时不必真的算出 ΔW，先算 Ax（r 维）再乘 B，更省。
【初始化】A 的每个元素 ~ N(0, 1/k)（σ = 1/√k，让 Ax 的每一维方差约等于 x 的平均方差），B = 0。
于是一开始 ΔW = 0，改装后的模型和预训练模型完全一样——微调从"原点"出发。
【合并（merge）】训练完以后可以把 ΔW 加回 W₀：W ← W₀ + (α/r)BA，之后就是一个普通线性层，
推理时没有任何额外计算（slides 的 "hot swapping"）；unmerge() 把它减回去，换上另一个任务的 B'A'。
【边界】
- 偏置 b 不参与 LoRA（slides：偏置本来就只有 d 个数，"已经是秩 1 了"，没什么可压缩的）。
- 合并 / 拆回用 float32 做加减，会有 1e-7 量级的舍入误差（01 号脚本量给你看）。
- 合并状态下再训练没有意义（A、B 不再参与前向），所以 merged 时 forward 不加增量。
"""
def __init__(self, base, r, alpha=None, gen=None, init="zero_B"):
super().__init__()
if r < 1:
raise ValueError("r 至少是 1")
if init not in ("zero_B", "both_random"):
raise ValueError("init 只能是 zero_B 或 both_random")
self.base = base
for p in base.parameters():
p.requires_grad_(False)
d, k = base.weight.shape
self.r, self.alpha = r, alpha if alpha is not None else r
self.scale = self.alpha / r
self.A = nn.Parameter(torch.randn(r, k, generator=gen) / math.sqrt(k))
if init == "zero_B":
self.B = nn.Parameter(torch.zeros(d, r))
else:
self.B = nn.Parameter(torch.randn(d, r, generator=gen) / math.sqrt(r))
self.merged = False
def delta_weight(self):
"""ΔW = (α/r)·B A，形状和 W₀ 相同 (d, k)。"""
return self.scale * (self.B @ self.A)
def forward(self, x):
out = self.base(x)
if not self.merged:
# 【为什么先乘 A 再乘 B】x @ Aᵀ 得到 r 维，再 @ Bᵀ 回到 d 维：计算量 r(k + d)，比先算出 d×k 的 ΔW 小得多
out = out + self.scale * ((x @ self.A.t()) @ self.B.t())
return out
@torch.no_grad()
def merge(self):
"""把增量加进 W₀（之后前向就是普通线性层）。重复调用不会加两次。"""
if not self.merged:
self.base.weight += self.delta_weight()
self.merged = True
@torch.no_grad()
def unmerge(self):
"""把增量从 W 里减掉，回到 W₀（slides：W ← W − BA = W₀）。"""
if self.merged:
self.base.weight -= self.delta_weight()
self.merged = False
ATTN_NAMES = ("W_q", "W_k", "W_v", "W_o")
FFN_NAMES = ("up", "down")
ALL_LINEAR = ATTN_NAMES + FFN_NAMES # 一个块里全部 6 个线性层
def apply_lora(model, r, targets=("W_q", "W_v"), alpha=None, seed=0, init="zero_B"):
"""
把模型里每个块的指定线性层换成 LoRALinear，并冻结其余全部参数。
【参数】
model : 第 2 讲的 GPT 或第 5 讲的 ViT（有 .blocks，每个块有 .attn（W_q/W_k/W_v/W_o）和 .ffn（up/down））
r、alpha、init : 见 LoRALinear
targets : 要加 LoRA 的线性层名字，取自 ATTN_NAMES 和 FFN_NAMES。
默认只加在 W_q、W_v 上——LoRA 论文对 GPT-3 的结论（slides："只加在 query 和 value 上最划算"）
seed : 初始化 A 用的随机种子
【返回】新加的 LoRALinear 模块列表（方便之后 merge / unmerge）
【边界】原地修改；模型的输出层（head）、嵌入、LayerNorm 全部冻结。分类任务需要的新输出层请调用者自己解冻。
"""
bad = [t for t in targets if t not in ATTN_NAMES + FFN_NAMES]
if bad:
raise ValueError(f"不认识的目标层：{bad}")
freeze(model)
gen = torch.Generator().manual_seed(seed)
added = []
for blk in model.blocks:
for name in targets:
owner = blk.attn if name in ATTN_NAMES else blk.ffn
lora = LoRALinear(getattr(owner, name), r, alpha=alpha, gen=gen, init=init)
# 【Python 小知识】setattr(对象, "名字", 值) 等价于 对象.名字 = 值；nn.Module 会把它登记成子模块
setattr(owner, name, lora)
added.append(lora)
return added
def lora_modules(model):
"""找出模型里所有的 LoRALinear（按 named_modules 的顺序）。"""
return [m for m in model.modules() if isinstance(m, LoRALinear)]
# =====================================================================
# 2. Adapter
# =====================================================================
class Adapter(nn.Module):
"""
瓶颈 adapter（Houlsby et al., 2019）：h ↦ h + W_up · GELU(W_down · h)
【参数】d：输入输出维度；r：中间的瓶颈维度（r << d）
【思路】一个只有一个隐藏层的前馈网络 + 残差连接（slides 的图：d → r → d）。
参数量 2dr + d + r，远小于一个 Transformer 层。
【初始化】W_up = 0、偏置 = 0：一开始 adapter 是恒等映射，插进去不改变原模型的输出（和 LoRA 的 B = 0 同一个道理）。
【边界】adapter 是**串行**插在原网络里的额外一层：推理时它永远要多算一次，没法像 LoRA 那样合并掉
——这就是 LoRA 论文说的"adapter 带来推理延迟"（05 号脚本量这个延迟）。
"""
def __init__(self, d, r, gen=None):
super().__init__()
self.down = nn.Linear(d, r)
self.up = nn.Linear(r, d)
with torch.no_grad():
self.down.weight.copy_(torch.randn(r, d, generator=gen) / math.sqrt(d))
self.down.bias.zero_()
self.up.weight.zero_()
self.up.bias.zero_()
def forward(self, h):
return h + self.up(F.gelu(self.down(h)))
class WithAdapter(nn.Module):
"""把一个子层（注意力或前馈）和它后面的 adapter 串起来：x ↦ Adapter(sublayer(x))。"""
def __init__(self, sublayer, adapter):
super().__init__()
self.sublayer, self.adapter = sublayer, adapter
def forward(self, x):
return self.adapter(self.sublayer(x))
def apply_adapters(model, r, where=("attn", "ffn"), seed=0):
"""
在每个块的注意力子层和 / 或前馈子层之后插一个 adapter（slides 的图：两处都插），冻结其余参数。
【思路】我们的 Block 是 pre-norm：x = x + Attn(LN(x))。把 blk.attn 换成 WithAdapter(attn, adapter) 后变成
x = x + Adapter(Attn(LN(x)))——adapter 位于子层输出、残差相加之前，和 Houlsby 的位置一致。
【返回】新加的 Adapter 列表
"""
freeze(model)
gen = torch.Generator().manual_seed(seed)
added = []
for blk in model.blocks:
d = blk.ln1.normalized_shape[0]
for name in where:
ad = Adapter(d, r, gen=gen)
setattr(blk, name, WithAdapter(getattr(blk, name), ad))
added.append(ad)
return added
# =====================================================================
# 3. 前缀微调（prefix tuning）
# =====================================================================
class PrefixAttention(nn.Module):
"""
前缀微调（Li & Liang, 2021）：在每一层的注意力里，假装真实序列前面还有 n 个"虚拟 token"，
它们的 key、value 是可训练的参数 P_k、P_v（形状 (n, d_model)），其余参数全部冻结。
【参数】
attn : 原来的 MultiHeadAttention（冻结，权重直接复用）
n_prefix : 虚拟 token 的个数 n
reparam : True 时按 slides 的第 3 步做重参数化：P = MLP(Q)，Q 的维度 d_small 比 P 小，
训练时更稳定；训练完以后只需要保存 P（MLP 可以扔掉）
d_small、hidden : 重参数化的 Q 的维度、MLP 隐藏层的宽度
gen : 初始化用的随机数发生器
【思路】第 t 个真实 token 的查询 q_t 现在能看到 n + t 个位置：n 个前缀（任何位置都能看）+ 自己和之前的真实 token。
前缀不需要被预测、也不需要位置编码；它们的作用相当于"每层都塞进一段学出来的上下文"。
【边界】和 LoRA / adapter 不同，前缀一加上，注意力的 softmax 分母就多了 n 项——**训练前**模型的输出就已经变了。
Li & Liang 发现随机初始化的前缀不太稳定（用真实词的激活初始化更好）；这里用很小的随机值，
05 号脚本会量"改装后、训练前"损失变了多少。
"""
def __init__(self, attn, n_prefix, reparam=True, d_small=16, hidden=64, gen=None):
super().__init__()
if not isinstance(attn, MultiHeadAttention):
raise TypeError("PrefixAttention 只能包住 MultiHeadAttention")
self.attn, self.n = attn, n_prefix
d = attn.d_model
self.reparam = reparam
if reparam:
self.Q = nn.Parameter(torch.randn(n_prefix, d_small, generator=gen))
self.mlp = nn.Sequential(nn.Linear(d_small, hidden), nn.Tanh(), nn.Linear(hidden, 2 * d))
with torch.no_grad():
for lin in (self.mlp[0], self.mlp[2]):
lin.weight.copy_(torch.randn(lin.weight.shape, generator=gen) * 0.02)
lin.bias.zero_()
else:
self.P = nn.Parameter(torch.randn(n_prefix, 2 * d, generator=gen) * 0.02)
def prefix_kv(self):
"""返回 (P_k, P_v)，各自形状 (n, d_model)。"""
P = self.mlp(self.Q) if self.reparam else self.P
d = self.attn.d_model
return P[:, :d], P[:, d:]
def forward(self, x):
a = self.attn
B, T, _ = x.shape
q, k, v = a.split_heads(a.W_q(x)), a.split_heads(a.W_k(x)), a.split_heads(a.W_v(x))
pk, pv = self.prefix_kv()
# (n, d) → (1, n, d) → 切头 (1, h, n, d_k) → 沿批维度广播成 (B, h, n, d_k)
pk = a.split_heads(pk.unsqueeze(0)).expand(B, -1, -1, -1)
pv = a.split_heads(pv.unsqueeze(0)).expand(B, -1, -1, -1)
k, v = torch.cat([pk, k], dim=2), torch.cat([pv, v], dim=2)
# 掩码 (T, n + T)：前 n 列（前缀）全部可见，后 T 列照常是因果掩码
mask = torch.cat([torch.zeros(T, self.n, dtype=x.dtype), causal_mask(T).to(x.dtype)], dim=1) if a.causal \
else None
out, _ = attention(q, k, v, mask)
return a.W_o(out.transpose(1, 2).reshape(B, T, a.d_model))
def apply_prefix(model, n_prefix, reparam=True, seed=0, d_small=16, hidden=64):
"""把每个块的注意力换成 PrefixAttention，冻结其余参数；返回新模块列表。"""
freeze(model)
gen = torch.Generator().manual_seed(seed)
added = []
for blk in model.blocks:
pa = PrefixAttention(blk.attn, n_prefix, reparam=reparam, d_small=d_small, hidden=hidden, gen=gen)
blk.attn = pa
added.append(pa)
return added
# =====================================================================
# 4. 只调一部分参数；通用小工具
# =====================================================================
def freeze(model):
"""把模型现有的全部参数设成不需要梯度。"""
for p in model.parameters():
p.requires_grad_(False)
def train_top_layers(model, K, head=True):
"""
slides 的最简单基线：冻结全部参数，只解冻最上面 K 个块（以及最后的 LayerNorm 和输出层）。
【思路】梯度只需要传过 K 层：更低的层既不需要存"参数的梯度"，反向传播也可以在第 K 层之下停住。
【边界】K = 0 且 head=True 时只训练输出层（第 4 讲的"线性探针"）；K 不能超过总层数。
"""
if K > len(model.blocks):
raise ValueError("K 超过了总层数")
freeze(model)
mods = list(model.blocks)[len(model.blocks) - K:] if K else []
if head:
mods += [m for m in (getattr(model, "ln_f", None), getattr(model, "norm", None), model.head) if m is not None]
for m in mods:
for p in m.parameters():
p.requires_grad_(True)
def count_trainable(model):
"""(可训练参数个数, 全部参数个数)。"""
total = sum(p.numel() for p in model.parameters())
train = sum(p.numel() for p in model.parameters() if p.requires_grad)
return train, total
def trainable_state(model):
"""取出所有可训练参数的副本（早停时保存"最好的那一刻"用）。"""
return {n: p.detach().clone() for n, p in model.named_parameters() if p.requires_grad}
def load_trainable_state(model, state):
with torch.no_grad():
params = dict(model.named_parameters())
for n, v in state.items():
params[n].copy_(v)
# =====================================================================
# 5. 随机子空间训练（内在维度，Li et al. 2018；Aghajanyan et al. 2020）
# =====================================================================
class Subspace(nn.Module):
"""
只在一个随机的 d 维子空间里训练模型：θ = θ₀ + P·θ_d，其中 θ₀ 是初始参数（D 维，冻结），
P 是固定的 D×d 随机矩阵，**只有 θ_d（d 个数）是可训练的**。
【参数】model：任意 nn.Module（不会被修改）；d：子空间维度；seed：生成 P 的种子
trainable_extra : 参数名的前缀列表，这些参数**不**放进子空间、而是照常完整训练（例如新任务的分类头）
【思路】Li et al. (2018) 的定义：从小到大试 d，第一次能达到"完整训练效果的 90%"时的 d 叫内在维度 d₉₀。
如果一个问题真的需要 D 个自由度，d << D 时应该学不会；如果 d 很小就够了，说明"真正需要的自由度"很少。
【P 怎么造】D 可能有十几万，d 上千时稠密的 P 要几亿个数。我们用"非常稀疏的随机投影"（Li, Hastie & Church, 2006）：
P 的每一列只有 nnz_per_col 个非零元素（默认 64 个；随机位置、随机正负号），再把每列归一化成单位长度
（Li et al. 2018 也做了这个归一化）。d 个这样的随机列彼此几乎正交，张成一个"随机的 d 维子空间"。
用 torch.sparse 存，乘法和反向传播都很快。Li et al. 用的是每列约 √D 个非零元素的稀疏投影或 Fastfood 变换，
我们取更稀疏的 64 个是为了在 CPU 上跑得动。
【Python / PyTorch 小知识】torch.func.functional_call(model, 参数字典, 输入) 用给定的参数值跑模型的前向，
而不动模型自己的参数——这样 θ 就可以是 θ_d 的函数，梯度自动传回 θ_d。
【边界】θ_d 初始化为 0，所以一开始 θ = θ₀，模型就是原来的模型。
"""
def __init__(self, model, d, seed=0, trainable_extra=(), nnz_per_col=64):
super().__init__()
self.model = model
freeze(model)
self.names, self.shapes, chunks = [], [], []
self.extra = nn.ParameterDict()
for n, p in model.named_parameters():
if any(n.startswith(e) for e in trainable_extra):
self.extra[n.replace(".", "__")] = nn.Parameter(p.detach().clone())
continue
self.names.append(n)
self.shapes.append(p.shape)
chunks.append(p.detach().reshape(-1))
self.register_buffer("theta0", torch.cat(chunks))
D = self.theta0.numel()
self.D, self.d = D, d
g = torch.Generator().manual_seed(seed)
k = min(nnz_per_col, D) # 每列的非零个数
rows = torch.randint(D, (d * k,), generator=g)
cols = torch.arange(d).repeat_interleave(k)
vals = (torch.randint(2, (d * k,), generator=g).float() * 2 - 1) / math.sqrt(k)
# 【Python 小知识】coalesce() 把同一位置的重复元素加起来（随机位置偶尔会撞上），得到规范的稀疏张量
self.P = torch.sparse_coo_tensor(torch.stack([rows, cols]), vals, (D, d), check_invariants=True).coalesce()
self.theta_d = nn.Parameter(torch.zeros(d))
def params(self):
theta = self.theta0 + torch.sparse.mm(self.P, self.theta_d.unsqueeze(1)).squeeze(1)
out, i = {}, 0
for n, s in zip(self.names, self.shapes):
m = math.prod(s)
out[n] = theta[i:i + m].view(s)
i += m
for key, p in self.extra.items():
out[key.replace("__", ".")] = p
return out
def forward(self, *args):
return torch.func.functional_call(self.model, self.params(), args)
# =====================================================================
# 6. 农场语言的两个下游任务
# =====================================================================
# "冬天的农场"：语法和用词完全不变，只是动物们的习惯变了——大多待在谷仓（barn）里，牛和马主要吃干草（hay）。
# 只改了"时间 / 地点短语"和两种动物的食物分布，是一个**内容上的小变化**。
WINTER = copy.deepcopy(PROFILES)
_WINTER_TIMES = {
"bat": {"barn": 7, "night": 3},
"cow": {"barn": 8, "morning": 2},
"dog": {"barn": 4, "morning": 4, "night": 2},
"cat": {"barn": 6, "night": 4},
"owl": {"barn": 5, "night": 5},
"hen": {"barn": 8, "morning": 2},
"fox": {"night": 8, "field": 2},
"horse": {"barn": 7, "day": 3},
}
for _a, _t in _WINTER_TIMES.items():
WINTER[_a]["times"] = _t
WINTER["cow"]["objects"]["eat"] = {"hay": 8, "corn": 2}
WINTER["horse"]["objects"]["eat"] = {"hay": 9, "grass": 1}
_TIME_PHRASES = {tuple(p): k for k, p in TIMES.items()}
@contextlib.contextmanager
def profiles_patched(profiles):
"""
临时把 toy_corpus 的 PROFILES 换成另一套习性，离开 with 块时一定换回来。
【为什么这样做】第 1 讲的 generate_sentence 在运行时才去读模块里的全局变量 PROFILES，
所以临时替换它，就能用**完全相同的语法和概率计算**生成"冬天"的句子，并得到正确的真实对数概率。
【Python 小知识】@contextlib.contextmanager 把一个带 yield 的函数变成可以用在 with 语句里的对象；
try/finally 保证即使中间出错，原来的 PROFILES 也会被恢复。
"""
old = toy_corpus.PROFILES
toy_corpus.PROFILES = profiles
try:
yield
finally:
toy_corpus.PROFILES = old
def time_phrase_start(tokens):
"""原来语序的句子里，句末时间短语从第几个词开始（下标从 0 算）。时间短语只有 2 个词或 3 个词两种长度。"""
body = tokens[:-1]
for n in (3, 2):
if tuple(body[-n:]) in _TIME_PHRASES:
return len(body) - n
raise ValueError(f"找不到时间短语：{tokens}")
def front_time(tokens):
"""
"时间前置的方言"：把句末的时间 / 地点短语挪到句首。
["the", "dogs", "chase", "the", "cats", "at", "night", "."] → ["at", "night", "the", "dogs", "chase", "the", "cats", "."]
【边界】这是一一对应的变换（时间短语只有 5 种，一定在句号前面），所以方言句子的真实概率和原句完全相同。
"""
i = time_phrase_start(tokens)
return tokens[i:-1] + tokens[:i] + ["."]
TASKS = {
"winter": "冬天的农场（只改了习性）",
"fronted": "时间前置的方言（改了语序）",
"original": "原来的农场语言",
}
def make_task(task, n, seed):
"""
生成某个任务的 n 句话。
【返回】列表，每个元素是 (词列表, 真实对数概率 log p_true)
【边界】"original" 就是第 1 讲的农场语言（用来量"遗忘"）；不同 seed 的样本互相独立。
"""
import random
rng = random.Random(seed)
out = []
if task == "winter":
with profiles_patched(WINTER):
for _ in range(n):
ex = generate_sentence(rng)
out.append((ex.tokens, ex.logprob_true))
elif task in ("fronted", "original"):
for _ in range(n):
ex = generate_sentence(rng)
toks = front_time(ex.tokens) if task == "fronted" else ex.tokens
out.append((toks, ex.logprob_true))
else:
raise ValueError(f"不认识的任务 {task}")
return out
# =====================================================================
# 7. 读取预训练模型、编码、评估、微调
# =====================================================================
def load_base(log=None):
"""读取第 2 讲缓存的农场语言 GPT（11 万参数）。【返回】(model, vocab)"""
model, vocab, _ = load_farm_gpt(log=log)
model.eval()
return model, vocab
def encode_task(data, vocab):
"""[(词列表, logp)] → ([, ..., ] 的 id 列表, 真实 NLL 的总和, 预测的 token 总数)。"""
seqs = [[vocab.bos_id] + vocab.encode(t) + [vocab.eos_id] for t, _ in data]
true_nll = -sum(lp for _, lp in data)
n_tok = sum(len(s) - 1 for s in seqs)
return seqs, true_nll, n_tok
@torch.no_grad()
def model_nll(model, seqs, pad_id, batch_size=500, skip_first=False):
"""
模型在这些序列上的总负对数似然（nats）。
【参数】skip_first=True 时不计第一个预测（ 之后的第一个词），见 excess_nll 的说明。
"""
model.eval()
total = 0.0
for i in range(0, len(seqs), batch_size):
inp, tgt = make_batches(seqs[i:i + batch_size], batch_size, pad_id)[0]
if skip_first:
tgt = tgt.clone()
tgt[:, 0] = IGNORE
logits = model(inp)
logits = logits[0] if isinstance(logits, tuple) else logits
total += F.cross_entropy(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1),
ignore_index=IGNORE, reduction="sum").item()
return total
def excess_nll(model, encoded, pad_id, skip_first=False):
"""
超出真实分布的损失（nats / 词）= (模型的 NLL − 真实分布的 NLL) / 预测的 token 数。
【直觉】这是"模型分布离真实分布还有多远"的平均每词 KL 的估计：0 表示和真实分布一样好，越小越好。
比直接报困惑度好比较：不同任务的真实熵不同，减掉它之后数字可以横着比。
【skip_first】量"遗忘"时用：在时间前置的方言里，句子从不以 the 开头，所以一个完美的方言模型必然认为
原来的句子"开头就错了"——这一步的损失是两种语言**本来就冲突**的地方，不算遗忘。去掉第一个预测以后，
剩下的部分（" the" 之后怎么接）方言数据完全没有约束，模型还记得多少，才是真正的遗忘。
原来的语言里第一个词一定是 the（真实概率 1，真实 NLL 为 0），所以真实 NLL 那一项不用改。
"""
seqs, true_nll, n_tok = encoded
if skip_first:
n_tok -= len(seqs)
return (model_nll(model, seqs, pad_id, skip_first=skip_first) - true_nll) / n_tok
def finetune(model, train_seqs, valid_seqs, pad_id, steps=600, lr=1e-3, batch=32, eval_every=50, seed=0,
params=None):
"""
带早停（early stopping）的微调：只更新 requires_grad=True 的参数（或显式给出的 params）。
【参数】
model : 已经改装好（冻结了该冻结的）的模型
train_seqs : 训练序列（[, ..., ] 的 id 列表）；valid_seqs：验证序列（用来早停）
steps、lr、batch、seed
eval_every : 每隔多少步在验证集上算一次损失；验证损失最低的那一刻的参数会被保留
【返回】字典 {"best_step", "best_valid", "valid_at_0", "curve"：[(步数, 验证 NLL/词)], "ms_per_step"}
【边界】
- 训练数据少（几百句）时，跑满 steps 步会过拟合；早停相当于"在验证集上挑最好的训练步数"。
- 验证 NLL 按"每个预测的 token"平均。
- 计时只算训练步（不含验证），标 ⏱ 的数字每次运行略有不同。
"""
import random
params = [p for p in model.parameters() if p.requires_grad] if params is None else list(params)
opt = torch.optim.Adam(params, lr=lr)
n_valid_tok = sum(len(s) - 1 for s in valid_seqs)
def valid_loss():
return model_nll(model, valid_seqs, pad_id) / n_valid_tok
best = valid_loss()
info = {"valid_at_0": best, "best_step": 0, "curve": [(0, best)]}
best_state = trainable_state(model)
rng = random.Random(seed)
batches, t_train, step = [], 0.0, 0
while step < steps:
if not batches:
batches = list(make_batches(train_seqs, batch, pad_id, rng))
inp, tgt = batches.pop()
model.train()
t0 = time.perf_counter()
logits = model(inp)
logits = logits[0] if isinstance(logits, tuple) else logits
loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1), ignore_index=IGNORE)
opt.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(params, 1.0)
opt.step()
t_train += time.perf_counter() - t0
step += 1
if step % eval_every == 0 or step == steps:
v = valid_loss()
info["curve"].append((step, v))
if v < best:
best, info["best_step"] = v, step
best_state = trainable_state(model)
load_trainable_state(model, best_state)
model.eval()
info["best_valid"] = best
info["ms_per_step"] = 1000 * t_train / max(step, 1)
return info