nanoGPT 是一个极简的 GPT 训练实现，由 Andrej Karpathy 开发，旨在用于学习和实验。整个模型定义在 model.py 中约 300 行，训练循环在 train.py 中约 300 行，代码干净、可修改，没有额外抽象，仅使用纯 PyTorch。
用途
nanoGPT 主要用于：学习 GPT 工作原理、实验 Transformer 变体、教学与教育、快速原型开发，以及在有限算力（甚至 CPU）上运行。
主要能力
- 字符级莎士比亚文本训练与生成
- 在 OpenWebText 上复现 GPT-2（124M）
- 微调预训练的 GPT-2 模型（gpt2、gpt2-medium、gpt2-large、gpt2-xl）
- 支持自定义数据集训练
- 支持多 GPU 分布式数据并行（DDP）训练
- 支持 PyTorch 2.0 编译加速和混合精度
使用方式
安装
pip install torch numpy transformers datasets tiktoken wandb tqdm
工作流 1：字符级莎士比亚
# 准备数据
python data/shakespeare_char/prepare.py
# 训练（CPU 约 5 分钟）
python train.py config/train_shakespeare_char.py
# 生成文本
python sample.py --out_dir=out-shakespeare-char
配置示例（config/train_shakespeare_char.py）：
n_layer = 6
n_head = 6
n_embd = 384
block_size = 256
batch_size = 64
learning_rate = 1e-3
max_iters = 5000
eval_interval = 500
device = 'cpu'
compile = False
工作流 2：复现 GPT-2（124M）
# 准备 OpenWebText（约 1 小时）
python data/openwebtext/prepare.py
# 使用 8 个 GPU 进行 DDP 训练
torchrun --standalone --nproc_per_node=8 \
train.py config/train_gpt2.py
# 采样
python sample.py --out_dir=out
配置示例（config/train_gpt2.py）：
n_layer = 12
n_head = 12
n_embd = 768
block_size = 1024
dropout = 0.0
batch_size = 12
gradient_accumulation_steps = 5 * 8
learning_rate = 6e-4
max_iters = 600000
lr_decay_iters = 600000
compile = True
工作流 3：微调预训练 GPT-2
init_from = 'gpt2'
python train.py config/finetune_shakespeare.py
配置示例（config/finetune_shakespeare.py）：
init_from = 'gpt2'
dataset = 'shakespeare_char'
batch_size = 1
block_size = 1024
learning_rate = 3e-5
max_iters = 2000
warmup_iters = 100
weight_decay = 1e-1
工作流 4：自定义数据集
# data/custom/prepare.py
import numpy as np
with open('my_data.txt', 'r') as f:
text = f.read()
chars = sorted(list(set(text)))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
data = np.array([stoi[ch] for ch in text], dtype=np.uint16)
n = len(data)
train_data = data[:int(n*0.9)]
val_data = data[int(n*0.9):]
train_data.tofile('data/custom/train.bin')
val_data.tofile('data/custom/val.bin')
python data/custom/prepare.py
python train.py --dataset=custom
适用场景
- 学习 GPT 和 Transformer 原理
- 实验 Transformer 变体
- 教学与教育
- 快速原型开发
- 算力有限（可在 CPU 上运行）
限制与替代方案
nanoGPT 适合学习和实验，但以下场景建议使用替代方案：
- 生产环境、需要多种模型：HuggingFace Transformers
- 大规模分布式训练：Megatron-LM
- 更多架构、生产就绪：LitGPT
- 需要高层框架：PyTorch Lightning
常见问题
- CUDA 内存不足：减小 batch_size 或 block_size，增加 gradient_accumulation_steps
- 训练太慢：启用 compile=True（约 2 倍加速），使用 dtype='bfloat16' 或 'float16'
- 生成质量差：增加 max_iters，降低 temperature，添加 top_k 采样
- 无法加载 GPT-2 权重：安装 transformers，检查模型名称（gpt2、gpt2-medium、gpt2-large、gpt2-xl）
硬件要求
- 莎士比亚（字符级）：CPU 约 5 分钟，GPU（T4）约 1 分钟，显存 <1GB
- GPT-2（124M）：1× A100 约 1 周，8× A100 约 4 天，每 GPU 显存约 16GB
- GPT-2 Medium（350M）：8× A100 约 2 周，每 GPU 显存约 40GB
性能：compile=True 约 2 倍加速；dtype=bfloat16 减少 50% 内存。
资源
- GitHub：https://github.com/karpathy/nanoGPT
- 视频："Let's build GPT" by Andrej Karpathy
- 论文："Attention is All You Need"（Vaswani et al.）
- OpenWebText：https://huggingface.co/datasets/Skylion007/openwebtext