docs: make Chinese README the default
This commit is contained in:
@@ -1,30 +1,36 @@
|
||||
<!-- WEHUB_ZH_README -->
|
||||
> [!NOTE]
|
||||
> 本文档由 WeHub 基于上游 README 翻译整理,属于社区翻译,非官方中文文档。
|
||||
> [English](./README.en.md) · [原始项目](https://github.com/FareedKhan-dev/train-llm-from-scratch) · [上游 README](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/HEAD/README.md)
|
||||
> 原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。
|
||||
|
||||

|
||||
|
||||
<div align="center">
|
||||
|
||||
<!-- omit in toc -->
|
||||
# Train LLM From Scratch
|
||||
# 从零训练大语言模型(LLM)
|
||||
|
||||
   [](https://fareedkhan-dev.github.io/train-llm-from-scratch/)
|
||||
|
||||
**I am Looking for a PhD position in AI**. [GitHub](https://github.com/FareedKhan-dev)
|
||||
**我正在寻找人工智能方向的博士职位。** [GitHub](https://github.com/FareedKhan-dev)
|
||||
|
||||
</div>
|
||||
|
||||
I implemented a transformer model from scratch using PyTorch, based on the paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). You can use my scripts to train your own **billion** or **million** parameter LLM using a single GPU.
|
||||
我基于论文 [Attention is All You Need](https://arxiv.org/abs/1706.03762). 使用 PyTorch 从零实现了一个 Transformer 模型。你可以用我的脚本在单张 GPU 上训练自己的 **十亿** 或 **百万** 参数的大语言模型。
|
||||
|
||||
This started as a pretraining tutorial. It now goes all the way from raw text to an aligned, reasoning style model, with every algorithm hand written in plain PyTorch (no `trl`, no `peft`, no `transformers`). The whole journey is one idea repeated: turn text into numbers, predict the next token, then keep changing the data and the loss until the model does what we want.
|
||||
这最初是一份预训练教程。现在它涵盖了从原始文本到对齐的、具备推理风格模型的完整流程,每个算法都是用纯 PyTorch 手写的(没有 `trl`,没有 `peft`,没有 `transformers`)。整个旅程重复着一个核心思路:把文本变成数字,预测下一个 token,然后不断调整数据和损失函数,直到模型按我们的期望行事。
|
||||
|
||||

|
||||

|
||||
|
||||
Here is the path we will walk, end to end:
|
||||
下面是我们将端到端走过的路径:
|
||||
|
||||
```
|
||||
raw text -> tokens -> a Transformer -> next-token loss -> a base model
|
||||
base model -> SFT -> Reward Model -> {PPO, DPO} -> GRPO -> evaluation and chat
|
||||
```
|
||||
|
||||
Below is the output of a trained 13 million parameter LLM, just so you can see where the small end of this starts:
|
||||
下面是一个训练好的 1300 万参数大语言模型的输出,让你了解这套方法在小型模型端的表现:
|
||||
|
||||
```
|
||||
In ***1978, The park was returned to the factory-plate that
|
||||
@@ -36,80 +42,80 @@ Odambinais is uncertain and fortune established in rural areas.
|
||||
```
|
||||
|
||||
<!-- omit in toc -->
|
||||
## Table of Contents
|
||||
- [Who this is for](#who-this-is-for)
|
||||
- [Prerequisites and Training Time](#prerequisites-and-training-time)
|
||||
- [Setup](#setup)
|
||||
- [Code Structure](#code-structure)
|
||||
- [Step 1: Preparing the Data](#step-1-preparing-the-data)
|
||||
- [Step 2: The Model, Built From Small Pieces](#step-2-the-model-built-from-small-pieces)
|
||||
- [Multi Layer Perceptron (MLP)](#multi-layer-perceptron-mlp)
|
||||
- [Single Head Attention](#single-head-attention)
|
||||
- [Multi Head Attention](#multi-head-attention)
|
||||
- [The Transformer Block](#the-transformer-block)
|
||||
- [The Full Transformer](#the-full-transformer)
|
||||
- [Step 3: Pretraining the Base Model](#step-3-pretraining-the-base-model)
|
||||
- [Step 4: Generating Text](#step-4-generating-text)
|
||||
- [Step 5: Post-Training, Turning a Base Model Into an Assistant](#step-5-post-training-turning-a-base-model-into-an-assistant)
|
||||
- [SFT (Supervised Fine-Tuning)](#sft-supervised-fine-tuning)
|
||||
- [The Reward Model](#the-reward-model)
|
||||
- [DPO, ORPO and KTO](#dpo-orpo-and-kto)
|
||||
## 目录
|
||||
- [适用读者](#who-this-is-for)
|
||||
- [前置要求与训练时间](#prerequisites-and-training-time)
|
||||
- [环境搭建](#setup)
|
||||
- [代码结构](#code-structure)
|
||||
- [步骤 1:准备数据](#step-1-preparing-the-data)
|
||||
- [步骤 2:由小组件构建模型](#step-2-the-model-built-from-small-pieces)
|
||||
- [多层感知机(MLP)](#multi-layer-perceptron-mlp)
|
||||
- [单头注意力(Single Head Attention)](#single-head-attention)
|
||||
- [多头注意力(Multi Head Attention)](#multi-head-attention)
|
||||
- [Transformer 块](#the-transformer-block)
|
||||
- [完整 Transformer](#the-full-transformer)
|
||||
- [步骤 3:预训练基座模型](#step-3-pretraining-the-base-model)
|
||||
- [步骤 4:生成文本](#step-4-generating-text)
|
||||
- [步骤 5:后训练——将基座模型变为助手](#step-5-post-training-turning-a-base-model-into-an-assistant)
|
||||
- [SFT(Supervised Fine-Tuning,监督微调)](#sft-supervised-fine-tuning)
|
||||
- [奖励模型(Reward Model)](#the-reward-model)
|
||||
- [DPO、ORPO 和 KTO](#dpo-orpo-and-kto)
|
||||
- [PPO](#ppo)
|
||||
- [GRPO / RLVR](#grpo--rlvr)
|
||||
- [Step 6: Evaluation](#step-6-evaluation)
|
||||
- [Step 7: Talking to the Model](#step-7-talking-to-the-model)
|
||||
- [The Streamlit Control Panel](#the-streamlit-control-panel)
|
||||
- [The Documentation Site](#the-documentation-site)
|
||||
- [Run the Whole Thing](#run-the-whole-thing)
|
||||
- [What's Next](#whats-next)
|
||||
- [步骤 6:评估](#step-6-evaluation)
|
||||
- [步骤 7:与模型对话](#step-7-talking-to-the-model)
|
||||
- [Streamlit 控制面板](#the-streamlit-control-panel)
|
||||
- [文档站点](#the-documentation-site)
|
||||
- [运行完整流程](#run-the-whole-thing)
|
||||
- [下一步](#whats-next)
|
||||
|
||||
## Who this is for
|
||||
## 适用读者
|
||||
|
||||
I tried to write this so one page works for very different readers:
|
||||
我尽量让这一页能同时满足不同背景的读者:
|
||||
|
||||
- If you are a **student**, read top to bottom. Every block of code comes after a plain explanation of what it does and why, and most blocks are followed by the output you should expect.
|
||||
- If you are a **developer**, the commands and file paths are all here. You can copy, run, and read the referenced source files directly.
|
||||
- If you are a **researcher**, the post-training half is the interesting part: SFT, a Bradley-Terry reward model, PPO with GAE, DPO/ORPO/KTO, and GRPO, all from scratch on the same small Transformer, trained on real public datasets.
|
||||
- 如果你是 **学生**,请从上到下阅读。每段代码之前都有通俗说明——它做什么、为什么这样做;大多数代码块之后还附有你应该看到的输出。
|
||||
- 如果你是 **开发者**,所需的命令和文件路径都在这里。你可以直接复制、运行,并阅读引用的源文件。
|
||||
- 如果你是 **研究者**,后训练部分更值得一看:SFT、Bradley-Terry 奖励模型、带 GAE 的 PPO、DPO/ORPO/KTO,以及 GRPO——全部在同一小型 Transformer 上从零实现,并在真实公开数据集上训练。
|
||||
|
||||
Every diagram in this README is colored the same way, so the colors mean something:
|
||||
本 README 中的每张图都采用统一配色,颜色含义如下:
|
||||
|
||||
- green is raw data
|
||||
- teal is stored, tokenized data on disk
|
||||
- blue is a plain processing step
|
||||
- yellow is the model or a training step
|
||||
- orange is the reinforcement learning and reward parts
|
||||
- red is a loss
|
||||
- grey is a saved checkpoint
|
||||
- purple is the final output or evaluation
|
||||
- 绿色表示原始数据
|
||||
- 青色表示磁盘上已存储、已分词的数据
|
||||
- 蓝色表示普通处理步骤
|
||||
- 黄色表示模型或训练步骤
|
||||
- 橙色表示强化学习与奖励相关部分
|
||||
- 红色表示损失(loss)
|
||||
- 灰色表示已保存的检查点(checkpoint)
|
||||
- 紫色表示最终输出或评估
|
||||
|
||||
## Prerequisites and Training Time
|
||||
## 前置要求与训练时间
|
||||
|
||||
You need a basic understanding of object oriented programming, neural networks, and PyTorch. Below are some resources to help you get started:
|
||||
你需要具备面向对象编程、神经网络和 PyTorch 的基础知识。以下是一些入门资源:
|
||||
|
||||
| Topic | Video Link |
|
||||
| 主题 | 视频链接 |
|
||||
|---------------------|-----------------------------------------------------------|
|
||||
| OOP | [OOP Video](https://www.youtube.com/watch?v=Ej_02ICOIgs) |
|
||||
| Neural Network | [Neural Network Video](https://www.youtube.com/watch?v=Jy4wM2X21u0) |
|
||||
| Pytorch | [Pytorch Video](https://www.youtube.com/watch?v=V_xro1bcAuA) |
|
||||
| OOP(面向对象编程) | [OOP Video](https://www.youtube.com/watch?v=Ej_02ICOIgs) |
|
||||
| Neural Network(神经网络) | [Neural Network Video](https://www.youtube.com/watch?v=Jy4wM2X21u0) |
|
||||
| Pytorch | [Pytorch Video](https://www.youtube.com/watch?v=V_xro1bcAuA) |
|
||||
|
||||
You will need a GPU to train. A free Colab or Kaggle T4 is enough for the 13 million parameter model, but it will not fit a billion parameter model. Here is a rough guide:
|
||||
训练需要 GPU。免费的 Colab 或 Kaggle T4 足以训练 1300 万参数模型,但无法容纳十亿参数模型。以下是粗略参考:
|
||||
|
||||
| GPU Name | Memory | 2B LLM Training | 13M LLM Training | Max Practical LLM Size (Training) |
|
||||
| GPU 名称 | 显存 | 20 亿参数 LLM 训练 | 1300 万参数 LLM 训练 | 最大实用 LLM 规模(训练) |
|
||||
|--------------------------|--------|-----------------|------------------|-----------------------------------|
|
||||
| NVIDIA A100 | 40 GB | ✔ | ✔ | ~6B to 8B |
|
||||
| NVIDIA V100 | 16 GB | ✘ | ✔ | ~2B |
|
||||
| NVIDIA RTX 4090 | 24 GB | ✔ | ✔ | ~4B |
|
||||
| NVIDIA RTX 5090 | 32 GB | ✔ | ✔ | 13M verified, larger configs TBD |
|
||||
| NVIDIA RTX 3090 | 24 GB | ✔ | ✔ | ~3.5B to 4B |
|
||||
| NVIDIA RTX 4080 | 16 GB | ✘ | ✔ | ~2B |
|
||||
| NVIDIA RTX 4060 | 8 GB | ✘ | ✔ | ~1B |
|
||||
| Tesla T4 | 16 GB | ✘ | ✔ | ~1.5B to 2B |
|
||||
| NVIDIA A100 | 40 GB | ✔ | ✔ | ~6B 至 8B |
|
||||
| NVIDIA V100 | 16 GB | ✘ | ✔ | ~2B |
|
||||
| NVIDIA RTX 4090 | 24 GB | ✔ | ✔ | ~4B |
|
||||
| NVIDIA RTX 5090 | 32 GB | ✔ | ✔ | 1300 万参数已验证,更大配置待定 |
|
||||
| NVIDIA RTX 3090 | 24 GB | ✔ | ✔ | ~3.5B 至 4B |
|
||||
| NVIDIA RTX 4080 | 16 GB | ✘ | ✔ | ~2B |
|
||||
| NVIDIA RTX 4060 | 8 GB | ✘ | ✔ | ~1B |
|
||||
| Tesla T4 | 16 GB | ✘ | ✔ | ~1.5B 至 2B |
|
||||
|
||||
If a large config runs out of memory, the pretraining script has opt-in flags (`--amp`, `--grad-checkpointing`, `--grad-accum`) that bring the memory down a lot. More on those later.
|
||||
若大型配置显存不足,预训练脚本提供可选标志(`--amp`、`--grad-checkpointing`、`--grad-accum`),可显著降低显存占用。后续会详细介绍。
|
||||
|
||||
## Setup
|
||||
## 环境搭建
|
||||
|
||||
Clone the repository and install it in editable mode. The editable install puts `config`, `src`, `data_loader`, and `ui` on your import path, so you do not need to set `PYTHONPATH` by hand anymore:
|
||||
克隆仓库并以可编辑模式安装。可编辑安装会把 `config`、`src`、`data_loader` 和 `ui` 加入导入路径,因此你不再需要手动设置 `PYTHONPATH`:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git
|
||||
@@ -117,7 +123,7 @@ cd train-llm-from-scratch
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
There are optional extras for the parts you want:
|
||||
可按需安装可选扩展:
|
||||
|
||||
```bash
|
||||
pip install -e ".[train]" # datasets + wandb, for downloading data and logging
|
||||
@@ -126,14 +132,14 @@ pip install -e ".[docs]" # mkdocs, for the documentation site
|
||||
pip install -e ".[all]" # everything
|
||||
```
|
||||
|
||||
There are two config systems, and it helps to know which is which from the start:
|
||||
项目有两套配置系统,一开始就分清它们会很有帮助:
|
||||
|
||||
- `config/config.py` is the original, simple config for the legacy pretraining script `scripts/train_transformer.py`. It is plain Python constants.
|
||||
- `config/post_training_config.py` plus the JSON files in `configs/` drive everything else (pretraining the bigger base, SFT, reward, DPO, PPO, GRPO). You edit a small JSON file per stage, and any field can also be overridden on the command line, for example `--lr 2e-5 --batch_size 16`.
|
||||
- `config/config.py` 是面向旧版预训练脚本 `scripts/train_transformer.py` 的原始简易配置,采用普通 Python 常量。
|
||||
- `config/post_training_config.py` 以及 `configs/` 中的 JSON 文件驱动其余所有流程(更大基座的预训练、SFT、奖励模型、DPO、PPO、GRPO)。每个阶段编辑一个小 JSON 文件即可,任何字段也可在命令行覆盖,例如 `--lr 2e-5 --batch_size 16`。
|
||||
|
||||
For fast checks there is a tiny `configs/smoke/` variant of every stage that shrinks the model so a full run finishes in seconds on a CPU or a single GPU.
|
||||
为快速验证,每个阶段都提供精简版 `configs/smoke/`,会缩小模型规模,使完整运行在 CPU 或单张 GPU 上数秒内完成。
|
||||
|
||||
## Code Structure
|
||||
## 代码结构
|
||||
|
||||
```bash
|
||||
train-llm-from-scratch/
|
||||
@@ -157,40 +163,40 @@ train-llm-from-scratch/
|
||||
└── pyproject.toml # pip install -e .
|
||||
```
|
||||
|
||||
## Step 1: Preparing the Data
|
||||
## 第 1 步:准备数据
|
||||
|
||||
A model only ever sees integers. So the first job is always the same: take text, turn it into token ids, and store those ids on disk in a format that is fast to read during training. We do this four times, once for each kind of training we will do later.
|
||||
模型看到的只有整数。因此第一步永远是相同的:把文本转成 token id,并以训练时便于快速读取的格式存到磁盘。我们会做四次,对应后面要进行的几类训练。
|
||||
|
||||

|
||||

|
||||
|
||||
The four streams are:
|
||||
四条数据流分别是:
|
||||
|
||||
1. **Pretraining text** from [The Pile](https://huggingface.co/datasets/monology/pile-uncopyrighted), stored as a flat array of token ids in an HDF5 file.
|
||||
2. **Instruction data** (Alpaca, Dolly, GSM8K) for SFT, packed into fixed length rows with a mask that says which tokens are the assistant's answer.
|
||||
3. **Preference pairs** (Anthropic HH-RLHF, UltraFeedback) for the reward model and DPO, stored as `{prompt, chosen, rejected}`.
|
||||
4. **RL prompts** (GSM8K and a small arithmetic warm-up) for PPO and GRPO, stored as `{prompt, gold}`.
|
||||
1. 来自 [The Pile](https://huggingface.co/datasets/monology/pile-uncopyrighted), 的 **预训练文本**,以 token id 扁平数组形式存储在 HDF5 文件中。
|
||||
2. 用于 SFT 的 **指令数据**(Alpaca、Dolly、GSM8K),打包成定长行,并附带掩码标明哪些 token 属于助手的回答。
|
||||
3. 用于奖励模型和 DPO 的 **偏好对**(Anthropic HH-RLHF、UltraFeedback),存储为 `{prompt, chosen, rejected}`。
|
||||
4. 用于 PPO 和 GRPO 的 **RL 提示**(GSM8K 及小规模算术热身),存储为 `{prompt, gold}`。
|
||||
|
||||
### Tokenization
|
||||
### 分词(Tokenization)
|
||||
|
||||
We use the `r50k_base` tokenizer from OpenAI's `tiktoken`, the same one GPT-3 used. Text becomes a list of integers, and we append a special `<|endoftext|>` token (id 50256) at the end of every document so the model learns where one piece of text stops and the next begins.
|
||||
我们使用 OpenAI 的 `tiktoken` 中的 `r50k_base` 分词器,与 GPT-3 所用相同。文本会变成整数列表,并在每篇文档末尾追加特殊的 `<|endoftext|>` token(id 50256),以便模型学会一段文本何处结束、下一段何处开始。
|
||||
|
||||

|
||||

|
||||
|
||||
For the legacy path, download a slice of The Pile and tokenize it into HDF5:
|
||||
旧路径:下载 The Pile 的一个切片并分词为 HDF5:
|
||||
|
||||
```bash
|
||||
python scripts/data_download.py # downloads the validation file + 1 training shard
|
||||
python scripts/data_preprocess.py # tokenizes to data/train/pile_train.h5 and data/val/pile_dev.h5
|
||||
```
|
||||
|
||||
The newer, faster path streams and batch-encodes the same data straight into a flat token array:
|
||||
更新、更快的路径会流式读取并批量编码同一数据,直接写入扁平 token 数组:
|
||||
|
||||
```bash
|
||||
python scripts/prepare_pretrain_data.py --split val --out data/pile_dev.h5
|
||||
python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out data/pile_train.h5
|
||||
```
|
||||
|
||||
Once tokenized, the data is just a long line of integers. Here is a real peek at the validation file I prepared for this README (8.76 million tokens), the first ten ids, and what they decode back to:
|
||||
分词完成后,数据就是一长串整数。下面是为此 README 准备的验证文件的真实片段(8.76 million token):前十个 id,以及它们解码回的内容:
|
||||
|
||||
```python
|
||||
#### OUTPUT ####
|
||||
@@ -201,11 +207,11 @@ decoded back to text:
|
||||
performance in patients with refractory focal epilepsy ...'
|
||||
```
|
||||
|
||||
That is the whole idea of tokenization in one output: text in, a flat array of integers out, and the integers decode straight back to the original words.
|
||||
这就是分词的核心:输入文本,输出整数扁平数组,而这些整数可直接解码回原文。
|
||||
|
||||
### The chat format and loss mask
|
||||
### 对话格式与损失掩码(loss mask)
|
||||
|
||||
For everything after pretraining the model has to know who is talking. The `r50k_base` tokenizer has only one special token, so instead of inventing new ones we use plain text role markers that the model simply learns during SFT. A single turn looks like this (see `src/post_training/chat_template.py`):
|
||||
预训练之后,模型必须知道谁在说话。`r50k_base` 分词器只有一个特殊 token,因此我们不再发明新 token,而是用纯文本角色标记,让模型在 SFT 中自行学习。单轮对话如下(见 `src/post_training/chat_template.py`):
|
||||
|
||||
```
|
||||
<|user|>
|
||||
@@ -213,13 +219,13 @@ For everything after pretraining the model has to know who is talking. The `r50k
|
||||
{assistant content}<|endoftext|>
|
||||
```
|
||||
|
||||
For math and reasoning we ask the assistant to show its work in a fixed structure, because the reinforcement learning reward later checks the number inside the answer tags:
|
||||
数学与推理任务中,我们要求助手按固定结构展示推导过程,因为后续强化学习奖励会检查答案标签内的数字:
|
||||
|
||||
```
|
||||
<think>step by step reasoning ...</think><answer>42</answer>
|
||||
```
|
||||
|
||||
The important trick is the **loss mask**. When we encode a conversation we also build a 0/1 mask that is 1 only on the assistant tokens (and the `<|endoftext|>` that ends the turn). That way SFT trains the model to write answers, not to parrot the prompt back. Here is the exact code that builds the ids and the aligned mask:
|
||||
关键技巧是 **loss mask**。编码对话时,我们同时构建 0/1 掩码:仅在助手 token(以及结束该轮的 `<|endoftext|>`)上为 1。这样 SFT 训练的是写回答,而不是复述提示。下面是构建 id 与对齐掩码的代码:
|
||||
|
||||
```python
|
||||
def encode_chat(messages, add_generation_prompt=False):
|
||||
@@ -241,7 +247,7 @@ def encode_chat(messages, add_generation_prompt=False):
|
||||
return ids, mask
|
||||
```
|
||||
|
||||
Here is a real rendered conversation and the verifier reward in action, printed from this repo:
|
||||
下面是本仓库打印的真实渲染对话与验证器奖励示例:
|
||||
|
||||
```python
|
||||
#### OUTPUT ####
|
||||
@@ -255,7 +261,7 @@ reward_gsm8k("<answer>42</answer>", 42.0) -> 1.2 # correct AND well format
|
||||
reward_gsm8k("<answer>7</answer>", 42.0) -> 0.2 # wrong, but it used the format
|
||||
```
|
||||
|
||||
And here is one real packed SFT row, showing how only the assistant tokens are trained (the mask is 1 on 48 of the 512 tokens in this row):
|
||||
下面是一条真实的打包 SFT 行,展示仅助手 token 参与训练(该行 512 个 token 中,掩码为 1 的有 48 个):
|
||||
|
||||
```python
|
||||
#### OUTPUT ####
|
||||
@@ -268,7 +274,7 @@ row 0 decoded:
|
||||
The Boston Marathon is the world's oldest annual marathon, beginning on April 19th 1897.
|
||||
```
|
||||
|
||||
The data prep scripts for these stages are:
|
||||
各阶段的数据准备脚本如下:
|
||||
|
||||
```bash
|
||||
python scripts/prepare_sft_data.py # Alpaca + Dolly + GSM8K -> sft_packed.h5
|
||||
@@ -276,13 +282,13 @@ python scripts/prepare_preference_data.py # HH-RLHF + UltraFeedback -> prefere
|
||||
python scripts/prepare_rl_prompts.py # GSM8K + arithmetic -> rl_prompts.jsonl
|
||||
```
|
||||
|
||||
## Step 2: The Model, Built From Small Pieces
|
||||
## 第 2 步:由小组件组装模型
|
||||
|
||||
A Transformer looks scary as one block of code, so we build it from four small pieces and then stack them. Each piece is a tiny `nn.Module`. We start at the bottom.
|
||||
Transformer 若写成一整块代码会显得吓人,因此我们拆成四个小组件再堆叠。每个组件都是一个小型 `nn.Module`。我们从最底层开始。
|
||||
|
||||
### Multi Layer Perceptron (MLP)
|
||||
### 多层感知机(Multi Layer Perceptron,MLP)
|
||||
|
||||
The MLP is the part of each block that does the per-token "thinking". It takes each token vector, expands it to four times its size, applies a `ReLU`, and squeezes it back down. The expansion gives the layer room to mix features before projecting back.
|
||||
MLP 是每个块中负责逐 token「思考」的部分。它接收每个 token 向量,先扩展到四倍维度,施加 `ReLU`,再压缩回原尺寸。扩展为投影回去之前混合特征留出空间。
|
||||
|
||||

|
||||
|
||||
@@ -301,13 +307,13 @@ class MLP(nn.Module):
|
||||
return x
|
||||
```
|
||||
|
||||
The `__init__` sets up the two linear layers and the activation. The `forward` runs them in order. Input and output shapes are the same, `(B, T, n_embed)`, so blocks can be stacked without any reshaping. The code is in `src/models/mlp.py`.
|
||||
`__init__` 配置两个线性层与激活函数。`forward` 按顺序执行它们。输入与输出形状相同,为 `(B, T, n_embed)`,因此块可直接堆叠而无需变形。代码位于 `src/models/mlp.py`。
|
||||
|
||||
### Single Head Attention
|
||||
### 单头注意力(Single Head Attention)
|
||||
|
||||
Attention is the part that lets a token look at other tokens. Each head builds three views of the input: a query (what am I looking for), a key (what do I contain), and a value (what I will pass on if chosen). We score every query against every key, scale the scores, hide the future with a causal mask, turn the scores into weights with a softmax, and take a weighted sum of the values.
|
||||
注意力让某个 token 能关注其他 token。每个头从输入构建三种视图:query(我在找什么)、key(我包含什么)、value(若被选中要传递什么)。我们对每个 query 与每个 key 打分,缩放分数,用因果掩码遮蔽未来,经 softmax 将分数变为权重,再对 value 做加权求和。
|
||||
|
||||

|
||||

|
||||
|
||||
```python
|
||||
class Head(nn.Module):
|
||||
@@ -333,13 +339,13 @@ class Head(nn.Module):
|
||||
return out
|
||||
```
|
||||
|
||||
The causal mask is what makes this a language model: position `t` can only attend to positions `0..t`, never to the future it is trying to predict. The code is in `src/models/attention.py`.
|
||||
因果掩码(causal mask)正是使其成为语言模型的关键:位置 `t` 只能关注位置 `0..t`,而绝不能关注它试图预测的未来。代码位于 `src/models/attention.py`。
|
||||
|
||||
### Multi Head Attention
|
||||
### 多头注意力(Multi Head Attention)
|
||||
|
||||
One head learns one kind of relationship. We want many, running in parallel, so the model can track several patterns at once (for example, a pronoun and the noun it refers to). We run `n_head` heads, concatenate their outputs, and pass the result through one more linear layer.
|
||||
单个注意力头学习一种关系。我们希望并行运行多个头,以便模型能同时追踪多种模式(例如,代词及其所指的名词)。我们运行 `n_head` 个注意力头,拼接它们的输出,并将结果再通过一个线性层。
|
||||
|
||||

|
||||

|
||||
|
||||
```python
|
||||
class MultiHeadAttention(nn.Module):
|
||||
@@ -356,13 +362,13 @@ class MultiHeadAttention(nn.Module):
|
||||
return x
|
||||
```
|
||||
|
||||
Each head works in a smaller subspace of size `n_embed // n_head`, so the concatenation lands right back at `n_embed`. The final projection lets the heads talk to each other.
|
||||
每个注意力头在大小为 `n_embed // n_head` 的较小子空间中工作,因此拼接后的结果恰好回到 `n_embed`。最终的投影层让各个头能够相互通信。
|
||||
|
||||
### The Transformer Block
|
||||
### Transformer 块
|
||||
|
||||
Now we combine attention and the MLP into one block. The block uses pre-norm residual connections: we normalize, run a sub-layer, and add the result back to the input. The "add back" (the residual) is what lets gradients flow through a deep stack without vanishing.
|
||||
现在我们将注意力与 MLP 合并为一个块。该块使用 pre-norm 残差连接:我们先做归一化,运行一个子层,再将结果加回输入。这种「加回去」(即残差)使梯度能够流经深层堆叠而不至于消失。
|
||||
|
||||

|
||||

|
||||
|
||||
```python
|
||||
class Block(nn.Module):
|
||||
@@ -379,13 +385,13 @@ class Block(nn.Module):
|
||||
return x
|
||||
```
|
||||
|
||||
Read `x = x + self.attn(self.ln1(x))` as "look at the other tokens, then add what you learned back onto yourself". The MLP line is the same idea for the per-token thinking. The code is in `src/models/transformer_block.py`.
|
||||
将 `x = x + self.attn(self.ln1(x))` 理解为「关注其他 token,然后将所学内容加回到自身」。MLP 那一行对逐 token 的思考采用同样的思路。代码位于 `src/models/transformer_block.py`。
|
||||
|
||||
### The Full Transformer
|
||||
### 完整 Transformer
|
||||
|
||||
Finally we wrap everything. Token ids become vectors through an embedding table, we add a position embedding so the model knows token order, we run the stack of blocks, normalize one last time, and project to vocabulary-sized scores called logits. If we pass targets, the model also returns the cross-entropy loss.
|
||||
最后我们将所有部分封装在一起。Token id 通过嵌入表变为向量,我们添加位置嵌入以使模型了解 token 顺序,运行块堆叠,最后再归一化一次,并投影到词汇表大小的分数(称为 logits)。如果传入 targets,模型还会返回交叉熵(cross-entropy)损失。
|
||||
|
||||

|
||||

|
||||
|
||||
```python
|
||||
class Transformer(nn.Module):
|
||||
@@ -413,9 +419,9 @@ class Transformer(nn.Module):
|
||||
return logits, loss
|
||||
```
|
||||
|
||||
One small detail worth pointing out: we use `.reshape` and not `.view` on the targets. The target batch is a non-contiguous slice of the data, and `.view` refuses to work on that on CPU. `.reshape` handles both cases and is identical in every other way. The full model, including `forward_hidden` (which the reward and value heads reuse later) and `generate`, lives in `src/models/transformer.py`.
|
||||
有一个小细节值得指出:我们在 targets 上使用 `.reshape` 而非 `.view`。target 批次是数据的一个非连续切片,而 `.view` 在 CPU 上无法处理这种情况。`.reshape` 两种情况都能处理,其他方面完全相同。完整模型(包括 `forward_hidden`——奖励头与价值头稍后也会复用——以及 `generate`)位于 `src/models/transformer.py`。
|
||||
|
||||
When we build the model it prints its parameter count. Here are the three sizes used in this repo:
|
||||
构建模型时会打印其参数量。本仓库使用的三种规模如下:
|
||||
|
||||
```python
|
||||
#### OUTPUT ####
|
||||
@@ -424,13 +430,13 @@ this tutorial's base (n_embed=512, n_head=8, n_blocks=8): 77,031,552 params
|
||||
post-training default (n_embed=1024, n_head=16, n_blocks=24): 406,359,168 params
|
||||
```
|
||||
|
||||
## Step 3: Pretraining the Base Model
|
||||
## 第 3 步:预训练基础模型
|
||||
|
||||
Pretraining is the long pole. We read random windows of tokens, ask the model to predict the next token at every position, measure how wrong it was with cross-entropy, and nudge the weights. We repeat that a few thousand times.
|
||||
预训练是最耗时的环节。我们读取随机的 token 窗口,要求模型在每个位置预测下一个 token,用交叉熵衡量其错误程度,并微调权重。我们将此重复数千次。
|
||||
|
||||

|
||||

|
||||
|
||||
The simplest version is the original `scripts/train_transformer.py`, which reads `config/config.py` and trains on one GPU. To train the 13 million parameter model, set these values in `config/config.py`:
|
||||
最简单的版本是原始的 `scripts/train_transformer.py`,它读取 `config/config.py` 并在单块 GPU 上训练。要训练 1300 万参数的模型,请在 `config/config.py` 中设置以下值:
|
||||
|
||||
```python
|
||||
VOCAB_SIZE = 50304
|
||||
@@ -446,20 +452,20 @@ then run:
|
||||
python scripts/train_transformer.py
|
||||
```
|
||||
|
||||
For long runs you can save periodic checkpoints and resume after an interruption:
|
||||
对于长时间训练,可以定期保存检查点并在中断后恢复:
|
||||
|
||||
```bash
|
||||
python scripts/train_transformer.py --checkpoint-every 1000 --keep-last 3
|
||||
python scripts/train_transformer.py --resume latest
|
||||
```
|
||||
|
||||
If a bigger config does not fit in memory, turn on the opt-in memory savers (all off by default, so default behavior never changes):
|
||||
如果更大的配置无法放入内存,可开启可选的内存节省选项(默认全部关闭,因此默认行为不会改变):
|
||||
|
||||
```bash
|
||||
python scripts/train_transformer.py --amp --grad-checkpointing --grad-accum 8
|
||||
```
|
||||
|
||||
The bigger, modern path is `scripts/pretrain_base.py`. It is the same recipe with the things you need to train a mid-size base: DistributedDataParallel across GPUs, bf16 autocast, gradient accumulation, a cosine learning-rate schedule with warmup, and periodic checkpoints. One GPU or many, same command shape:
|
||||
更大规模的现代路径是 `scripts/pretrain_base.py`。配方相同,但加入了训练中等规模基础模型所需的要素:跨 GPU 的 DistributedDataParallel、bf16 autocast、梯度累积、带 warmup 的余弦学习率调度,以及定期保存检查点。单 GPU 或多 GPU,命令形式相同:
|
||||
|
||||
```bash
|
||||
# one GPU
|
||||
@@ -468,7 +474,7 @@ python scripts/pretrain_base.py
|
||||
torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py
|
||||
```
|
||||
|
||||
The core of the loop is small. Each step pulls a batch, runs the forward pass under bf16, scales the loss for gradient accumulation, backpropagates, clips the gradient, and steps the optimizer:
|
||||
循环的核心很简单。每一步拉取一个批次,在 bf16 下运行前向传播,为梯度累积缩放损失,反向传播,裁剪梯度,并执行优化器步进:
|
||||
|
||||
```python
|
||||
for micro in range(cfg.grad_accum):
|
||||
@@ -481,7 +487,7 @@ torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
|
||||
optimizer.step()
|
||||
```
|
||||
|
||||
While it trains it prints the step, the loss, the learning rate, and the throughput, and at eval time it also prints the dev loss and peak GPU memory:
|
||||
训练过程中会打印步数、损失、学习率和吞吐量;在评估时还会打印 dev 损失和峰值 GPU 内存:
|
||||
|
||||
```
|
||||
#### OUTPUT ####
|
||||
@@ -502,17 +508,17 @@ step 1900 | loss 3.7725 | lr 6.36e-05 | 151,488 tok/s
|
||||
Done. Final checkpoint -> /ephemeral/ckpts/base_pretrained.pt
|
||||
```
|
||||
|
||||
### The loss curve
|
||||
### 损失曲线
|
||||
|
||||
Here is the real training and dev loss for the 77 million parameter base I trained for this README on 2x L40 GPUs. The loss starts near `ln(vocab_size)`, which is about 10.8 (the loss of a model that guesses uniformly), and drops as the model learns the statistics of the text:
|
||||
以下是我为本 README 在 2×L40 GPU 上训练的 7700 万参数基础模型的真实训练损失与 dev 损失。损失起始于接近 `ln(vocab_size)`,约为 10.8(均匀猜测的模型的损失),并随模型学习文本统计规律而下降:
|
||||
|
||||

|
||||

|
||||
|
||||
This run started at a loss of **11.14** (just above the uniform-guess line at `ln(50304) = 10.83`) and came down to about **3.73 train / 3.76 dev** after 2000 steps, training at roughly **130,000 to 150,000 tokens per second** across the two L40s. The curve going down is the whole story of pretraining: the model is slowly compressing the patterns of language into its weights.
|
||||
本次运行初始损失为 **11.14**(略高于均匀猜测基线 `ln(50304) = 10.83`),在 2000 步后降至约 **3.73 训练 / 3.76 验证**,在两张 L40 上训练速度约为每秒 **130,000 至 150,000** 个 token。损失曲线下降,正是预训练(pretraining)的全部故事:模型正慢慢将语言模式压缩进其权重之中。
|
||||
|
||||
## Step 4: Generating Text
|
||||
## 步骤 4:生成文本
|
||||
|
||||
A trained model predicts a distribution over the next token. To generate, we sample one token from that distribution, append it, and feed the longer sequence back in. We repeat until we have enough tokens.
|
||||
训练好的模型会预测下一个 token 的概率分布。要生成文本,我们从该分布中采样一个 token,将其追加到序列末尾,再把更长的序列送回模型。如此重复,直到得到足够多的 token。
|
||||
|
||||
```python
|
||||
def generate(self, idx, max_new_tokens):
|
||||
@@ -526,25 +532,25 @@ def generate(self, idx, max_new_tokens):
|
||||
return idx
|
||||
```
|
||||
|
||||
Run it from a saved checkpoint:
|
||||
从已保存的检查点运行:
|
||||
|
||||
```bash
|
||||
python scripts/generate_text.py --model_path models/transformer_B.pt --input_text "The" --max_new_tokens 100
|
||||
```
|
||||
|
||||
The 13 million parameter model already produces real words and roughly correct grammar, which is the encouraging part of starting small.
|
||||
这个 1300 万参数的模型已经能产出真实词汇和大体正确的语法,这正是从小规模起步令人鼓舞之处。
|
||||
|
||||
## Step 5: Post-Training, Turning a Base Model Into an Assistant
|
||||
## 步骤 5:后训练——将基座模型变成助手
|
||||
|
||||
A base model can continue text, but it cannot follow instructions or reason on purpose. That takes post-training. The good news is that the model never changes. We reuse the exact same Transformer backbone and only change two things at each stage: the data, and the loss.
|
||||
基座模型可以续写文本,却无法按指令行事或有意识地推理。这需要后训练(post-training)。好消息是模型本身始终不变。我们复用完全相同的 Transformer 骨干网络,每个阶段只改两件事:数据和损失函数。
|
||||
|
||||

|
||||

|
||||
|
||||
The one design idea that makes all of this fit in a small repo is **wrap, do not rewrite**. The educational `Transformer` gains a single extra method, `forward_hidden`, which returns the hidden states right before `lm_head`. The reward head, the value head, and all the log-probability math compose around that one method. Nothing in `src/models/` had to be rewritten.
|
||||
让所有这一切都能装进一个小仓库的关键设计思路是:**包装,而非重写**。教学用的 `Transformer` 只增加一个额外方法 `forward_hidden`,该方法返回 `lm_head` 之前的隐状态。奖励头、价值头以及所有对数概率计算,都围绕这一个方法组合而成。`src/models/` 中无需重写任何内容。
|
||||
|
||||
### SFT (Supervised Fine-Tuning)
|
||||
### SFT(监督微调,Supervised Fine-Tuning)
|
||||
|
||||
SFT teaches the base model to answer in the chat format. It is still next-token prediction, with one change: the loss is only counted on the assistant tokens, using the mask we built back in Step 1.
|
||||
SFT 教会基座模型以对话格式作答。它仍是下一个 token 预测,唯有一点不同:损失只对助手 token 计算,使用的是步骤 1 中构建的掩码。
|
||||
|
||||

|
||||
|
||||
@@ -560,20 +566,20 @@ def sft_loss(logits, tokens, loss_mask):
|
||||
return ce.sum() / mask.sum().clamp(min=1.0) # average over assistant tokens only
|
||||
```
|
||||
|
||||
Prepare the data and train:
|
||||
准备数据并训练:
|
||||
|
||||
```bash
|
||||
python scripts/prepare_sft_data.py --context_length 1024
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
|
||||
```
|
||||
|
||||
The loss code is in `src/post_training/sft.py`, the trainer in `scripts/train_sft.py`.
|
||||
损失函数代码在 `src/post_training/sft.py`,训练器在 `scripts/train_sft.py`。
|
||||
|
||||
### The Reward Model
|
||||
### 奖励模型
|
||||
|
||||
To do reinforcement learning we need a number that says how good an answer is. One way to get it is to train a reward model on human preference pairs. We put a small linear head on top of the SFT backbone that reads one scalar off the last real token, and we train it with the Bradley-Terry loss so the chosen answer always scores higher than the rejected one.
|
||||
要做强化学习,就需要一个数字来衡量答案好坏。获取它的方法之一,是在人类偏好对上训练奖励模型。我们在 SFT 骨干之上加一个小型线性头,从最后一个真实 token 读出一个标量,并用 Bradley-Terry 损失训练,使被选答案的得分始终高于被拒答案。
|
||||
|
||||

|
||||

|
||||
|
||||
```python
|
||||
def bradley_terry_loss(chosen_rewards, rejected_rewards):
|
||||
@@ -586,7 +592,7 @@ python scripts/prepare_preference_data.py --source both
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_reward.py
|
||||
```
|
||||
|
||||
The headline metric is preference accuracy, the fraction of held-out pairs where the model scores the chosen answer higher. In this run on 7974 real preference pairs it reached **0.574** (above the 0.5 chance line); with a larger model and more data this climbs toward 0.65 to 0.75.
|
||||
核心指标是偏好准确率(preference accuracy),即留出配对中被模型判为更高分的被选答案所占比例。本次在 7974 对真实偏好对上运行,达到 **0.574**(高于 0.5 的随机基线);若使用更大模型和更多数据,该值可升至 0.65 至 0.75。
|
||||
|
||||
```
|
||||
#### OUTPUT ####
|
||||
@@ -596,11 +602,11 @@ Reward model from sft.pt | 7974 pairs | total_steps=996
|
||||
Done RM. test_acc 0.574 margin 0.063 -> reward.pt
|
||||
```
|
||||
|
||||
The code is in `src/post_training/reward_model.py` and `src/post_training/reward_train.py`.
|
||||
代码位于 `src/post_training/reward_model.py` 与 `src/post_training/reward_train.py`。
|
||||
|
||||
### DPO, ORPO and KTO
|
||||
### DPO、ORPO 与 KTO
|
||||
|
||||
DPO skips the reward model and the RL loop entirely. It works directly on preference pairs by comparing how much more likely the policy makes the chosen answer (relative to a frozen reference copy of the SFT model) than the rejected one.
|
||||
DPO 完全跳过奖励模型与 RL 循环。它直接作用于偏好对,比较策略使被选答案(相对于冻结的 SFT 参考副本)比被拒答案更可能的程度。
|
||||
|
||||

|
||||
|
||||
@@ -620,15 +626,15 @@ torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo
|
||||
# --loss_type kto works from an unpaired desirable / undesirable signal
|
||||
```
|
||||
|
||||
In this run, DPO reached an implicit-reward accuracy of **0.574** on the held-out pairs (the fraction where the policy prefers the chosen response more than the frozen reference does). All three objectives are in `src/post_training/dpo.py`.
|
||||
本次运行中,DPO 在留出配对上达到 **0.574** 的隐式奖励准确率(implicit-reward accuracy),即策略比冻结参考更偏好被选回复的比例。三种目标函数均在 `src/post_training/dpo.py` 中。
|
||||
|
||||
### PPO
|
||||
|
||||
PPO is the classic RLHF loop. For each prompt the model writes an answer (a rollout), we score it (with the reward model or with a GSM8K answer checker), add a small per-token penalty for drifting away from the reference model, estimate how good each token was with GAE, and take a few clipped gradient steps.
|
||||
PPO 是经典的 RLHF 循环。对每个提示,模型生成一个答案(rollout),我们对其评分(用奖励模型或 GSM8K 答案校验器),对偏离参考模型施加按 token 计的小惩罚,用 GAE 估计每个 token 的好坏,再执行若干步裁剪梯度更新。
|
||||
|
||||

|
||||
|
||||
The two load-bearing pieces, the advantage estimate and the clipped policy loss, are short:
|
||||
两个关键组件——优势估计与裁剪策略损失——都很简短:
|
||||
|
||||
```python
|
||||
def ppo_policy_loss(new_logp, old_logp, advantages, mask, clip=0.2):
|
||||
@@ -645,11 +651,11 @@ torchrun --standalone --nproc_per_node=2 scripts/train_ppo.py --reward_source ve
|
||||
# --reward_source rm to use the trained reward model instead of the answer checker
|
||||
```
|
||||
|
||||
The actor-critic shares the backbone through a small value head (`src/post_training/value_head.py`), and the GAE, clipped policy loss, and clipped value loss are in `src/post_training/ppo.py`.
|
||||
Actor-critic 通过小型价值头(`src/post_training/value_head.py`)共享骨干,GAE、裁剪策略损失与裁剪价值损失均在 `src/post_training/ppo.py` 中。
|
||||
|
||||
### GRPO / RLVR
|
||||
|
||||
GRPO is the 2025, DeepSeek-R1 style method. It throws away the value network. For each prompt it samples a whole group of answers, scores them with a verifiable reward (did the final number match the gold answer), and uses the group's own mean and standard deviation as the baseline. The advantage of an answer is just how much better it did than its group.
|
||||
GRPO 是 2025 年、DeepSeek-R1 风格的方法。它舍弃价值网络。对每个提示,它采样一整组答案,用可验证奖励(最终数字是否与标准答案一致)评分,并以该组自身的均值与标准差作为基线。某答案的优势,就是它比同组其他答案好多少。
|
||||
|
||||

|
||||
|
||||
@@ -666,11 +672,11 @@ def group_advantages(rewards, group_size, eps=1e-4):
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_grpo.py --group_size 8
|
||||
```
|
||||
|
||||
A short arithmetic curriculum runs first, so the model gets some non-zero reward to learn from before it faces full GSM8K. The group-relative advantage, the clipped surrogate, and the k3 KL penalty are in `src/post_training/grpo.py`.
|
||||
先运行一段简短算术课程,让模型在面对完整 GSM8K 之前能获得一些非零奖励来学习。组相对优势、裁剪替代目标与 k3 KL 惩罚均在 `src/post_training/grpo.py` 中。
|
||||
|
||||
## Step 6: Evaluation
|
||||
## 步骤 6:评估
|
||||
|
||||
The single number that ties all the stages together is greedy GSM8K accuracy. We give the model a math question, let it generate, pull the number out of the `<answer>` tags, and check it against the gold answer.
|
||||
贯穿所有阶段的单一指标是 greedy(贪心)GSM8K 准确率。我们给模型一道数学题,让它生成回答,从 `<answer>` 标签中提取数字,并与标准答案核对。
|
||||
|
||||

|
||||
|
||||
@@ -681,15 +687,15 @@ done
|
||||
python scripts/eval_post_training.py --table logs/table.jsonl
|
||||
```
|
||||
|
||||
This builds the headline comparison so you can track every stage on one axis, Base, SFT, DPO, PPO, and GRPO, all with the same greedy decoding and the same verifiable reward: parse the number inside the `<answer>` tags and check it against the gold answer. The same command runs on any checkpoint, so you fill in the table for your own run, and the bigger the base model and the more pretraining compute you give it, the higher these scores go.
|
||||
这会构建核心对比,让你能在同一坐标轴上追踪每个阶段:Base、SFT、DPO、PPO 和 GRPO,全部使用相同的 greedy decoding(贪心解码)和相同的可验证奖励:解析 `<answer>` 标签内的数字并与标准答案核对。同一命令可在任意 checkpoint 上运行,因此你可以填入自己运行的表格;基础模型越大、预训练算力越多,这些分数就越高。
|
||||
|
||||
### What changes after SFT
|
||||
### SFT 之后的变化
|
||||
|
||||
The clearest effect of SFT is a change in behavior. The base model only knows how to continue text, so when you give it a question it just writes more text. After SFT the model has learned the chat format: it answers inside the `<think>...</think><answer>...</answer>` structure it was trained on, which is exactly the structure the reward model and the RL stages then optimize against. That learned format is the foundation every stage after it builds on. You can talk to any stage's checkpoint with `scripts/chat.py` and watch this directly, which is the next section.
|
||||
SFT 最明显的效果是行为变化。基础模型只知道如何续写文本,因此当你给它一个问题时,它只会继续写更多文字。SFT 之后,模型学会了对话格式:它会在训练时使用的 `<think>...</think><answer>...</answer>` 结构内作答,而这正是奖励模型和后续 RL 阶段所优化的结构。这种学到的格式是之后每个阶段所依赖的基础。你可以用 `scripts/chat.py` 与任意阶段的 checkpoint 对话,直接观察这一点——这正是下一节的内容。
|
||||
|
||||
## Step 7: Talking to the Model
|
||||
## 第 7 步:与模型对话
|
||||
|
||||
`scripts/chat.py` loads any checkpoint, reads the model dimensions from the checkpoint itself, and lets you talk to it. It applies the chat template for instruction models, or treats the base model as raw continuation.
|
||||
`scripts/chat.py` 可加载任意 checkpoint,从 checkpoint 自身读取模型维度,并让你与之对话。它对指令模型应用 chat template(对话模板),或将基础模型视为原始续写。
|
||||
|
||||

|
||||
|
||||
@@ -703,59 +709,59 @@ python scripts/chat.py --ckpt models/base_pretrained.pt --raw --prompt "Once upo
|
||||
python scripts/chat.py --ckpt models/sft.pt
|
||||
```
|
||||
|
||||
Generation reuses the same tested core as training and eval, so what you see in chat is exactly what the RL stages optimized. Sampling is controlled by `--temperature`, `--top_p`, `--top_k`, or `--greedy`. The code is in `src/post_training/inference.py`.
|
||||
生成复用与训练和评估相同的经过验证的核心,因此你在对话中看到的内容正是 RL 阶段所优化的。采样由 `--temperature`、`--top_p`、`--top_k` 或 `--greedy` 控制。代码位于 `src/post_training/inference.py`。
|
||||
|
||||
## The Streamlit Control Panel
|
||||
## Streamlit 控制面板
|
||||
|
||||
If you would rather click than type, there is a small control panel that can launch every stage, watch the loss live, run evaluation, and chat with a checkpoint:
|
||||
如果你更愿意点击而非输入命令,有一个小型控制面板可以启动每个阶段、实时查看 loss、运行评估,并与 checkpoint 对话:
|
||||
|
||||
```bash
|
||||
pip install -e ".[ui]"
|
||||
streamlit run ui/app.py
|
||||
```
|
||||
|
||||
It has one page per stage (Data, Pretrain, SFT, Reward, DPO, PPO, GRPO, Evaluate, Chat), each one a form over the same JSON configs you would edit by hand.
|
||||
每个阶段对应一页(Data、Pretrain、SFT、Reward、DPO、PPO、GRPO、Evaluate、Chat),每页都是基于你可手动编辑的同一套 JSON 配置的表单。
|
||||
|
||||
## The Documentation Site
|
||||
## 文档站点
|
||||
|
||||
Every stage has a deeper write-up, with the theory, the diagram, the real code, and what each metric means, on the documentation site:
|
||||
文档站点为每个阶段提供更深入的说明,包括理论、图示、真实代码以及各指标的含义:
|
||||
|
||||
**https://fareedkhan-dev.github.io/train-llm-from-scratch/**
|
||||
|
||||
To run it locally:
|
||||
本地运行:
|
||||
|
||||
```bash
|
||||
pip install -e ".[docs]"
|
||||
mkdocs serve
|
||||
```
|
||||
|
||||
There is also a Foundations section that explains the ideas this code assumes you know (tokenization, the decoder-only Transformer, attention, objectives, optimization, and generation).
|
||||
还有一个 Foundations(基础)章节,讲解本代码假定你已了解的概念(tokenization、decoder-only Transformer、attention、objectives、optimization 和 generation)。
|
||||
|
||||
## Run the Whole Thing
|
||||
## 运行完整流程
|
||||
|
||||
Once the base is pretrained and the data is prepared, one script runs the entire post-training chain and prints the across-stages table:
|
||||
基础模型完成预训练且数据准备就绪后,一个脚本即可运行整个 post-training(后训练)流程,并打印跨阶段对比表:
|
||||
|
||||
```bash
|
||||
bash scripts/run_posttraining.sh # uses both GPUs via torchrun
|
||||
NPROC=1 bash scripts/run_posttraining.sh # single GPU
|
||||
```
|
||||
|
||||
For a fast end-to-end check on a tiny model that finishes in seconds, every stage has a smoke config:
|
||||
若要对秒级完成的小模型进行快速端到端检查,每个阶段都提供了 smoke config(冒烟配置):
|
||||
|
||||
```bash
|
||||
python tests/test_post_training_smoke.py # core math, on CPU
|
||||
python scripts/train_sft.py --config configs/smoke/sft.json # a real (tiny) training run
|
||||
```
|
||||
|
||||
## What's Next
|
||||
## 下一步
|
||||
|
||||
I recommend you start by training the 13 million parameter model, see it produce sensible words, then scale `n_embed` and `n_blocks` up with the memory flags until you hit your GPU limit. After that, walk the post-training chain one stage at a time and watch the GSM8K number move. Every stage is small enough to read in one sitting, and they all share the same model.
|
||||
我建议你先训练 1300 万参数模型,观察它能产出通顺的词,然后用 memory flags 逐步放大 `n_embed` 和 `n_blocks`,直到触及 GPU 上限。之后逐阶段走完 post-training 流程,观察 GSM8K 分数的变化。每个阶段的代码都足够精简,一次通读即可,且它们共用同一套模型。
|
||||
|
||||
If you want to go deeper on any single stage, the documentation site has a focused page for each one.
|
||||
若想深入了解某个特定阶段,文档站点为每个阶段都设有专题页面。
|
||||
|
||||
<hr>
|
||||
|
||||
Wanna chat on something? [My Linkedin](https://www.linkedin.com/in/fareed-khan-dev/)
|
||||
想聊点什么?[My Linkedin](https://www.linkedin.com/in/fareed-khan-dev/)
|
||||
|
||||
## Star History
|
||||
|
||||
|
||||
Reference in New Issue
Block a user