train-llm-from-scratch 项目深度分析

分析日期：2026-06-01
项目地址：https://github.com/FareedKhan-dev/train-llm-from-scratch
开源协议：MIT License
作者：Fareed Khan

---

一、项目介绍

1.1 项目概述

train-llm-from-scratch 是一个从零开始实现 Transformer 模型并训练大语言模型的教学项目。基于经典论文《Attention is All You Need》，使用 PyTorch 从头实现了完整的 Transformer 架构，用户可在单块 GPU 上训练自己的百万到十亿参数级别的 LLM。

项目不依赖 HuggingFace 等高层框架，纯 PyTorch 手写所有组件，覆盖了从数据下载、预处理、训练到文本推理的完整全流程。README 长达 1240 行，包含极其详细的逐步代码解释和架构图。

13M 参数模型输出示例：

In 1978, The park was returned to the factory-plate that 
the public share to the lower of the electronic fence that 
follow from the Station's cities. The Canal of ancient Western 
nations were confined to the city spot...

1.2 项目地址

类型	地址
GitHub 仓库	https://github.com/FareedKhan-dev/train-llm-from-scratch
作者 GitHub	https://github.com/FareedKhan-dev
Medium 技术文章	https://levelup.gitconnected.com/building-a-million-parameter-llm-from-scratch-using-python-f612398f06c2
训练数据集	https://huggingface.co/datasets/monology/pile-uncopyrighted

1.3 项目示意图

Transformer 架构图（来自作者 Medium 文章）：

README 中还包含以下架构示意图（由作者手绘）：

• MLP（前馈神经网络）结构图
• Single Head Attention（单头注意力）结构图
• Multi Head Attention（多头注意力）结构图
• Transformer Block（Transformer 块）结构图
• Transformer Class（完整模型）结构图
• Training Loss Comparison（训练损失对比）图

---

二、项目亮点

2.1 完全从零实现

不依赖 HuggingFace Transformers 等高层框架，纯 PyTorch 手写所有组件：MLP、单头注意力、多头注意力、Transformer Block、完整模型。每个组件都有详细的代码解释和架构图。

2.2 极低硬件门槛

13M 参数模型仅需消费级 GPU 即可训练。作者提供了 30 款 GPU 的详细兼容性对比表，涵盖 NVIDIA 和 AMD 显卡，从 GTX 1660 Ti 到 A100 全覆盖。

GPU	显存	13M 模型	2B 模型	最大可训练
RTX 4090	24 GB	✅	✅	~4B
RTX 3090	24 GB	✅	✅	~3.5-4B
RTX 4080	16 GB	✅	❌	~2B
RTX 4060	8 GB	✅	❌	~1B
T4	16 GB	✅	❌	~1.5-2B
A100	40 GB	✅	✅	~6-8B

2.3 完整训练流水线

覆盖从数据到推理的全流程，5 步即可完成：

数据下载 → 预处理（JSONL.ZST 解压 + tokenize）→ 训练 → 生成文本

2.4 灵活的模型规模

通过配置文件调整参数，支持从 13M 到 30 亿参数的模型：

模型规模	CONTEXT_LENGTH	N_EMBED	N_HEAD	N_BLOCKS
13M（入门）	128	128	8	1
1B	256	1024	16	16
2B	512	2048	16	64
3B（默认）	512	2048	16	64

2.5 详尽的文档

README 长达 1240 行，包含：

• Transformer 架构原理详解
• 每个组件的代码逐步解释
• 30 款 GPU 兼容性对比表
• 训练损失曲线对比
• 13M 和十亿参数模型的输出示例

2.6 使用 The Pile 数据集

使用 The Pile 数据集（825GB，22 个子数据集），包含书籍、论文、代码、网页等多样化文本，是 LLM 训练的优质开源数据集。

2.7 附带 SFT/RLHF 指南

提供 sft_rlhf_guide.ipynb Jupyter Notebook，指导用户在基础模型上进行监督微调和人类反馈强化学习。

2.8 模块化代码设计

代码已重构为模块化结构，各组件职责清晰：

• src/models/ — 模型定义（MLP、Attention、Block、Transformer）
• config/ — 集中配置管理
• data_loader/ — 数据加载器
• scripts/ — 训练、推理、数据预处理脚本

---

三、项目运行环境介绍和运行条件介绍

3.1 系统要求

项目	最低配置（13M 模型）	推荐配置（1B+ 模型）
操作系统	Windows / Linux / macOS	Linux（推荐）
Python	3.8+	3.10+
GPU	6 GB 显存（如 RTX 4050）	16 GB+ 显存（如 RTX 4080）
内存	8 GB	16 GB+
存储	~5 GB（小数据集）	50 GB+（大数据集）
CUDA	11.8+	12.1+

GPU 非必需但强烈推荐。CPU 训练速度极慢，仅适合代码验证。

3.2 Python 依赖

包名	用途
torch	深度学习框架（CUDA 11.8）
torchvision	图像处理（辅助）
torchaudio	音频处理（辅助）
numpy	数值计算
h5py	HDF5 数据格式存储
tiktoken	OpenAI 分词器（r50k_base，与 GPT-2 相同）
zstandard	Zstandard 解压缩（.jsonl.zst 格式）
requests	HTTP 下载
tqdm	进度条

3.3 快速安装步骤

# 1. 克隆项目
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git
cd train-llm-from-scratch

# 2. 设置 Python 路径
export PYTHONPATH="$PYTHONPATH:."

# 3. 安装依赖
pip install -r requirements.txt

# 4. 下载训练数据（默认下载 1 个文件，约 11GB）
python scripts/data_download.py

# 5. 预处理数据（tokenize 并保存为 HDF5）
python scripts/data_preprocess.py

# 6. 修改配置（编辑 config/config.py，改为 13M 配置）

# 7. 训练模型
python scripts/train_transformer.py

# 8. 生成文本
python scripts/generate_text.py --model_path models/your_model.pth --input_text hi

3.4 前置知识要求

需要具备以下基础知识：

• 面向对象编程（OOP）
• 神经网络基础（前向传播、反向传播、损失函数）
• PyTorch 基础（Tensor 操作、autograd、nn.Module）

README 中提供了 3 个 YouTube 视频教程链接作为补充学习资源。

---

四、项目代码介绍

4.1 代码架构图

train-llm-from-scratch/
├── src/models/                      # 模型定义层
│   ├── mlp.py                       #   前馈神经网络（FFN）
│   ├── attention.py                 #   注意力机制（单头 + 多头）
│   ├── transformer_block.py         #   Transformer 块（LN + Attn + MLP）
│   └── transformer.py               #   完整 Transformer 模型
│
├── config/
│   └── config.py                    # 集中配置（模型参数 + 训练超参数）
│
├── data_loader/
│   └── data_loader.py               # HDF5 数据加载器（无限批次迭代器）
│
├── scripts/                         # 可执行脚本
│   ├── data_download.py             #   Pile 数据集下载
│   ├── data_preprocess.py           #   数据预处理（解压 + tokenize → HDF5）
│   ├── train_transformer.py         #   模型训练脚本
│   └── generate_text.py             #   文本生成/推理脚本
│
├── data/                            # 数据存储
│   ├── train/                       #   训练数据（.jsonl.zst / .h5）
│   └── val/                         #   验证数据
│
├── models/                          # 训练好的模型保存目录
├── sft_rlhf_guide.ipynb             # SFT 和 RLHF 指南
├── requirements.txt                 # 依赖列表
└── LICENSE                          # MIT 许可证

4.2 核心模块介绍

模块一：transformer.py — 完整 Transformer 模型

组合所有子模块，构成完整的语言模型。

核心类：Transformer(nn.Module)

组件	说明
token_embed	词嵌入层（vocab_size → n_embed）
position_embed	位置嵌入层（context_length → n_embed）
attn_blocks	N 个 Transformer Block 的 ModuleList
layer_norm	全局 LayerNorm
lm_head	语言模型头（n_embed → vocab_size）

核心方法：

• forward(idx, targets=None) — 前向传播，计算 logits 和交叉熵损失
• generate(idx, max_new_tokens) — 自回归文本生成（multinomial 采样）

模块二：attention.py — 注意力机制

实现单头和多头注意力。

Head(nn.Module) — 单头注意力：

q = query(x)   # (B, T, head_size)
k = key(x)     # (B, T, head_size)
attn = softmax((q @ k^T) / sqrt(head_size) + causal_mask)
out = attn @ value(x)

MultiHeadAttention(nn.Module) — 多头注意力：

• 将 n_embed 维度均匀分配给 n_head 个头
• 所有头并行计算后拼接输出

模块三：mlp.py — 前馈神经网络

x → Linear(n_embed, 4*n_embed) → ReLU → Linear(4*n_embed, n_embed) → out

隐藏层扩展 4 倍，使用 ReLU 激活函数。

模块四：transformer_block.py — Transformer 块

x = x + attn(ln1(x))    # Pre-Norm + 多头注意力 + 残差
x = x + mlp(ln2(x))     # Pre-Norm + MLP + 残差

采用 Pre-Norm 架构（现代 LLM 主流做法，训练更稳定）。

模块五：data_loader.py — 数据加载器

从 HDF5 文件创建无限批次迭代器：

• 每个样本取 context_length + 1 个连续 token
• 输入 xb：前 context_length 个 token
• 目标 yb：后 context_length 个 token（偏移 1 位）
• 自动循环 epoch，适合按步数控制的训练

模块六：train_transformer.py — 训练脚本

完整的训练流程：模型初始化 → AdamW 优化器 → 训练循环（按步数）→ 学习率衰减 → 评估 → 模型保存。

4.3 Transformer 架构流程图

输入 token 索引 (B, T)
       │
       ▼
  Token Embedding (vocab_size → n_embed)
       +  Position Embedding (context_length → n_embed)
       │
       ▼
  ┌─────────────────────────────────┐
  │     Transformer Block × N       │  ← 重复 N_BLOCKS 次
  │  ┌───────────────────────────┐ │
  │  │ LayerNorm                  │ │
  │  │ Multi-Head Attention       │ │  ← n_head 个 Head 并行
  │  │ + Residual Connection      │ │
  │  │ LayerNorm                  │ │
  │  │ MLP (4x expand + ReLU      │ │  ← 前馈网络
  │  │      + project back)       │ │
  │  │ + Residual Connection      │ │
  │  └───────────────────────────┘ │
  └─────────────────────────────────┘
       │
       ▼
  LayerNorm (全局)
       │
       ▼
  Linear (n_embed → vocab_size)  ← 语言模型头
       │
       ▼
  Logits (B, T, vocab_size)
       │
       ▼
  Cross-Entropy Loss (与 targets 比较)

4.4 核心代码解析：训练循环

# scripts/train_transformer.py 核心流程

# 1. 初始化模型和优化器
model = Transformer(**config).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=config['t_lr'])

# 2. 训练循环（按步数控制）
for step in range(config['t_train_steps']):
    xb, yb = next(batch_iterator)        # 从 HDF5 获取批次
    logits, loss = model(xb, yb)         # 前向传播 + 计算损失
    optimizer.zero_grad()
    loss.backward()                       # 反向传播
    optimizer.step()                      # 更新参数

    # 学习率阶梯衰减
    if step == config['t_lr_decay_step']:
        for g in optimizer.param_groups:
            g['lr'] = config['t_lr_decayed']

    # 定期评估
    if step % config['t_eval_steps'] == 0:
        losses = estimate_loss(model)     # 在 train/dev 上评估
        print(f"step {step}: train loss {losses['train']:.4f}, dev loss {losses['dev']:.4f}")

# 3. 保存模型
torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss_history': loss_history,
}, config['t_out_path'])

4.5 模块间调用关系

scripts/train_transformer.py (训练入口)
    │
    ├──→ config/config.py (读取模型和训练参数)
    │
    ├──→ src/models/transformer.py (创建模型)
    │       ├──→ src/models/transformer_block.py (N 个 Block)
    │       │       ├──→ src/models/attention.py (多头注意力)
    │       │       │       └──→ src/models/attention.py (单头注意力)
    │       │       └──→ src/models/mlp.py (前馈网络)
    │       └──→ src/models/attention.py (嵌入层)
    │
    ├──→ data_loader/data_loader.py (获取训练/验证批次)
    │       └──→ data/train/pile_train.h5 (HDF5 数据)
    │
    └──→ models/transformer_B.pt (保存训练好的模型)

scripts/generate_text.py (推理入口)
    │
    ├──→ config/config.py (读取参数)
    ├──→ src/models/transformer.py (重建模型 + 加载权重)
    └──→ tiktoken (编码/解码文本)

scripts/data_preprocess.py (数据预处理)
    │
    ├──→ tiktoken (r50k_base 分词器)
    └──→ data/train/*.jsonl.zst → data/train/pile_train.h5

---

五、项目的应用、优点和不足

5.1 应用场景

场景	说明
LLM 教学和学习	最主要的用途，帮助开发者从底层理解 Transformer 原理到完整训练流程
自然语言生成	训练模型生成连贯文本
特定领域微调	先训练 13M 基础模型，再用领域数据微调
私有数据训练	在私有数据上安全地训练小型专用模型
AI 研究实验	快速验证新的模型架构或训练技巧
面试准备	深入理解 Transformer 有助于 AI 岗位面试

5.2 项目优点

1. 真正的从零实现：不依赖 HuggingFace 等高层框架，纯 PyTorch 手写所有组件
2. 极低硬件门槛：13M 模型仅需 6GB 显存的消费级 GPU
3. 文档极其详尽：1240 行 README，包含逐步代码解释、架构图、GPU 兼容表
4. 完整训练流水线：数据下载 → 预处理 → 训练 → 推理，5 步完成
5. 灵活的模型规模：通过配置文件支持 13M 到 30 亿参数
6. 模块化代码设计：各组件职责清晰，代码可读性高
7. 使用优质数据集：The Pile（825GB，22 个子数据集）
8. 附带 SFT/RLHF 指南：提供微调和强化学习的进阶学习路径
9. MIT 开源协议：完全自由使用
10. 社区认可度高：2.7K Star，CSDN/今日头条等多篇中文教程推荐

5.3 项目不足

1. MultiHeadAttention 缺少输出投影层：标准 Transformer 在多头拼接后应有 W_o 线性变换，该项目省略了
2. MLP 使用 ReLU 而非 GELU：现代 LLM 普遍使用 GELU 或 SwiGLU，ReLU 可能限制模型表达能力
3. 缺少学习率 warmup：现代 LLM 训练通常需要 warmup 阶段防止初始训练不稳定
4. 缺少梯度裁剪：大模型训练中可能导致梯度爆炸
5. 未实现混合精度训练（AMP）：内存效率有较大优化空间
6. 权重未共享：lm_head 和 token_embed 独立参数，增加了参数量（GPT-2 采用权重共享）
7. 默认配置面向 3B 模型：新手直接运行会 OOM，需手动改为 13M 配置
8. 无 Release 版本管理：45 次提交但 0 个 Release，版本管理不够规范
9. 训练效率较低：未实现 Flash Attention、数据并行等现代训练加速技术
10. 生成质量有限：13M 模型仅能生成语法正确但语义不连贯的文本