nanogpt

nanoGPT 是一个极简的 GPT 训练实现，由 Andrej Karpathy 开发，主要用于学习和实验。整个模型定义在 model.py 中约 300 行，训练循环在 train.py 中约 300 行，代码干净、可修改，没有额外抽象，直接使用纯 PyTorch。

用途

nanoGPT 适合以下场景：学习 GPT 工作原理、实验 Transformer 变体、教学与教育目的、快速原型开发、计算资源有限（可在 CPU 上运行）。如果需要生产环境使用、多种模型支持或大规模分布式训练，建议考虑 HuggingFace Transformers、Megatron-LM、LitGPT 或 PyTorch Lightning 等替代方案。

主要能力

• 字符级莎士比亚训练：在 CPU 上约 5 分钟完成训练，GPU（T4）约 1 分钟，显存占用小于 1GB。
• 复现 GPT-2（124M）：在 OpenWebText 数据集上进行多 GPU 训练，使用 8×A100 约 4 天，单卡 A100 约 1 周，每 GPU 显存约 16GB。
• 微调预训练 GPT-2：可从 OpenAI 检查点（gpt2、gpt2-medium、gpt2-large、gpt2-xl）开始微调。
• 自定义数据集训练：支持用户准备自己的文本数据，生成 train.bin 和 val.bin 后进行训练。
• 性能优化：支持 PyTorch 2.0 的 compile=True（约 2 倍加速）和 dtype=bfloat16（约 50% 显存减少）。

使用方式

安装依赖： pip install torch numpy transformers datasets tiktoken wandb tqdm

工作流 1：字符级莎士比亚 准备数据：python data/shakespeare_char/prepare.py 训练：python train.py config/train_shakespeare_char.py 生成文本：python sample.py --out_dir=out-shakespeare-char 配置示例：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 使用 DDP 在 8 GPU 上训练：torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py 采样：python sample.py --out_dir=out 配置示例：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 在 train.py 或配置中设置 init_from='gpt2'（可选 gpt2、gpt2-medium、gpt2-large、gpt2-xl），模型会自动加载 OpenAI 权重。 运行：python train.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，读取文本、创建字符映射、分词、划分训练/验证集并保存为 train.bin 和 val.bin。 运行：python data/custom/prepare.py 然后 python train.py --dataset=custom

适用场景

• 学习 GPT 内部结构，如 GPT 块、多头注意力、MLP 层。
• 实验 Transformer 变体。
• 教学演示。
• 快速原型验证。
• 在有限算力（CPU）上运行小规模训练。

限制与常见问题

• CUDA 显存不足：减小 batch_size 或 block_size，增加 gradient_accumulation_steps 以维持有效批次。
• 训练太慢：启用 compile=True（PyTorch 2.0+）或使用 dtype='bfloat16' 或 'float16'。
• 生成质量差：增加 max_iters，降低 temperature（如 0.7），添加 top_k 采样（如 200）。
• 无法加载 GPT-2 权重：安装 transformers，并确认 init_from 使用有效模型名。

硬件要求

• 莎士比亚字符级：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。

资源

GitHub：https://github.com/karpathy/nanoGPT （48,000+ 星） 视频：Andrej Karpathy 的 "Let's build GPT" 论文："Attention is All You Need"（Vaswani 等） OpenWebText：https://huggingface.co/datasets/Skylion007/openwebtext

更多高级主题可参考 references/architecture.md（GPT 块结构、多头注意力、MLP 层）、references/training.md（学习率调度、梯度累积、分布式数据并行）、references/data.md（字符级与 BPE 分词策略、二进制格式细节）。