Train LLM From Scratch 从零训练大语言模型

项目简介

Train LLM From Scratch 是一个用纯 PyTorch 从零实现 Transformer 的开源项目，基于论文《Attention is All You Need》。它把大语言模型训练的完整流程拆解成可运行的脚本，从下载原始文本、分词、预训练，一直走到 SFT、奖励模型、DPO/PPO/GRPO 对齐，最终生成对话回复。

整个项目不依赖 transformers、trl、peft 等高层封装，所有算法都手写在 PyTorch 中，适合想真正搞懂 LLM 训练细节的开发者。

核心流程

项目把训练路径总结为一条清晰的流水线：

raw text -> tokens -> Transformer -> next-token loss -> base model
base model -> SFT -> Reward Model -> {PPO, DPO} -> GRPO -> evaluation and chat

主要阶段包括：

• 数据准备：下载公开数据集并完成分词与磁盘存储
• 模型搭建：从 MLP、单头注意力、多头注意力到完整 Transformer Block
• 预训练：训练基础模型，支持 AMP、梯度检查点、梯度累积
• 文本生成：用训练好的模型采样生成文本
• 后训练对齐：SFT、Bradley-Terry 奖励模型、DPO/ORPO/KTO、PPO、GRPO/RLVR
• 评估与对话：对模型进行评估并通过聊天界面交互

适用人群

• 学生：按顺序阅读，每段代码前都有通俗解释，并附预期输出
• 开发者：命令与文件路径齐全，可直接复制运行并阅读源码
• 研究者：后训练部分是重点，包含 SFT、奖励模型、PPO with GAE、DPO/ORPO/KTO、GRPO 等从零实现

硬件与配置

项目对显存需求做了说明：免费 Colab 或 Kaggle 的 T4 足以训练 1300 万参数的小模型，但十亿参数模型需要更大显存。常见 GPU 参考：

预训练脚本提供 --amp、--grad-checkpointing、--grad-accum 等选项，可显著降低显存占用。

安装与运行

克隆仓库后以可编辑模式安装，即可把 config、src、data_loader、ui 加入导入路径：

pip install -e .

按需安装可选依赖：

pip install -e ".[train]"   # datasets + wandb
pip install -e ".[ui]"      # streamlit + pandas + altair
pip install -e ".[docs]"    # mkdocs
pip install -e ".[all]"     # 全部

项目提供两套配置系统：config/config.py 用于旧的预训练脚本，config/post_training_config.py 与 configs/ 下的 JSON 文件驱动其余阶段，命令行也可覆盖字段，例如 --lr 2e-5 --batch_size 16。此外 configs/smoke/ 提供每个阶段的极简配置，可在 CPU 或单卡上秒级跑通验证。

项目亮点

• 全流程手写，原理透明，适合深入理解 Transformer 与对齐算法
• 覆盖预训练到 RLHF 的完整链路，教学价值高
• 提供 Streamlit 控制面板与 MkDocs 文档站，交互与查阅方便
• 支持从小模型到大模型的渐进式实验，硬件门槛可调