第 4 章 组装 GPT-2

struct GPTConfig {
    int64_t vocab_size = 50257;       // BPE 词表（第 2 章）
    int64_t context_length = 1024;    // 最大可观察的词元数
    int64_t emb_dim = 768;            // 每个词元的嵌入向量维度
    int64_t n_heads = 12;             // 注意力头数（第 3 章）
    int64_t n_layers = 12;            // Transformer 块数
    double drop_rate = 0.1;           // dropout 比例
    bool qkv_bias = false;            // Q/K/V 投影是否加偏置（GPT-2 官方为 true，本教程取 false）
};

class LayerNormImpl : public torch::nn::Module {
public:
    explicit LayerNormImpl(int64_t emb_dim) {
        scale_ = register_parameter("scale", torch::ones(emb_dim));
        shift_ = register_parameter("shift", torch::zeros(emb_dim));
    }
    torch::Tensor forward(torch::Tensor x) {
        auto mean = x.mean(-1, /*keepdim=*/true);
        auto var  = x.var({-1}, /*unbiased=*/false, /*keepdim=*/true);  // 有偏方差！
        auto norm_x = (x - mean) / torch::sqrt(var + eps_);             // eps_ = 1e-5
        return scale_ * norm_x + shift_;
    }
    torch::Tensor scale_, shift_;
    double eps_ = 1e-5;
};

层           无快捷连接      有快捷连接     提升倍数
layers.0     0.0002017      0.2217        ≈ 1099×
layers.1     0.0001201      0.2069        ≈ 1723×
layers.2     0.0007152      0.3290        ≈ 460×
layers.4     0.0050496      1.3259        ≈ 263×

torch::Tensor forward(torch::Tensor x) {
    // 注意力子块：Pre-LayerNorm + 快捷连接
    auto shortcut = x;
    x = norm1->forward(x);
    x = att->forward(x);            // 第 3 章的 MultiHeadAttention
    x = drop_shortcut->forward(x);
    x = x + shortcut;               // 残差

    // 前馈子块：同样一套
    shortcut = x;
    x = norm2->forward(x);
    x = ff->forward(x);
    x = drop_shortcut->forward(x);
    x = x + shortcut;
    return x;
}

torch::Tensor forward_hidden(torch::Tensor in_idx) {   // [batch, seq_len] 词元 ID
    auto tok_embeds = tok_emb->forward(in_idx);        // (b, T, 768)  词元嵌入
    auto pos_embeds = pos_emb->forward(torch::arange(seq_len));  // (T, 768) 位置嵌入
    auto x = tok_embeds + pos_embeds;                  // 广播相加（第 2 章 2.7）
    x = drop_emb->forward(x);
    x = trf_blocks->forward(x);                        // 12 个 TransformerBlock
    x = final_norm->forward(x);
    return x;                                          // (b, T, 768) 隐藏表示
}
torch::Tensor forward(torch::Tensor in_idx) {
    return out_head->forward(forward_hidden(in_idx));  // (b, T, 50257) logits
}

for (int64_t i = 0; i < max_new_tokens; ++i) {
    auto idx_cond = idx.index({Slice(), Slice(-context_size, None)});  // 截断上下文
    torch::Tensor logits;
    { torch::NoGradGuard no_grad;
      logits = model->forward(idx_cond); }
    logits  = logits.index({Slice(), -1, Slice()});      // 只取最后一个位置的 logits
    auto idx_next = torch::argmax(torch::softmax(logits, -1), -1, /*keepdim=*/true);
    idx = torch::cat({idx, idx_next}, 1);                // 拼回序列，进入下一轮
}

Total number of parameters: 163,009,536            ← 全部参数（含输出头）
Number of trainable parameters considering weight tying: 124,412,160
Total size of the model: 621.83 MB                 ← 163,009,536 × 4 字节

[练习 4.1] FeedForward 参数量 = 4,722,432, MultiHeadAttention 参数量 = 2,360,064
[练习 4.2] 公式估算 vs 实际（124M）: 163009536 (实际 163009536)
  GPT-2 medium (emb=1024, layers=24, heads=16): 总参数 ...，weight tying 后 ...
  GPT-2 large (emb=1280, layers=36, heads=20):  ...
  GPT-2 xl (emb=1600, layers=48, heads=25):    ...

./scripts/run.sh chapter04_gpt

encoded_tensor.shape: [1, 4]
Output length: 10
Decoded text: "Hello, I am <乱码词元>"      ← 未训练模型的预期行为