第 3 章 组装训练

这一章我们补齐训练流程，和上一章介绍的语言模型组件一起构成整体训练系统，当然也包括简单的推理，如果你对工程设计和代码不感兴趣，可以跳过本章。

3.1 数据长什么样

磁盘上的普通文本

hello world
hello torch

编码之后得到一串整数。训练时从这串整数上切长度为 block_size 的窗口。设 block_size = 8，某次切到 ID：

idx    = [h e l l o   w o]     形状 [B, 8]，送进模型
目标   = [e l l o   w o r]     形状 [B, 8]，与 idx 错开一位

实现上不必另存一份已经错开的张量：idx 与 labels 取同一窗口，在 loss 里做 logits[:, :-1] 对 labels[:, 1:]。位置 0 的分布去拟合 e，位置 6 的分布去拟合 o，最后一个位置的 logits 没有下一位标签，直接丢掉。

有 padding 时，补出来的位置在 labels 里写成 -100。F.cross_entropy(..., ignore_index=-100) 不给这些位置梯度。本例的随机切片没有 pad，标签里全是真实字符。

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


class CharTokenizer:
    """从语料统计字符表。正式训练换 BPE 时，只替换 encode/decode 与 vocab_size。"""

    def __init__(self, text: str):
        chars = sorted(set(text))
        self.stoi = {ch: i for i, ch in enumerate(chars)}
        self.itos = {i: ch for ch, i in self.stoi.items()}

    @property
    def vocab_size(self) -> int:
        return len(self.stoi)

    def encode(self, s: str) -> list[int]:
        return [self.stoi[c] for c in s]

    def decode(self, ids: list[int]) -> str:
        return "".join(self.itos[i] for i in ids)


def sample_batch(token_ids: torch.Tensor, batch_size: int, block_size: int, device: torch.device):
    """token_ids: 整份语料的一维 LongTensor。返回 idx, labels，形状都是 [B, block_size]。"""
    # 窗口右端还要能取到「下一位」，所以起点上限是 len - block_size - 1
    hi = token_ids.numel() - block_size - 1
    if hi <= 0:
        raise ValueError("语料短于 block_size+1，无法切出一个带标签的窗口")
    starts = torch.randint(0, hi, (batch_size,))           # 每条样本一个随机起点
    windows = torch.stack([token_ids[s : s + block_size] for s in starts.tolist()])
    return windows.to(device), windows.to(device)          # idx 与 labels 同一窗口；错位在 loss 里做

语料至少要比 block_size 长出 1 个 token，否则没有「下一个字」可监督。这份代码 过拟合演示用几 KB 文本即可；拿它衡量生成质量是没有意义的，只用来确认前向反向走得通。

3.2 模型

因果掩码注册成 buffer，换设备时跟着走，避免每步在 GPU 上重新 torch.ones。scaled_dot_product_attention(..., is_causal=True) 与手写上三角 -inf 一致；PyTorch 2 在 CUDA 上会走到更快的核。这里仍把掩码显式留下，方便对照第 2 章。

@dataclass
class GPTConfig:
    vocab_size: int
    block_size: int = 128
    n_layer: int = 4
    n_head: int = 4
    n_embd: int = 128
    dropout: float = 0.0
    tie_embeddings: bool = True


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        assert cfg.n_embd % cfg.n_head == 0
        self.n_head = cfg.n_head
        self.d_h = cfg.n_embd // cfg.n_head
        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False)
        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
        self.drop = nn.Dropout(cfg.dropout)
        # 上三角（不含对角线）为 True：这些位置在 softmax 前填 -inf
        mask = torch.triu(torch.ones(cfg.block_size, cfg.block_size, dtype=torch.bool), diagonal=1)
        self.register_buffer("causal", mask, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, n, _ = x.shape                                    # x: [B, n, d]
        qkv = self.qkv(x).view(B, n, 3, self.n_head, self.d_h).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]                   # 各 [B, h, n, d_h]；无 cache，用完即丢
        att = (q @ k.transpose(-2, -1)) * (self.d_h ** -0.5)  # [B, h, n, n]
        att = att.masked_fill(self.causal[:n, :n], float("-inf"))  # 看不见未来
        att = self.drop(att.softmax(dim=-1))
        y = (att @ v).transpose(1, 2).contiguous().view(B, n, -1)  # 拼回 [B, n, d]
        return self.proj(y)


class MLP(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=False)
        self.proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=False)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.drop(self.proj(F.gelu(self.fc(x))))


class Block(nn.Module):
    """Post-LN：x = LN(x + Attn(x))，再 x = LN(x + MLP(x))。"""

    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.attn = CausalSelfAttention(cfg)
        self.mlp = MLP(cfg)
        self.ln1 = nn.LayerNorm(cfg.n_embd)
        self.ln2 = nn.LayerNorm(cfg.n_embd)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.ln1(x + self.attn(x))
        x = self.ln2(x + self.mlp(x))
        return x


class GPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer))
        self.ln_f = nn.LayerNorm(cfg.n_embd)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.tok_emb.weight
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def forward(self, idx: torch.Tensor, labels: torch.Tensor | None = None):
        B, n = idx.shape                                     # idx: [B, n]
        if n > self.cfg.block_size:
            raise ValueError(f"n={n} 超过 block_size={self.cfg.block_size}")
        pos = torch.arange(n, device=idx.device)             # 位置 0..n-1
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))  # → [B, n, d]
        for block in self.blocks:                            # L 层，形状一直 [B, n, d]
            x = block(x)
        logits = self.lm_head(self.ln_f(x))                  # → [B, n, V]
        loss = None
        if labels is not None:                               # 训练才算 loss；生成不传 labels
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.size(-1)),  # 预测：前 n-1 个位置
                labels[:, 1:].reshape(-1),                    # 标签：后 n-1 个 token（错开一位）
                ignore_index=-100,
            )
        return logits, loss

我们看下网络中的张量形状

• 入口需要：给定 id \(i\)，变成向量 \(e_i\in\mathbb R^d\)。做法是一张 \(V\times d\) 的表 tok_emb，第 \(i\) 行就是 \(e_i\)，查表即可。
• 出口需要：最后一层给出 \(h\in\mathbb R^d\)，要在 \(V\) 个词里选出下一个 id。做法是给每个词一个 \(d\) 维「探针」向量，算 \(h\) 与每个探针的点积，得到 \(V\) 个分数（logits），再 softmax / 取样。这 \(V\) 个探针排成矩阵，形状仍是 \(V\times d\)，就是 lm_head.weight（无 bias 的 Linear(d, V)）。

所以会出现两张同形的表：一张把 id 读成向量，一张把隐状态 比成「更像哪个 id」。它们也可以各学各的。tie_embeddings=True 时令二者共用同一张表——词 \(i\) 的输入向量 \(e_i\)，同时当作出口上与 \(h\) 比较的那一根探针。含义是：下一个词该是 \(i\)，当且仅当 \(h\) 靠近「\(i\) 被读进网络时用的那个向量」。少存一份 \(Vd\) 参数；输入、输出被迫共用同一套词向量几何。实现就是：

self.lm_head.weight = self.tok_emb.weight

GPT-2 一类默认打开。关掉则入口、出口各一份矩阵，多占显存，多一点自由度。

与 apply(_init_weights) 的顺序：上面代码先绑定、再 apply——共享后同一块权重可能被初始化写两次（Embedding 分支一次、Linear 分支一次），分布都是 \(N(0,0.02)\)，留下最后一次。若先 apply、再绑定，会先初始化两块，绑定后丢掉 lm_head 私有的那份。两种最终都是一份共享权重；书里固定为「先绑定、再统一初始化」。

3.3 训练细节

一步里的顺序是固定的：取 batch → 前向出 loss → zero_grad → backward → 看梯度范数 → 超过上限则裁剪 → optimizer.step。学习率写在参数组上，AdamW 下一步就用新值。

配置项	这里的取值	作用
AdamW lr	3e-4，前 100 步线性抬升	小模型字符级过拟合的常用起点
betas	(0.9, 0.95)	二阶矩衰减比 CV 里常见的 0.999 更短，大模型训练里更常见
weight_decay	0.1	只衰减矩阵权重；AdamW 把衰减从梯度更新里拆开
clip	1.0	全局梯度 L2 范数上限
dropout	0	先过拟合一小份文本，确认实现；正则留到数据变大之后

def learning_rate(step: int, warmup: int, max_lr: float) -> float:
    # step 从 0 计；第 0 步也给一个非零 lr，避免整步梯度被乘成 0
    if step < warmup:
        return max_lr * (step + 1) / warmup
    return max_lr


def train(model: GPT, token_ids: torch.Tensor, steps: int, batch_size: int,
          max_lr: float = 3e-4, warmup: int = 100, clip: float = 1.0) -> None:
    device = next(model.parameters()).device
    block = model.cfg.block_size
    # 矩阵/嵌入做 weight decay；LayerNorm 的一维 γ、β 不做
    decay_params, nodecay_params = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        (decay_params if p.ndim >= 2 else nodecay_params).append(p)
    opt = torch.optim.AdamW(
        [
            {"params": decay_params, "weight_decay": 0.1},
            {"params": nodecay_params, "weight_decay": 0.0},
        ],
        lr=max_lr,
        betas=(0.9, 0.95),
    )
    model.train()
    for step in range(steps):                              # ① 每个 step 更新一次参数
        lr = learning_rate(step, warmup, max_lr)           # ② warmup 内抬升学习率
        for group in opt.param_groups:
            group["lr"] = lr
        idx, labels = sample_batch(token_ids, batch_size, block, device)  # ③ 切一批 [B, T]
        _, loss = model(idx, labels=labels)                # ④ 前向 + 交叉熵（内部已错开一位）
        opt.zero_grad(set_to_none=True)                    # ⑤ 清梯度
        loss.backward()                                    # ⑥ 反传
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), clip)  # ⑦ 裁剪
        opt.step()                                         # ⑧ AdamW 更新权重
        if step % 50 == 0 or not math.isfinite(loss.item()):
            print(f"step {step:4d}  lr {lr:.2e}  loss {loss.item():.4f}  grad {grad_norm:.2f}")
        if not math.isfinite(loss.item()):
            raise RuntimeError("loss 变成 NaN/Inf，停止。先把 max_lr 降一个数量级再查掩码与 shift")

二维及以上的参数进 weight decay，一维参数（LayerNorm 的 \(\gamma,\beta\)）不进，避免把归一化增益往 0 拉。

日志里同时看 loss 和 grad。loss 在 1e2 以上长期不动，多半是学习率过小或标签全被 ignore；grad 长期顶在 clip 上，多半是学习率偏大或有一批样本异常。

3.4 训练时会碰到什么

loss 一开始就是 NaN。 学习率过大（例如直接 1e-2 且没有 warmup），或注意力分数在 softmax 前溢出。先把 max_lr 降到 3e-4，确认 warmup 步数大于 0。bf16 训练一般不需要 fp16 那套 loss scaling；本章默认 fp32，单卡小模型足够。

loss 降得异常快，生成却在背文。 数据只有几 KB、模型又记得住整窗，过拟合是预期结果。它证明反传通了，不证明模型会写新句子。要看泛化，得换一份没进训练窗的文本测 loss。

loss 几乎不降。 常见三处：labels 又在外部错开了一次，监督对到了「下下个 token」；因果掩码用了下三角，把左边挡住了；block_size 大于语料，sample_batch 直接报错或你改成了重复 padding 且全部标成 -100，有效 token 数为 0，交叉熵返回 NaN。

生成重复同一个字。 温度过低（接近贪心）加上模型还没训够，会黏在高频字符上。先看训练 loss 是否已经低于「均匀猜字符」的 \(\ln V\)：字符表大约 50 时，\(\ln 50\approx 3.9\)，loss 还在 3.5 附近就不要指望连贯句子。

位置一对不上。 n > block_size 时 pos_emb 没有对应行，代码里直接抛错。把长文截到 block_size 再送；想训更长窗口，先加大 block_size 并重新初始化位置表，旧位置权重不能假设能外推。

权重绑定之后优化器建了两次同一参数。 named_parameters() 对绑定的嵌入只出现一次（lm_head.weight 是同一 Parameter）。若手写参数列表时既加入 tok_emb.weight 又加入 lm_head.weight，AdamW 会为同一张量存两份状态并更新两次。用 named_parameters() 遍历可以避开。

3.5 生成与入口

前面的 train(...) 只负责把 loss 压下去。真正往外写新 token 的是下面的 generate；main 里先 train，再调用一次 generate 打印续写。

@torch.no_grad()
def generate(model: GPT, idx: torch.Tensor, max_new: int, temperature: float = 0.8) -> torch.Tensor:
    model.eval()
    for _ in range(max_new):                          # ① 循环 max_new 次 → 写出 max_new 个 token
        ctx = idx[:, -model.cfg.block_size :]          # ② 当前前缀（超长则只留最后 T 个）
        logits, _ = model(ctx)                         # ③ 整段 ctx 再跑一遍完整前向
        logits = logits[:, -1, :] / temperature        # ④ 只要最后一个位置的词表分数
        probs = logits.softmax(dim=-1)
        nxt = torch.multinomial(probs, num_samples=1)  # ⑤ 抽出 1 个新 id
        idx = torch.cat([idx, nxt], dim=1)             # ⑥ 拼回前缀，下一轮再用
    return idx

是一个 token 一个 token 生成的。 for 每转一圈只执行一次 multinomial(..., num_samples=1)，只得到形状 \([B,1]\) 的 nxt，再 cat 进 idx。本章词表是字符级，一个 token 往往就是一个字符；换成 BPE 后仍是「每次一个 token」，只是一个 token 可能对应多个字符。

哪里看出没有 KV Cache。 对照三处：

1. model 的 forward(self, idx, labels=None) 只有输入 id，没有 past_key_values / cache 之类参数；注意力里每次现场算 \(Q,K,V\)，算完就丢，没有写成成员变量存下来。
2. 生成循环里每一步都是 logits, _ = model(ctx)：把当前整段 ctx（最长 \(T\)）再送进去。上一步已经为前缀算过的 K、V，这一步全部重算一遍。
3. 因此代价随前缀变长而涨：第 \(t\) 步 ctx 长度约为 \(\min(n_0+t,\,T)\)，注意力分是 \([B,h,n',n']\)。若有 KV Cache，新步通常只算长度为 1 的新位置，并拼上旧 K、V；本章代码没有这条路径。

@torch.no_grad() 只表示生成不算梯度，和有没有缓存无关。model.eval() 在 dropout=0 时与 train() 数值相同；一旦打开 dropout，生成前必须 eval()。

def main() -> None:
    raw = open("input.txt", encoding="utf-8").read()
    tok = CharTokenizer(raw)
    token_ids = torch.tensor(tok.encode(raw), dtype=torch.long)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    cfg = GPTConfig(vocab_size=tok.vocab_size, block_size=64, n_layer=4, n_head=4, n_embd=128)
    model = GPT(cfg).to(device)
    n_param = sum(p.numel() for p in model.parameters())
    print(f"device={device}  vocab={cfg.vocab_size}  params={n_param/1e6:.2f}M")
    train(model, token_ids, steps=300, batch_size=16)   # 训练
    prompt_ids = tok.encode(raw[:16])
    prompt = torch.tensor([prompt_ids], dtype=torch.long, device=device)
    out = generate(model, prompt, max_new=120, temperature=0.8)  # 自回归续写 120 个 token
    print(tok.decode(out[0].tolist()))


if __name__ == "__main__":
    main()

input.txt 放在工作目录。第 4 章会把「③ 整段重算」改成「只算新 token，K/V 留下复用」。

3.6 这一份代码里的形状

\(d\) 就是配置字段 n_embd。训练时 sample_batch 切出的窗口长度等于 block_size，故下面训练路径里常有 \(n=T\)；生成时 \(n\) 从短前缀涨到至多 \(T\)。

记号	配置 / 变量	main() 例子	含义
\(B\)	batch_size	16	batch
\(n\)	idx.shape[1]	训练时 64；生成时 \(1\sim 64\)	当前序列长
\(T\)	block_size	64	最大窗口；pos_emb 与 causal 的边长
\(V\)	vocab_size	视语料	词表大小
\(d\)	n_embd	128	隐状态宽度
\(h\)	n_head	4	头数
\(d_h\)	n_embd // n_head	32	每头宽度；必须 \(d=h\cdot d_h\)
\(L\)	n_layer	4	Block 个数
\(N\)	token_ids.numel()	语料长度	一维语料 token 数

本章没有 KV 缓存：没有 k_cache / v_cache / past_key_values。训练时自动求导会暂存激活供 backward，那是框架行为，不是手写缓存张量。

1. 数据

张量	形状	代码位置
token_ids	\([N]\)，long	语料 encode
starts	\([B]\)	sample_batch 里每个样本的起点
idx	\([B,n]\)，训练时 \(n=T\)	送进 model
labels	\([B,n]\)	本代码与 idx 同一块；-100 表示 ignore（本例没有）

2. GPT.forward 入口

张量	形状	说明
idx	\([B,n]\)	要求 \(n\le T\)
pos = arange(n)	\([n]\)	窗口内位置下标 \(0..n-1\)
tok_emb(idx)	\([B,n,d]\)	按 id 查行
pos_emb(pos)	\([n,d]\)	按位置查行；与上一行相加时广播成 \([B,n,d]\)
tok_emb + pos_emb	\([B,n,d]\)
x（进 blocks 前）	\([B,n,d]\)	再过 self.drop

3. 一层里的注意力 CausalSelfAttention.forward

记输入为 \(x\in\mathbb R^{B\times n\times d}\)。

步骤	张量	形状
线性	qkv(x)	\([B,n,3d]\)
view		\([B,n,3,h,d_h]\)
permute(2,0,3,1,4)	qkv 五维	\([3,B,h,n,d_h]\)
拆开	\(Q,K,V\)	各 \([B,h,n,d_h]\)
\(K\) 转置	k.transpose(-2,-1)	\([B,h,d_h,n]\)
分数	att = (Q @ Kᵀ) * d_h^{-1/2}	\([B,h,n,n]\)
掩码	causal[:n,:n]	\([n,n]\)，从缓冲 \([T,T]\) 切开；True 处填 -inf
softmax + drop	att	仍是 \([B,h,n,n]\)
加权	att @ V	\([B,h,n,d_h]\)
transpose(1,2)		\([B,n,h,d_h]\)
view	拼头	\([B,n,d]\)（\(d=h\cdot d_h\)）
输出投影	proj(...)	\([B,n,d]\)

缓冲 self.causal 常驻 \([T,T]\)，与 batch 无关；每个注意力模块各有一份（本代码每层一个 CausalSelfAttention）。

4. 一层里的 MLP 与 Block

步骤	张量	形状
fc(x)		\([B,n,4d]\)
gelu(...)		\([B,n,4d]\)
proj(...)		\([B,n,d]\)
drop(...)	MLP 输出	\([B,n,d]\)

Post-LN 残差（形状全程 \([B,n,d]\)）：

步骤	表达式	形状
注意力支路	attn(x)	\([B,n,d]\)
残差 + LN	ln1(x + attn(x))	\([B,n,d]\)
MLP 支路	mlp(...)	\([B,n,d]\)
残差 + LN	ln2(... + mlp(...))	\([B,n,d]\) → 本层输出

\(L\) 层串行：第 \(\ell\) 层出入都是 \([B,n,d]\)，宽度不变。

5. 出口与 loss

张量	形状	说明
最后一层输出 \(x\)	\([B,n,d]\)
ln_f(x)	\([B,n,d]\)
logits = lm_head(...)	\([B,n,V]\)	每个位置一份词表分数
logits[:, :-1]	\([B,n-1,V]\)	丢掉最后一列（没有下一位标签）
labels[:, 1:]	\([B,n-1]\)	丢掉第一列，与上面对齐
reshape 后预测	\([B(n-1),\,V]\)
reshape 后标签	\([B(n-1)]\)
loss	标量	cross_entropy 平均

labels is None（生成路径）时不算 loss，只返回 logits。

6. 生成 generate（仍无 KV 缓存）

设 batch=1。第 \(t\) 步（已写出 \(t\) 个新 token 之后）：

张量	形状
idx（当前前缀）	\([1,\,n_0+t]\)，再截断 [:, -T:] 得 \(n'\le T\)
ctx	\([1,n']\)
整网 forward(ctx) 的 logits	\([1,n',V]\)
logits[:, -1, :]	\([1,V]\)
/ temperature 再 softmax	probs \([1,V]\)
multinomial → nxt	\([1,1]\)
cat 后新 idx	\([1,\,n_0+t+1]\)

每步注意力分是 \([1,h,n',n']\)，随 \(n'\) 变大；历史 K/V 不另存，下一步整段重算。

7. 权重（参数）形状

nn.Linear(in_f, out_f) 的 .weight 为 \([out_f,\,in_f]\)；无 bias 时没有 .bias。

模块	参数	形状	个数（模块数）
词嵌入	tok_emb.weight	\([V,d]\)	1
位置嵌入	pos_emb.weight	\([T,d]\)	1
注意力	qkv.weight	\([3d,\,d]\)	\(L\)
注意力	proj.weight	\([d,d]\)	\(L\)
MLP	fc.weight	\([4d,\,d]\)	\(L\)
MLP	proj.weight	\([d,\,4d]\)	\(L\)
Block 内 LN	ln1/ln2 的 weight,bias	各 \([d]\)	每层 2 个 LN
末尾 LN	ln_f.weight/bias	各 \([d]\)	1
输出头	lm_head.weight	\([V,d]\)	与 tok_emb 同一对象（tie_embeddings=True）；关掉则多一份独立 \([V,d]\)
缓冲（非参数）	causal	\([T,T]\)，bool	每层注意力 1 份；不计进 parameters()

绑定打开时，可训练参数量粗算：

\[ Vd + Td + L\bigl(3d\cdot d + d\cdot d + 4d\cdot d + d\cdot 4d\bigr) + \underbrace{(2L+1)\cdot 2d}_{\text{各 LN 的 }\gamma,\beta} = Vd + Td + 12 L d^2 + (4L+2)d. \]

（注意力 \(3d^2+d^2\)，MLP \(4d^2+4d^2\)，合计每层 \(12d^2\)。）main 打印的 params 应与此一致。

8. 代入 main() 数字看一眼

取 \(B=16,\,n=T=64,\,d=128,\,h=4,\,d_h=32,\,L=4\)，并设 \(V=65\)（示意）：

激活	形状	元素个数（约）
\(x\)	\([16,64,128]\)	\(1.3\times 10^5\)
\(Q\)（单层）	\([16,4,64,32]\)	同量级
att（单层）	\([16,4,64,64]\)	\(2.6\times 10^5\)
MLP 中间	\([16,64,512]\)	\(5.2\times 10^5\)
logits	\([16,64,65]\)	\(6.7\times 10^4\)

隐状态宽度在残差干线上始终是 128；只有注意力头内是 32、MLP 里临时 512。若某步打印出第三维既不是 128 也不是 512，多半是 view/reshape 写错。

---

核心问答

Q1：idx 和 labels 为什么可以是同一块窗口？ 因为错位发生在 loss 内部：用前 \(n-1\) 个位置的 logits 去拟合后 \(n-1\) 个 token。若 DataLoader 已经给出 labels = idx 向左移一位并在末尾补 -100，模型里就不能再切一次 [:, 1:]。

Q2：交叉熵里的 reshape(-1, V) 在做什么？ logits[:, :-1] 是 [B, n-1, V]。cross_entropy 要求预测为 [N, V]、目标为 [N]。把 batch 和长度折到同一维，等价于对所有有效位置求平均，不需要手写双重循环。

Q3：ignore_index 会让 loss 的分母变小吗？ 会。被忽略的位置不进入平均。若一个 batch 里有效标签一个都没有，PyTorch 的交叉熵得到 NaN。padding 比例极高时要检查是不是整批都被标成了 -100。

Q4：为什么 LayerNorm 参数不做 weight decay？ 衰减把参数往 0 拉。矩阵权重上这是正则；LayerNorm 的缩放若被拉向 0，这一层的输出幅度会被人为压扁，和「把激活拉回稳定尺度」的用途相反。所以一维参数放进 weight_decay=0 的那一组。

Q5：梯度范数长期等于 clip，要不要紧？ 偶尔顶到上限是裁剪在起作用。若几百步里几乎步步都顶满，更新方向一直被同一比例缩小，等价于实际学习率远小于你设的 max_lr，应降学习率或查是不是有输入异常（全 0、错位的极大 ID）。

Q6：loss 降了，为什么续写仍是乱码？ 字符级、几百 step、几万参数，模型只记住了局部 n-gram。loss 低于 \(\ln V\) 只说明比「均匀乱猜」好。连贯句子依赖词表、数据规模和训练步数，不由这一章的环是否写对单独决定。

Q7：pos_emb 用下标 0..n-1，和 token 的 ID 有关吗？ 无关。位置表按窗口内的序号查行，token 表按字符 ID 查行，两者相加。同一字符出现在窗口的不同列，拿到的位置向量不同。

Q8：打开 dropout 之后，训练 loss 和重新前向的 loss 对不上？ 训练态下 dropout 每次掩码不同，同一条 idx 两次前向 loss 可以差一截。对照实现时用 model.eval()，或把 dropout 设为 0。

Q9：zero_grad(set_to_none=True) 和置零有什么差别？ set_to_none 把 .grad 设成 None，下一步 backward 直接写入，少一次把梯度缓冲填 0 的读写。语义上本步梯度不与上步相加。若你依赖梯度累积（多步 backward 再一次 step），中途不能 zero_grad。

Q10：绑定词嵌入之后，参数量怎么数？ lm_head.weight 与 tok_emb.weight 是同一对象，parameters() 只计一次。手算是 \(V \times d\) 而不是 \(2Vd\)，再加上位置表 \(block\_size \times d\) 和各层投影。

Q11：温度设成 0 会怎样？ logits / temperature 会除零。贪心应改成 logits.argmax，不要把温度写成 0。温度远大于 1 时分布趋近均匀，续写接近从词表里乱抽。

Q12：这份代码和 HuggingFace GPT2LMHeadModel 差在哪？ 差在层的细节：这里是 Post-LN、可学习绝对位置、GELU 两层 FFN、字符词表。训练环（shift 交叉熵、AdamW、裁剪、按步改学习率）同一结构。读开源 forward 时，先找 labels 是否在内部 shift、pad 是否为 -100。