超越 NanoGPT：从 LLM 初学者到 AI 研究者！

Beyond-NanoGPT 是一个极简且具有教育意义的代码库，旨在 连接 NanoGPT 与研究级别的深度学习。
该仓库包含了近 100 种前沿深度学习中关键技术的注释版和从零开始实现的代码，帮助初学者掌握足够的知识，从而能够独立开展实验。

仓库涵盖了从 LLM 中的 KV 缓存和推测解码，到视觉 Transformer 和 MLP 混合器等架构；从线性注意力或多潜在空间注意力等变体，到去噪扩散模型和流匹配算法等生成模型；以及 PPO、A3C 和 AlphaZero 等强化学习领域的里程碑式论文，再到 GPU 通信算法和数据/张量并行等系统基础内容。

由于所有内容都是手写实现的，代码注释详细解释了那些在论文和生产级代码库中常常被忽略的微妙细节。

LESSONS.md 记录了我在编写这个代码库的几个月里所学到的一些经验教训。

快速入门

1. 克隆仓库：

   git clone https://github.com/tanishqkumar/beyond-nanogpt.git
   

2. 安装最小依赖：

   pip install torch numpy torchvision wandb tqdm transformers datasets diffusers matplotlib pillow jupyter gym 
   

3. 开始学习吧！
   这些代码的设计目的是供你仔细阅读、动手修改，然后尝试从头自行实现，并与原版进行对比。
   你可以使用原生 Python 直接运行 .py 文件，例如：

   cd architectures/
   python train_dit.py
   

   或者：

   cd rl/fundamentals/
   python train_reinforce.py --verbose --wandb
   

   所有代码都设计为在单个 GPU 上运行。代码本身带有注释，帮助理解背后的直觉，并详细阐述我在实现过程中遇到的难点。
   每个文件底部都列出了可配置的参数。Jupyter 笔记本则适合逐步浏览。

当前实现与路线图

星号 (*) 表示特别复杂的实现。

架构

• 基础 Transformer language-models/transformer.py 和 train_naive.py [论文]
• 视觉 Transformer (ViT) architectures/train_vit.py [论文]
• 扩散 Transformer (DiT) architectures/train_dit.py [论文]
• 循环神经网络 (RNN) architectures/train_rnn.py [论文]
• 残差网络 (ResNet) architectures/train_resnet.py [论文]
• MLP 混合器 architectures/train_mlp_mixer.py [论文]
• LSTM architectures/train_lstm.py [论文]
• 专家混合模型* (MoE) architectures/train_moe.py [论文]
• Mamba* architectures/train_mamba.py [论文]

注意力变体

• 原始自注意力 attention-variants/vanilla_attention.ipynb [论文]
• 多头自注意力 attention-variants/mhsa.ipynb [论文]
• 分组查询注意力 attention-variants/gqa.ipynb [论文]
• 线性注意力* attention-variants/linear_attention.ipynb [论文]
• 稀疏注意力 attention-variants/sparse_attention.ipynb [论文]
• 交叉注意力 attention-variants/cross_attention.ipynb [论文]
• 多潜在空间注意力* attention-variants/mla.ipynb [论文]

语言模型

• 优化的数据加载 language-models/dataloaders [参考]
  • 生产者-消费者异步数据加载
  • 序列打包
• 字节对编码 language-models/bpe.ipynb [论文]
• KV 缓存 language-models/KV_cache.ipynb [参考]
• 推测解码 language-models/speculative_decoding.ipynb [论文]
• RoPE 嵌入* language-models/rope.ipynb [论文]
• 多标记预测 language-models/train_mtp.py [论文]

强化学习

• 深度强化学习
  • 基础知识 rl/fundamentals
    • DQN train_dqn.py [论文]
    • REINFORCE train_reinforce.py [论文]
    • PPO train_ppo.py [论文]
  • 策略-价值网络及其主要变体 rl/actor-critic
    • 优势策略-价值网络 (A2C) train_a2c.py [论文]
    • 异步优势策略-价值网络 (A3C) train_a3c.py [论文]
    • IMPALA*（分布式强化学习）train_impala.py [论文]
    • 深度确定性策略梯度 (DDPG) train_ddpg.py [论文]
    • 软策略-价值网络* (SAC) train_sac.py [论文]
  • 基于模型的强化学习 rl/model-based
    • 模型预测控制 (MPC) train_mpc.py[参考文献]
    • 专家迭代 (MCTS) train_expert_iteration.py [论文]
    • 轨迹采样概率集成 (PETS)
  • 神经国际象棋引擎 (AlphaZero) rl/chess [论文]
    • 定义架构和环境 model.py 和 env.py
    • MCTS 用于走法搜索 mcts.py
    • 自对弈 train.py
    • 动态批处理和多进程 mcts.py
• 大语言模型 rl/llms
  • RLHF：使用 UltraFeedback 对基础模型进行微调
  • DPO：使用 UltraFeedback 对基础模型进行微调
  • GRPO 用于推理 train_grpo_gsm.py [论文]
  • GRPO 用于幽默生成 train_grpo_humor.py

生成模型

• 生成对抗网络 (GAN) generative-models/train_gan.py [论文]
• Pix2Pix（条件 GAN）generative-models/train_pix2pix.py [论文]
• 变分自编码器 (VAE) generative-models/train_vae.py [论文]
  • 训练用于重建的自编码器 generative-models/train_autoencoder.py
• 神经辐射场 (NeRF)
• 去噪扩散概率模型* (DDPM) generative-models/train_ddpm.py [论文]
• 基于分类器的扩散引导 generative-models/ddpm_classifier_guidance.py [论文]
  • 无分类器扩散引导 generative-models/ddpm_classifier_free_guidance.py [论文]
• 流匹配 generative-models/train_flow_matching.py [论文]

MLSys

• GPU 通信算法*（scatter、gather、ring/tree allreduce）mlsys/comms.py [参考文献]
• 分布式数据并行 mlsys/train_ddp.py [论文]
• 张量并行* mlsys/train_tp.py [论文]
• 环形注意力（上下文并行）
• 分页注意力
• 连续批处理
• Triton 内核 mlsys/kernels
  • 向量加法 vector_add.py
  • 一维卷积（朴素 + 分块）conv1d.py
  • 二维步进复制 copy2d.py
  • 层归一化前向传播 layernorm.py
  • 数组反转 reverse_array.py
  • SwiGLU 激活函数 swiglu.py
  • 分块 GEMM tiled_gemm.py
  • FlashAttention 前向传播

评估

• BERT 在 SST-2 数据集上的表现（传统 NLP）
• GSM8k（生成任务）evals/eval_gsm8k.py [论文]
• MMLU（多项选择题）evals/eval_mmlu.py [论文]
• SimpleQA（LLM 评判）evals/eval_simpleqa.py [论文]
• 设计我们自己的评估标准（“良好品味”）

RAG

• 训练小型嵌入模型和重排序模型
• RAG 入门：基于问答答案的检索 rag/intro_rag.py
• 多跳分解 RAG
• 稀疏与稠密检索
• 图 RAG

代理

• 让 LLM 使用互联网搜索进行问答 agents/basic-search-use
• 编码代理 agents/coding-agent
  • 工具使用（搜索、运行代码、读写文件）及强大工具的沙箱隔离 [/tools]
  • ReAct（迭代思维链，中间穿插工具使用）agent.py
  • 短期与长期记忆的区分及记忆/上下文管理 memory.py
  • 评估：它能否针对 GitHub 问题端到端地提交一个正确的 PR？[演示]
• 模拟由语言模型构成的社会
• 思维树深度研究代理
• 并行多代理深度研究

---

注释

• 该代码库通常可以在 CPU 或 GPU 上运行，但大多数实现实际上需要 GPU，否则速度会慢得难以忍受。建议使用带有 GPU 的消费级笔记本电脑，或者付费使用 Colab/Runpod，如果预算不足，也可以向计算资源提供商或当地大学申请计算资助（这种方法往往效果很好，人们通常非常慷慨）。显而易见的例外情况，如数据并行或张量并行，则需要多 GPU 节点。
• 所有 .py 脚本在运行时都接受 --verbose 和 --wandb 作为命令行参数。您可以根据需要随意修改这些参数。
• 如有任何反馈、实现或功能请求，或发现任何错误，请随时通过电子邮件 tanishq@stanford.edu 与我联系，并将其作为 GitHub 问题提交。如果您觉得这个代码库对您有所帮助，请分享并给它一颗星！您可以在您的工作中按如下方式引用该仓库。

@misc{kumar2025beyond,
  author = {Tanishq Kumar},
  title = {超越 NanoGPT：从 LLM 初学者到 AI 研究者},
  year = {2025},
  howpublished = {\url{https://github.com/tanishqkumar/beyond-nanogpt}},
  note = {访问日期：2025-01-XX}
}

祝编程愉快，愿您的梯度永不消失！

Beyond-NanoGPT 快速上手指南

Beyond-NanoGPT 是一个极简的教育性开源项目，旨在填补 nanoGPT 与前沿深度学习研究之间的空白。它从零开始实现了近 100 种现代关键技术（包括 LLM 优化、注意力变体、强化学习、生成模型及系统底层原理），并通过详细的代码注释解释了许多在论文和生产代码中常被忽略的细节。适合希望从 LLM 初学者进阶为 AI 研究者的开发者。

环境准备

本项目主要基于 PyTorch 构建，虽然部分代码可在 CPU 上运行，但绝大多数实现（尤其是训练任务）强烈建议使用 GPU，否则速度将极慢。

• 操作系统: Linux, macOS, Windows (需配置 CUDA 环境)
• 硬件要求:
  • 单张 NVIDIA GPU (推荐显存 ≥ 8GB)
  • 部分分布式训练示例（如 Tensor Parallelism）需要多卡环境
• 软件依赖:
  • Python 3.8+
  • CUDA Toolkit (根据显卡型号安装对应版本)
  • Git

安装步骤

1. 克隆仓库

首先将项目代码克隆到本地：

git clone https://github.com/tanishqkumar/beyond-nanogpt.git
cd beyond-nanogpt

2. 安装依赖

使用 pip 安装所需的最小化依赖包。

国内加速建议：如果您在中国大陆，建议使用清华或阿里镜像源以加快下载速度：

# 使用默认源
pip install torch numpy torchvision wandb tqdm transformers datasets diffusers matplotlib pillow jupyter gym 

# 或使用清华镜像源 (推荐国内用户)
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch numpy torchvision wandb tqdm transformers datasets diffusers matplotlib pillow jupyter gym

基本使用

本项目的核心理念是阅读代码、修改实验、并尝试从头复现。所有 .py 文件均包含自文档化的注释，参数定义在文件底部。

运行单个训练脚本

你可以直接使用原生 Python 运行任意脚本。以下是两个典型示例：

示例 1：训练扩散 Transformer (DiT) 进入架构目录并运行：

cd architectures/
python train_dit.py

示例 2：运行强化学习 (REINFORCE 算法) 进入强化学习基础目录，并开启详细日志与 WandB 监控：

cd rl/fundamentals/
python train_reinforce.py --verbose --wandb

探索 Jupyter Notebook

对于注意力机制变体（如 Linear Attention, MLA）和数据加载优化等内容，项目提供了交互式 Notebook。请在 Jupyter Lab 或 VS Code 中逐步运行以下文件以可视化结果：

# 例如查看注意力变体的速度对比图
jupyter lab attention-variants/linear_attention.ipynb

自定义参数

每个脚本底部的 argparse 部分定义了可调整的参数。您可以直接编辑代码或在命令行中传递参数（如 --batch_size, --lr 等）来进行自己的实验。

---

提示：更多学习心得和实现细节请参阅项目根目录下的 LESSONS.md 文件。