nanoGPT 代码解析
项目概况
nanoGPT 是 Andrej Karpathy 写的最简洁 GPT 实现:
> "The simplest, fastest repository for training/finetuning medium-sized GPTs."
---
model.py 模块结构
LayerNorm # 带可选偏置的层归一化
CausalSelfAttention # 因果多头自注意力
MLP # 两层前馈网络
Block # 一个完整的 Transformer 块
GPTConfig # 配置数据类
GPT # 主模型类
---
GPT 类的方法节选(省略上下文,不是可独立执行的完整代码)
`__init__` — 模型初始化
class GPT(nn.Module):
def __init__(self, config):
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)]), # N个Block
ln_f = LayerNorm(config.n_embd, bias=config.bias), # 最终LayerNorm
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
# 权重共享
self.transformer.wte.weight = self.lm_head.weight
`forward` — 前向传播
def forward(self, idx, targets=None):
device = idx.device
b, t = idx.size()
# 位置序列
pos = torch.arange(0, t, dtype=torch.long, device=device)
# 嵌入层
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)
# 逐层处理
for block in self.transformer.h:
x = block(x)
# 最终归一化
x = self.transformer.ln_f(x)
# 计算 logits 和损失
if targets is not None:
logits = self.lm_head(x)
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
else:
logits = self.lm_head(x[:, [-1], :]) # 推理时只算最后一个位置
loss = None
return logits, loss
`configure_optimizers` — 优化器配置
def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
# 区分有权重衰减和无权重衰减的参数
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
# 二维嵌入矩阵参与权重衰减;偏置和一维LayerNorm参数不衰减
# 使用 AdamW 优化器
optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=use_fused)
`from_pretrained` — 加载 GPT-2 预训练权重
@classmethod
def from_pretrained(cls, model_type, override_args=None):
# 支持: 'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'
# 从 HuggingFace 下载 OpenAI 权重
from transformers import GPT2LMHeadModel
model_hf = GPT2LMHeadModel.from_pretrained(model_type)
# 手动映射权重名称(OpenAI 命名 vs nanoGPT 命名)
`estimate_mfu` — 模型利用率估算
def estimate_mfu(self, fwdbwd_per_iter, dt):
"""估算 Model FLOPs Utilization (MFU),以 A100 bfloat16 峰值 FLOPS 为基准"""
# 用于监控 GPU 利用效率
---
关键设计决策
| 设计 | 选择 | 原因 |
|------|------|------|
| 归一化位置 | Pre-Norm(先归一化再处理) | 训练更稳定 |
| 激活函数 | GELU(非 ReLU) | GPT-2 原始选择,性能更好 |
| 注意力实现 | SDPA接口或手工回退 | 实际内核与性能依运行条件 |
| 偏置项 | 可配置(默认 True) | 可选择去掉 bias 减少参数 |
| 权重初始化 | 按 GPT-2 论文规范 | 0.02 正态分布,残差层缩放 |
---
参数量计算
def get_num_params(self, non_embedding=True):
n_params = sum(p.numel() for p in self.parameters())
if non_embedding:
n_params -= self.transformer.wpe.weight.numel()
return n_params
# 默认统计扣除位置嵌入;完整参数量还取决于词表、bias等配置
---
相关页面
来源
本轮核查来源与限制(2026-09-17)
只修本轮明确列出的陈述,不将旧页整篇升级为已验证。themis_only为本地保存和禁止外发的兼容标记,不代表本地模型推理。