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