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 机制