上一关 第 8 章 · SFT 是全量微调 —— 124M 参数全更新。
这一关换手段不换任务:还是同一件 SFT,但用
LoRA
—— 冻住底座,只在每层旁边挂一条低秩旁路 B·A,只训那 ~1% 的新参数,存盘只是个几 MB 的
adapter。
配套代码 phase2-sft-lora/07_lora.py。
第 8 章的全量 SFT 把 124M 个参数全部更新了一遍,训练要为每个参数都存梯度/优化器状态,存盘是一份完整的 ~498MB 新模型。可我们要教的只是"答题的姿态"这点小改动 —— 真有必要动到每一根钢筋吗? LoRA 的想法:底座一个数都不动,只在旁边搭一条很细的"改装管道",改动量压到 1% 级别。
原来一层是 y = W·x,W 是个大方阵(例如 768×768 ≈ 59 万个数)。
LoRA 不去改 W,而是在旁边并上 y = W·x + (B·A)·x ——
A 把 768 维压到 r 维(瓶颈),B 再升回 768 维。
r 很小,所以 A、B 加起来比 W 小得多。拨动 r 看整模型的账:
ΔW,其实用一个很"扁"的矩阵就能近似(秩很低),
不需要一个满秩的大方阵。所以拿 B·A(秩最多 r)去逼近它就够了。r 越大表达力越强、参数也越多 —— 是个旋钮。
lora_A: in→r,lora_B: r→out)· 实测 r=8 → 可训练 1,179,648
注入 LoRA 时只动两件事:① 把原层 W 整段冻结(requires_grad=False,永不更新);
② 旁路里 A 用小随机数、B 全初始化成 0。于是训练第 0 步旁路输出恒为 0,
"装了 LoRA 的模型"和原始 base 一模一样 —— 不会一上来就把 base 搅乱。点下面切换看:
normal_(A) + zeros_(B))与 inject_lora 前的全量 requires_grad=FalseA、B 有一个是 0,乘积 B·A 就是 0,初始旁路都为 0 —— 这点上两者等价。
但若两个都置零,反向传播时梯度也会卡住(谁都推不动谁),旁路永远学不起来。
所以标准做法是一个随机(A)、一个置零(B):既保证"初始 ΔW=0",又留出梯度通路让它能被训起来。
06_sft.py 逐字相同,只是优化器只拿到 A、B
zero_grad → forward(算 loss) → backward → step —— 一步没变。
loss mask、对话模板、EOS 全照搬 06。
唯一区别:交给 AdamW 的只有 requires_grad=True 的 A、B,
backward 的梯度也只流到旁路,底座纹丝不动。点"跑一个 epoch"看真实 loss:
1e-3 vs 全量 2e-5)、
多过几遍数据。底座本来就冻着,不怕学崩,可以大胆点。
lora_params=[p for p in model.parameters() if p.requires_grad] 只把它们交给优化器
下面是 07_lora.py 在 124M 上的真实采样。"训练前"是装了 LoRA 但 B=0(等价 base),
"训练后"只更新了 A、B。重点和第 8 章一样:有没有套上格式、有没有在 EOS 处停。
python 07_lora.py --ckpt ../phase1-124m/ckpt10b/latest.pt --epochs 60Wx + scaling·BAx(本关代码就是这样,一份底座可挂多个 adapter 切换任务);
② 合并 —— 把 scaling·B·A 直接加进 W 得到新权重(W' = W + scaling·BA),推理零额外开销,但就回到一份完整大模型、失去"可插拔"。
inject_lora 给每个 Block 的 注意力 c_attn / c_proj 和 MLP c_fc / c_proj 各挂一条(12 层 × 4 = 48 条)。
lm_head 和 wte 权重绑定、属底座,不加。