nanogpt

nanoGPT - 极简 GPT 训练

快速开始

nanoGPT 是一个简化的 GPT 实现，旨在用于学习和实验。

安装：

pip install torch numpy transformers datasets tiktoken wandb tqdm

在莎士比亚文本上训练（适合 CPU）：

# Prepare data
python data/shakespeare_char/prepare.py

# Train (5 minutes on CPU)
python train.py config/train_shakespeare_char.py

# Generate text
python sample.py --out_dir=out-shakespeare-char

输出：

ROMEO:
What say'st thou? Shall I speak, and be a man?

JULIET:
I am afeard, and yet I'll speak; for thou art
One that hath been a man, and yet I know not
What thou art.

常见工作流

工作流 1：字符级莎士比亚文本

完整训练流程：

# Step 1: Prepare data (creates train.bin, val.bin)
python data/shakespeare_char/prepare.py

# Step 2: Train small model
python train.py config/train_shakespeare_char.py

# Step 3: Generate text
python sample.py --out_dir=out-shakespeare-char

配置（config/train_shakespeare_char.py）：

# Model config
n_layer = 6          # 6 transformer layers
n_head = 6           # 6 attention heads
n_embd = 384         # 384-dim embeddings
block_size = 256     # 256 char context

# Training config
batch_size = 64
learning_rate = 1e-3
max_iters = 5000
eval_interval = 500

# Hardware
device = 'cpu'  # Or 'cuda'
compile = False # Set True for PyTorch 2.0

训练时间：约 5 分钟（CPU），约 1 分钟（GPU）

工作流 2：复现 GPT-2（124M）

在 OpenWebText 上进行多 GPU 训练：

# Step 1: Prepare OpenWebText (takes ~1 hour)
python data/openwebtext/prepare.py

# Step 2: Train GPT-2 124M with DDP (8 GPUs)
torchrun --standalone --nproc_per_node=8 \
  train.py config/train_gpt2.py

# Step 3: Sample from trained model
python sample.py --out_dir=out

配置（config/train_gpt2.py）：

# GPT-2 (124M) architecture
n_layer = 12
n_head = 12
n_embd = 768
block_size = 1024
dropout = 0.0

# Training
batch_size = 12
gradient_accumulation_steps = 5 * 8  # Total batch ~0.5M tokens
learning_rate = 6e-4
max_iters = 600000
lr_decay_iters = 600000

# System
compile = True  # PyTorch 2.0

训练时间：约 4 天（8× A100）

工作流 3：微调预训练的 GPT-2

从 OpenAI 检查点开始：

# In train.py or config
init_from = 'gpt2'  # Options: gpt2, gpt2-medium, gpt2-large, gpt2-xl

# Model loads OpenAI weights automatically
python train.py config/finetune_shakespeare.py

配置示例（config/finetune_shakespeare.py）：

# Start from GPT-2
init_from = 'gpt2'

# Dataset
dataset = 'shakespeare_char'
batch_size = 1
block_size = 1024

# Fine-tuning
learning_rate = 3e-5  # Lower LR for fine-tuning
max_iters = 2000
warmup_iters = 100

# Regularization
weight_decay = 1e-1

工作流 4：自定义数据集

使用你自己的文本进行训练：

# data/custom/prepare.py
import numpy as np

# Load your data
with open('my_data.txt', 'r') as f:
    text = f.read()

# Create character mappings
chars = sorted(list(set(text)))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

# Tokenize
data = np.array([stoi[ch] for ch in text], dtype=np.uint16)

# Split train/val
n = len(data)
train_data = data[:int(n*0.9)]
val_data = data[int(n*0.9):]

# Save
train_data.tofile('data/custom/train.bin')
val_data.tofile('data/custom/val.bin')

训练：

python data/custom/prepare.py
python train.py --dataset=custom

何时使用以及与其他方案的比较

适合使用 nanoGPT 的场景：

• 学习 GPT 的工作原理
• 尝试不同的 Transformer 变体
• 教学用途
• 快速原型开发
• 计算资源有限（可在 CPU 上运行）

简洁性的优势：

• 约 300 行：整个模型都在 model.py 中
• 约 300 行：训练循环位于 train.py 中
• 易于修改：方便进行修改
• 没有抽象层：纯 PyTorch

其他替代方案：

• HuggingFace Transformers：用于生产环境，提供多种模型
• Megatron-LM：用于大规模分布式训练
• LitGPT：支持更多架构，可用于生产环境
• PyTorch Lightning：需要高层框架时使用

常见问题

问题：CUDA 显存不足

减小批次大小或上下文长度：

batch_size = 1  # Reduce from 12
block_size = 512  # Reduce from 1024
gradient_accumulation_steps = 40  # Increase to maintain effective batch

问题：训练速度太慢

启用编译（PyTorch 2.0+）：

compile = True  # 2× speedup

使用混合精度：

dtype = 'bfloat16'  # Or 'float16'

问题：生成质量不佳

延长训练时间：

max_iters = 10000  # Increase from 5000

降低温度：

# In sample.py
temperature = 0.7  # Lower from 1.0
top_k = 200       # Add top-k sampling

问题：无法加载 GPT-2 权重

安装 transformers：

pip install transformers

检查模型名称：

init_from = 'gpt2'  # Valid: gpt2, gpt2-medium, gpt2-large, gpt2-xl

进阶主题

模型架构：参见 references/architecture.md，了解 GPT 模块结构、多头注意力和 MLP 层的简单说明。

训练循环：参见 references/training.md，了解学习率调度、梯度累积和分布式数据并行设置。

数据准备：参见 references/data.md，了解分词策略（字符级与 BPE）和二进制格式详情。

硬件要求

• 莎士比亚（字符级）：

  • CPU：5 分钟
  • GPU（T4）：1 分钟
  • VRAM：<1GB

• GPT-2（124M）：

  • 1× A100：约 1 周
  • 8× A100：约 4 天
  • VRAM：每个 GPU 约 16GB

• GPT-2 Medium（350M）：

  • 8× A100：约 2 周
  • VRAM：每个 GPU 约 40GB

性能：

• 使用 compile=True：速度提升 2 倍
• 使用 dtype=bfloat16：内存减少 50%

资源

• 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
• 教育用途：最适合从头理解 Transformer

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 确认输入。 根据数据或用户提供的信息确认格式、标识符和参数。凡是会影响结果的值，都要询问，不要猜测。
3. 先运行小规模版本。 进行冒烟测试：使用小型模型或子集、少量步骤，并确认损失下降且检查点成功保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件信息。
6. 汇报。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果发生这种情况	请执行以下操作
CUDA 内存不足	减小批次大小，启用梯度累积/检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否有无效值，并启用梯度裁剪。
这些说明中的函数、标志或端点在已安装版本中不存在	查看已安装版本自己的文档（help()、--help、官方文档），进行适配并告知用户。切勿捏造 API。
必需的输入、标识符或参数含义不明确	询问用户，或在运行前明确说明所采用的假设。

诚信守则

• 切勿捏造结果、参数、标识符、引用或统计数据。如果无法运行或验证某项内容，请如实说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前，应先估算计算成本。
• 请将此处与特定版本相关的细节视为可能过时，并根据已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征得同意。

相关技能

• ray-data：为 ML 工作负载提供可扩展的数据处理。
• optimizing-attention-flash：使用 PyTorch 原生 scaled_dot_product_attention（PyTorch 2.2+）或 flash-attn 库，为 Transformer 模型启用 Flash Attention，…
• huggingface-accelerate：使用 HuggingFace Accelerate（Accelerator 类、accelerate config、accelerate launch）封装现有的 PyTorch 训练脚本，以便相同的 c…