一、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 Nonedef forward(self, input):return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)作用层归一化让每层的输出分布更稳定关键点- weight 和 bias 是可学习参数- PyTorch 的 F.layer_norm 默认不支持 biasFalse这里做了封装- 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, Vself.c_attn nn.Linear(config.n_embd, 3 * config.n_embd, biasconfig.bias)# 输出投影self.c_proj nn.Linear(config.n_embd, config.n_embd, biasconfig.bias)# 正则化self.attn_dropout nn.Dropout(config.dropout)self.resid_dropout nn.Dropout(config.dropout)self.n_head config.n_headself.n_embd config.n_embdself.dropout config.dropout# Flash Attention 检查PyTorch 2.0self.flash hasattr(torch.nn.functional, scaled_dot_product_attention)if not self.flash:# 因果掩码确保只能看到左边的tokenself.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, Vq, k, v self.c_attn(x).split(self.n_embd, dim2)# 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.0y torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_maskNone,dropout_pself.dropout if self.training else 0,is_causalTrue)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 → 4Cself.c_fc nn.Linear(config.n_embd, 4 * config.n_embd, biasconfig.bias)self.gelu nn.GELU() # 激活函数# 输出投影4C → Cself.c_proj nn.Linear(4 * config.n_embd, config.n_embd, biasconfig.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, biasconfig.bias)self.attn CausalSelfAttention(config)self.ln_2 LayerNorm(config.n_embd, biasconfig.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行dataclassclass 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 # 是否使用 bias6. 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)]), # 多层 Blockln_f LayerNorm(config.n_embd, biasconfig.bias), # 最终归一化))# 语言模型头self.lm_head nn.Linear(config.n_embd, config.vocab_size, biasFalse)# 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, mean0.0, std0.02/math.sqrt(2 * config.n_layer))**前向传播第170-193行**pythondef forward(self, idx, targetsNone):device idx.deviceb, t idx.size()# 位置编码pos torch.arange(0, t, dtypetorch.long, devicedevice)# 1. Token Embedding Position Embeddingtok_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. 输出 logitsif 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 Nonereturn 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, temperature1.0, top_kNone):for _ in range(max_new_tokens):# 1. 如果序列太长只保留最后 block_size 个 tokenidx_cond idx if idx.size(1) self.config.block_size else idx[:, -self.config.block_size:]# 2. 前向传播logits, _ self(idx_cond)# 3. 取最后一个位置的 logits除以 temperaturelogits 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_samples1)# 7. 追加到序列idx torch.cat((idx, idx_next), dim1)return idx采样策略| 参数 | 作用 | 推荐值 ||--------------------|-------------------------------------------------|-----------|| temperature | 控制随机性越低越保守 | 0.8-1.0 || top_k | 只从概率最高的 k 个 token 中采样 | 50-200 |二、train.py 完整解析train.py 是训练循环约 330 行。1. 配置参数第32-78行# 数据dataset openwebtextbatch_size 12 # 每个 GPU 的 batch sizeblock_size 1024 # 序列长度gradient_accumulation_steps 5 * 8 # 梯度累积步数# 模型n_layer 12n_head 12n_embd 768dropout 0.0# 优化器learning_rate 6e-4max_iters 600000weight_decay 1e-1beta1 0.9beta2 0.95grad_clip 1.0# 学习率调度warmup_iters 2000lr_decay_iters 600000min_lr 6e-52. 数据加载第114-131行def get_batch(split):# 使用 memmap 避免内存泄漏if split train:data np.memmap(os.path.join(data_dir, train.bin), dtypenp.uint16, moder)else:data np.memmap(os.path.join(data_dir, val.bin), dtypenp.uint16, moder)# 随机选择 batch_size 个起始位置ix torch.randint(len(data) - block_size, (batch_size,))# 提取输入 x 和目标 yx torch.stack([torch.from_numpy((data[i:iblock_size]).astype(np.int64)) for i in ix])y torch.stack([torch.from_numpy((data[i1:i1block_size]).astype(np.int64)) for i in ix])# 移到 GPUif device_type cuda:x, y x.pin_memory().to(device, non_blockingTrue), y.pin_memory().to(device, non_blockingTrue)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_locationdevice)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. 线性 warmupif 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│└─────────────────────────── iter0 2000 6000006. 训练循环第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(fstep {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_noneTrue)# 7. 计时和日志t1 time.time()dt t1 - t0# ... 打印 loss, time, mfu ...iter_num 1if 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.py2. 读代码逐行理解 model.py 的每个类3. 改参数尝试修改 n_layer, n_head, n_embd 观察变化4. 加功能尝试加入新的层或修改 attention 机制