Skip to content

程式碼逐行解讀 — microgpt.py

檔案標頭(第 1–7 行)

python
"""The most atomic way to train and run inference for a GPT in pure, dependency-free Python.
This file is the complete algorithm.
Everything else is just efficiency.
@karpathy"""

作者開宗明義:這份檔案是完整的演算法,其他的(PyTorch、CUDA、分散式訓練)只是效率問題。

導入與亂數種子(第 9–12 行)

python
import os, math, random
random.seed(42)

僅使用 Python 標準函式庫。設定亂數種子確保可重現性。

資料集載入(第 15–21 行)

python
if not os.path.exists('input.txt'):
    import urllib.request
    names_url = 'https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt'
    urllib.request.urlretrieve(names_url, 'input.txt')
docs = [line.strip() for line in open('input.txt') if line.strip()]
random.shuffle(docs)

使用 makemore 專案的 names.txt——一個包含約 32,000 個英文名字的資料集。每行一個名字,讀取後打亂順序。

Tokenizer(第 24–27 行)

python
uchars = sorted(set(''.join(docs)))
BOS = len(uchars)
vocab_size = len(uchars) + 1

字元級 tokenizer:

  • uchars:資料集中所有不重複字元的排序列表
  • BOS:特殊的序列開始(Beginning of Sequence)token,ID = 字元總數
  • vocab_size:字元數 + 1(BOS token)

Value 類別 — 自製 Autograd(第 30–72 行)

python
class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads')

__slots__ 優化

限制實例只能有這四個屬性,減少每個物件的記憶體開銷。

__init__(第 33–37 行)

python
def __init__(self, data, children=(), local_grads=()):
    self.data = data       # 前向傳播的數值結果
    self.grad = 0          # 反向傳播的梯度(對 loss 的導數)
    self._children = children    # 計算圖中的子節點
    self._local_grads = local_grads  # 局部導數

運算子多載

加法(第 39–41 行):

python
def __add__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    return Value(self.data + other.data, (self, other), (1, 1))

局部導數為 (1, 1),因為 d(a+b)/da = 1, d(a+b)/db = 1

乘法(第 43–45 行):

python
def __mul__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    return Value(self.data * other.data, (self, other), (other.data, self.data))

局部導數為 (other.data, self.data),因為 d(a*b)/da = b, d(a*b)/db = a

Power(第 47 行):

python
def __pow__(self, other):
    return Value(self.data**other, (self,), (other * self.data**(other-1),))

支援 x**n,局部導數為 n * x^(n-1)

LogExp(第 48–49 行):

python
def log(self):  return Value(math.log(self.data), (self,), (1/self.data,))
def exp(self):  return Value(math.exp(self.data), (self,), (math.exp(self.data),))

ReLU(第 50 行):

python
def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))

局部導數:data > 0 時為 1,否則為 0。

反向傳播(第 59–72 行)

python
def backward(self):
    topo = []
    visited = set()
    def build_topo(v):
        if v not in visited:
            visited.add(v)
            for child in v._children:
                build_topo(child)
            topo.append(v)
    build_topo(self)
    self.grad = 1
    for v in reversed(topo):
        for child, local_grad in zip(v._children, v._local_grads):
            child.grad += local_grad * v.grad

流程:

  1. 拓撲排序:透過 DFS 找出計算圖中從 loss 到所有葉節點的拓撲順序
  2. 初始化self.grad = 1dloss/dloss = 1
  3. 反向傳播:反向遍歷拓撲順序,對每個節點將其梯度乘上局部梯度後累加到子節點

超參數與參數初始化(第 74–90 行)

python
n_layer = 1       # Transformer 深度
n_embd = 16       # 嵌入維度
block_size = 16   # 上下文視窗大小
n_head = 4        # 注意力頭數
head_dim = n_embd // n_head  # 每個頭的維度 = 4

參數矩陣

  • wte:token embedding (vocab_size × n_embd)
  • wpe:position embedding (block_size × n_embd)
  • lm_head:語言模型投影頭 (vocab_size × n_embd)
  • 每層的:attn_wq, attn_wk, attn_wv, attn_wo, mlp_fc1, mlp_fc2

matrix 輔助函式使用 Value(random.gauss(0, 0.08)) 初始化每個參數。

模型架構(第 92–144 行)

Linear(第 94–95 行)

python
def linear(x, w):
    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

純 Python 實作的矩陣乘法:對權重矩陣的每一列(wo),計算 wo · x

Softmax(第 97–101 行)

python
def softmax(logits):
    max_val = max(val.data for val in logits)
    exps = [(val - max_val).exp() for val in logits]
    total = sum(exps)
    return [e / total for e in exps]

先減去最大值再取 exp,這是數值穩定的標準作法。

RMSNorm(第 103–106 行)

python
def rmsnorm(x):
    ms = sum(xi * xi for xi in x) / len(x)
    scale = (ms + 1e-5) ** -0.5
    return [xi * scale for xi in x]

RMSNorm 比 LayerNorm 更簡單:不需減去均值,只需計算 RMS 後縮放。1e-5 是防止除以零。

gpt() 主函式(第 108–144 行)

Token + Position Embedding(第 109–111 行):

python
tok_emb = state_dict['wte'][token_id]
pos_emb = state_dict['wpe'][pos_id]
x = [t + p for t, p in zip(tok_emb, pos_emb)]

Pre-attention RMSNorm(第 112 行): 先在殘差連接前做一次 RMSNorm,與標準 Pre-LN Transformer 一致。

注意力區塊(第 114–134 行):

python
x_residual = x
x = rmsnorm(x)
q = linear(x, state_dict[f'layer{li}.attn_wq'])
k = linear(x, state_dict[f'layer{li}.attn_wk'])
v = linear(x, state_dict[f'layer{li}.attn_wv'])
keys[li].append(k)    # 儲存 K,供後續位置使用
values[li].append(v)  # 儲存 V

每個頭的計算(第 124–132 行):

python
for h in range(n_head):
    hs = h * head_dim
    q_h = q[hs:hs+head_dim]
    k_h = [ki[hs:hs+head_dim] for ki in keys[li]]
    v_h = [vi[hs:hs+head_dim] for vi in values[li]]
    attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]
    attn_weights = softmax(attn_logits)
    head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]
    x_attn.extend(head_out)

這裡特別注意:由於是自迴歸生成,keysvalues 會在每個時間步累積——這實作了因果注意力(causal attention),因為當前的 Q 只能看到過去的 K 和 V。

MLP 區塊(第 135–141 行):

python
x_residual = x
x = rmsnorm(x)
x = linear(x, state_dict[f'layer{li}.mlp_fc1'])
x = [xi.relu() for xi in x]
x = linear(x, state_dict[f'layer{li}.mlp_fc2'])
x = [a + b for a, b in zip(x, x_residual)]

標準的兩層 FFN:n_embd → 4*n_embd → n_embd,中間用 ReLU 激活。

Adam 優化器(第 147–149 行)

python
learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8
m = [0.0] * len(params)  # 一階動量
v = [0.0] * len(params)  # 二階動量

訓練迴圈(第 152–184 行)

資料準備(第 156–157 行)

python
doc = docs[step % len(docs)]
tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]

每個 step 取一個名字,前後加上 BOS token。

前向傳播(第 161–169 行)

python
keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
losses = []
for pos_id in range(n):
    token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
    logits = gpt(token_id, pos_id, keys, values)
    probs = softmax(logits)
    loss_t = -probs[target_id].log()
    losses.append(loss_t)
loss = (1 / n) * sum(losses)

對序列中的每個位置:輸入當前 token,預測下一個 token,計算交叉熵損失,最後取平均值。

反向傳播(第 172 行)

python
loss.backward()

觸發完整的反向傳播,計算所有參數的梯度。

參數更新(第 175–182 行)

python
lr_t = learning_rate * (1 - step / num_steps)
for i, p in enumerate(params):
    m[i] = beta1 * m[i] + (1 - beta1) * p.grad
    v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
    m_hat = m[i] / (1 - beta1 ** (step + 1))
    v_hat = v[i] / (1 - beta2 ** (step + 1))
    p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
    p.grad = 0

標準 Adam 更新,加上線性學習率衰減。

推論(第 186–200 行)

python
temperature = 0.5
for sample_idx in range(20):
    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
    token_id = BOS
    sample = []
    for pos_id in range(block_size):
        logits = gpt(token_id, pos_id, keys, values)
        probs = softmax([l / temperature for l in logits])
        token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
        if token_id == BOS:
            break
        sample.append(uchars[token_id])

從 BOS token 開始,逐字元生成。遇到 BOS 即停止。溫度控制用於調整生成的多樣性。