"""第 23 章：词元加权的梯度累积，以及 CPU 上可精确恢复的训练。"""

import argparse
from hashlib import sha256
from math import cos, pi
from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn
from tiny_transformer import TinyTransformer
from small_language_model import PAD, corpus, vocabulary, make_windows

PLAN = dict(total=12, warmup=3, peak=0.003, floor=0.0003, batch=16, micro=8)


class DropoutLanguageModel(TinyTransformer):
    def __init__(self, dropout=0.1, **kwargs):
        super().__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
        self.config["dropout"] = dropout

    def forward(self, ids):
        length = ids.shape[1]
        if not 1 <= length <= self.config["max_context"]:
            raise ValueError("输入长度超出位置表范围")
        positions = torch.arange(length, device=ids.device)
        x = self.dropout(self.tokens(ids)+self.positions(positions))
        return self.output(self.final_norm(self.blocks(x)))


def learning_rate(step, plan):
    if step < plan["warmup"]:
        return plan["peak"]*(step+1)/plan["warmup"]
    fraction = (step-plan["warmup"]+1)/(plan["total"]-plan["warmup"])
    return plan["floor"] + (plan["peak"]-plan["floor"])*(1+cos(pi*fraction))/2


def dataset():
    train, _, _ = corpus()
    vocab = vocabulary(train)
    # 16 条全文，16 条只留第一句，使每个样本的有效目标数不同。
    texts = [text if i % 2 == 0 else text[:14] for i, text in enumerate(train[:32])]
    fingerprint = sha256("\n".join(texts).encode()).hexdigest()
    return make_windows(texts, vocab, 32), vocab, fingerprint


def fresh(vocab):
    torch.manual_seed(2026)
    model = DropoutLanguageModel(vocab_size=len(vocab), width=16).double()
    optimizer = torch.optim.AdamW(model.parameters(), lr=PLAN["peak"], weight_decay=0.01)
    sampler = torch.Generator().manual_seed(2027)
    return model, optimizer, sampler


def accumulate(model, x, y, micro):
    # 每个 microbatch 都除以同一个分母，最后相加才是整批词元平均。
    total_targets = (y != PAD).sum()
    if not total_targets:
        raise ValueError("整个批次没有有效目标")
    loss_sum = 0.0
    for start in range(0, len(x), micro):
        logits = model(x[start:start+micro])
        numerator = nn.functional.cross_entropy(
            logits.flatten(0, 1), y[start:start+micro].flatten(),
            ignore_index=PAD, reduction="sum")
        loss = numerator/total_targets
        if not torch.isfinite(loss):
            raise RuntimeError("损失出现非有限值")
        loss.backward()
        loss_sum += loss.detach().item()
    return loss_sum


def advance(model, optimizer, sampler, data, start, stop, plan):
    x, y, _ = data
    history = []
    model.train()
    for step in range(start, stop):
        batch = torch.randint(len(x), (plan["batch"],), generator=sampler)
        optimizer.zero_grad(set_to_none=True)
        for group in optimizer.param_groups:
            group["lr"] = learning_rate(step, plan)
        loss = accumulate(model, x[batch], y[batch], plan["micro"])
        nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        history.append(loss)
    return history


def save(path, model, optimizer, sampler, completed, vocab, fingerprint, plan):
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(dict(model=model.state_dict(), config=model.config,
                    optimizer=optimizer.state_dict(), completed=completed, plan=plan,
                    sampler=sampler.get_state(), rng=torch.get_rng_state(),
                    vocab=vocab, data_sha256=fingerprint, torch_version=str(torch.__version__)), path)


def restore(path, vocab, fingerprint, plan, restore_optimizer=True):
    saved = torch.load(path, map_location="cpu", weights_only=True)
    if saved["vocab"] != vocab or saved["data_sha256"] != fingerprint or saved["plan"] != plan:
        raise ValueError("词表、数据或训练计划与检查点不同")
    if saved["torch_version"] != str(torch.__version__):
        raise ValueError("精确续训示例要求相同 PyTorch 版本")
    model = DropoutLanguageModel(**saved["config"]).double()
    model.load_state_dict(saved["model"])
    optimizer = torch.optim.AdamW(model.parameters(), lr=plan["peak"], weight_decay=0.01)
    if restore_optimizer:
        optimizer.load_state_dict(saved["optimizer"])
    sampler = torch.Generator()
    sampler.set_state(saved["sampler"])
    # 初始化模型会消耗随机数，因此必须在它之后恢复全局 RNG。
    torch.set_rng_state(saved["rng"])
    return model, optimizer, sampler, saved["completed"]


def packing():
    # 输入是 [BOS, a, EOS, BOS, b]；目标向右错开一位。
    segments = torch.tensor([0, 0, 0, 1, 1, 1])
    input_segments = segments[:-1]
    allowed = torch.ones(5, 5, dtype=torch.bool).tril()
    allowed &= input_segments[:, None] == input_segments[None, :]
    targets_used = input_segments == segments[1:]
    return allowed, targets_used, torch.tensor([0, 1, 2, 0, 1])


def demonstration():
    data, vocab, fingerprint = dataset()
    full, optimizer, sampler = fresh(vocab)
    continuous = advance(full, optimizer, sampler, data, 0, PLAN["total"], PLAN)
    partial, optimizer, sampler = fresh(vocab)
    first = advance(partial, optimizer, sampler, data, 0, 6, PLAN)
    with TemporaryDirectory() as directory:
        path = Path(directory)/"resume.pt"
        save(path, partial, optimizer, sampler, 6, vocab, fingerprint, PLAN)
        resumed, optimizer, sampler, step = restore(path, vocab, fingerprint, PLAN)
        second = advance(resumed, optimizer, sampler, data, step, PLAN["total"], PLAN)
        reset, optimizer, sampler, step = restore(path, vocab, fingerprint, PLAN, False)
        advance(reset, optimizer, sampler, data, step, PLAN["total"], PLAN)
    same = lambda a, b: all(torch.equal(x, y) for x, y in zip(a.parameters(), b.parameters()))
    print("连续训练与恢复训练：")
    print(f"  12 步损失逐项完全相同：{continuous == first+second}")
    print(f"  最终参数逐项完全相同：{same(full, resumed)}")
    print(f"  不恢复优化器时仍完全相同：{same(full, reset)}")
    print(f"首步损失={continuous[0]:.6f}；末步损失={continuous[-1]:.6f}")
    counts = sorted(set((data[1] != PAD).sum(1).tolist()))
    print(f"每条样本的真实目标数：{counts}；微批次 8，有效批次 16")
    allowed, targets_used, positions = packing()
    print("分组因果掩码：")
    for row in allowed.int().tolist():
        print(" ".join(map(str, row)))
    print(f"参与损失的位置：{targets_used.int().tolist()}")
    print(f"重置后的输入位置：{positions.tolist()}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--save", type=Path, help="训练到第 6 步并保存")
    group.add_argument("--resume", type=Path, help="从检查点继续到第 12 步")
    args = parser.parse_args()
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    if args.save or args.resume:
        data, vocab, fingerprint = dataset()
        if args.save:
            model, optimizer, sampler = fresh(vocab)
            advance(model, optimizer, sampler, data, 0, 6, PLAN)
            save(args.save, model, optimizer, sampler, 6, vocab, fingerprint, PLAN)
            print("已保存第 6 步的完整训练状态。")
        else:
            model, optimizer, sampler, step = restore(args.resume, vocab, fingerprint, PLAN)
            history = advance(model, optimizer, sampler, data, step, PLAN["total"], PLAN)
            print(f"从第 {step} 步继续至第 {PLAN['total']} 步；末步损失={history[-1]:.6f}")
    else:
        demonstration()


if __name__ == "__main__":
    main()

1. 并行计算描述运算的组织方式；因果掩码描述信息能流到哪里。矩阵同时算，不代表每个元素能访问全部输入。
2. $(3\times4+9\times2)/12=2.5$，不是两个均值的简单平均 3。
3. 不能。EOS 是符号，跨文档可见性由注意力掩码决定。
4. 我们希望每篇独立建模，不训练从一篇结尾跳到下一篇开头的连接任务。
5. 需要当前随机数状态而非最初种子；还缺优化器历史、样本顺序、调度进度等。
6. 网络初始化消耗随机数，会使后续 Dropout 序列改变。恢复顺序是实验的一部分。
7. 先清空梯度，四个微批次累积完再裁剪和更新，随后清空。所有微批次仍除以同一个有效目标总数。
8. 这是修改训练计划后的继续训练。它可以是合理实验，但已经改变了学习率轨迹，不能期待与原 12 步连续运行完全相同。