"""第 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()

1. $8(512+256)=6144$ 个；原矩阵有 131,072 个权重，仍需保留。
2. 基础权重仍在，旁边又加了 $A$、$B$；减少的主要是冻结参数对应的梯度和优化器状态。
3. 两个因子相乘，某一因子的梯度依赖另一因子；都为零时两侧可能都得不到启动信号。
4. 冻结指定参数仍允许其他参数所需的梯度传递；关闭整个前向的梯度记录会切断 LoRA 的学习路径。
5. 实际输出由基础路径加增量决定，权重冻结不等于函数不变。
6. 增量只对配套的权重、形状和词表约定有意义。
7. 本例目标层不变时，新增参数翻倍为 8,792。状态字节可按数据类型和优化器估算；激活峰值、速度和效果需测量。
8. 基础权重、激活、量化元数据、临时工作区、算子效率、额外路径与版本维护，也需要考虑。