Appearance
microgpt.py — 中文註解版
python
"""
以最原子化的方式,在純 Python(無依賴)環境中訓練與執行 GPT 推論。
這份檔案包含了完整的演算法。
其他一切(PyTorch、CUDA、flash attention)只是效率問題。
@karpathy
"""
import os # os.path.exists — 檢查檔案是否存在
import math # math.log, math.exp — 數學運算
import random # random.seed, random.choices, random.gauss, random.shuffle — 隨機操作
random.seed(42) # 設定亂數種子,確保可重現性
# === 第一步:資料集 ===
# docs: list[str] — 文件列表(例如名字列表)
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)
print(f"num docs: {len(docs)}")
# === 第二步:Tokenizer(標記器)===
# 將字串轉換為整數序列(token),以及反向轉換
uchars = sorted(set(''.join(docs))) # 所有不重複字元,排序後作為 token ID 0..n-1
BOS = len(uchars) # 特殊 token:序列開始(Beginning of Sequence)
vocab_size = len(uchars) + 1 # 總 token 數 = 字元數 + 1(BOS)
print(f"vocab size: {vocab_size}")
# === 第三步:Autograd(自動微分)===
# 透過 chain rule 在計算圖上遞迴計算梯度
class Value:
__slots__ = ('data', 'grad', '_children', '_local_grads') # 減少記憶體用量
def __init__(self, data, children=(), local_grads=()):
self.data = data # 前向傳播的數值結果
self.grad = 0 # 對 loss 的梯度(反向傳播時計算)
self._children = children # 計算圖的子節點
self._local_grads = local_grads # 對子節點的局部導數
def __add__(self, other):
other = other if isinstance(other, Value) else Value(other)
return Value(self.data + other.data, (self, other), (1, 1))
# d(a+b)/da = 1, d(a+b)/db = 1
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))
# d(a*b)/da = b, d(a*b)/db = a
def __pow__(self, other):
return Value(self.data**other, (self,), (other * self.data**(other-1),))
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),))
def relu(self):
return Value(max(0, self.data), (self,), (float(self.data > 0),))
def __neg__(self): return self * -1
def __radd__(self, other): return self + other
def __sub__(self, other): return self + (-other)
def __rsub__(self, other): return other + (-self)
def __rmul__(self, other): return self * other
def __truediv__(self, other): return self * other**-1
def __rtruediv__(self, other): return other * self**-1
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 # dloss/dloss = 1
for v in reversed(topo):
for child, local_grad in zip(v._children, v._local_grads):
child.grad += local_grad * v.grad
# === 第四步:參數初始化 ===
n_layer = 1 # Transformer 深度
n_embd = 16 # 嵌入維度
block_size = 16 # 最大上下文長度(最長名字為 15 字元)
n_head = 4 # 注意力頭數
head_dim = n_embd // n_head # 每個頭的維度
matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]
state_dict = {
'wte': matrix(vocab_size, n_embd), # token embedding
'wpe': matrix(block_size, n_embd), # position embedding
'lm_head': matrix(vocab_size, n_embd) # 語言模型投影頭
}
for i in range(n_layer):
state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd) # Q 投影
state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd) # K 投影
state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd) # V 投影
state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd) # 輸出投影
state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd) # FFN 第一層
state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd) # FFN 第二層
# 將所有參數攤平成單一列表
params = [p for mat in state_dict.values() for row in mat for p in row]
print(f"num params: {len(params)}")
# === 第五步:模型架構 ===
# 將 token 和參數映射為下一個 token 的 logits
# 遵循 GPT-2 設計,但有兩個差異:RMSNorm 取代 LayerNorm,無 bias,ReLU 取代 GeLU
def linear(x, w):
"""線性層:y = x @ W.T"""
return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]
def softmax(logits):
"""穩定的 softmax:先減去最大值再取 exp"""
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]
def rmsnorm(x):
"""RMS 層歸一化(比 LayerNorm 更簡單,無需減去均值)"""
ms = sum(xi * xi for xi in x) / len(x)
scale = (ms + 1e-5) ** -0.5
return [xi * scale for xi in x]
def gpt(token_id, pos_id, keys, values):
"""GPT 前向傳播:輸入 token → 輸出 logits"""
# Token + Position Embedding
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)]
x = rmsnorm(x)
for li in range(n_layer):
# 1) 多頭注意力區塊
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
x_attn = []
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]]
# Attention(Q,K,V) = softmax(QK^T / sqrt(d)) V
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)
x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])
x = [a + b for a, b in zip(x, x_residual)] # 殘差連接
# 2) MLP 區塊
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)] # 殘差連接
logits = linear(x, state_dict['lm_head'])
return logits
# === 第六步:Adam 優化器 ===
learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8
m = [0.0] * len(params) # 一階動量(梯度移動平均)
v = [0.0] * len(params) # 二階動量(梯度平方移動平均)
# === 第七步:訓練迴圈 ===
num_steps = 1000
for step in range(num_steps):
# 取一個文件,tokenize,前後加上 BOS
doc = docs[step % len(docs)]
tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]
n = min(block_size, len(tokens) - 1)
# 前向傳播:建立計算圖直至 loss
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) # 序列平均損失
# 反向傳播:計算所有參數的梯度
loss.backward()
# Adam 參數更新
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 # 重置梯度
print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}", end='\r')
# === 第八步:推論 ===
temperature = 0.5 # 溫度參數,控制生成的多樣性
print("\n--- inference (new, hallucinated names) ---")
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 # 遇到 BOS 停止生成
sample.append(uchars[token_id])
print(f"sample {sample_idx+1:2d}: {''.join(sample)}")