nanoGPT 逐行讲解

一、model.py 完整解析

model.py 是整个项目的核心，只有 330 行代码，却实现了完整的 GPT 模型。

1. LayerNorm（第18-27行）

class LayerNorm(nn.Module):
def __init__(self, ndim, bias):
super().__init__()
self.weight = nn.Parameter(torch.ones(ndim))
self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None

def forward(self, input):
return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)

作用：层归一化，让每层的输出分布更稳定

关键点：
- `weight` 和 `bias` 是可学习参数
- PyTorch 的 `F.layer_norm` 默认不支持 bias=False，这里做了封装
- `1e-5` 是 epsilon，防止除零

2. CausalSelfAttention（第29-76行）—— 最重要！

class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0

# QKV 投影：一个线性层同时计算 Q, K, V
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
# 输出投影
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
# 正则化
self.attn_dropout = nn.Dropout(config.dropout)
self.resid_dropout = nn.Dropout(config.dropout)

self.n_head = config.n_head
self.n_embd = config.n_embd
self.dropout = config.dropout

# Flash Attention 检查（PyTorch >= 2.0）
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
if not self.flash:
# 因果掩码：确保只能看到左边的token
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size))

def forward(self, x):
B, T, C = x.size() # Batch, Time, Channel

# 1. 投影得到 Q, K, V
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)

# 2. 多头重塑：[B, T, C] → [B, heads, T, head_dim]
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)

# 3. 计算注意力
if self.flash:
# 高效实现（PyTorch 2.0+）
y = torch.nn.functional.scaled_dot_product_attention(
q, k, v, attn_mask=None,
dropout_p=self.dropout if self.training else 0,
is_causal=True
)
else:
# 手动实现
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
y = att @ v

# 4. 合并多头输出
y = y.transpose(1, 2).contiguous().view(B, T, C)

# 5. 输出投影
y = self.resid_dropout(self.c_proj(y))
return y

数据流详解：
输入: [B, T, C] (Batch, Sequence, Embedding Dim)
│
▼
┌─────────────────────────────────────┐
│ c_attn: Linear(C → 3*C) │
│ 一次投影同时得到 Q, K, V │
└─────────────────────────────────────┘
│
▼
┌─────────────────────────────────────┐
│ 重塑为多头 │
│ [B, T, 3*C] → [B, heads, T, hs] │
│ 其中 hs = C / heads │
└─────────────────────────────────────┘
│
▼
┌─────────────────────────────────────┐
│ Attention(Q, K, V) │
│ = softmax(QK^T / √hs) @ V │
└─────────────────────────────────────┘
│
▼
┌─────────────────────────────────────┐
│ 合并多头 + 输出投影 │
│ [B, heads, T, hs] → [B, T, C] │
└─────────────────────────────────────┘

为什么叫"Causal"（因果）？

因为掩码矩阵 `self.bias` 是下三角矩阵：

[[1, 0, 0, 0],
[1, 1, 0, 0],
[1, 1, 1, 0],
[1, 1, 1, 1]]

这确保位置 i 只能看到位置 0~i，看不到未来。

3. MLP / Feed-Forward Network（第78-92行）

class MLP(nn.Module):
def __init__(self, config):
super().__init__()
# 输入投影：C → 4C
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
self.gelu = nn.GELU() # 激活函数
# 输出投影：4C → C
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
self.dropout = nn.Dropout(config.dropout)

def forward(self, x):
x = self.c_fc(x) # [B, T, C] → [B, T, 4C]
x = self.gelu(x) # GELU 激活
x = self.c_proj(x) # [B, T, 4C] → [B, T, C]
x = self.dropout(x)
return x

作用：逐位置处理，每个 token 独立变换

为什么是 4 倍？
- GPT-2 论文中的标准设置
- 扩大容量后再压缩，学习更复杂的特征

4. Block（第94-106行）—— Transformer 层

class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
self.attn = CausalSelfAttention(config)
self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
self.mlp = MLP(config)

def forward(self, x):
# 预归一化 + 残差连接
x = x + self.attn(self.ln_1(x)) # 注意力子层
x = x + self.mlp(self.ln_2(x)) # FFN 子层
return x

结构图示：
x
│
├─────────────────────────┐
│ │
▼ ▼
┌───────┐ ┌───────┐
│ Layer │ │ Layer │
│ Norm │ │ Norm │
└───────┘ └───────┘
│ │
▼ ▼
┌─────────┐ ┌───────┐
│ Attention │ │ MLP │
│ (Multi- │ │(FFN) │
│ Head) │ └───────┘
└─────────┘ │
│ │
└──────────┬───────────┘
│
▼
x (输出)

残差连接的作用：
- 梯度可以直通，缓解消失梯度问题
- 保留原始信息，神经网络只学"增量"

5. GPTConfig（第108-116行）

@dataclass
class GPTConfig:
block_size: int = 1024 # 最大序列长度
vocab_size: int = 50304 # GPT-2 词表大小 (50257 补齐到 64 的倍数)
n_layer: int = 12 # Transformer 层数
n_head: int = 12 # 注意力头数
n_embd: int = 768 # 嵌入维度
dropout: float = 0.0 # dropout 率
bias: bool = True # 是否使用 bias

6. GPT 模型（第118-330行）

class GPT(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config

# Transformer 主体
self.transformer = nn.ModuleDict(dict(
wte = nn.Embedding(config.vocab_size, config.n_embd), # 词嵌入
wpe = nn.Embedding(config.block_size, config.n_embd), # 位置嵌入
drop = nn.Dropout(config.dropout),
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), # 多层 Block
ln_f = LayerNorm(config.n_embd, bias=config.bias), # 最终归一化
))

# 语言模型头
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

# Weight Tying：输出投影复用词嵌入权重
self.transformer.wte.weight = self.lm_head.weight

# 权重初始化
self.apply(self._init_weights)
for pn, p in self.named_parameters():
if pn.endswith('c_proj.weight'):
torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
```

**前向传播（第170-193行）：**

```python
def forward(self, idx, targets=None):
device = idx.device
b, t = idx.size()

# 位置编码
pos = torch.arange(0, t, dtype=torch.long, device=device)

# 1. Token Embedding + Position Embedding
tok_emb = self.transformer.wte(idx) # [B, T, C]
pos_emb = self.transformer.wpe(pos) # [T, C]
x = self.transformer.drop(tok_emb + pos_emb)

# 2. 通过所有 Transformer 层
for block in self.transformer.h:
x = block(x)

# 3. 最终归一化
x = self.transformer.ln_f(x)

# 4. 输出 logits
if targets is not None:
# 训练模式：计算所有位置的损失
logits = self.lm_head(x)
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
else:
# 推理模式：只取最后一个位置的 logits（效率优化）
logits = self.lm_head(x[:, [-1], :])
loss = None

return logits, loss

关键优化：推理时只取最后一个位置

logits = self.lm_head(x[:, [-1], :]) # 只取最后一步

因为 GPT 是自回归模型，我们只需要预测下一个 token，前面的 position 的 logits 不用计算。

7. 生成函数（第305-330行）

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
for _ in range(max_new_tokens):
# 1. 如果序列太长，只保留最后 block_size 个 token
idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]

# 2. 前向传播
logits, _ = self(idx_cond)

# 3. 取最后一个位置的 logits，除以 temperature
logits = logits[:, -1, :] / temperature

# 4. Top-k 裁剪
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')

# 5. Softmax 转概率
probs = F.softmax(logits, dim=-1)

# 6. 采样
idx_next = torch.multinomial(probs, num_samples=1)

# 7. 追加到序列
idx = torch.cat((idx, idx_next), dim=1)

return idx

采样策略：

| 参数 | 作用 | 推荐值 |
|--------------------|-------------------------------------------------|-----------|
| `temperature` | 控制随机性，越低越保守 | 0.8-1.0 |
| `top_k` | 只从概率最高的 k 个 token 中采样 | 50-200 |

二、train.py 完整解析

train.py 是训练循环，约 330 行。

1. 配置参数（第32-78行）

# 数据
dataset = 'openwebtext'
batch_size = 12 # 每个 GPU 的 batch size
block_size = 1024 # 序列长度
gradient_accumulation_steps = 5 * 8 # 梯度累积步数

# 模型
n_layer = 12
n_head = 12
n_embd = 768
dropout = 0.0

# 优化器
learning_rate = 6e-4
max_iters = 600000
weight_decay = 1e-1
beta1 = 0.9
beta2 = 0.95
grad_clip = 1.0

# 学习率调度
warmup_iters = 2000
lr_decay_iters = 600000
min_lr = 6e-5

2. 数据加载（第114-131行）

def get_batch(split):
# 使用 memmap 避免内存泄漏
if split == 'train':
data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
else:
data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')

# 随机选择 batch_size 个起始位置
ix = torch.randint(len(data) - block_size, (batch_size,))

# 提取输入 x 和目标 y
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])

# 移到 GPU
if device_type == 'cuda':
x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
else:
x, y = x.to(device), y.to(device)

return x, y

关键设计：
- `memmap`：内存映射文件，不占用实际内存
- `x` 和 `y` 的关系：y 是 x 右移一位（预测下一个 token）
- `pin_memory()`：加速 CPU→GPU 数据传输

3. 模型初始化（第146-193行）

# 三种初始化方式
if init_from == 'scratch':
# 从零开始
gptconf = GPTConfig(**model_args)
model = GPT(gptconf)
elif init_from == 'resume':
# 从检查点恢复
checkpoint = torch.load(ckpt_path, map_location=device)
model.load_state_dict(checkpoint['model'])
elif init_from.startswith('gpt2'):
# 从 OpenAI 预训练权重加载
model = GPT.from_pretrained(init_from, override_args)

4. 损失估算（第215-228行）

@torch.no_grad()
def estimate_loss():
out = {}
model.eval()
for split in ['train', 'val']:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X, Y = get_batch(split)
with ctx:
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out

为什么用 `@torch.no_grad()`？
- 评估时不需要计算梯度，节省内存
- 速度更快

5. 学习率调度（第230-242行）

def get_lr(it):
# 1. 线性 warmup
if it < warmup_iters:
return learning_rate * (it + 1) / (warmup_iters + 1)

# 2. 超过 decay_iters 后使用最小学习率
if it > lr_decay_iters:
return min_lr

# 3. Cosine 衰减
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (learning_rate - min_lr)

学习率曲线：

lr
│ ╭─── warmup
│ ╱
│ ╱
│ ╱
╭─╯
│ ╲_________ cosine decay
│
└─────────────────────────── iter
0 2000 600000

6. 训练循环（第249-333行）

# 获取第一批数据
X, Y = get_batch('train')
t0 = time.time()

while True:
# 1. 设置学习率
lr = get_lr(iter_num)
for param_group in optimizer.param_groups:
param_group['lr'] = lr

# 2. 定期评估和保存
if iter_num % eval_interval == 0 and master_process:
losses = estimate_loss()
print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
# 保存最佳模型
if losses['val'] < best_val_loss:
best_val_loss = losses['val']
torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))

# 3. 前向 + 反向 + 更新
for micro_step in range(gradient_accumulation_steps):
if ddp:
model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)

with ctx:
logits, loss = model(X, Y)
loss = loss / gradient_accumulation_steps # 梯度累积

# 异步预取下一批数据
X, Y = get_batch('train')

# 反向传播
scaler.scale(loss).backward()

# 4. 梯度裁剪
if grad_clip != 0.0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)

# 5. 优化器步进
scaler.step(optimizer)
scaler.update()

# 6. 清空梯度
optimizer.zero_grad(set_to_none=True)

# 7. 计时和日志
t1 = time.time()
dt = t1 - t0
# ... 打印 loss, time, mfu ...

iter_num += 1
if iter_num > max_iters:
break

三、整体数据流图

训练流程
═══════════════════════════════════════════════════

train.bin / val.bin (uint16 token IDs)
│
▼
┌─────────────────┐
│ get_batch() │ ← 随机采样 block_size 个连续 token
└────────┬────────┘
│
▼
┌─────────────────┐
│ GPT Model │
│ ┌───────────┐ │
│ │ wte + wpe │ │ ← Token + Position Embedding
│ ├───────────┤ │
│ │ Block × N │ │ ← N 层 Transformer
│ │ - Attn │ │
│ │ - FFN │ │
│ └───────────┘ │
│ │ │
│ ▼ │
│ lm_head │ ← 投影到词表大小
└────────┬────────┘
│
▼
┌─────────────────┐
│ CrossEntropy │ ← 计算 loss
│ Loss │
└────────┬────────┘
│
▼
┌─────────────────┐
│ Backward + │ ← 反向传播
│ Optimizer │
└────────┬────────┘
│
▼
更新权重

四、关键数字总结

| 参数 | GPT-2 (124M) | 说明 |
|-------------------|-------------|-------------------------|
| `vocab_size` | 50304 | 词表大小 |
| `n_embd` | 768 | 嵌入维度 |
| `n_layer` | 12 | Transformer 层数 |
| `n_head` | 12 | 注意力头数 |
| `block_size` | 1024 | 最大序列长度 |
| 参数量 | ~124M | 1.24 亿参数 |

计算量估算：

每 token FLOPs ≈ 6 × 参数量 = 6 × 124M ≈ 744M FLOPs

五、学习建议

1. 先跑通：`python train.py config/train_shakespeare_char.py`
2. 读代码：逐行理解 model.py 的每个类
3. 改参数：尝试修改 n_layer, n_head, n_embd 观察变化
4. 加功能：尝试加入新的层或修改 attention 机制