nanoGPT

坚持就是胜利。——游吉祥

nanoGPT：用最简单最快的方式训练/微调中等规模 GPT

在大模型成为开发“标配”的今天，很多人都希望能亲手训练或微调一个可用的 GPT。但现实往往是：工程复杂、脚本繁多、动辄几十个开关，令人望而却步。Andrej Karpathy 的开源项目——nanoGPT，则用一句话击中痛点：The simplest, fastest repository for training/finetuning medium-sized GPTs.

• 项目地址：karpathy/nanoGPT
• 项目描述：The simplest, fastest repository for training/finetuning medium-sized GPTs.
• 语言：Python
• 许可证：MIT
• Stars：50,711+
• Forks：8,492+

它的目标不是“最全功能”，而是“最小可用且高效”——把训练中等规模 GPT 的路径压缩到极简，既能上手学习，也能作为小团队/个人项目的工程基座。

---

为什么选择 nanoGPT？

• 简洁直接：围绕训练/微调核心流程设计，尽可能减少“无关复杂度”。
• 性能敏捷：在合理硬件下快速得到效果，适合实验和迭代。
• 教学友好：作为“动手理解 GPT 训练”的入口非常适合。
• 工程基座：可以在此基础上替换数据、分词、配置，搭建你的定制化训练管线。

对于想“把训练流程跑起来”的开发者而言，nanoGPT就像一条捷径：你可以在最短学习曲线内完成从数据到模型的闭环，然后再逐步扩展。

---

核心能力概览

• 支持中等规模 GPT 的训练与微调（适合自有领域数据）
• 简洁的训练脚本与配置，便于快速修改和对比实验
• 可与主流分词器、数据处理方式结合（根据你的项目需要替换）
• 适合单机或小规模资源环境，便于复现与分享

---

训练流程基本思路

nanoGPT的设计哲学是“抓住本质，删繁就简”。把训练过程拆成四步：

1. 准备数据

• 将文本数据整理为模型可读的序列（通常是 token 序列）
• 按固定长度切块（block size）以便批处理训练

2. 选择分词器

• 采用已有工具（如 BPE、SentencePiece、tiktoken 等），或复用项目内置方案
• 让训练数据的 token 化稳定且可复用

3. 配置模型与训练

• 配置 GPT 模型的层数、头数、宽度、上下文长度等参数
• 选择优化器（如 AdamW）、学习率策略、精度（AMP）、梯度裁剪

4. 监控与评估

• 打印 loss、保存 checkpoint
• 小规模验证生成质量（温度、top-k/top-p）以及过拟合情况

---

代码示例：从数据到训练的最小骨架

以下是“概念示例”，帮助你理解 nanoGPT 的工程结构与最小可用流程。具体实现与API以仓库文档为准。

1. 数据准备与数据集封装

import torch
from torch.utils.data import Dataset, DataLoader

class TextDataset(Dataset):
    def __init__(self, token_ids, block_size=256):
        self.token_ids = token_ids
        self.block_size = block_size

    def __len__(self):
        return max(0, len(self.token_ids) - self.block_size - 1)

    def __getitem__(self, idx):
        x = torch.tensor(self.token_ids[idx:idx+self.block_size], dtype=torch.long)
        y = torch.tensor(self.token_ids[idx+1:idx+1+self.block_size], dtype=torch.long)
        return x, y

def build_loader(token_ids, block_size=256, batch_size=32, shuffle=True):
    ds = TextDataset(token_ids, block_size)
    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle)

说明：真实环境下请使用成熟分词器与更规范的数据清洗流程。这里强调“训练用的样本对 (x → y)”。

---

2. GPT 模型最小实现（Causal Mask）

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout=0.1):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.scale = self.head_dim ** -0.5

        self.qkv = nn.Linear(n_embd, 3 * n_embd)
        self.proj = nn.Linear(n_embd, n_embd)
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)

        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)).unsqueeze(0).unsqueeze(0))

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * self.scale
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v
        y = y.transpose(1, 2).reshape(B, T, C)
        y = self.resid_drop(self.proj(y))
        return y

class MLP(nn.Module):
    def __init__(self, n_embd, dropout=0.1):
        super().__init__()
        self.fc1 = nn.Linear(n_embd, 4 * n_embd)
        self.fc2 = nn.Linear(4 * n_embd, n_embd)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = F.gelu(x)
        x = self.fc2(x)
        return self.drop(x)

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd, dropout)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class MiniGPT(nn.Module):
    def __init__(self, vocab_size, block_size=256, n_layer=8, n_head=8, n_embd=512, dropout=0.1):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        x = self.drop(x)
        for blk in self.blocks:
            x = blk(x)
        x = self.ln_f(x)
        logits = self.head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

这个骨架帮助你理解 GPT 的训练核心：嵌入、位置编码、自注意力（带因果mask）、前馈网络、LayerNorm及输出头。

---

3. 训练循环（AMP + AdamW）

import torch
from torch.optim import AdamW
from torch.cuda.amp import autocast, GradScaler

def train(model, dataloader, epochs=1, lr=3e-4, weight_decay=0.01, grad_clip=1.0, device="cuda"):
    model.to(device)
    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scaler = GradScaler(enabled=(device == "cuda"))

    model.train()
    step = 0
    for epoch in range(epochs):
        for x, y in dataloader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad(set_to_none=True)
            with autocast(enabled=(device == "cuda")):
                _, loss = model(x, y)
            scaler.scale(loss).backward()
            if grad_clip is not None:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            scaler.step(optimizer)
            scaler.update()

            if step % 50 == 0:
                print(f"epoch {epoch} step {step} loss {loss.item():.4f}")
            step += 1

这段代码体现了 nanoGPT 倡导的“有效即足够”：精简但关键的训练循环，搭配 AMP 与常见优化器即可跑起来。

---

4. 推理（温度与Top-k）

import torch
import torch.nn.functional as F

@torch.no_grad()
def generate(model, idx, max_new_tokens=100, temperature=1.0, top_k=None):
    model.eval()
    device = next(model.parameters()).device
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -model.block_size:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :] / max(temperature, 1e-5)
        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = -float('Inf')
        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_id], dim=1)
    return idx

---

如何在你的场景中用好 nanoGPT？

• 领域微调：用你自己的文档、代码、对话数据微调模型，提升特定任务效果。
• 小规模预训练：做一个迷你语料的预训练实验，建立直觉与验证管线。
• 项目模板：把训练脚本做成团队模板，统一实验参数、日志与落盘规则。
• 教学用途：用最短路径让学生或同事“看懂并跑通”GPT训练。

---

实用建议与常见坑

• 数据清洗重于一切：重复、乱码、超短/超长样本都会影响训练稳定性与效果。
• 合理的上下文长度：过长会极大消耗显存与时间；先从 256/512 开始。
• 监控与断点：loss、学习率、梯度裁剪、checkpoint保存要常规化。
• 推理策略要调：温度、top-k/top-p能显著改变文本质量与多样性。
• 先跑通再扩展：遵循 nanoGPT 的“简洁哲学”，优先完成从 0 到 1。

---

总结

nanoGPT 的魅力在于“把复杂问题做得很简单”。它不是“功能最全”，但它是“最简单最快”，适合任何想要快速理解并实践 GPT 训练的人。你可以用它作为脚手架，跑起来、看清楚、再根据场景做加法——这恰恰是中等规模 GPT 训练最需要的工程思维。

项目主页：https://github.com/karpathy/nanoGPT

如果你想把“训练 GPT”变成一件可随手尝试的事，nanoGPT 是一个绝佳的起点。把它克隆下来，准备好数据，按需调整配置，立刻开始你的小而美的 GPT 训练之旅吧！