程式碼逐行解讀

整體流程

prepare.py (一次性) → train.py (每輪訓練)
    ↓                     ↓
下載資料 + 訓練          GPT 模型定義
BPE tokenizer            Muon + AdamW 最佳化器
                         5 分鐘訓練迴圈
                         val_bpb 評估

train.py 重點解讀

1. 模型架構（GPTConfig + GPT）

@dataclass
class GPTConfig:
    sequence_len: int = 2048
    vocab_size: int = 32768
    n_layer: int = 12
    n_head: int = 6
    n_kv_head: int = 6
    n_embd: int = 768
    window_pattern: str = "SSSL"

Agent 可修改的範圍包含：層數、注意力頭數、嵌入維度、視窗模式等。

2. Value Residual（ResFormer）

def has_ve(layer_idx, n_layer):
    return layer_idx % 2 == (n_layer - 1) % 2

# Value embedding 每兩層交替出現
if ve is not None:
    gate = 2 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
    v = v + gate.unsqueeze(-1) * ve

Value Residual 機制：將 token 的 value embedding 以 input-dependent gate 混合到 attention 的 value 中。來自 ResFormer 論文的概念。

3. QK-Normalization

q, k = norm(q), norm(k)

在 attention 計算前對 query 與 key 做 RMS Norm，穩定訓練。

4. Logit Softcapping

softcap = 15
logits = self.lm_head(x)
logits = softcap * torch.tanh(logits / softcap)

限制 logits 範圍在 [-softcap, softcap]，避免極端值影響 softmax。

5. Muon 最佳化器

# Polar Express orthogonalization (Newton-Schulz iterations)
X = g.bfloat16()
X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
for a, b, c in polar_express_coeffs[:ns_steps]:
    A = X.mT @ X
    B = b * A + c * (A @ A)
    X = a * X + X @ B

Muon 最佳化器比 AdamW 更適合訓練 LLM，核心思想是對梯度矩陣做正交化（近似 SGD 的極致版本）。

6. 固定時間訓練迴圈

while True:
    # ... 訓練步驟 ...
    progress = min(total_training_time / TIME_BUDGET, 1.0)
    lrm = get_lr_multiplier(progress)
    # ...
    if step > 10 and total_training_time >= TIME_BUDGET:
        break

使用 wall-clock time 而非 step count 決定何時停止。前 10 步為 warmup 階段（用於 torch.compile 編譯）。