从零开始构建大语言模型训练框架 —— 纯 PyTorch 实现，面向教学与生产级 LLM 训练
功能特性 • 版本能力 • 快速开始 • 项目结构 • 配置说明 • 路线图
llm-from-zero-to-one 是一个从零手写实现的大语言模型（LLM）训练框架，旨在帮助开发者深入理解 LLM 训练的每一个核心环节。项目采用纯 PyTorch 实现，不依赖 HuggingFace Transformers 的模型代码，完整覆盖了从数据预处理、模型定义、训练循环到评估推理的全流程。
本项目适合：
- 学习者：想理解 Transformer/GPT 模型底层原理和训练细节的开发者
- 研究者：需要灵活定制训练流程的研究人员
- 工程师：希望掌握 LLM 工程化实践的技术人员
本项目的设计理念和架构深受 MiniMind 项目的启发。MiniMind 是一个优秀的开源项目，致力于以极低的成本（3 元 GPU 费用 + 2 小时训练时间）从零训练超小语言模型（~25.8M 参数），其"大道至简"的哲学——使用纯 PyTorch 原生实现所有核心算法、不依赖第三方框架封装——与本项目的初衷高度一致。
MiniMind 项目亮点：
- 🧠 全阶段开源复现：Pretrain → SFT → LoRA → DPO → RLHF (PPO/GRPO) → 模型蒸馏
- 🔬 核心代码从零手写，帮助理解 LLM 每一个细节
- 📚 完整教程导向，适合 LLM 入门学习
本项目示例配置中使用的预训练数据集来自 minimind_dataset（HuggingFace），分词器采用 minimind-3。感谢 MiniMind 社区开放的高质量数据资源。
| 模块 | 能力 | 说明 |
|---|---|---|
| 模型架构 | GPT-2 从零实现 | 完整的 Decoder-only Transformer，包含因果自注意力、MLP、Pre-Norm、位置编码 |
| 预训练 | 全量训练循环 | 支持 AdamW/Adam 优化器、余弦学习率调度、梯度累积、梯度裁剪 |
| 混合精度 | AMP (fp16/bf16) | CUDA 上支持自动混合精度训练，提升训练效率与显存利用率 |
| 编译加速 | torch.compile | 自动启用 PyTorch 2.x 编译加速（CUDA/MPS） |
| 数据加载 | 双策略支持 | Padding（边训边分词）/ Megatron（预先分词 .bin/.idx）两种模式 |
| 检查点管理 | 断点续训 | 支持按步保存、自动恢复最新检查点、配置兼容性校验 |
| 评估系统 | 困惑度评估 | 训练中可周期性评估 Loss / Perplexity / Token 吞吐量 |
| 实验跟踪 | SwanLab 集成 | 可选接入 SwanLab 进行实验可视化与参数记录 |
| 数据处理工具 | 数据预处理 | 提供文本分词 → .bin + .idx 文件生成的完整流水线 |
- 注册表模式：模型与配置通过装饰器注册，扩展新架构只需几行代码
- YAML 驱动配置：所有超参统一通过 YAML 管理，支持环境变量替换
${VAR:-default} - 类型安全：使用
jaxtyping进行张量形状注解，dataclass统一配置管理 - 确定性训练：完整的种子控制（全局/Dataloader Worker），保证实验可复现
| 能力维度 | 支持状态 | 详细说明 |
|---|---|---|
| 模型架构 | ✅ 已支持 | GPT-2 (Decoder-only Transformer) |
| 训练模式 | ✅ 已支持 | Pre-training（语言模型预训练） |
| 优化器 | ✅ 已支持 | AdamW, Adam |
| 学习率调度 | ✅ 已支持 | Cosine Annealing + Linear Warmup |
| 混合精度训练 (AMP) | ✅ 已支持 | fp16, bf16 (CUDA) |
| torch.compile 加速 | ✅ 已支持 | CUDA / MPS 自动启用 |
| 梯度累积 | ✅ 已支持 | 可配置 accumulation_steps |
| 梯度裁剪 | ✅ 已支持 | 全局范数裁剪 |
| 数据加载 - Padding 策略 | ✅ 已支持 | 边训练边分词，HuggingFace Datasets 集成 |
| 数据加载 - Megatron 策略 | ✅ 已支持 | 预先分词 (.bin + .idx 内存映射) |
| 检查点保存/恢复 | ✅ 已支持 | 按步保存 + 断点续训 + 配置兼容校验 |
| 训练中评估 | ✅ 已支持 | Loss / Perplexity / 吞吐量 |
| 独立评估模式 | ✅ 已支持 | 对已有 checkpoint 执行评估 |
| SwanLab 实验跟踪 | ✅ 已支持 | 参数/指标可视化 |
| YAML 配置管理 | ✅ 已支持 | 环境变量替换 + 校验 + 默认生成 |
| 数据处理工具链 | ✅ 已支持 | 分词 → .bin/.idx 生成 |
| 单元测试 | ✅ 已覆盖 | 核心模块 pytest 测试 |
| 能力维度 | 计划版本 | 说明 |
|---|---|---|
| 分布式训练 (DDP) | v0.2.0 | 单机多卡数据并行 |
| SFT / DPO 微调 | v0.3.0 | 监督微调 + 偏好对齐 |
| Qwen / DeepSeek 模型架构 | v0.3.0 | Dense & MoE 架构支持 |
| HuggingFace 格式互转 | v0.3.0 | 导出/导入 HF 格式权重 |
| DeepSpeed / FSDP | v0.4.0 | 大规模分布式训练加速 |
- Python >= 3.10
- PyTorch >= 2.10.0
- CUDA (推荐) 或 MPS (Apple Silicon)
# 克隆项目
git clone <repo-url>
cd llm_from_zero_to_one
# 创建虚拟环境并安装依赖
uv sync
# 激活虚拟环境
source .venv/bin/activate# 使用内置配置文件启动预训练
uv run python3 -m main --config configs/pretrain_padding_example.yaml
# 生成默认配置模板
uv run python3 -m main --config my_config.yaml --generate-config# 对已有 checkpoint 执行评估
uv run python3 -m main --config configs/eval_minimind_unseen_1024.yaml# 将原始文本数据集处理为 .bin + .idx 格式
uv run python3 tools/llm_data_processor.pyllm_from_zero_to_one/
├── configs/ # YAML 配置文件
│ ├── pretrain_padding_example.yaml # Padding 策略预训练示例
│ ├── pretrain_megatron_example.yaml # Megatron 策略预训练示例
│ └── eval_*.yaml # 评估配置
├── src/
│ ├── models/ # 模型模块
│ │ ├── gpt2/ # GPT-2 模型实现
│ │ │ ├── model.py # CausalSelfAttention, MLP, Block, GPT2
│ │ │ └── config.py # GPT2Config
│ │ ├── base.py # BaseModel 基类（含 generate/save/load）
│ │ ├── common_module.py # LayerNorm 等公共组件
│ │ ├── registry.py # 模型注册表（装饰器）
│ │ └── model_factory.py # 模型工厂
│ ├── trainer/ # 训练模块
│ │ ├── pretrain/pretrain.py # PreTrainTrainer 核心训练逻辑
│ │ ├── train.py # CLI 入口 & 模式分发
│ │ └── train_args.py # YAML 配置 dataclass 定义
│ ├── dataset/ # 数据加载模块
│ │ ├── dataset_factory.py # 数据策略工厂
│ │ ├── padding_dataset.py # Padding 策略数据集
│ │ └── simple_megatron_dataset/ # Megatron 策略数据集 (.bin/.idx)
│ ├── checkpoint_manager/ # 检查点管理
│ │ └── checkpoint_manager.py # 保存/加载/恢复
│ ├── evaluator/ # 评估模块
│ │ └── checkpoint_evaluator.py # Perplexity 评估器
│ └── logger/ # 日志模块
├── tools/ # 辅助工具脚本
│ ├── llm_data_processor.py # 数据预处理（生成 .bin/.idx）
│ ├── check_idx.py # 索引文件检查
│ └── select_unseen_pretrain_samples.py # 未见过样本筛选
├── tests/ # 测试套件
├── main.py # 程序入口
└── pyproject.toml # 项目配置
项目通过 YAML 配置文件 驱动所有训练参数，核心配置项：
mode: pretrain # 运行模式: pretrain | eval
name: experiment_name # 实验名称（用于 checkpoint 目录）
model:
name: gpt2 # 模型名称
config:
vocab_size: 6400 # 词表大小
n_layer: 8 # Transformer 层数
n_head: 8 # 注意力头数
n_embd: 768 # 嵌入维度
training:
batch_size: 16 # 批大小
seq_len: 1024 # 序列长度
epoch_num: 1 # 训练轮数
learning_rate: 3e-4 # 学习率
warmup_steps: 100 # 预热步数
grad_clip: 1.0 # 梯度裁剪阈值
accumulation_steps: 4 # 梯度累积步数
amp: true # 是否启用混合精度
amp_dtype: bf16 # 混合精度类型: fp16 | bf16
data:
data_strategy: 'padding' # 数据策略: padding | megatron
dataset_config:
dataset_path: 'your/dataset' # 数据集路径
tokenizer_path: 'your/tokenizer' # 分词器路径
dataloader_config:
num_workers: 4 # 数据加载进程数
shuffle: true # 是否打乱数据
checkpoint:
save_steps: 1000 # 每 N 步保存一次
checkpoint_dir: checkpoints/pretrain # 检查点目录
eval:
steps: 1000 # 每 N 步执行一次评估（0=关闭）
dataset_path: ... # 评估数据集路径
max_samples: 256 # 最大评估样本数
optimizer:
name: adamw # 优化器: adamw | adam
weight_decay: 0.01 # 权重衰减
swanlab:
enabled: false # 是否启用 SwanLab
project: llm-pretrain # 项目名称
experiment_name: experiment # 实验名称💡 提示：配置文件支持环境变量替换，例如
dataset_path: ${DATA_PATH:-default/path}
# 运行全部测试
uv run python3 -m pytest tests/ -v
# 运行指定测试文件
uv run python3 -m pytest tests/test_models.py -v# 代码格式化 (ruff + isort)
ruff format .
ruff check --fix .
# 运行测试
uv run python3 -m pytest tests/ -vMIT License