從零實作一個能生成文字的 GPT 模型Implementing a GPT model from scratch to generate text

attention 只是引擎，要組成一台能吐字的 GPT 還需要一批「看似配角、缺一不可」的零件。這是本章的積木清單——它會隨你往下讀自動打勾（也可以點擊跳到該節）：

暖身：tensor 就是「多層的 array」

模型看不懂文字，只吃數字。一個字（token）→ 一串數字（vector）；一句話 → 一疊向量（matrix）；一批句子 → 再疊一層（3 維 tensor）。以單字 moves 為例：

「參數（parameters）」就是模型裡可被訓練調整的權重（weights）。一層 2048×2048 的權重矩陣就有 2048×2048 = 4,194,304 個參數。我們瞄準 GPT-2 small（約 124M 參數）：它的預訓練權重公開、筆電跑得動；GPT-3 架構幾乎相同，只是規模放大到 175B。

所有超參數集中在一份設定字典。點任一行看說明：

佔位骨架 DummyGPTModel

在刻好每個零件前，先用「什麼都不做、原樣回傳」的假模組把骨架搭起來，看清楚資料怎麼從輸入流到輸出：

class DummyGPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])   # token 查表成向量
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"]) # 位置查表成向量
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        self.trf_blocks = nn.Sequential(                                  # 12 個（假的）block
            *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        self.final_norm = DummyLayerNorm(cfg["emb_dim"])                  # （假的）最終正規化
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(in_idx.shape[1], device=in_idx.device))
        x = tok_embeds + pos_embeds     # 「是什麼字」+「在第幾位」
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        return self.out_head(x)         # 投影回 50257 維 → logits

動畫：一個英文字如何走進模型

在餵整個 batch 之前，先看單一個字的完整旅程：文字 → 查 vocabulary 得到整數 ID → 到 embedding 表取出那一列向量。選一個字、按播放：

實際餵資料：兩句話的 batch

用 tiktoken 把兩句話編碼成一個 batch。滑過下面的字或數字，看文字 ↔ token ID 的對應（注意兩句開頭都是 Every → 同一個 ID 6109）：

輸出形狀 [2, 4, 50257]

把 batch 餵進模型（torch.manual_seed(123)），輸出 logits 的形狀是 torch.Size([2, 4, 50257])。點下面三個維度，看每一維在 array 裡對應什麼：

深入：為什麼能一次算 4 個位置，又不會「偷看未來」？（causal mask）🏋️ 訓練

雖然整句 token 是一起送進 transformer 做一次矩陣運算，但 causal mask 讓每個位置的 hidden state 只能利用自己與前面的 token——所以它產生的 logits 從未看過未來。點一個位置，看它的視野和它在回答的問題：

深層網路常遇到 vanishing / exploding gradient（梯度消失／爆炸），訓練極不穩定。Layer normalization 的核心想法：把一層的輸出調整成 mean 0、variance 1（unit variance），加速收斂、穩定訓練。

互動①：訊號穿過 30 層，會發生什麼事？

把一個訊號丟進 30 層的網路（真實矩陣相乘）。第 l 顆點＝訊號過完第 l 層時的大小——精確說是該層輸出向量的 RMS（均方根：√mean(x²)，24 個數字「典型有多大」；不用平均值是因為正負會抵銷）。點縮小到看不見＝訊號消失、點變紅＝數值爆掉；滑鼠停在點上可看實際數值。

互動②：那 loss 會怎樣？（瀏覽器裡真的訓練 300 步）🏋️ 訓練

同一個 6 層小網路，只差在初始權重的倍率 s（就是互動①那顆旋鈕，訓練開始時權重的整體大小）：s=0.3（太小）、s=3.0（太大）、s=3.0 但每層加 LayerNorm。在你的瀏覽器裡實際跑 300 步訓練（真實 forward + backward + 權重更新），把 loss 畫出來：

先看 dim 往哪個方向算

書中的例子：torch.randn(2,5) 過一層 Linear(5,6)+ReLU，得到一個 2×6 的輸出（ReLU 把負值壓成 0，所以沒有負數）。切換 dim，看平均值是「沿哪個方向」算的：

一步步做正規化（數值即時計算，和書中輸出一致）

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5                                  # 防止除以零
        self.scale = nn.Parameter(torch.ones(emb_dim))   # 可訓練縮放
        self.shift = nn.Parameter(torch.zeros(emb_dim))  # 可訓練平移

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)              # keepdim 保形狀好做廣播
        var  = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

LLM 不用傳統 ReLU，改用更平滑的 GELU（Gaussian error linear unit）。精確定義是 GELU(x) = x·Φ(x)（Φ 是標準常態的 CDF），實務用便宜的近似式（原始 GPT-2 也用它）：

GELU(x) ≈ 0.5·x·(1 + tanh[ √(2/π)·(x + 0.044715·x³) ])

互動：GELU 的好處——神經元不會「死掉」

FeedForward：768 → 3072 → 768

transformer block 裡的前饋網路＝兩個 Linear 夾一個 GELU。先放大 4 倍到更寬的空間探索特徵、過非線性、再收縮回原維度。點「下一步」看形狀流動（用書中的輸入 [2, 3, 768]）：

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),  # 768 → 3072（放大 4 倍）
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),  # 3072 → 768（收縮回來）
        )

下一節會用「各層梯度」證明 shortcut 的威力——先花兩分鐘搞懂梯度從哪來。forward pass 把輸入算成預測、算出 loss；backward pass（backpropagation）用連鎖律從 loss 往回，算出每個權重「該往哪調」。

Shortcut connection（又稱 skip / residual connection，源自電腦視覺的 ResNet）作法樸素到不可思議：把某一層的輸入，直接加到那一層的輸出上。這等於替反向傳播的梯度開一條「抄近路」的旁路，繞過中間運算、避免梯度層層縮小。

書中用一個 5 層小網路（每層 Linear+GELU）實測。forward 裡可選擇要不要加捷徑：

def forward(self, x):
    for layer in self.layers:
        layer_output = layer(x)
        if self.use_shortcut and x.shape == layer_output.shape:
            x = x + layer_output      # 加回輸入＝捷徑
        else:
            x = layer_output
    return x

切換開關，比較各層權重的平均絕對梯度（右側是書中實際的 print 輸出）：

把 MultiHeadAttention、LayerNorm、FeedForward、dropout、shortcut 全部組進一個 transformer block——GPT-2 small 把它重複 12 次。兩個子模組分工明確：

• multi-head attention：「跨位置」的資訊交流——讓每個 token 看到並吸收其他 token 的資訊。
• feed forward：「逐位置」的加工——對每個位置的向量獨立做非線性變換。

逐行走一遍 forward（注意兩個子層都是同一個範式：存 shortcut → LayerNorm → 子層 → dropout → 加回）：

把骨架裡的假模組換成真貨，就是完整的 GPTModel：embedding → 12× TransformerBlock → final LayerNorm → linear output head。輸出形狀依然 [2, 4, 50257]——複習一下這三個數字：2＝batch（一次餵了 2 句話）、4＝每句 4 個 token、50257＝每個 token 位置對詞彙表全部 50257 個字各給一個分數（logits）。也就是「2 句 × 每句 4 個位置 × 每個位置一整排『下一個字是誰』的評分」（回 4.1 的維度互動）。

self.trf_blocks = nn.Sequential(
    *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])   # 這次是真的 12 層
self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

參數量：163M 與 124M 的差別（weight tying）

• 省多少：50257 × 768 ≈ 38.6M 個參數（float32 約 147 MB）——tying 後 163M → 124M、權重檔 622 MB → 約 475 MB。
• 代價：輸入與輸出被迫共用同一套表徵，少了一份可以分工的權重。實務上兩派都有：原始 GPT-2 綁（省參數）；本書實作不綁——經驗上各自獨立的訓練表現通常略好。
• 之後會再遇到：載入 OpenAI 的 GPT-2 預訓練權重時，因為官方是綁的，要把同一份 embedding 權重同時填進 tok_emb 與 out_head。

動手感受規模：參數計算器

拖動滑桿即時計算總參數與記憶體。preset 是書中列的四種 GPT-2 尺寸（medium 1024/24 層、large 1280/36、XL 1600/48）。計算器用的公式如下（帶入 GPT-2 small 的數值驗算）：

# 記 E = emb_dim(768)、V = vocab_size(50257)、C = context_length(1024)、L = n_layers(12)

Token embedding      = V × E                       = 50257×768          =  38,597,376
Positional embedding = C × E                       = 1024×768           =     786,432

# 每個 transformer block：
  attention  = 3E²（Q,K,V，無 bias）+ (E²+E)（輸出投影含 bias）
             = 3×768² + (768²+768)                 = 1,769,472+590,592  =   2,360,064
  feed forward = (E·4E + 4E) + (4E·E + E)          # 768→3072 再 3072→768，皆含 bias
             = (768×3072+3072) + (3072×768+768)                         =   4,722,432
  2× LayerNorm = 2×(E + E)（scale 與 shift）        = 4×768              =       3,072
  → 每 block 小計                                                        =   7,085,568

Transformer ×L       = 12 × 7,085,568                                   =  85,026,816
Final LayerNorm      = 2E                          = 2×768              =       1,536
── 到這裡的小計（＝weight tying 後的總數）                                 124,412,160
Output head（不綁時）= V × E                        = 50257×768          =  38,597,376
── 全部（我們的實作）                                                     163,009,536

# 記憶體：163,009,536 × 4 bytes (float32) ≈ 621.83 MB

生成式模型一次產生一個 token、接龍式寫句子。單步機制：模型輸出 logits → 只取最後一個位置那條 50257 維向量 → softmax 轉機率 → argmax 挑最大 → 把新 token ID 接回輸入 → 再跑一輪。

def generate_text_simple(model, idx, max_new_tokens, context_size):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]        # 超長就裁掉前面
        with torch.no_grad():                    # 推論不用算梯度，省時省記憶體
            logits = model(idx_cond)
        logits = logits[:, -1, :]                # (batch, n_token, vocab) → 只取最後一格
        probas = torch.softmax(logits, dim=-1)
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # greedy：挑機率最高
        idx = torch.cat((idx, idx_next), dim=1)  # 接回序列尾巴
    return idx

實跑：模型還沒訓練，會發生什麼事？

把 "Hello, I am" 編碼成 [15496, 11, 314, 716]、模型切到 model.eval()（關掉 dropout），生成 6 個新 token。按「生成下一個 token」逐步看 ID array 長大——這些是書中用 seed 123 跑出的真實輸出：

• Layer normalization 沿特徵維（dim=-1）把輸出調成 mean 0 / var 1，穩定訓練；scale/shift 可訓練、eps 防除零、unbiased=False 相容 GPT-2 權重。
• GELU 平滑、負值保留非零梯度，比 ReLU 更適合深層 LLM。
• FeedForward：768 → 3072 → 768，「膨脹—非線性—收縮」探索更豐富的表徵。
• Shortcut connection 把輸入加回輸出；5 層實測：無捷徑第 1 層梯度 0.0002（消失）、有捷徑 0.22（穩定），差三個數量級。
• Transformer block＝attention（跨位置）＋ feed forward（逐位置），Pre-LayerNorm＋雙捷徑；形狀不變、內容重新編碼。
• GPTModel＝embedding ＋ 12×block ＋ final norm ＋ output head；我們 163M，GPT-2 靠 weight tying 記為 124M，權重約 622 MB。
• 生成＝取最後一格 logits → softmax → argmax（greedy）→ 接回輸入循環；未訓練 → 亂碼。

// 六題，作答後立即看解析