"""第 27 章：手写 LoRA，比较训练状态，并保存、加载和合并低秩增量。""" import argparse from copy import deepcopy from hashlib import sha256 import json from pathlib import Path from tempfile import TemporaryDirectory import torch from torch import nn from tiny_transformer import TinyTransformer from small_language_model import corpus from instruction_tuning import (prepare, setup, vocabulary, instruction_records, finetune, metrics, supervised_data) class LoRALinear(nn.Module): def __init__(self, base, rank=4, alpha=4): super().__init__() if not 1 <= rank <= min(base.in_features, base.out_features): raise ValueError("示例要求 rank 不超过矩阵较短的一边") self.base = base self.base.requires_grad_(False) self.rank, self.alpha = rank, alpha self.A = nn.Parameter(torch.randn(rank, base.in_features, dtype=base.weight.dtype)*.02) self.B = nn.Parameter(torch.zeros(base.out_features, rank, dtype=base.weight.dtype)) def forward(self, x): return self.base(x)+(x @ self.A.T @ self.B.T)*(self.alpha/self.rank) def merged(self): result = deepcopy(self.base) with torch.no_grad(): result.weight.add_(self.B @ self.A*(self.alpha/self.rank)) return result def attach_lora(base, rank=4, alpha=4): model = deepcopy(base) model.zero_grad(set_to_none=True) model.requires_grad_(False) torch.manual_seed(2030) for module in list(model.modules()): for name, child in list(module.named_children()): if isinstance(child, nn.Linear): setattr(module, name, LoRALinear(child, rank, alpha)) return model def merge_lora(model): result = deepcopy(model) for module in list(result.modules()): for name, child in list(module.named_children()): if isinstance(child, LoRALinear): setattr(module, name, child.merged()) return result def state_fingerprint(model): digest = sha256() for name, tensor in sorted(model.state_dict().items()): digest.update(name.encode()) digest.update(tensor.detach().cpu().contiguous().numpy().tobytes()) return digest.hexdigest() def save_adapter(path, model, base, vocab): adapters = {name: p.detach().cpu() for name, p in model.named_parameters() if p.requires_grad} modules = {name: dict(rank=module.rank, alpha=module.alpha) for name, module in model.named_modules() if isinstance(module, LoRALinear)} path.parent.mkdir(parents=True, exist_ok=True) torch.save(dict(adapters=adapters, modules=modules, vocab=vocab, base_sha256=state_fingerprint(base), base_config=base.config), path) def load_adapter(path, base, vocab): saved = torch.load(path, map_location="cpu", weights_only=True) if (saved["base_sha256"] != state_fingerprint(base) or saved["vocab"] != vocab or saved["base_config"] != base.config): raise ValueError("基础权重或词表与适配器不匹配") settings = list(saved["modules"].values()) if not settings or any(item != settings[0] for item in settings): raise ValueError("本例要求所有适配器采用同一配置") model = attach_lora(base, **settings[0]) current_modules = {name for name, module in model.named_modules() if isinstance(module, LoRALinear)} if current_modules != set(saved["modules"]): raise ValueError("LoRA 注入位置不同") state = model.state_dict() expected = {name for name, p in model.named_parameters() if p.requires_grad} if expected != set(saved["adapters"]): raise ValueError("适配器参数缺失或多余") state.update(saved["adapters"]) model.load_state_dict(state) return model def run(base_path=None, adapter_path=None): setup() if base_path: train, validation, test = corpus() vocab = vocabulary(train) saved = torch.load(base_path, map_location="cpu", weights_only=True) if (saved["experiment"] != "instruction-base-v1" or saved["vocab"] != vocab or saved["train_sha256"] != sha256("\n".join(train).encode()).hexdigest()): raise ValueError("需要第 26 章生成的基础模型") base = TinyTransformer(**saved["config"]).double() base.load_state_dict(saved["state"]) else: base, (train, validation, test), vocab = prepare() records = instruction_records(train[:96]) full, full_cost = finetune(base, records, train, vocab) initial = attach_lora(base) x, _ = supervised_data(instruction_records(validation), vocab) with torch.no_grad(): initially_equal = torch.equal(base(x), initial(x)) adapted, adapter_cost = finetune(initial, records, train, vocab) if adapter_path: save_adapter(adapter_path, adapted, base, vocab) frozen_equal = all(torch.equal(p, initial.state_dict()[name]) for name, p in adapted.named_parameters() if not p.requires_grad) with TemporaryDirectory() as directory: path = Path(directory)/"adapter.pt" save_adapter(path, adapted, base, vocab) restored = load_adapter(path, base, vocab) merged = merge_lora(restored) with torch.no_grad(): reference = adapted(x) restored_equal = torch.equal(reference, restored(x)) merge_error = (reference-merged(x)).abs().max().item() results = [] for name, model, cost in [("全量微调", full, full_cost), ("LoRA rank=4", adapted, adapter_cost)]: results.append(dict(name=name, cost=cost, validation=metrics(model, instruction_records(validation), validation, vocab), test=metrics(model, instruction_records(test), test, vocab))) return dict(results=results, initially_equal=initially_equal, frozen_equal=frozen_equal, restored_equal=restored_equal, merge_error=merge_error, base_parameters=sum(p.numel() for p in base.parameters())) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--base", type=Path, help="可选：复用第 26 章的 --save-base 文件") parser.add_argument("--out", type=Path) parser.add_argument("--save-adapter", type=Path) args = parser.parse_args() report = run(args.base, args.save_adapter) print(f"基础参数={report['base_parameters']:,}；两组均用 160 次更新、相同指令批次") for result in report["results"]: cost, val, test = result["cost"], result["validation"], result["test"] print(f"{result['name']}：可训练参数={cost['trainable']:,}；" f"验证回答损失={val['answer_loss']:.3f}；测试回答={test['correct']}/64；" f"旧续写={test['old_exact']}/32") state = cost["state_bytes"] print(f" 权重={state['weights']/1024:.1f} KiB；梯度={state['gradients']/1024:.1f} KiB；" f"优化器={state['optimizer']/1024:.1f} KiB；合计={sum(state.values())/1024:.1f} KiB") print(f"初始输出与基础模型完全相同：{report['initially_equal']}") print(f"训练后冻结参数完全不变：{report['frozen_equal']}") print(f"保存/加载适配器后输出完全相同：{report['restored_equal']}") print(f"合并后输出误差小于 1e-12：{report['merge_error'] < 1e-12}") if args.save_adapter: print(f"适配器文件：{args.save_adapter}") if args.out: args.out.parent.mkdir(parents=True, exist_ok=True) args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2)) print(f"实验记录：{args.out}") if __name__ == "__main__": main()