""" ============================================================================= 第 6 讲 · 代码 3：从零实现 LoRA，并检验它背后的"低秩"假设 ============================================================================= 对应讲义：lectures/lecture06_adaptation.md 第 4.6–4.8 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- LoRA（Hu et al. 2021）的全部想法只有一行：W ← W + (α/r)·B·A。这个文件把它拆开看： 演示 1：零件。参数怎么数；B 初始化为 0 时输出和原模型逐位相同；训练完把 B·A 合并进 W， 合并前后的输出差多少；一个 LoRA 插件存成文件有多大。 演示 2：秩 r 和缩放 α/r。r = 1、4、16、64，两种缩放（α = r，以及论文的做法：α 固定）； 预算相同时，LoRA 放在注意力上还是放在所有线性层上。 演示 3：全量微调的更新量 ΔW 真的是低秩的吗？对每个权重矩阵的 ΔW 做奇异值分解（SVD）， 看能量集中在前几个方向；再把 ΔW 截断到秩 r，看准确率还剩多少。 演示 4：内在维度（intrinsic dimension，Aghajanyan et al. 2020）：只允许参数在一个随机的 d 维子空间里移动， d 要多大才能做到全量微调九成的准确率？ 演示 5：一个基座、多个插件：两个任务各训练一个 LoRA，换插件就换任务；比较存储。 任务和代码 01 相同（评论的新格式 + 新标签词），另加一个"服务好不好"的新任务（演示 5）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 03_lora_from_scratch.py # 约 50 秒 需要 PyTorch、numpy。依赖本目录的 adapt_lab.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py。 【PyTorch 小知识】 - torch.linalg.svd(M, full_matrices=False)：M = U · diag(S) · Vh，S 从大到小排列。 奇异值的平方 S² 是矩阵在各个方向上的"能量"，前 r 个奇异值就是 M 的最佳秩 r 近似（Eckart–Young 定理）。 - torch.func.functional_call(model, params, (x,))：用另一组参数（一个 dict）跑同一个模型的前向， 不改模型本身。演示 4 用它让参数 = θ0 + P·z，梯度只流向 z。 - io.BytesIO + torch.save：把张量"存"进内存里的一个文件，量出它有多大，不用真的写磁盘。 """ import copy import io import math import os import random import sys import time import numpy as np import torch import torch.nn.functional as F from torch.func import functional_call HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from adapt_lab import (VERDICT_WORDS, LoRALinear, choice_accuracy, encode_pairs, load_base, # noqa: E402 make_adapted, train_on_pairs, verdict_items) from console import pad, print_header # noqa: E402 from fm_world import FOOD_NEG, FOOD_POS, make_review # noqa: E402 STEPS = 150 def make_data(): """和代码 01 完全相同的数据（同样的种子 11），这样两边的数字可以对照。""" rng = random.Random(11) seen, unseen = (FOOD_POS[:4], FOOD_NEG[:4]), (FOOD_POS[4:], FOOD_NEG[4:]) return dict(pool=verdict_items(400, rng, seen), dev=verdict_items(200, rng, seen), test_id=verdict_items(300, rng, seen), test_ood=verdict_items(300, rng, unseen)) def sample_k(pool, k, seed): """ 和代码 01 相同的抽样（同一个 k、同一个种子抽到同样的样本），直接返回训练用的 (提示, " 标签词.") 对。 【参数】pool：标注池；k：样本数；seed：种子 【返回】list[(str, str)] """ rng = np.random.default_rng(1000 + 10 * k + seed) return [(pool[i][0], f" {pool[i][1]}.") for i in rng.permutation(len(pool))[:k]] def lora_modules(model): """ 模型里所有的 LoRALinear，按模块的遍历顺序返回（每次调用顺序相同，所以可以和保存的 A、B 列表一一对应）。 【Python 小知识】model.modules() 递归地遍历模型里的每一个子模块（包括它自己）。 """ return [m for m in model.modules() if isinstance(m, LoRALinear)] def nbytes(tensors): """ 一组张量（dict）用 torch.save 存下来有多少字节。 【边界】包含少量文件头和键名，所以比"参数个数 × 4 字节"略多一点。 """ buf = io.BytesIO() torch.save(tensors, buf) return buf.getbuffer().nbytes # ============================================================================= # 演示 1：零件 # ============================================================================= def demo_parts(base, data): print_header("演示 1：LoRA 的零件 —— 参数怎么数、为什么从 0 开始、合并之后还剩什么") d_in, d_out = 64, 192 print(f"一个线性层 W：{d_in} → {d_out}（每个块的 qkv 矩阵）。全量微调要更新 {d_in * d_out:,} 个数；") print("LoRA 只训练 A（r × d_in）和 B（d_out × r）：") for r in [1, 4, 8, 64]: print(f" r = {r:<3d} → {r * (d_in + d_out):6,} 个参数（{r * (d_in + d_out) / (d_in * d_out):.1%}）") ad = make_adapted(base, "lora4", seed=0) test = data["test_id"] x, _, _ = encode_pairs([(p, f" {a}.") for p, a in test[:50]]) with torch.no_grad(): d0 = (ad.fwd(x) - base(x)).abs().max().item() print(f"\n刚初始化（A 随机、B = 0）：50 条评论上，输出与 base model 的最大差 = {d0:.1e}") train_on_pairs(ad, sample_k(data["pool"], 16, 0), STEPS, 1e-3, seed=0) with torch.no_grad(): before = ad.fwd(x).clone() acc_before = choice_accuracy(ad.fwd, test, VERDICT_WORDS) lora_state = {n: p.detach().clone() for n, p in ad.model.named_parameters() if n.endswith((".A", ".B"))} norms = [(m.delta_weight().norm() / m.base.weight.norm()).item() for m in lora_modules(ad.model)] for m in lora_modules(ad.model): m.merge() with torch.no_grad(): d1 = (ad.fwd(x) - before).abs().max().item() acc_after = choice_accuracy(ad.fwd, test, VERDICT_WORDS) full_state = {n: p for n, p in base.state_dict().items()} print(f"训练 {STEPS} 步（r = 4，k = 16）：同分布测试 {acc_before:.3f}；" f"ΔW 的大小只有 W 的 {min(norms):.1%}–{max(norms):.1%}（Frobenius 范数之比）") print(f"把 B·A 合并进 W 之后：输出的最大差 = {d1:.1e}，准确率 {acc_after:.3f}（合并前 {acc_before:.3f}）") print(f"存一个 LoRA 插件：{nbytes(lora_state) / 1024:.1f} KB（{sum(t.numel() for t in lora_state.values()):,} 个数）；" f"存一整份模型：{nbytes(full_state) / 1024:.1f} KB") print(""" 观察 Observations： · B = 0 让训练从"原模型"出发：第 0 步的输出和 base model 逐位相同。如果 A、B 都随机初始化， 一开始就给模型加了一个随机扰动，等于先把它弄坏一点再修。 · 合并后的输出和合并前只差浮点舍入（1e-5 量级），准确率完全相同：部署时把 W + B·A 当成新的 W， 推理的计算量和原模型一模一样——这是 LoRA 相对 adapter 的一个实际优势（adapter 串在计算路径上，去不掉）。 · ΔW 只有 W 的百分之几到十几：适配是在原权重上做一个小的修正，而不是把它推翻。""") # ============================================================================= # 演示 2：秩与缩放 # ============================================================================= def demo_rank(base, data): print_header("演示 2：秩 r 与缩放 α/r（k = 16，学习率 1e-3，150 步，2 个种子平均；同分布 / 分布外）") ranks = [1, 4, 16, 64] print(pad("r", 6) + pad("可训练参数", 12, ">") + pad("α = r（缩放 1）", 22, ">") + pad("α = 8 固定（缩放 8/r）", 26, ">")) for r in ranks: cells = [] n = 0 for alpha in ["r", 8]: accs = [] for seed in [0, 1]: ad = make_adapted(base, f"lora{r}", seed=seed) n = ad.n_trainable if alpha != "r": for m in lora_modules(ad.model): m.alpha = float(alpha) train_on_pairs(ad, sample_k(data["pool"], 16, seed), STEPS, 1e-3, seed=seed) accs.append((choice_accuracy(ad.fwd, data["test_id"], VERDICT_WORDS), choice_accuracy(ad.fwd, data["test_ood"], VERDICT_WORDS))) a = np.mean(accs, axis=0) cells.append(f"{a[0]:.3f} / {a[1]:.3f}") print(pad(r, 6) + pad(f"{n:,}", 12, ">") + pad(cells[0], 22, ">") + pad(cells[1], 26, ">")) print("\n同样约 2,048 个参数，放在哪里？（α = r，学习率 1e-3，2 个种子平均）") for name, label in [("loraqv4", "只放在注意力的 qkv 上，r = 4"), ("lora1", "4 个线性层都放，r = 1")]: accs = [] for seed in [0, 1]: ad = make_adapted(base, name, seed=seed) train_on_pairs(ad, sample_k(data["pool"], 16, seed), STEPS, 1e-3, seed=seed) accs.append((choice_accuracy(ad.fwd, data["test_id"], VERDICT_WORDS), choice_accuracy(ad.fwd, data["test_ood"], VERDICT_WORDS))) a = np.mean(accs, axis=0) print(f" {pad(label, 30)} 参数 {ad.n_trainable:,}：{a[0]:.3f} / {a[1]:.3f}") print(""" 观察 Observations： · α = r 时（缩放恒为 1），r 越小学得越慢：同样的学习率、同样 150 步，秩 1 的 ΔW 走不了那么远。 · 按论文的做法把 α 固定、缩放取 α/r，秩 1 的成绩明显上升；秩大的时候两种缩放差别不大。 缩放 α/r 本质上是在给 LoRA 的参数换一个学习率——论文引入它，就是为了换秩时不用重新调学习率。 · 在 k = 16、150 步这个设定里，秩从 4 加到 16 还有好处，从 16 到 64（已经等于模型宽度，满秩）就不再明显提高。 代码 01 里 k = 64 时 r = 1 也到了 0.970：数据多一些，小秩就够了。LoRA 论文在 GPT-3 上看到的是同一个方向—— 只调 W_q、W_v 时 r = 1 就基本够用（WikiSQL 73.4、MultiNLI 91.3）。 · 同样约 2,048 个参数，"4 个线性层各放一点"明显好于"全部放在注意力的 qkv 上"。LoRA 论文只调注意力 （它的预算表里 W_q 和 W_v 一起调最好）；后来 QLoRA（2023）发现，要追上全量微调，LoRA 得放在**所有**线性层上。""") # ============================================================================= # 演示 3：全量微调的 ΔW 是低秩的吗？ # ============================================================================= def demo_full_ft_rank(base, data): print_header("演示 3：全量微调的更新量 ΔW 是低秩的吗？（k = 64，lr 3e-4，150 步，种子 0）") ad = make_adapted(base, "full", seed=0) train_on_pairs(ad, sample_k(data["pool"], 64, 0), STEPS, 3e-4, seed=0) acc_full = choice_accuracy(ad.fwd, data["test_id"], VERDICT_WORDS) print(f"全量微调后同分布测试 {acc_full:.3f}。每个权重矩阵的 ΔW = W_微调后 − W_原来：") bp = dict(base.named_parameters()) fp = dict(ad.model.named_parameters()) mats = [n for n, p in bp.items() if p.dim() == 2 and "emb" not in n] g = torch.Generator().manual_seed(0) print(pad("矩阵", 26) + pad("形状", 12) + pad("‖ΔW‖/‖W‖", 11, ">") + pad("前 1 个方向", 12, ">") + pad("前 4 个", 9, ">") + pad("前 16 个", 10, ">") + pad("随机矩阵：前 4 个", 18, ">")) for n in mats: dW = (fp[n] - bp[n]).detach() s2 = torch.linalg.svdvals(dW) ** 2 e = s2.cumsum(0) / s2.sum() rs2 = torch.linalg.svdvals(torch.randn(dW.shape, generator=g)) ** 2 re = rs2.cumsum(0) / rs2.sum() print(pad(n.replace(".weight", ""), 26) + pad(str(tuple(dW.shape)), 12) + pad(f"{(dW.norm() / bp[n].norm()).item():.3f}", 11, ">") + pad(f"{e[0]:.2f}", 12, ">") + pad(f"{e[3]:.2f}", 9, ">") + pad(f"{e[15]:.2f}", 10, ">") + pad(f"{re[3]:.2f}", 18, ">")) print("\n把每个矩阵的 ΔW 截断到秩 r（其余参数——嵌入、偏置、LayerNorm——的改动原样保留）：") print(pad("秩 r", 8) + pad("同分布", 10, ">") + pad("分布外", 10, ">")) for rk in [1, 2, 4, 8, 16, 64]: m = copy.deepcopy(base) with torch.no_grad(): for n, p in m.named_parameters(): d = (fp[n] - bp[n]).detach() if n in mats: U, S, Vh = torch.linalg.svd(d, full_matrices=False) d = U[:, :rk] @ torch.diag(S[:rk]) @ Vh[:rk] p.copy_(bp[n] + d) print(pad(rk, 8) + pad(f"{choice_accuracy(m, data['test_id'], VERDICT_WORDS):.3f}", 10, ">") + pad(f"{choice_accuracy(m, data['test_ood'], VERDICT_WORDS):.3f}", 10, ">")) print(""" 观察 Observations： · 全量微调并没有被要求"低秩"，可它的更新量自己就集中在少数几个方向上：前 4 个奇异方向占了大约一半到七成的能量， 同样形状的随机矩阵只有一两成。这就是 LoRA 论文的出发点："适配时权重的变化量有一个低的'内在秩'"。 · 截断实验更直接：每个 ΔW 只保留前 4 个方向就有 0.96，前 8 个就和全量微调一样；只留 1、2 个方向只剩六成多。 所以在这个任务上，"有效的秩"是个位数——比矩阵的宽度 64 小得多，但也不是 1。 · 注意这只是一个任务、一个小模型。Biderman et al. 2024 在真实的 7B 模型上发现，学编程、数学这类 "新知识多"的任务，全量微调的更新秩比常用的 LoRA 设定高 10–100 倍——低秩假设对"格式类"适配最成立。""") # ============================================================================= # 演示 4：内在维度 # ============================================================================= def demo_intrinsic_dimension(base, data): print_header("演示 4：内在维度 —— 参数只能在一个随机的 d 维子空间里动：θ = θ0 + P·z（k = 64，150 步）") theta0 = {n: p.detach().clone() for n, p in base.named_parameters()} names = list(theta0) D = sum(t.numel() for t in theta0.values()) pairs = sample_k(data["pool"], 64, 0) x, y, mask = encode_pairs(pairs) print(f"全部参数 D = {D:,}。P 是一个随机投影：每个参数随机对应 z 的一个分量、带一个随机正负号，" f"再乘 √(d/D)（让 P 的每一列长度为 1）。") print(pad("d", 8) + pad("占 D 的比例", 12, ">") + pad("训练损失", 10, ">") + pad("同分布", 10, ">") + pad("分布外", 10, ">")) rows = {} for d in [16, 64, 256, 1024, 4096]: res = [] for seed in [0, 1]: g = torch.Generator().manual_seed(seed) idx = {n: torch.randint(0, d, t.shape, generator=g) for n, t in theta0.items()} sign = {n: torch.randint(0, 2, t.shape, generator=g).float() * 2 - 1 for n, t in theta0.items()} scale = math.sqrt(d / D) z = torch.zeros(d, requires_grad=True) # 学习率：让每个原始参数每步的变化量和全量微调同一个量级（≈ 3e-3） opt = torch.optim.Adam([z], lr=3e-3 / scale) def fwd(xx, _z=z, _idx=idx, _sign=sign, _scale=scale): params = {n: theta0[n] + _scale * _sign[n] * _z[_idx[n]] for n in names} return functional_call(base, params, (xx,)) gb = torch.Generator().manual_seed(seed) for _ in range(STEPS): bi = torch.randint(0, len(pairs), (16,), generator=gb) logits = fwd(x[bi]) per = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y[bi].reshape(-1), reduction="none") loss = (per * mask[bi].reshape(-1)).sum() / mask[bi].sum() opt.zero_grad() loss.backward() opt.step() res.append((loss.item(), choice_accuracy(fwd, data["test_id"], VERDICT_WORDS), choice_accuracy(fwd, data["test_ood"], VERDICT_WORDS))) a = np.mean(res, axis=0) rows[d] = a print(pad(d, 8) + pad(f"{d / D:.2%}", 12, ">") + pad(f"{a[0]:.3f}", 10, ">") + pad(f"{a[1]:.3f}", 10, ">") + pad(f"{a[2]:.3f}", 10, ">")) print("（对照：代码 01 里全量微调在 k = 64 时同分布 1.000；LoRA r = 1 有 2,048 个可训练参数，0.970。）") print(""" 观察 Observations： · 在 10 万维的参数空间里随便挑一个 d 维的随机子空间，只在里面找解：d = 16、64 不够，几百维开始像样， 一千维左右（不到全部参数的 1%）同分布就做到了全量微调的九成以上。 · 但分布外一直落后于同分布（0.79–0.85）：在随机子空间里硬凑出来的解，比全量微调、LoRA 更依赖训练例子里的形容词。 · Aghajanyan et al. 2020 用同样的思路量真实模型：RoBERTa-Large 在 MRPC 上只要约 200 个参数 （随机投影回全空间）就能达到全量微调九成的水平；而且模型越大、预训练越久，内在维度越低。 他们的解释：预训练把模型带到了一个"只需在很少几个方向上调整、就能适配各种任务"的位置。 · 随机子空间比 LoRA 笨得多（LoRA 的方向是学出来的，而且按矩阵结构组织），所以 LoRA 用 2,048 个参数就能做到的事， 随机子空间要上千维。两者共同说明的是：**适配所需的自由度远小于模型的参数量**。""") return rows # ============================================================================= # 演示 5：一个基座，多个插件 # ============================================================================= def staff_items(n, rng): """ 第二个任务的样本：服务好不好。 【思路】"staff check: <评论> staff:" → " nice." / " mean."。格式和标签词同样从没出现在语料里； 和 verdict 任务用同一种评论，但要看的是另一个形容词——两个插件得学不同的东西。 【参数】n：条数；rng：random.Random 【返回】list[(str, str)] """ out = [] for _ in range(n): text, _, _, service_pos = make_review(rng) out.append((f"\nstaff check: {text[len('review: '):]} staff:", "nice" if service_pos else "mean")) return out def demo_many_adapters(base, data): print_header("演示 5：一个基座、多个插件 —— 两个任务各训练一个 LoRA（r = 4，k = 64），换插件就换任务") rng = random.Random(23) staff_pool, staff_test = staff_items(400, rng), staff_items(300, rng) tasks = { "评论 verdict（3 选 1）": (sample_k(data["pool"], 64, 0), data["test_id"], VERDICT_WORDS), "服务 staff（2 选 1）": ([(p, f" {a}.") for p, a in staff_pool[:64]], staff_test, ["nice", "mean"]), } adapters = {} host = make_adapted(base, "lora4", seed=0) # 同一个"宿主"模型：基座权重只有一份 zero = {id(m): (m.A.detach().clone(), m.B.detach().clone()) for m in lora_modules(host.model)} for name, (train, _, _) in tasks.items(): with torch.no_grad(): # 每个任务从同一个初始化（B = 0）开始 for m in lora_modules(host.model): m.A.copy_(zero[id(m)][0]) m.B.copy_(zero[id(m)][1]) train_on_pairs(host, train, STEPS, 1e-3, seed=0) adapters[name] = [(m.A.detach().clone(), m.B.detach().clone()) for m in lora_modules(host.model)] def load_adapter(weights): with torch.no_grad(): for m, (A, B) in zip(lora_modules(host.model), weights): m.A.copy_(A) m.B.copy_(B) print(pad("装上的插件", 26) + "".join(pad(t, 26, ">") for t in tasks)) zero_w = [(A, torch.zeros_like(B)) for A, B in adapters[list(tasks)[0]]] for label, w in [("不装（B = 0，就是 base）", zero_w)] + [(f"插件：{n}", adapters[n]) for n in tasks]: load_adapter(w) print(pad(label, 26) + "".join(pad(f"{choice_accuracy(host.fwd, test, cands):.3f}", 26, ">") for _, test, cands in tasks.values())) n_base = sum(p.numel() for p in base.parameters()) n_lora = sum(A.numel() + B.numel() for A, B in adapters[list(tasks)[0]]) print(f"\n存储：两份全量微调的模型 = 2 × {n_base:,} 个数；一个基座 + 两个插件 = {n_base:,} + 2 × {n_lora:,} 个数。" f"\n100 个任务时：全量微调 {100 * n_base:,} 个数，LoRA {n_base + 100 * n_lora:,} 个数（约 {100 * n_base / (n_base + 100 * n_lora):.1f} 倍的差别）。") print(""" 观察 Observations： · 每个插件只在自己的任务上有用：装错插件，另一个任务的成绩回到不装插件时的水平，甚至更差。 基座只有一份，换任务就是换掉几 KB 的 A、B——W22 笔记说的"lightweight fine-tuning 可以为很多用户各存一份"。 · 插件没有合并时，同一个 batch 里不同的请求甚至可以用不同的插件（基座的计算共享，每行加上自己的 B·A·x）； S-LoRA（2023）把这个想法做成了在一张 GPU 上同时服务上千个 LoRA 的系统。合并进 W 之后就失去了这种灵活性 （LoRA 论文自己也指出了这个限制）。""") if __name__ == "__main__": T0 = time.time() world, base = load_base() data = make_data() demo_parts(base, data) demo_rank(base, data) demo_full_ft_rank(base, data) demo_intrinsic_dimension(base, data) demo_many_adapters(base, data) print(f"\n（总用时约 {time.time() - T0:.0f} 秒；计时每次运行略有不同）")