""" ============================================================================= 第 2 讲 · 代码 7：一个完整的最小训练循环 ============================================================================= 对应讲义：lectures/lecture02_pytorch_resource_accounting.md 第 8、10 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 把前面所有零件组装成一个真正能跑的语言模型训练程序。它包含 CS336 作业 1 训练部分要求的全部要素： 1. 数据：把 token 存成 uint16 二进制文件，用 np.memmap 按需读取（可以处理比内存还大的数据集） 2. 批次：get_batch —— 随机截取若干段长度为 context_length 的序列 3. 模型：参数初始化为什么重要；一个"固定窗口"的神经语言模型（Bengio 2003 的思路） 4. 损失：数值稳定的交叉熵 5. 训练：AdamW + 预热余弦学习率 + 梯度裁剪 + 定期在验证集上评估 6. 可复现性：为什么同样的随机种子、同样的代码，两次运行结果却不一样？ 7. 检查点：保存 / 恢复，并验证"中断后继续训练"与"一口气训练完"逐比特一致 8. 生成：从训练好的模型采样文本 模型刻意选得很简单（不是 Transformer），因为本讲的重点是"训练的基础设施"； 第 3 讲会把模型换成 Transformer，其余部分几乎不用改。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 07_training_loop.py （CPU 上约 10~30 秒） 数据使用第 1 讲的小语料 code/lecture01/data/sample_corpus.txt，按字节（0~255）作为 token。 """ import math import tempfile import time from pathlib import Path import numpy as np import torch import torch.nn.functional as F CORPUS_PATH = Path(__file__).parent.parent / "lecture01" / "data" / "sample_corpus.txt" def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：数据 —— uint16 二进制文件 + memmap # ============================================================================= def write_token_file(text, path): """ 把文本按 UTF-8 字节编码成 token（0~255），存成 uint16 二进制文件。 【思路】 真实训练中，我们会先用 BPE 分词器把整个语料编码一次，把 token id 存成一个大的二进制文件， 训练时直接读这个文件，不再重复分词。 - 为什么用 uint16？CS336 作业的词表是 1 万 ~ 3.2 万，不超过 65536，uint16（2 字节）刚好够； 比 Python 的 int 或 int64（8 字节）省 4 倍空间。这里用字节级 token（<256）， 本来 uint8 就够，但我们保持和真实场景一致。 - arr.tofile(path) 把数组的原始字节直接写进文件，没有任何格式头。 【参数】 text : str path : 文件路径 【返回】 int，写入的 token 数 """ arr = np.frombuffer(text.encode("utf-8"), dtype=np.uint8).astype(np.uint16) arr.tofile(path) return len(arr) def load_token_file(path): """ 用内存映射（memmap）打开 token 文件。 【思路】 np.memmap 不会把整个文件读进内存，而是让操作系统"按需"加载用到的部分。 即使文件有 100 GB，打开它也几乎不花时间和内存； get_batch 每次只读取几十 KB，真正从磁盘读入的数据就只有这么多。 【边界】 dtype 必须和写文件时一致（uint16），否则读出来的数全是错的，而且不会报错！ 二进制文件本身不记录 dtype，这是一个常见的坑。 """ return np.memmap(path, dtype=np.uint16, mode="r") def get_batch(data, batch_size, context_length, device, rng): """ 从 token 序列中随机取一个批次（CS336 作业 1 的 get_batch）。 【思路】 随机选 batch_size 个起点 i，每个起点取： 输入 x = data[i : i + context_length] 目标 y = data[i + 1 : i + context_length + 1] （x 整体右移一位） 这样模型在每个位置 t 都要预测"下一个 token"y[t] = x[t+1]。 一个长度为 L 的样本同时提供了 L 个预测任务，非常高效。 【参数】 data : 一维整数数组（np.memmap 或 np.ndarray） batch_size : int，>= 1 context_length : int，>= 1 device : torch.device 或字符串 rng : np.random.Generator，用它抽样，保证可复现、可保存到检查点 【返回】 (x, y)，两个形状为 (batch_size, context_length) 的 torch.long 张量 【边界】 - 数据长度必须 > context_length，否则连一个样本都取不出来，抛 ValueError。 - 起点的上限是 len(data) - context_length - 1（包含），这样 y 的最后一个元素不会越界。 """ n = len(data) if n <= context_length: raise ValueError(f"数据长度 {n} 必须大于 context_length {context_length}") starts = rng.integers(0, n - context_length, size=batch_size) # 上界不包含 x = np.stack([data[i:i + context_length] for i in starts]).astype(np.int64) y = np.stack([data[i + 1:i + context_length + 1] for i in starts]).astype(np.int64) # 【PyTorch 小知识】torch.from_numpy 与 numpy 数组共享内存，不复制；.to(device) 再搬到目标设备 return torch.from_numpy(x).to(device), torch.from_numpy(y).to(device) def split_documents(text): """ 按 <|endoftext|> 把语料切成文档，留出最后一篇英文和最后一篇中文作为验证集。 【思路】 验证集必须是训练时**没见过**的文本。如果简单地取"文件最后 10%"， 我们的语料末尾恰好全是中文故事，验证集就和训练集的分布差别很大，评估结果会误导人。 所以按文档划分，并让验证集同时包含英文和中文（与第 1 讲的做法一致）。 【返回】 (train_text, val_text) """ docs = [d.strip() for d in text.split("<|endoftext|>") if d.strip()] is_zh = lambda d: any("\u4e00" <= ch <= "\u9fff" for ch in d) english = [d for d in docs if not is_zh(d)] chinese = [d for d in docs if is_zh(d)] return "\n".join(english[:-1] + chinese[:-1]), "\n".join([english[-1], chinese[-1]]) def demo_data(tmpdir): """演示 1：写入、memmap 读取、取批次。返回 (训练数据, 验证数据)。""" print_header("演示 1：数据 —— uint16 文件 + memmap + get_batch") train_text, val_text = split_documents(CORPUS_PATH.read_text(encoding="utf-8")) train_path, val_path = Path(tmpdir) / "train.bin", Path(tmpdir) / "val.bin" n_train = write_token_file(train_text, train_path) n_val = write_token_file(val_text, val_path) train_data, val_data = load_token_file(train_path), load_token_file(val_path) print(f" 训练集 {n_train:,} 个 token（文件 {train_path.stat().st_size:,} 字节，每个 token 2 字节）；" f"验证集 {n_val:,} 个 token") print(f" memmap 读取：类型 {type(train_data).__name__}，dtype {train_data.dtype}，" f"前 10 个 token {train_data[:10].tolist()}") rng = np.random.default_rng(0) x, y = get_batch(train_data, batch_size=2, context_length=12, device="cpu", rng=rng) to_text = lambda t: bytes(t.tolist()).decode("utf-8", errors="replace") print(f" 一个批次：x.shape = {tuple(x.shape)}") for i in range(2): print(f" x[{i}] = {to_text(x[i])!r}") print(f" y[{i}] = {to_text(y[i])!r} ← 右移一位") print(" （起点是随机的字节位置，可能切在一个汉字的中间，于是开头出现 �。对字节级模型这很正常。）") # 【边界】用错 dtype 读文件 wrong = np.memmap(train_path, dtype=np.uint8, mode="r") print(f" 【边界】用 uint8 读 uint16 文件，前 6 个数：{wrong[:6].tolist()}（每个 token 被拆成了两半）") return train_data, val_data # ============================================================================= # 第 2 部分：参数初始化 # ============================================================================= def demo_initialization(): """演示 2：为什么初始化的尺度这么重要。""" print_header("演示 2：参数初始化 —— 信号在深层网络里会爆炸或消失") torch.manual_seed(0) d, num_layers = 512, 10 x = torch.randn(64, d) schemes = { "N(0, 1)": lambda: torch.randn(d, d), "N(0, 0.01²)": lambda: torch.randn(d, d) * 0.01, "N(0, 1/d_in)": lambda: torch.randn(d, d) / math.sqrt(d), "截断正态 std=√(2/(d_in+d_out))": lambda: torch.nn.init.trunc_normal_( torch.empty(d, d), std=math.sqrt(2 / (d + d)), a=-3 * math.sqrt(2 / (d + d)), b=3 * math.sqrt(2 / (d + d))), } for name, make in schemes.items(): h = x for _ in range(num_layers): h = h @ make() print(f" {name:<32} 经过 {num_layers} 个线性层后，输出的标准差 = {h.std().item():.3g}") print(f""" 直觉：y = x W 的每个输出是 d_in 个乘积之和。若 x 和 W 的元素都是独立的、标准差分别为 1 和 σ， 则 y 的方差 = d_in × σ²。想让方差保持为 1，就要 σ = 1/√d_in。 - σ = 1：每层方差放大 d = {d} 倍，10 层后标准差达到 √{d}^10 ≈ {math.sqrt(d) ** 10:.1e}（爆炸）； - σ 太小：每层都缩小，信号消失； - σ = 1/√d_in：保持稳定。CS336 作业 1 用的是 std = √(2/(d_in+d_out)) 的截断正态分布 （Xavier 初始化的一种形式，d_in = d_out 时两者相同），并把超过 ±3σ 的值截断，避免极端值。""") # ============================================================================= # 第 3 部分：模型 # ============================================================================= class Linear(torch.nn.Module): """ 不带偏置的线性层 y = x W^T（与 CS336 作业 1 的约定一致）。 【思路】 权重 W 的形状是 (d_out, d_in)，这是 PyTorch 的惯例（与 nn.Linear 一致）。 用 einsum 写成 "... d_in, d_out d_in -> ... d_out"，任意前导维度都可以。 【PyTorch 小知识】 - 继承 nn.Module，在 __init__ 里把 nn.Parameter 赋给 self 的属性，PyTorch 会自动登记它为参数， model.parameters() 就能找到它，优化器才能更新它。 - forward 定义前向计算；调用 layer(x) 时会自动调用 forward(x)。 """ def __init__(self, d_in, d_out): super().__init__() std = math.sqrt(2 / (d_in + d_out)) w = torch.empty(d_out, d_in) torch.nn.init.trunc_normal_(w, mean=0.0, std=std, a=-3 * std, b=3 * std) self.weight = torch.nn.Parameter(w) def forward(self, x): return torch.einsum("...i,oi->...o", x, self.weight) class Embedding(torch.nn.Module): """ 嵌入层：一个 (vocab_size, d) 的查找表，token id → 向量。 【思路】 embedding(ids) 就是 weight[ids]：按下标取行。它没有乘法，FLOPs 为 0， 但它的参数量 vocab_size × d 在小模型里占很大比例（第 1 讲实验 5）。 【边界】ids 必须在 [0, vocab_size) 范围内，否则下标越界报错。 """ def __init__(self, vocab_size, d): super().__init__() w = torch.empty(vocab_size, d) torch.nn.init.trunc_normal_(w, mean=0.0, std=1.0, a=-3.0, b=3.0) self.weight = torch.nn.Parameter(w) def forward(self, ids): return self.weight[ids] class WindowMLPLM(torch.nn.Module): """ 固定窗口的神经语言模型：用前 k 个 token 预测下一个 token。 【思路】（Bengio 等人 2003 年第一个神经语言模型的思路） 对每个位置 t： 1. 取位置 t-k+1 … t 这 k 个 token 的嵌入向量，拼成一个长向量（k × d_embed 维）； 2. 经过一个隐藏层（ReLU）； 3. 输出层给出词表上每个 token 的分数（logits）。 开头不足 k 个 token 的位置，用全 0 向量补齐（左侧填充），保证模型"看不到未来"。 【形状】 输入 ids : (batch, seq) 输出 logits: (batch, seq, vocab_size) 【边界】 这个模型只能看到最近 k 个 token，更早的信息完全丢失 —— 这正是 Transformer 要解决的问题（第 3 讲）。 """ def __init__(self, vocab_size=256, window=8, d_embed=16, d_hidden=128): super().__init__() self.window = window self.embed = Embedding(vocab_size, d_embed) self.hidden = Linear(window * d_embed, d_hidden) self.out = Linear(d_hidden, vocab_size) def forward(self, ids): e = self.embed(ids) # (B, S, E) e = F.pad(e, (0, 0, self.window - 1, 0)) # 在序列维左侧补 k-1 个零向量 → (B, S+k-1, E) # 【PyTorch 小知识】unfold(dim, size, step) 沿 dim 取滑动窗口： # (B, S+k-1, E) → (B, S, E, k)，第 t 个窗口正好是原位置 t-k+1 … t win = e.unfold(1, self.window, 1) win = win.transpose(-1, -2).reshape(ids.shape[0], ids.shape[1], -1) # → (B, S, k*E) h = torch.relu(self.hidden(win)) return self.out(h) # ============================================================================= # 第 4 部分：数值稳定的交叉熵 # ============================================================================= def cross_entropy(logits, targets): """ 交叉熵损失：-log softmax(logits)[target]，对所有位置取平均。 【推导】 softmax(o)_i = exp(o_i) / Σ_j exp(o_j) -log softmax(o)_y = -o_y + log Σ_j exp(o_j) 直接算 exp(o_j) 很危险：o_j = 1000 时 exp 会上溢成 inf。 技巧：减去最大值 m = max_j o_j 不改变结果： log Σ_j exp(o_j) = m + log Σ_j exp(o_j - m) 此时所有 o_j - m ≤ 0，exp 不会上溢；而且至少有一项是 exp(0) = 1，求和不会下溢成 0。 【参数】 logits : (..., vocab_size) 的浮点张量 targets : (...) 的整数张量，取值在 [0, vocab_size) 【返回】 标量张量，平均损失（单位：nats，即以 e 为底） 【边界】 内部把 logits 转成 float32 计算：即使模型用 BF16，损失也应该用 FP32 算（见 02 号文件）。 """ logits = logits.float() m = logits.max(dim=-1, keepdim=True).values log_sum_exp = m.squeeze(-1) + torch.log(torch.exp(logits - m).sum(dim=-1)) # 【PyTorch 小知识】gather 按下标取值：从每个位置的 logits 中取出正确答案对应的那个分数 target_logit = logits.gather(-1, targets.unsqueeze(-1)).squeeze(-1) return (log_sum_exp - target_logit).mean() def naive_cross_entropy(logits, targets): """【错误示范】不减最大值的交叉熵，logits 很大时会得到 nan/inf。""" p = torch.exp(logits) / torch.exp(logits).sum(dim=-1, keepdim=True) return -torch.log(p.gather(-1, targets.unsqueeze(-1))).mean() def demo_cross_entropy(): """演示 3：交叉熵的正确性与数值稳定性。""" print_header("演示 3：数值稳定的交叉熵") torch.manual_seed(0) logits = torch.randn(4, 8, 256) targets = torch.randint(0, 256, (4, 8)) ours = cross_entropy(logits, targets).item() ref = F.cross_entropy(logits.view(-1, 256), targets.view(-1)).item() print(f" 普通 logits：我们的 {ours:.6f}，PyTorch 的 {ref:.6f}") print(f" 随机猜测的损失 = ln(256) = {math.log(256):.6f}（模型刚初始化时应该接近这个值）") big = logits * 1000 print(f" logits × 1000：稳定版 {cross_entropy(big, targets).item():.2f}，" f"朴素版 {naive_cross_entropy(big, targets).item()} ← exp 上溢") # ============================================================================= # 第 5 部分：训练循环 # ============================================================================= def lr_cosine_schedule(t, lr_max, lr_min, warmup_iters, cosine_cycle_iters): """预热 + 余弦退火（与 06 号文件相同，这里复制一份以便本文件独立运行）。""" if t < warmup_iters: return t / warmup_iters * lr_max if t <= cosine_cycle_iters: progress = (t - warmup_iters) / max(1, cosine_cycle_iters - warmup_iters) return lr_min + 0.5 * (1 + math.cos(math.pi * progress)) * (lr_max - lr_min) return lr_min CONFIG = dict( batch_size=32, context_length=64, lr_max=3e-3, lr_min=3e-4, warmup_iters=50, max_iters=1000, weight_decay=0.1, grad_clip=1.0, eval_interval=100, eval_batches=10, ) @torch.no_grad() def estimate_loss(model, data, cfg, rng): """ 在若干个随机批次上估计平均损失。 【PyTorch 小知识】 @torch.no_grad() 让函数内部不记录计算图：评估时不需要反向传播，这样更省内存、更快。 model.eval() / model.train() 切换评估/训练模式（影响 dropout 等层，本模型没有，但这是好习惯）。 """ model.eval() losses = [] for _ in range(cfg["eval_batches"]): x, y = get_batch(data, cfg["batch_size"], cfg["context_length"], "cpu", rng) losses.append(cross_entropy(model(x), y).item()) model.train() return sum(losses) / len(losses) def train_steps(model, optimizer, train_data, cfg, rng, start_iter, end_iter): """ 从 start_iter 训练到 end_iter（不含），返回最后一步的训练损失。 【每一步做的事】 1. 按调度设置学习率； 2. 取一个批次； 3. 前向，算损失； 4. 清空旧梯度 → 反向传播； 5. 梯度裁剪； 6. 优化器更新参数。 【参数】 rng : np.random.Generator，用于取批次（它的状态会被检查点保存） """ loss = None for it in range(start_iter, end_iter): lr = lr_cosine_schedule(it, cfg["lr_max"], cfg["lr_min"], cfg["warmup_iters"], cfg["max_iters"]) for group in optimizer.param_groups: # 【PyTorch 小知识】修改学习率的标准方式 group["lr"] = lr x, y = get_batch(train_data, cfg["batch_size"], cfg["context_length"], "cpu", rng) loss = cross_entropy(model(x), y) optimizer.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["grad_clip"]) optimizer.step() return None if loss is None else loss.item() def build(cfg, seed=0): """按固定随机种子创建模型和优化器。""" torch.manual_seed(seed) model = WindowMLPLM() optimizer = torch.optim.AdamW(model.parameters(), lr=cfg["lr_max"], betas=(0.9, 0.95), weight_decay=cfg["weight_decay"]) return model, optimizer # ============================================================================= # 第 6 部分：可复现性 # ============================================================================= def run_short_training(train_data, steps=50): """用固定种子训练 steps 步，返回模型（用于比较两次运行是否一致）。""" cfg = dict(CONFIG, max_iters=steps) model, opt = build(cfg) train_steps(model, opt, train_data, cfg, np.random.default_rng(0), 0, steps) return model def same_params(m1, m2): """两个模型的参数是否逐比特相同。""" return all(torch.equal(a, b) for a, b in zip(m1.parameters(), m2.parameters())) def demo_determinism(train_data): """演示 4：同样的种子、同样的代码，两次运行却不一样？""" print_header("演示 4：可复现性 —— 固定随机种子还不够") a, b, c = 0.1, 0.2, 0.3 print(f" 浮点加法不满足结合律：(0.1 + 0.2) + 0.3 = {(a + b) + c!r}") print(f" 0.1 + (0.2 + 0.3) = {a + (b + c)!r}") torch.use_deterministic_algorithms(False) r1 = same_params(run_short_training(train_data), run_short_training(train_data)) torch.use_deterministic_algorithms(True) r2 = same_params(run_short_training(train_data), run_short_training(train_data)) print(f"\n 同样的种子训练 50 步，两次结果是否逐比特相同：") print(f" 默认设置：{r1}") print(f" torch.use_deterministic_algorithms(True)：{r2}") print(""" 原因：为了快，很多运算会用多个线程并行地把一堆数加起来（比如嵌入层的反向传播要把 同一个 token 在不同位置的梯度累加到同一行）。线程完成的先后每次不同，加法顺序就不同， 而浮点加法的结果依赖顺序 —— 于是出现最后几位的差异，训练几十步后差异会被放大。 【PyTorch 小知识】torch.use_deterministic_algorithms(True) 强制使用确定性的实现（可能稍慢）； 没有确定性实现的运算会直接报错，提醒你。 【大模型联系】大规模训练一般不追求逐比特可复现（太慢），但调试时，能复现是定位 bug 的前提。 本文件接下来的演示都开启确定性模式。""") # ============================================================================= # 第 7 部分：训练 + 检查点 # ============================================================================= def save_checkpoint(model, optimizer, iteration, rng, path): """ 保存训练状态（CS336 作业 1 的 save_checkpoint 的扩展版）。 【要保存什么】 - 模型参数 model.state_dict() - 优化器状态 optimizer.state_dict()（AdamW 的 m、v 和步数！只存模型的话，恢复后 Adam 要重新"热身"） - 当前步数（学习率调度要用） - 随机数生成器的状态（保证恢复后抽到的批次和不中断时完全一样） 【Python 小知识】 torch.save 可以保存任意嵌套的字典，里面可以有张量、数字、字符串。 """ torch.save({ "model": model.state_dict(), "optimizer": optimizer.state_dict(), "iteration": iteration, "rng_state": rng.bit_generator.state, }, path) def load_checkpoint(path, model, optimizer): """ 从检查点恢复模型、优化器，返回 (步数, 随机数生成器)。 【边界】 - model 和 optimizer 必须是用相同结构、相同超参数新建的对象；load_state_dict 会覆盖它们的状态。 - PyTorch 2.6 起 torch.load 默认 weights_only=True，只允许加载张量和基本类型，更安全 （加载来路不明的检查点可能执行恶意代码）。我们保存的内容都是基本类型，可以直接加载。 """ ckpt = torch.load(path) model.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) rng = np.random.default_rng() rng.bit_generator.state = ckpt["rng_state"] return ckpt["iteration"], rng def demo_train(train_data, val_data, tmpdir): """演示 5：完整训练，定期评估，保存验证损失最低的检查点。返回最佳模型。""" print_header("演示 5：训练一个字节级的固定窗口语言模型") cfg = CONFIG model, optimizer = build(cfg) n_params = sum(p.numel() for p in model.parameters()) tokens_per_step = cfg["batch_size"] * cfg["context_length"] print(f" 模型参数 {n_params:,} 个；每步 {cfg['batch_size']} × {cfg['context_length']} = " f"{tokens_per_step:,} 个 token；共 {cfg['max_iters']} 步") rng = np.random.default_rng(0) best_val, best_step = float("inf"), 0 best_path = Path(tmpdir) / "best.pt" t0 = time.perf_counter() for start in range(0, cfg["max_iters"], cfg["eval_interval"]): end = min(start + cfg["eval_interval"], cfg["max_iters"]) train_steps(model, optimizer, train_data, cfg, rng, start, end) # 用固定的评估随机数，让每次评估看到同样的批次，数字之间才可比 tr = estimate_loss(model, train_data, cfg, np.random.default_rng(1)) va = estimate_loss(model, val_data, cfg, np.random.default_rng(1)) mark = "" if va < best_val: best_val, best_step = va, end save_checkpoint(model, optimizer, end, rng, best_path) mark = " ← 目前最好，保存检查点" print(f" 第 {end:>5} 步：训练损失 {tr:.3f}，验证损失 {va:.3f}（{va / math.log(2):.2f} 比特/字节）{mark}") dt = time.perf_counter() - t0 tokens = cfg["max_iters"] * tokens_per_step epochs = tokens / len(train_data) print(f"\n 用时 {dt:.1f} 秒（含评估），训练吞吐量约 {tokens / dt:,.0f} tokens/秒") print(f" 一共训练了 {tokens:,} 个 token = 把 {len(train_data):,} 字节的训练集重复看了约 {epochs:.0f} 遍") print(f""" - 刚开始的损失应接近 ln(256) ≈ 5.55（均匀猜测），训练后明显下降。 - "比特/字节" = 损失 / ln 2。第 1 讲说过"预测就是压缩"：这个数就是用该模型压缩文本时， 每个字节平均需要的比特数（原始 UTF-8 是 8 比特/字节）。 - 验证损失在第 {best_step} 步左右最低，之后训练损失继续下降、验证损失却回升：这就是**过拟合**。 原因是训练集太小，模型把同样的数据看了 {epochs:.0f} 遍，开始"背诵"训练集。 真实的大模型训练数据极多，通常每个 token 只看一遍（第 1 讲 2.6 节），几乎不会遇到这个问题。 - 应对小数据过拟合的办法：保存验证损失最低的检查点（"早停"）、更小的模型、更多数据。""") best_model, best_opt = build(cfg, seed=123) load_checkpoint(best_path, best_model, best_opt) print(f" 已从检查点加载第 {best_step} 步的最佳模型（验证损失 {best_val:.3f}）。") return best_model def demo_checkpoint(train_data, tmpdir): """演示 6：中断后恢复训练，结果与不中断完全一致。""" print_header("演示 6：检查点 —— 中断再恢复，结果逐比特一致") cfg = dict(CONFIG, max_iters=200) # 方案 A：一口气训练 200 步 model_a, opt_a = build(cfg) train_steps(model_a, opt_a, train_data, cfg, np.random.default_rng(0), 0, 200) # 方案 B：训练 100 步 → 保存 → 用全新的对象恢复 → 再训练 100 步 model_b, opt_b = build(cfg) rng_b = np.random.default_rng(0) train_steps(model_b, opt_b, train_data, cfg, rng_b, 0, 100) path = Path(tmpdir) / "ckpt.pt" save_checkpoint(model_b, opt_b, 100, rng_b, path) del model_b, opt_b, rng_b # 模拟"程序被中断" model_c, opt_c = build(cfg, seed=999) # 故意用不同的种子新建，证明状态全部来自检查点 it, rng_c = load_checkpoint(path, model_c, opt_c) train_steps(model_c, opt_c, train_data, cfg, rng_c, it, 200) same = same_params(model_a, model_c) print(f" 检查点文件大小：{path.stat().st_size / 1e3:.0f} KB（参数 + AdamW 的 m 和 v ≈ 3 倍参数）") print(f" 一口气训练 vs 中断后恢复：参数完全相同 = {same}") assert same # 边界：如果只恢复模型、不恢复优化器状态会怎样？ model_d, opt_d = build(cfg) rng_d = np.random.default_rng(0) train_steps(model_d, opt_d, train_data, cfg, rng_d, 0, 100) fresh_opt = torch.optim.AdamW(model_d.parameters(), lr=cfg["lr_max"], betas=(0.9, 0.95), weight_decay=cfg["weight_decay"]) train_steps(model_d, fresh_opt, train_data, cfg, rng_d, 100, 200) diff = max((pa - pd).abs().max().item() for pa, pd in zip(model_a.parameters(), model_d.parameters())) print(f" 【边界】只恢复模型、用新的优化器继续：与正确结果的最大参数差异 = {diff:.3f}（不再一致）") print(" 因为 AdamW 的 m、v 被清零，偏差修正也从头开始，相当于训练中途又'重启'了一次优化器。") # ============================================================================= # 第 8 部分：生成 # ============================================================================= @torch.no_grad() def generate(model, prompt, max_new_tokens=200, temperature=0.8, seed=0): """ 从模型中采样生成文本。 【思路】 重复：用当前序列算出下一个 token 的 logits → 除以温度 → softmax 得到概率 → 按概率抽一个 token → 接到末尾。 - 温度 < 1：分布更"尖"，更保守、更重复； - 温度 > 1：分布更"平"，更随机、更容易胡言乱语； - 温度 → 0：总是选概率最大的（贪心解码）。 【参数】 prompt : str，开头文本 temperature : float，> 0 【返回】 str（按 UTF-8 解码，非法字节替换为 �） """ if temperature <= 0: raise ValueError("温度必须大于 0") g = torch.Generator().manual_seed(seed) ids = list(prompt.encode("utf-8")) for _ in range(max_new_tokens): x = torch.tensor([ids[-model.window:]]) # 这个模型只需要最近 window 个 token logits = model(x)[0, -1] / temperature probs = torch.softmax(logits, dim=-1) ids.append(torch.multinomial(probs, 1, generator=g).item()) return bytes(ids).decode("utf-8", errors="replace") def demo_generate(model): """演示 7：让训练好的模型写故事。""" print_header("演示 7：从模型中采样") for temp in [0.5, 1.0]: print(f"\n 温度 {temp}：") print(" " + generate(model, "Once upon a time", temperature=temp).replace("\n", "\n ")) print(f""" 模型只能看到最近 {model.window} 个字节：它能拼出一些真实的单词（and、the、said、Mom）， 也会造出很多"似是而非"的伪单词，句子更是毫无逻辑。 第 3 讲换成 Transformer 后，模型可以看到整个上下文。""") def main(): torch.manual_seed(0) # 【Python 小知识】TemporaryDirectory：临时文件夹，with 块结束后自动删除 with tempfile.TemporaryDirectory() as tmpdir: train_data, val_data = demo_data(tmpdir) demo_initialization() demo_cross_entropy() demo_determinism(train_data) # 这之后开启了确定性模式 model = demo_train(train_data, val_data, tmpdir) demo_checkpoint(train_data, tmpdir) demo_generate(model) print("\n全部演示运行完毕。") if __name__ == "__main__": main()