nanoGPT 代码解析

项目概况

nanoGPT 是 Andrej Karpathy 写的最简洁 GPT 实现:

  • `model.py`:约 300 行,完整的 GPT 模型定义
  • `train.py`:约 300 行,完整的训练循环
  • 可复现 GPT-2 (124M) 的性能
  • > "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等配置
    

    ---

    相关页面

  • 01-概念/Transformer架构 — 架构原理
  • 01-概念/注意力机制 — 注意力细节
  • 02-技术/训练流程与超参数 — train.py 解析
  • 02-技术/Flash-Attention优化 — 性能优化
  • 05-实体/Andrej_Karpathy — 作者信息
  • 来源

  • [nanoGPT model.py](https://github.com/karpathy/nanoGPT/blob/master/model.py)
  • [build-nanoGPT 讲解视频](https://github.com/karpathy/build-nanogpt)
  • 本轮核查来源与限制(2026-09-17)

  • https://raw.githubusercontent.com/karpathy/nanoGPT/master/model.py
  • https://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html
  • 只修本轮明确列出的陈述,不将旧页整篇升级为已验证。themis_only为本地保存和禁止外发的兼容标记,不代表本地模型推理。