搭起整台机器 · GPT 的 init 与 forward · nanoGPT 源码解读

逐行读 GPT 类——__init__ 怎么把词嵌入、位置嵌入、一摞 Block、输出层拼起来，forward 怎么从 token 编号一路算到 loss。外加三个精妙细节：权重绑定、推理期只算最后一个位置、以及那个 0.02/√(2·n_layer) 的残差缩放初始化。

搭起整台机器 · GPT 的 init 与 forward

零件集齐了，这一章把它们装成整台机器。GPT 类是 model.py 的总装车间：__init__ 负责搭骨架，forward 负责跑数据流。我们逐段读，顺带拆三个容易被略过、却很精妙的细节。

`init`：总装车间

class GPT(nn.Module):

    def __init__(self, config):
        super().__init__()
        ...
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = LayerNorm(config.n_embd, bias=config.bias),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

一个 ModuleDict 把主体零件收在一起，名字都眼熟（对应全景图）：

wte(word token embedding)：词嵌入表，vocab_size 个 token，每个一个 n_embd 维向量——「查字典」的入口。
wpe(word position embedding)：位置嵌入表，block_size 个位置各一个向量——告诉模型「你在第几位」。
drop：输入 dropout。
h：n_layer 个 Block 串成的列表——主体就是这一摞。
ln_f：最后一道归一化(final)。
lm_head：输出头，把 n_embd 维投影成 vocab_size 个分数——「做预测」的出口。注意 bias=False。

细节一：权重绑定(weight tying)

self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying

这一行看着不起眼，却很妙：让「输入词嵌入表」wte 和「输出投影」lm_head 共用同一个权重矩阵。

为什么合理？wte 是「token 编号 → 向量」，lm_head 是「向量 → token 分数」，两者恰好是一件事的正反两面。绑在一起，不但省下一大块参数（词表 × 嵌入维，是模型最大的矩阵之一），语义上也更自洽。这招叫 weight tying(Press & Wolf， 2016)，几乎所有 GPT 类模型都用。

细节二：残差缩放初始化

self.apply(self._init_weights)
# apply special scaled init to the residual projections, per GPT-2 paper
for pn, p in self.named_parameters():
    if pn.endswith('c_proj.weight'):
        torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

self.apply(self._init_weights)：给所有层做基础初始化（下面看 _init_weights）。
然后单独把所有名字以 c_proj.weight 结尾的权重（即每个子层输出投影：注意力的 c_proj 和 MLP 的 c_proj）的初始化标准差，乘一个 1/√(2·n_layer) 缩小。

为什么？ 回忆残差：每个 Block 都往 x 上累加两个子层的输出。层数 n_layer 越多，加的次数越多，残差流的方差会越滚越大、数值失控。GPT-2 论文的办法是：按层数把每个子层的输出「调小」一点（除以 √(2·n_layer)，那个 2 是因为每层有注意力、MLP 两个子层），让累加 N 次之后，整体方差还稳得住。一个用来「压住深度带来的数值膨胀」的初始化技巧——正是《大模型解构》归一化那章提到的「Pre-Norm 残差会随深度变大」的另一种应对。

再看基础初始化 _init_weights：

def _init_weights(self, module):
    if isinstance(module, nn.Linear):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            torch.nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

朴素得很：所有 Linear 和 Embedding 的权重，用均值 0、标准差 0.02 的正态分布初始化；bias 清零。0.02 是 GPT-2 传下来的经验值。

`forward`：数据从编号流到 loss

def forward(self, idx, targets=None):
    device = idx.device
    b, t = idx.size()
    assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
    pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)

    # forward the GPT model itself
    tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
    pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
    x = self.transformer.drop(tok_emb + pos_emb)
    for block in self.transformer.h:
        x = block(x)
    x = self.transformer.ln_f(x)
    ...

对着全景图那张流水线逐行走：

idx 是一批 token 编号，形状 (b, t)。
assert t <= block_size：序列不能超过上下文长度（位置表就那么长）。
pos = arange(0, t)：生成位置索引 [0, 1, ..., t-1]。
tok_emb = wte(idx)：查词嵌入表，(b, t, n_embd)。
pos_emb = wpe(pos)：查位置嵌入表，(t, n_embd)。
x = drop(tok_emb + pos_emb)：词向量 + 位置向量，过 dropout。这一步之后，每个 token 的向量既带「我是什么词」又带「我在第几位」。
for block in h: x = block(x)：流过每一层 Block——整个模型的重头戏，都在这个循环里。
x = ln_f(x)：最后归一化。

细节三：训练算 loss，推理只算最后一个位置

    if targets is not None:
        # if we are given some desired targets also calculate the loss
        logits = self.lm_head(x)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
    else:
        # inference-time mini-optimization: only forward the lm_head on the very last position
        logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
        loss = None

    return logits, loss

forward 有两种模式，由给不给 targets 决定：

训练时（给了 targets）：对每一个位置都算 logits(lm_head(x))，再和真实的下一个 token 比，算交叉熵损失。logits.view(-1, vocab_size) 和 targets.view(-1) 把 batch 和时间维摊平，一口气算所有位置的 loss。ignore_index=-1 是说「标成 -1 的位置不计入 loss」（留给需要 padding 的场景）。
推理时（没给 targets）：这里有个小优化——生成下一个 token，只需要最后一个位置的预测，前面那些位置的 logits 算了也没用。所以 x[:, [-1], :] 只取最后一个位置喂给 lm_head，省一大笔计算。注释特意说明：用 [-1]（列表）而不是 -1（标量），是为了保住时间维，让形状仍是 (b, 1, n_embd)。

交叉熵(cross entropy)在做什么？ 一句话：模型对「下一个 token 是谁」输出了一组分数(logits)，交叉熵衡量「这组分数和真实答案差多远」——真实那个 token 的预测概率越高，loss 越低。训练就是不停地把这个 loss 往下压。你在什么是 nanoGPT 那章的流水线图里见过它：模型出 logits → 和真实答案比 → 得 loss → 反向传播调参数。

模型部分，通关

到这里，model.py 的五个类——LayerNorm、CausalSelfAttention、MLP、Block、GPT——你已经一行不落地读完了。一个能前向传播、能算 loss 的 GPT，就这样搭起来了。

但它现在还是个「什么都不会」的新生儿：权重是随机初始化的，吐出来的全是乱码。要让它学会写莎士比亚，得把它扔进训练循环里，拿数据反复喂、拿梯度反复调。这就是第三站——训练。我们从「怎么把训练脚本配起来」开始。👉 训练的起点：配置与初始化

参考文献

[1] Press & Wolf. “Using the Output Embedding to Tie Word Embeddings.” EACL 2017. arXiv:1608.05859. —— 权重绑定。
[2] Radford et al. “Language Models are Unsupervised Multitask Learners (GPT-2).” 2019. —— 0.02 初始化与残差缩放初始化的出处。
[3] Vaswani et al. “Attention Is All You Need.” NeurIPS 2017. arXiv:1706.03762. —— Transformer 整体结构。
源码：Andrej Karpathy, nanoGPT, model.py. https://github.com/karpathy/nanoGPT

搭起整台机器 · GPT 的 init 与 forward

本页目录