Add Day 1 learning notes and model Chinese comments
- learning/day01: prereading, deep dive, and exam notes - model_minimind.py: added comprehensive Chinese comments for study - .gitignore: minor update Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>pull/792/head
parent
3f1a7cc25b
commit
b23a05868d
|
|
@ -4,3 +4,4 @@ __pycache__/
|
||||||
out
|
out
|
||||||
website/
|
website/
|
||||||
docs-minimind/
|
docs-minimind/
|
||||||
|
.gstack/
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,23 @@
|
||||||
|
# MiniMind 7-Day Training Learning Plan
|
||||||
|
|
||||||
|
目标:用 7 天、约 60 小时,亲手训练并理解 MiniMind 的核心训练链路。
|
||||||
|
|
||||||
|
每天包含 3 个文件:
|
||||||
|
|
||||||
|
- `01_prereading.md`:实战前必须掌握的基础概念、面试考点、工业要求、训练注意事项。
|
||||||
|
- `02_deep_dive.md`:结合本项目代码深入拆解一个或多个关键训练问题。
|
||||||
|
- `03_exam.md`:当天考核,包含代码阅读题、手写推导题、实验题和复盘题。
|
||||||
|
|
||||||
|
推荐顺序:
|
||||||
|
|
||||||
|
1. 先读 `01_prereading.md`,补齐概念。
|
||||||
|
2. 再按 `02_deep_dive.md` 打开项目代码,边读边做小实验。
|
||||||
|
3. 最后完成 `03_exam.md`,把答案写到自己的训练日志里。
|
||||||
|
|
||||||
|
最终交付:
|
||||||
|
|
||||||
|
- `out/pretrain_768.pth`
|
||||||
|
- `out/full_sft_768.pth`
|
||||||
|
- 训练日志和 loss 曲线记录
|
||||||
|
- 评测报告
|
||||||
|
- 至少一份 RL 或蒸馏算法理解报告
|
||||||
|
|
@ -0,0 +1,88 @@
|
||||||
|
# Day 1 Prereading: 环境、项目地图、Decoder-only 基础
|
||||||
|
|
||||||
|
## 今日目标
|
||||||
|
|
||||||
|
你今天要建立项目全局地图,知道 MiniMind 的模型结构、训练脚本、数据集、评测入口分别在哪里。不要急着长时间训练,先确认环境、GPU、依赖和最小 forward 能跑通。
|
||||||
|
|
||||||
|
## 必须掌握的基础概念
|
||||||
|
|
||||||
|
- Decoder-only Transformer:输入 token 序列,预测下一个 token。
|
||||||
|
- Causal LM loss:第 `t` 个位置的 logits 预测第 `t+1` 个 token。
|
||||||
|
> **笔记:** Causal = 只能看前面的 token(上三角 mask)。Loss = `CrossEntropy(logits[:-1], labels[1:])`,即位置 t 预测 t+1。与 BERT (Masked LM, 双向) 相对。
|
||||||
|
- Tokenizer:把文本变成 token id,模型只看 token id。
|
||||||
|
|
||||||
|
> **笔记:** Tokenizer 由两个文件构成:`tokenizer.json`(词表+切词规则)和 `tokenizer_config.json`(特殊 token 定义 + chat_template 对话格式模板)。关键特殊 token:id0=PAD, id1=BOS(`<|im_start|>`), id2=EOS(`<|im_end|>`), id21-22=tool_call, id25-26=think。
|
||||||
|
> 切词方式:BPE + ByteLevel。先把文本转 UTF-8 bytes(任何语言都不会 UNK),再按 6108 条训练好的合并规则把高频 byte 组合合并成 token。词表 6400,中文约 1.5-1.7 字符/token。
|
||||||
|
- Embedding / LM Head:输入词表映射到 hidden,输出 hidden 映射回 vocab logits。
|
||||||
|
> **笔记:** 接力关系:Tokenizer (文本→id, 模型外部) → Embedding (id→向量, 入口) → LM Head (向量→logits, 出口)。MiniMind 中 Embedding 和 LM Head 共享权重 (`tie_word_embeddings=True`)。
|
||||||
|
> Hidden = Embedding 只是起点(查表),经过 8 层 Transformer 融合上下文后的向量才是真正的 hidden (768维)。Vocab logits = hidden @ LM_Head矩阵(768×6400) 得到的得分向量 (6400维),softmax 后变概率。
|
||||||
|
> **Weight Tying 原理:** Embedding `[6400,768]` 和 LM Head `[768,6400]` 互为逆操作(id→向量 vs 向量→得分),共享同一份权重(LM Head = Embedding 转置)。省 490 万参数(占 MiniMind 约 20%),且让编码/解码共用语义空间,效果更好。GPT-2/LLaMA/Qwen 均采用。
|
||||||
|
- Pre-Norm:attention/MLP 前先做 RMSNorm,训练更稳定。
|
||||||
|
- Flash Attention:优化注意力计算的算法,通过分块(tiling)在 GPU SRAM 内完成计算,避免将 O(n²) 注意力矩阵反复搬运到慢速显存(HBM),数学等价但速度快 2-4x、显存降至 O(n)。PyTorch 2.0+ 通过 `F.scaled_dot_product_attention` 内置支持。
|
||||||
|
- RoPE:旋转位置编码,把位置信息注入 Q/K。
|
||||||
|
> **笔记:** RoPE 不是专为 Causal LM 发明,而是解决通用问题:Transformer 本身无位置感知。演进:Learned Absolute (固定长度) → Sinusoidal → RoPE (旋转 Q/K,天然编码相对距离,可外推)。Causal LM 对长序列需求大,所以 RoPE 在此场景尤其流行。
|
||||||
|
> **YaRN 外推:** `max_position_embeddings`(32K) > 训练长度 `max_seq_len`(8K),靠 YaRN 实现。原理:把 RoPE 频率分三区——低频(长距离)压缩适配、高频(短距离)保持不变、中频渐进过渡。避免简单等比缩放丢失近距离精度。对应参数:`yarn_scale`、`yarn_beta_fast/slow`。
|
||||||
|
- GQA:query heads 多于 key/value heads,降低 KV cache 和计算成本。
|
||||||
|
> **笔记:** MHA 8Q:8KV → GQA 8Q:4KV → MQA 8Q:1KV。MiniMind 用 GQA (每2个Q头共享1对KV),KV cache 省50%。
|
||||||
|
- SwiGLU:现代 LLM 常用 FFN 结构。
|
||||||
|
> **笔记:** `Down(SiLU(Gate(x)) * Up(x))`。比标准 FFN 多一个 Gate 矩阵 (参数+50%),Gate 起可学习门控作用,选择性保留信息。LLaMA/Qwen/DeepSeek 均采用。
|
||||||
|
|
||||||
|
## 面试考点
|
||||||
|
|
||||||
|
- 为什么 causal LM 训练时 logits 和 labels 要 shift?
|
||||||
|
> **答:** 自回归目标是"用位置 t 预测 t+1"。logits 的最后一个位置没有下一个 token 可预测,labels 的第一个 token 没有上文来预测它,所以 `logits[:-1]` 对齐 `labels[1:]`。不 shift 的话预测目标就错位了。
|
||||||
|
- Decoder-only 和 Encoder-Decoder 在训练目标上有什么区别?
|
||||||
|
> **答:** Decoder-only(GPT):所有 token 统一做从左到右的 next-token prediction,输入和输出共享同一个序列。Encoder-Decoder(T5):Encoder 双向编码输入(无因果 mask),Decoder 自回归生成输出,两者是分开的序列。Decoder-only 更简单,规模扩展更直接,是当前主流。
|
||||||
|
- RMSNorm 和 LayerNorm 的区别是什么?
|
||||||
|
> **答:** LayerNorm: `(x - mean) / std * γ + β`,需要均值和方差,有偏置 β。RMSNorm: `x / RMS(x) * γ`,只用均方根,无偏置。RMSNorm 省掉 mean 计算和 β 参数,速度快 ~10-15%,效果基本无损,所以现代 LLM 都用 RMSNorm。
|
||||||
|
- RoPE 相比 learned absolute position embedding 的优势是什么?
|
||||||
|
> **答:** ① 天然编码相对距离(Q·K 只依赖 |i-j|,而非绝对位置);② 可外推到训练时没见过的长度(配合 YaRN 更好);③ 不增加可学习参数;④ Learned absolute 长度固定,超出即崩,且只编码绝对位置,对相对关系的建模隐式且弱。
|
||||||
|
|
||||||
|
- GQA/MQA/MHA 的区别和工程取舍是什么?
|
||||||
|
> **答:** MHA: Q=K=V 头数相同,表达力最强但 KV cache 最大。MQA: 所有 Q 头共享 1 个 KV 头,KV cache 最小但质量可能下降。GQA: 折中,多个 Q 头分组共享 KV 头(MiniMind: 8Q:4KV)。工程取舍核心是推理时 KV cache 显存 vs 模型质量——大规模部署时 KV cache 是瓶颈,GQA 是当前最佳平衡点。
|
||||||
|
|
||||||
|
- 为什么小模型里 vocab size 会显著影响参数量?
|
||||||
|
> **答:** Embedding 层参数 = `vocab_size × hidden_size`。MiniMind: 6400×768 = 490 万。如果用 Qwen2 的 151K 词表: 151643×768 = 1.16 亿,**仅 Embedding 就超过整个模型 (64M)**。小模型总参数少,Embedding 占比极高,所以词表必须精简。MiniMind 用 6400 词表正是这个原因。
|
||||||
|
|
||||||
|
## 工业要求
|
||||||
|
- 训练前先确定可复现环境:Python、CUDA、PyTorch、GPU 显存。
|
||||||
|
- 所有实验记录命令、git 状态、数据版本、模型参数、随机种子。
|
||||||
|
- 第一天不追求效果,只追求可运行、可定位问题。
|
||||||
|
- 任何训练前先做小 batch smoke test,避免长时间跑错配置。
|
||||||
|
|
||||||
|
## 项目中要看的文件
|
||||||
|
- `README.md`:快速开始、数据、训练说明。
|
||||||
|
- `model/model_minimind.py`:模型结构。
|
||||||
|
- `trainer/train_pretrain.py`:预训练入口。
|
||||||
|
- `trainer/train_full_sft.py`:SFT 入口。
|
||||||
|
- `dataset/lm_dataset.py`:数据到 tensor 的转换。
|
||||||
|
- `eval_llm.py`:训练后推理评测。
|
||||||
|
|
||||||
|
## 今日动手任务
|
||||||
|
|
||||||
|
1. 确认环境:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python - <<'PY'
|
||||||
|
import torch
|
||||||
|
print(torch.__version__)
|
||||||
|
print(torch.cuda.is_available())
|
||||||
|
print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "no cuda")
|
||||||
|
PY
|
||||||
|
```
|
||||||
|
|
||||||
|
2. 阅读 `model/model_minimind.py` 中这些类:
|
||||||
|
|
||||||
|
- `MiniMindConfig`
|
||||||
|
- `Attention`
|
||||||
|
- `FeedForward`
|
||||||
|
- `MOEFeedForward`
|
||||||
|
- `MiniMindForCausalLM`
|
||||||
|
|
||||||
|
3. 写一段自己的结构笔记:输入 `[B, T]` 如何变成 logits `[B, T, vocab]`。
|
||||||
|
|
||||||
|
## 训练注意事项
|
||||||
|
|
||||||
|
- 今天不要直接跑完整训练。
|
||||||
|
- 如果没有 CUDA,后续训练计划必须缩小模型或只做代码理解。
|
||||||
|
- 如果用 Mac/MPS,混合精度、性能和部分算子兼容性要单独确认。
|
||||||
|
|
@ -0,0 +1,107 @@
|
||||||
|
# Day 1 Deep Dive: MiniMind 模型结构拆解
|
||||||
|
|
||||||
|
## 深挖问题 1:MiniMind 的最小 forward 路径是什么?
|
||||||
|
|
||||||
|
打开 `model/model_minimind.py`:
|
||||||
|
|
||||||
|
1. `MiniMindForCausalLM.forward`
|
||||||
|
2. `MiniMindModel.forward`
|
||||||
|
3. `MiniMindBlock.forward`
|
||||||
|
4. `Attention.forward`
|
||||||
|
5. `FeedForward.forward`
|
||||||
|
|
||||||
|
你要能追踪:
|
||||||
|
|
||||||
|
```text
|
||||||
|
input_ids
|
||||||
|
-> embed_tokens
|
||||||
|
-> dropout
|
||||||
|
-> N x MiniMindBlock
|
||||||
|
-> RMSNorm
|
||||||
|
-> lm_head
|
||||||
|
-> logits
|
||||||
|
-> cross_entropy loss
|
||||||
|
```
|
||||||
|
|
||||||
|
重点观察:
|
||||||
|
|
||||||
|
- `input_ids.shape == [batch, seq_len]`
|
||||||
|
- `hidden_states.shape == [batch, seq_len, hidden_size]`
|
||||||
|
- `logits.shape == [batch, seq_len, vocab_size]`
|
||||||
|
|
||||||
|
## 深挖问题 2:为什么 forward 里是这样算 loss?
|
||||||
|
|
||||||
|
代码位置:`MiniMindForCausalLM.forward`
|
||||||
|
|
||||||
|
核心逻辑:
|
||||||
|
|
||||||
|
```python
|
||||||
|
x = logits[..., :-1, :]
|
||||||
|
y = labels[..., 1:]
|
||||||
|
loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
|
||||||
|
```
|
||||||
|
|
||||||
|
理解方式:
|
||||||
|
|
||||||
|
- 位置 0 的 logits 预测 token 1。
|
||||||
|
- 位置 1 的 logits 预测 token 2。
|
||||||
|
- 最后一个位置没有下一个 token,所以丢掉。
|
||||||
|
- label 的第一个 token 没有被任何上文预测,所以丢掉。
|
||||||
|
|
||||||
|
`ignore_index=-100` 是训练中最重要的 mask 约定之一。pretrain 用它忽略 padding;SFT 用它忽略 user/system/prompt 部分。
|
||||||
|
|
||||||
|
## 深挖问题 3:Attention 里 GQA 是怎么实现的?
|
||||||
|
|
||||||
|
看 `Attention.__init__`:
|
||||||
|
|
||||||
|
- `num_attention_heads`
|
||||||
|
- `num_key_value_heads`
|
||||||
|
- `n_rep = q_heads // kv_heads`
|
||||||
|
|
||||||
|
看 `repeat_kv`:
|
||||||
|
|
||||||
|
- K/V heads 先少算。
|
||||||
|
- attention 前把 K/V 复制到和 Q heads 对齐。
|
||||||
|
|
||||||
|
工业理解:
|
||||||
|
|
||||||
|
- 训练时省一部分参数和计算。
|
||||||
|
- 推理时 KV cache 更小。
|
||||||
|
- 对小模型通常是合理取舍。
|
||||||
|
|
||||||
|
## 深挖问题 4:MoE 为什么有 aux loss?
|
||||||
|
|
||||||
|
看 `MOEFeedForward.forward`:
|
||||||
|
|
||||||
|
- gate 给每个 token 分配 expert。
|
||||||
|
- `topk_idx` 决定 token 去哪个 expert。
|
||||||
|
- `aux_loss` 鼓励 expert 负载更均衡。
|
||||||
|
|
||||||
|
工业风险:
|
||||||
|
|
||||||
|
- 没有负载均衡,部分 expert 可能被过度使用。
|
||||||
|
- 原生 PyTorch MoE 训练不一定快,因为 expert 分派有额外调度开销。
|
||||||
|
|
||||||
|
## 今日小实验
|
||||||
|
|
||||||
|
写一个临时 Python 片段,不修改项目文件:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python - <<'PY'
|
||||||
|
import torch
|
||||||
|
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||||
|
|
||||||
|
config = MiniMindConfig(hidden_size=128, num_hidden_layers=2)
|
||||||
|
model = MiniMindForCausalLM(config)
|
||||||
|
input_ids = torch.randint(0, config.vocab_size, (2, 16))
|
||||||
|
labels = input_ids.clone()
|
||||||
|
out = model(input_ids, labels=labels)
|
||||||
|
print(out.logits.shape)
|
||||||
|
print(out.loss.item())
|
||||||
|
PY
|
||||||
|
```
|
||||||
|
|
||||||
|
你要解释:
|
||||||
|
|
||||||
|
- 为什么随机模型也有 loss?
|
||||||
|
- loss 大概应接近 `log(vocab_size)`,为什么?
|
||||||
|
|
@ -0,0 +1,48 @@
|
||||||
|
# Day 1 Exam
|
||||||
|
|
||||||
|
## A. 代码阅读题
|
||||||
|
|
||||||
|
1. `MiniMindForCausalLM.forward` 中 `logits[..., :-1, :]` 的含义是什么?
|
||||||
|
2. `labels[..., 1:]` 为什么要右移?
|
||||||
|
3. `ignore_index=-100` 在 pretrain 和 SFT 中分别解决什么问题?
|
||||||
|
4. `Attention.forward` 中 `past_key_value` 的作用是什么?
|
||||||
|
5. `repeat_kv` 为什么只重复 K/V,不重复 Q?
|
||||||
|
|
||||||
|
## B. 手写题
|
||||||
|
|
||||||
|
给定:
|
||||||
|
|
||||||
|
```text
|
||||||
|
input_ids = [BOS, 我, 爱, AI, EOS, PAD, PAD]
|
||||||
|
```
|
||||||
|
|
||||||
|
请写出 causal LM 中每个位置预测的目标 token。
|
||||||
|
|
||||||
|
再写出 padding 后的 labels,其中 PAD 对应 `-100`。
|
||||||
|
|
||||||
|
## C. 面试题
|
||||||
|
|
||||||
|
1. RoPE 为什么能支持比训练长度更长的外推?
|
||||||
|
2. RMSNorm 为什么比 LayerNorm 更轻?
|
||||||
|
3. 小模型中为什么 vocab size 不能无限大?
|
||||||
|
4. GQA 相比 MHA 的主要收益是什么?
|
||||||
|
|
||||||
|
## D. 实验题
|
||||||
|
|
||||||
|
运行 Day 1 deep dive 的随机模型 forward 实验,记录:
|
||||||
|
|
||||||
|
- logits shape
|
||||||
|
- loss 数值
|
||||||
|
- `log(6400)` 的值
|
||||||
|
|
||||||
|
回答:随机初始化模型的 loss 为什么和 `log(vocab_size)` 接近?
|
||||||
|
|
||||||
|
## E. 通过标准
|
||||||
|
|
||||||
|
你能不看代码讲出:
|
||||||
|
|
||||||
|
```text
|
||||||
|
input_ids -> embeddings -> blocks -> lm_head -> logits -> shifted CE loss
|
||||||
|
```
|
||||||
|
|
||||||
|
并能解释 `-100`、shift、RoPE、GQA。
|
||||||
|
|
@ -0,0 +1,65 @@
|
||||||
|
# Day 2 Prereading: Tokenizer、数据格式、Label Mask
|
||||||
|
|
||||||
|
## 今日目标
|
||||||
|
|
||||||
|
你要理解数据怎样进入训练。MiniMind 里 pretrain 和 SFT 的训练循环很像,真正关键差异在 dataset 和 label mask。
|
||||||
|
|
||||||
|
## 必须掌握的基础概念
|
||||||
|
|
||||||
|
- JSONL 数据:一行一个样本。
|
||||||
|
- Pretrain 样本:`{"text": "..."}`
|
||||||
|
- SFT 样本:`{"conversations": [{"role": "...", "content": "..."}]}`
|
||||||
|
- Chat template:把多轮消息序列变成模型输入文本。
|
||||||
|
- Label mask:哪些 token 参与 loss,哪些 token 不参与。
|
||||||
|
- `-100`:PyTorch cross entropy 中忽略 label 的约定。
|
||||||
|
|
||||||
|
## 面试考点
|
||||||
|
|
||||||
|
- SFT 为什么通常只对 assistant answer 计算 loss?
|
||||||
|
- 如果 user prompt 也计算 loss,会有什么副作用?
|
||||||
|
- `max_seq_len` 截断会造成什么训练问题?
|
||||||
|
- 为什么 tool call 要通过 chat template 统一成特殊标签?
|
||||||
|
- padding side 在训练和生成中有什么影响?
|
||||||
|
|
||||||
|
## 工业要求
|
||||||
|
|
||||||
|
- 数据质量比模型小改动更重要。
|
||||||
|
- 训练前必须抽样查看 tokenized 后的样本。
|
||||||
|
- SFT 必须检查 assistant span mask 是否正确。
|
||||||
|
- 长样本截断不能只看字符数,要看 token 数。
|
||||||
|
- Tool call / reasoning / normal chat 混合数据要保持模板一致。
|
||||||
|
|
||||||
|
## 项目中要看的文件
|
||||||
|
|
||||||
|
- `dataset/lm_dataset.py`
|
||||||
|
- `model/tokenizer.json`
|
||||||
|
- `model/tokenizer_config.json`
|
||||||
|
- `README.md` 的数据章节
|
||||||
|
|
||||||
|
重点类:
|
||||||
|
|
||||||
|
- `PretrainDataset`
|
||||||
|
- `SFTDataset`
|
||||||
|
- `RLAIFDataset`
|
||||||
|
- `AgentRLDataset`
|
||||||
|
|
||||||
|
## 今日动手任务
|
||||||
|
|
||||||
|
1. 下载或确认数据:
|
||||||
|
|
||||||
|
```text
|
||||||
|
dataset/pretrain_t2t_mini.jsonl
|
||||||
|
dataset/sft_t2t_mini.jsonl
|
||||||
|
```
|
||||||
|
|
||||||
|
2. 打印一条 pretrain 样本和一条 SFT 样本。
|
||||||
|
|
||||||
|
3. 用 tokenizer 对一条 SFT 样本做 `apply_chat_template`,观察模板输出。
|
||||||
|
|
||||||
|
4. 手动判断哪些 token 应该参与 loss。
|
||||||
|
|
||||||
|
## 训练注意事项
|
||||||
|
|
||||||
|
- 今天仍然不需要完整训练。
|
||||||
|
- 如果数据文件不存在,先只用 README 样例做理解。
|
||||||
|
- 任何 SFT 训练异常,优先怀疑 label mask 和 chat template。
|
||||||
|
|
@ -0,0 +1,105 @@
|
||||||
|
# Day 2 Deep Dive: Dataset 到 Loss 的路径
|
||||||
|
|
||||||
|
## 深挖问题 1:PretrainDataset 做了什么?
|
||||||
|
|
||||||
|
代码位置:`dataset/lm_dataset.py` 的 `PretrainDataset`
|
||||||
|
|
||||||
|
流程:
|
||||||
|
|
||||||
|
```text
|
||||||
|
sample["text"]
|
||||||
|
-> tokenizer(..., add_special_tokens=False)
|
||||||
|
-> [BOS] + tokens + [EOS]
|
||||||
|
-> pad 到 max_length
|
||||||
|
-> labels = input_ids.clone()
|
||||||
|
-> labels[PAD] = -100
|
||||||
|
```
|
||||||
|
|
||||||
|
关键点:
|
||||||
|
|
||||||
|
- pretrain 让模型学习所有非 PAD token 的 next-token prediction。
|
||||||
|
- BOS/EOS 也参与训练,帮助模型学习开始和结束。
|
||||||
|
|
||||||
|
## 深挖问题 2:SFTDataset 为什么更复杂?
|
||||||
|
|
||||||
|
代码位置:`SFTDataset.create_chat_prompt` 和 `SFTDataset.generate_labels`
|
||||||
|
|
||||||
|
SFT 输入包含 system/user/assistant/tool 等角色。训练目标不是复述完整对话,而是让模型在看到上下文后生成 assistant 输出。
|
||||||
|
|
||||||
|
因此:
|
||||||
|
|
||||||
|
- user/system/tool context 作为条件。
|
||||||
|
- assistant 内容作为监督目标。
|
||||||
|
- 非 assistant token label 设为 `-100`。
|
||||||
|
|
||||||
|
## 深挖问题 3:`generate_labels` 如何找到 assistant span?
|
||||||
|
|
||||||
|
代码里构造:
|
||||||
|
|
||||||
|
```python
|
||||||
|
self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
|
||||||
|
self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
|
||||||
|
```
|
||||||
|
|
||||||
|
然后扫描 input_ids:
|
||||||
|
|
||||||
|
```text
|
||||||
|
找到 assistant 起点
|
||||||
|
-> 从 assistant 内容开始
|
||||||
|
-> 直到 eos
|
||||||
|
-> 这些 token 的 label = input_ids[j]
|
||||||
|
-> 其他位置 label = -100
|
||||||
|
```
|
||||||
|
|
||||||
|
你要特别注意:模板变化会直接影响这段扫描逻辑。如果 tokenizer template 改了,但 `bos_id/eos_id` 匹配没改,SFT loss mask 就可能错。
|
||||||
|
|
||||||
|
## 深挖问题 4:Tool Call 样本为什么要完整保留?
|
||||||
|
|
||||||
|
`pre_processing_chat` 里:
|
||||||
|
|
||||||
|
```python
|
||||||
|
if any(conv.get('tools') for conv in conversations): return conversations
|
||||||
|
```
|
||||||
|
|
||||||
|
原因:
|
||||||
|
|
||||||
|
- tool use 样本依赖 system 中的 tools schema。
|
||||||
|
- 随机加 system prompt 可能破坏工具调用上下文。
|
||||||
|
- tool call 训练要求模板结构严格。
|
||||||
|
|
||||||
|
## 今日小实验
|
||||||
|
|
||||||
|
不修改文件,运行:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python - <<'PY'
|
||||||
|
from transformers import AutoTokenizer
|
||||||
|
from dataset.lm_dataset import SFTDataset
|
||||||
|
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained("model")
|
||||||
|
sample = {
|
||||||
|
"conversations": [
|
||||||
|
{"role": "user", "content": "你好"},
|
||||||
|
{"role": "assistant", "content": "你好,我是 MiniMind。"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
ds = SFTDataset("dataset/sft_t2t_mini.jsonl", tokenizer, max_length=128)
|
||||||
|
prompt = ds.create_chat_prompt(sample["conversations"])
|
||||||
|
ids = tokenizer(prompt).input_ids[:128]
|
||||||
|
labels = ds.generate_labels(ids)
|
||||||
|
print(prompt)
|
||||||
|
for i, (tid, lab) in enumerate(zip(ids, labels)):
|
||||||
|
tok = tokenizer.decode([tid])
|
||||||
|
flag = "LOSS" if lab != -100 else "MASK"
|
||||||
|
print(i, repr(tok), flag)
|
||||||
|
PY
|
||||||
|
```
|
||||||
|
|
||||||
|
如果数据文件不存在,可以先等下载后再做。
|
||||||
|
|
||||||
|
## 复盘问题
|
||||||
|
|
||||||
|
- assistant 的 eos 是否参与 loss?为什么?
|
||||||
|
- 如果 `max_length` 在 assistant 回答中间截断,会发生什么?
|
||||||
|
- 多轮对话中每个 assistant span 是否都参与 loss?
|
||||||
|
|
@ -0,0 +1,52 @@
|
||||||
|
# Day 2 Exam
|
||||||
|
|
||||||
|
## A. 代码阅读题
|
||||||
|
|
||||||
|
1. `PretrainDataset.__getitem__` 为什么要手动加 BOS/EOS?
|
||||||
|
2. `SFTDataset.create_chat_prompt` 中 `tools` 是从哪里来的?
|
||||||
|
3. `SFTDataset.generate_labels` 如何定位 assistant 内容?
|
||||||
|
4. `post_processing_chat` 为什么会概率性移除空 `<think>`?
|
||||||
|
5. `RLAIFDataset.__getitem__` 为什么只返回 prompt,不返回标准 answer?
|
||||||
|
|
||||||
|
## B. 手写题
|
||||||
|
|
||||||
|
给定 SFT 对话:
|
||||||
|
|
||||||
|
```json
|
||||||
|
[
|
||||||
|
{"role": "user", "content": "2+2等于几?"},
|
||||||
|
{"role": "assistant", "content": "等于4。"}
|
||||||
|
]
|
||||||
|
```
|
||||||
|
|
||||||
|
请说明:
|
||||||
|
|
||||||
|
- 哪些部分是条件上下文?
|
||||||
|
- 哪些部分参与 loss?
|
||||||
|
- 如果 assistant 回答被截断到只剩 `"等于"`,训练信号有什么问题?
|
||||||
|
|
||||||
|
## C. 工业题
|
||||||
|
|
||||||
|
1. 如何发现 SFT mask 写错了?
|
||||||
|
2. 如果训练 loss 很低但模型推理不会聊天,可能是什么原因?
|
||||||
|
3. 如果 SFT 数据中 tool call JSON 经常不合法,会导致什么问题?
|
||||||
|
4. 为什么数据抽样检查比直接看 loss 更重要?
|
||||||
|
|
||||||
|
## D. 实验题
|
||||||
|
|
||||||
|
抽 3 条真实 SFT 样本,打印:
|
||||||
|
|
||||||
|
- chat template 后的文本
|
||||||
|
- token 数量
|
||||||
|
- loss token 数量
|
||||||
|
- loss token 占比
|
||||||
|
|
||||||
|
写出你的观察:样本长度分布和 loss token 占比是否合理?
|
||||||
|
|
||||||
|
## E. 通过标准
|
||||||
|
|
||||||
|
你能解释:
|
||||||
|
|
||||||
|
```text
|
||||||
|
同样是 cross entropy,pretrain 和 SFT 的差异主要来自 labels。
|
||||||
|
```
|
||||||
|
|
@ -0,0 +1,75 @@
|
||||||
|
# Day 3 Prereading: Pretrain 实战
|
||||||
|
|
||||||
|
## 今日目标
|
||||||
|
|
||||||
|
今天正式完成预训练。你要理解训练循环、学习率、梯度累积、混合精度、checkpoint,并产出 `out/pretrain_768.pth`。
|
||||||
|
|
||||||
|
## 必须掌握的基础概念
|
||||||
|
|
||||||
|
- Epoch:完整遍历一次数据。
|
||||||
|
- Step:一个 batch 的训练迭代。
|
||||||
|
- Gradient accumulation:多个 micro-batch 累积后再 optimizer step。
|
||||||
|
- Effective batch size:`batch_size * accumulation_steps * world_size`。
|
||||||
|
- AdamW:主流 LLM 优化器。
|
||||||
|
- Grad clip:限制梯度范数,降低训练爆炸风险。
|
||||||
|
- Mixed precision:用 bf16/fp16 降低显存和提升吞吐。
|
||||||
|
- Cosine LR schedule:学习率从高到低平滑衰减。
|
||||||
|
|
||||||
|
## 面试考点
|
||||||
|
|
||||||
|
- 为什么 LLM pretrain 通常需要较大的学习率?
|
||||||
|
- 梯度累积和直接大 batch 是否完全等价?
|
||||||
|
- bf16 相比 fp16 的优势是什么?
|
||||||
|
- 为什么要做 gradient clipping?
|
||||||
|
- checkpoint 应该保存哪些状态,为什么只保存模型不够?
|
||||||
|
|
||||||
|
## 工业要求
|
||||||
|
|
||||||
|
- 必须记录完整训练命令。
|
||||||
|
- 必须记录 loss、lr、吞吐或耗时。
|
||||||
|
- 必须测试断点恢复。
|
||||||
|
- 训练中如果 loss NaN,先检查学习率、精度、数据异常、grad clip。
|
||||||
|
- 每次保存权重时确认输出文件名,避免覆盖误解。
|
||||||
|
|
||||||
|
## 项目中要看的文件
|
||||||
|
|
||||||
|
- `trainer/train_pretrain.py`
|
||||||
|
- `trainer/trainer_utils.py`
|
||||||
|
- `dataset/lm_dataset.py::PretrainDataset`
|
||||||
|
|
||||||
|
## 今日推荐命令
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd trainer
|
||||||
|
python train_pretrain.py \
|
||||||
|
--epochs 2 \
|
||||||
|
--batch_size 32 \
|
||||||
|
--accumulation_steps 8 \
|
||||||
|
--max_seq_len 340 \
|
||||||
|
--learning_rate 5e-4 \
|
||||||
|
--log_interval 100 \
|
||||||
|
--save_interval 1000
|
||||||
|
```
|
||||||
|
|
||||||
|
如果显存不足:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python train_pretrain.py --batch_size 8 --accumulation_steps 32
|
||||||
|
```
|
||||||
|
|
||||||
|
## 今日记录模板
|
||||||
|
|
||||||
|
```text
|
||||||
|
命令:
|
||||||
|
GPU:
|
||||||
|
数据:
|
||||||
|
hidden_size / layers:
|
||||||
|
batch_size:
|
||||||
|
accumulation_steps:
|
||||||
|
max_seq_len:
|
||||||
|
初始 loss:
|
||||||
|
中段 loss:
|
||||||
|
最终 loss:
|
||||||
|
是否发生 OOM:
|
||||||
|
checkpoint 是否成功:
|
||||||
|
```
|
||||||
|
|
@ -0,0 +1,105 @@
|
||||||
|
# Day 3 Deep Dive: Pretrain 训练循环
|
||||||
|
|
||||||
|
## 深挖问题 1:训练循环每一步发生了什么?
|
||||||
|
|
||||||
|
代码位置:`trainer/train_pretrain.py::train_epoch`
|
||||||
|
|
||||||
|
流程:
|
||||||
|
|
||||||
|
```text
|
||||||
|
for input_ids, labels in loader:
|
||||||
|
move to device
|
||||||
|
update lr
|
||||||
|
autocast forward
|
||||||
|
loss = res.loss + res.aux_loss
|
||||||
|
loss /= accumulation_steps
|
||||||
|
backward
|
||||||
|
if accumulation boundary:
|
||||||
|
unscale
|
||||||
|
grad clip
|
||||||
|
optimizer step
|
||||||
|
scaler update
|
||||||
|
zero grad
|
||||||
|
log
|
||||||
|
save
|
||||||
|
```
|
||||||
|
|
||||||
|
你要能解释每一行为什么存在。
|
||||||
|
|
||||||
|
## 深挖问题 2:学习率怎么变?
|
||||||
|
|
||||||
|
代码位置:`trainer/trainer_utils.py::get_lr`
|
||||||
|
|
||||||
|
```python
|
||||||
|
return lr * (0.1 + 0.45 * (1 + cos(pi * current_step / total_steps)))
|
||||||
|
```
|
||||||
|
|
||||||
|
含义:
|
||||||
|
|
||||||
|
- 初始接近 `lr`。
|
||||||
|
- 最终接近 `0.1 * lr`。
|
||||||
|
- 中间 cosine 平滑下降。
|
||||||
|
|
||||||
|
工业注意:
|
||||||
|
|
||||||
|
- 这个实现没有 warmup。
|
||||||
|
- 小模型和小数据通常能跑,但大规模训练一般会加 warmup。
|
||||||
|
- 如果 early loss 抖动严重,可以考虑降低 lr 或加 warmup。
|
||||||
|
|
||||||
|
## 深挖问题 3:checkpoint 保存了什么?
|
||||||
|
|
||||||
|
代码位置:`trainer/trainer_utils.py::lm_checkpoint`
|
||||||
|
|
||||||
|
保存两类文件:
|
||||||
|
|
||||||
|
- `out/pretrain_768.pth`:主要用于推理/下游训练。
|
||||||
|
- `checkpoints/pretrain_768_resume.pth`:用于断点续训。
|
||||||
|
|
||||||
|
resume checkpoint 包含:
|
||||||
|
|
||||||
|
- model
|
||||||
|
- optimizer
|
||||||
|
- epoch
|
||||||
|
- step
|
||||||
|
- world_size
|
||||||
|
- wandb id
|
||||||
|
- scaler 等额外状态
|
||||||
|
|
||||||
|
为什么 optimizer 也要保存?
|
||||||
|
|
||||||
|
因为 AdamW 有动量状态。如果只恢复模型权重,优化器状态丢失,训练动态会突然改变。
|
||||||
|
|
||||||
|
## 深挖问题 4:MoE aux_loss 为什么加在 pretrain loss 上?
|
||||||
|
|
||||||
|
如果 `use_moe=1`,模型输出 `aux_loss`。训练总 loss:
|
||||||
|
|
||||||
|
```python
|
||||||
|
loss = res.loss + res.aux_loss
|
||||||
|
```
|
||||||
|
|
||||||
|
这表示主目标仍是 next-token prediction,但额外鼓励 expert 负载均衡。
|
||||||
|
|
||||||
|
## 今日实验
|
||||||
|
|
||||||
|
1. 先跑 200 step smoke test:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd trainer
|
||||||
|
python train_pretrain.py --epochs 1 --save_interval 200 --log_interval 20
|
||||||
|
```
|
||||||
|
|
||||||
|
2. 中断训练。
|
||||||
|
|
||||||
|
3. 恢复:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python train_pretrain.py --epochs 2 --from_resume 1
|
||||||
|
```
|
||||||
|
|
||||||
|
4. 检查日志是否显示跳过 step。
|
||||||
|
|
||||||
|
## 复盘问题
|
||||||
|
|
||||||
|
- 为什么 `loss / accumulation_steps` 后再 backward?
|
||||||
|
- 为什么 `scaler.unscale_` 后再 clip grad?
|
||||||
|
- 为什么保存权重前切 `model.eval()`,保存后再 `model.train()`?
|
||||||
|
|
@ -0,0 +1,55 @@
|
||||||
|
# Day 3 Exam
|
||||||
|
|
||||||
|
## A. 代码阅读题
|
||||||
|
|
||||||
|
1. `train_pretrain.py` 中 `autocast_ctx` 是如何根据 device 决定的?
|
||||||
|
2. `scaler` 什么时候启用?bf16 下是否启用?
|
||||||
|
3. `optimizer.zero_grad(set_to_none=True)` 的好处是什么?
|
||||||
|
4. `SkipBatchSampler` 在 resume 时解决什么问题?
|
||||||
|
5. `torch.save({k: v.half().cpu() ...})` 为什么要转 half 和 cpu?
|
||||||
|
|
||||||
|
## B. 计算题
|
||||||
|
|
||||||
|
假设:
|
||||||
|
|
||||||
|
```text
|
||||||
|
batch_size = 32
|
||||||
|
accumulation_steps = 8
|
||||||
|
world_size = 1
|
||||||
|
max_seq_len = 340
|
||||||
|
```
|
||||||
|
|
||||||
|
回答:
|
||||||
|
|
||||||
|
- effective batch size 是多少条样本?
|
||||||
|
- 每次 optimizer step 理论上看过多少 token?
|
||||||
|
- 如果 world_size 变成 4,effective batch size 如何变化?
|
||||||
|
|
||||||
|
## C. 故障排查题
|
||||||
|
|
||||||
|
1. 训练第一个 step 就 OOM,你优先改哪些参数?
|
||||||
|
2. loss 变成 NaN,可能原因有哪些?
|
||||||
|
3. resume 后 loss 突然跳变,可能是什么状态没恢复?
|
||||||
|
4. 训练很慢,如何判断瓶颈在数据加载还是 GPU 计算?
|
||||||
|
|
||||||
|
## D. 实验题
|
||||||
|
|
||||||
|
完成一次 pretrain,并记录:
|
||||||
|
|
||||||
|
- loss 起点
|
||||||
|
- loss 终点
|
||||||
|
- 每 100 step 的 lr
|
||||||
|
- 每个 epoch 耗时
|
||||||
|
- 最终权重路径
|
||||||
|
|
||||||
|
然后运行:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python eval_llm.py --weight pretrain
|
||||||
|
```
|
||||||
|
|
||||||
|
评价 pretrain 模型是否已经具备聊天能力。为什么?
|
||||||
|
|
||||||
|
## E. 通过标准
|
||||||
|
|
||||||
|
你能解释完整 pretrain loop,并成功得到 `out/pretrain_768.pth`。
|
||||||
|
|
@ -0,0 +1,73 @@
|
||||||
|
# Day 4 Prereading: SFT 实战
|
||||||
|
|
||||||
|
## 今日目标
|
||||||
|
|
||||||
|
今天从 `pretrain` 权重继续训练 `full_sft`,得到能按对话模板回答的模型。
|
||||||
|
|
||||||
|
## 必须掌握的基础概念
|
||||||
|
|
||||||
|
- SFT:Supervised Fine-Tuning,用高质量指令/对话数据训练助手行为。
|
||||||
|
- Instruction following:理解用户意图并按格式回答。
|
||||||
|
- Chat template:训练和推理必须一致。
|
||||||
|
- Assistant-only loss:只让模型学习 assistant 输出,不学习复述 user。
|
||||||
|
- Catastrophic forgetting:微调过窄数据可能损伤原有能力。
|
||||||
|
|
||||||
|
## 面试考点
|
||||||
|
|
||||||
|
- Pretrain 和 SFT 的本质区别是什么?
|
||||||
|
- SFT 能不能注入新知识?
|
||||||
|
- 为什么 SFT 学到的不只是格式,也可能改变知识分布?
|
||||||
|
- 为什么学习率比 pretrain 小很多?
|
||||||
|
- 为什么 SFT 后模型更会聊天,但不一定更真实?
|
||||||
|
|
||||||
|
## 工业要求
|
||||||
|
|
||||||
|
- SFT 数据必须检查格式一致性。
|
||||||
|
- 训练和推理的 chat template 必须一致。
|
||||||
|
- SFT 评测不能只看 loss,要看真实生成。
|
||||||
|
- 小模型 SFT 容易学表面风格,但事实性仍弱。
|
||||||
|
- 如果领域微调,要混合通用数据避免能力坍缩。
|
||||||
|
|
||||||
|
## 项目中要看的文件
|
||||||
|
|
||||||
|
- `trainer/train_full_sft.py`
|
||||||
|
- `dataset/lm_dataset.py::SFTDataset`
|
||||||
|
- `eval_llm.py`
|
||||||
|
- `scripts/eval_toolcall.py`
|
||||||
|
|
||||||
|
## 今日推荐命令
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd trainer
|
||||||
|
python train_full_sft.py \
|
||||||
|
--epochs 2 \
|
||||||
|
--from_weight pretrain \
|
||||||
|
--batch_size 16 \
|
||||||
|
--accumulation_steps 1 \
|
||||||
|
--max_seq_len 768 \
|
||||||
|
--learning_rate 1e-5 \
|
||||||
|
--log_interval 100 \
|
||||||
|
--save_interval 1000
|
||||||
|
```
|
||||||
|
|
||||||
|
训练后:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd ..
|
||||||
|
python eval_llm.py --weight full_sft
|
||||||
|
```
|
||||||
|
|
||||||
|
## 今日记录模板
|
||||||
|
|
||||||
|
```text
|
||||||
|
SFT 命令:
|
||||||
|
from_weight:
|
||||||
|
数据:
|
||||||
|
max_seq_len:
|
||||||
|
初始 loss:
|
||||||
|
最终 loss:
|
||||||
|
pretrain 回答样例:
|
||||||
|
full_sft 回答样例:
|
||||||
|
明显提升:
|
||||||
|
明显问题:
|
||||||
|
```
|
||||||
|
|
@ -0,0 +1,100 @@
|
||||||
|
# Day 4 Deep Dive: SFT 为什么能让模型会聊天
|
||||||
|
|
||||||
|
## 深挖问题 1:SFT 训练循环和 pretrain 有什么不同?
|
||||||
|
|
||||||
|
打开:
|
||||||
|
|
||||||
|
- `trainer/train_pretrain.py`
|
||||||
|
- `trainer/train_full_sft.py`
|
||||||
|
|
||||||
|
你会发现训练循环几乎一样。核心差异:
|
||||||
|
|
||||||
|
```text
|
||||||
|
PretrainDataset -> 所有正文 token 学 next token
|
||||||
|
SFTDataset -> 只有 assistant span 学 next token
|
||||||
|
```
|
||||||
|
|
||||||
|
这说明:训练算法变化不大,监督信号变了。
|
||||||
|
|
||||||
|
## 深挖问题 2:`from_weight=pretrain` 为什么重要?
|
||||||
|
|
||||||
|
SFT 默认:
|
||||||
|
|
||||||
|
```python
|
||||||
|
parser.add_argument('--from_weight', default='pretrain')
|
||||||
|
```
|
||||||
|
|
||||||
|
如果从 `none` 开始 SFT,模型只看到问答格式,缺少基础语言建模能力,效果会差很多。
|
||||||
|
|
||||||
|
工程理解:
|
||||||
|
|
||||||
|
- pretrain 建立语言和世界知识底座。
|
||||||
|
- SFT 对齐交互格式和回答风格。
|
||||||
|
- 小数据 SFT 不能替代 pretrain。
|
||||||
|
|
||||||
|
## 深挖问题 3:为什么 SFT 学习率低?
|
||||||
|
|
||||||
|
默认:
|
||||||
|
|
||||||
|
- pretrain:`5e-4`
|
||||||
|
- SFT:`1e-5`
|
||||||
|
|
||||||
|
原因:
|
||||||
|
|
||||||
|
- SFT 是在已有模型上微调。
|
||||||
|
- 学习率太大会破坏 pretrain 已学到的分布。
|
||||||
|
- SFT 数据更窄,过大学习率容易过拟合风格或格式。
|
||||||
|
|
||||||
|
## 深挖问题 4:Tool Call 能力来自哪里?
|
||||||
|
|
||||||
|
README 说明当前 `sft_t2t_mini` 已混入 tool call 样本。
|
||||||
|
|
||||||
|
数据格式大致是:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{"role": "system", "tools": "..."}
|
||||||
|
{"role": "user", "content": "..."}
|
||||||
|
{"role": "assistant", "tool_calls": "..."}
|
||||||
|
{"role": "tool", "content": "..."}
|
||||||
|
{"role": "assistant", "content": "..."}
|
||||||
|
```
|
||||||
|
|
||||||
|
模板会展开为:
|
||||||
|
|
||||||
|
```text
|
||||||
|
<tool_call>...</tool_call>
|
||||||
|
<tool_response>...</tool_response>
|
||||||
|
```
|
||||||
|
|
||||||
|
因此 full_sft 权重已经具备基础 tool call 格式能力。
|
||||||
|
|
||||||
|
## 今日实验
|
||||||
|
|
||||||
|
对比:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python eval_llm.py --weight pretrain
|
||||||
|
python eval_llm.py --weight full_sft
|
||||||
|
```
|
||||||
|
|
||||||
|
固定 prompt:
|
||||||
|
|
||||||
|
```text
|
||||||
|
你是谁?
|
||||||
|
解释什么是机器学习。
|
||||||
|
请用 Python 写一个斐波那契函数。
|
||||||
|
推荐一些中国美食。
|
||||||
|
```
|
||||||
|
|
||||||
|
记录:
|
||||||
|
|
||||||
|
- 回答是否像 assistant。
|
||||||
|
- 是否能按指令组织结构。
|
||||||
|
- 是否重复或胡编。
|
||||||
|
- 是否出现模板标签。
|
||||||
|
|
||||||
|
## 复盘问题
|
||||||
|
|
||||||
|
- 为什么 SFT loss 下降不代表模型事实性更强?
|
||||||
|
- 如果 SFT 后模型回答变短,可能是什么数据分布导致?
|
||||||
|
- 如果模型总输出 `<think>` 标签,可能和哪些训练样本有关?
|
||||||
|
|
@ -0,0 +1,53 @@
|
||||||
|
# Day 4 Exam
|
||||||
|
|
||||||
|
## A. 代码阅读题
|
||||||
|
|
||||||
|
1. `train_full_sft.py` 和 `train_pretrain.py` 的训练循环有哪些完全相同的地方?
|
||||||
|
2. SFT 的 `from_weight` 默认值是什么?为什么?
|
||||||
|
3. `max_seq_len=768` 相比 pretrain 默认 `340` 的意义是什么?
|
||||||
|
4. `SFTDataset.generate_labels` 是否会训练多轮对话里的多个 assistant?
|
||||||
|
5. `eval_llm.py` 中 pretrain 权重和 SFT 权重的 prompt 构造有什么区别?
|
||||||
|
|
||||||
|
## B. 对比题
|
||||||
|
|
||||||
|
用同一个 prompt 比较 `pretrain` 和 `full_sft`:
|
||||||
|
|
||||||
|
```text
|
||||||
|
解释什么是机器学习。
|
||||||
|
```
|
||||||
|
|
||||||
|
回答:
|
||||||
|
|
||||||
|
- 哪个更像聊天助手?
|
||||||
|
- 哪个更可能只是续写文本?
|
||||||
|
- 这种差异来自模型结构还是训练数据?
|
||||||
|
|
||||||
|
## C. 工业题
|
||||||
|
|
||||||
|
1. SFT 数据中如果有大量低质量长回答,会对模型造成什么影响?
|
||||||
|
2. 为什么 SFT 前后必须固定一组评测 prompt?
|
||||||
|
3. 如何判断模型是在学能力,还是只是在学模板?
|
||||||
|
4. 领域 SFT 为什么要混合通用数据?
|
||||||
|
|
||||||
|
## D. 实验题
|
||||||
|
|
||||||
|
完成 SFT 后,运行:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python eval_llm.py --weight full_sft
|
||||||
|
cd scripts && python eval_toolcall.py --weight full_sft
|
||||||
|
```
|
||||||
|
|
||||||
|
记录:
|
||||||
|
|
||||||
|
- 普通对话 8 个 prompt 的表现。
|
||||||
|
- tool call 自动测试中成功和失败样例。
|
||||||
|
- 失败原因是格式、工具选择、参数、还是最终回答。
|
||||||
|
|
||||||
|
## E. 通过标准
|
||||||
|
|
||||||
|
你能解释:
|
||||||
|
|
||||||
|
```text
|
||||||
|
SFT 训练循环不是魔法,真正关键是 chat template 和 assistant-only labels。
|
||||||
|
```
|
||||||
|
|
@ -0,0 +1,56 @@
|
||||||
|
# Day 5 Prereading: 评测、推理参数、失败分析
|
||||||
|
|
||||||
|
## 今日目标
|
||||||
|
|
||||||
|
今天不追求继续训练,而是建立评测习惯。你要能判断自己的模型哪里变好了,哪里只是看起来变好了。
|
||||||
|
|
||||||
|
## 必须掌握的基础概念
|
||||||
|
|
||||||
|
- Offline eval:固定测试集,离线比较。
|
||||||
|
- Human eval:人工主观评分。
|
||||||
|
- Regression eval:新权重不能明显退化旧能力。
|
||||||
|
- Sampling parameters:temperature、top_p、top_k、repetition_penalty。
|
||||||
|
- Hallucination:模型生成看似合理但错误的信息。
|
||||||
|
- Exposure bias:训练时看标准前缀,推理时看自己生成的前缀。
|
||||||
|
|
||||||
|
## 面试考点
|
||||||
|
|
||||||
|
- 为什么 loss 不能完全代表生成质量?
|
||||||
|
- temperature 和 top_p 分别控制什么?
|
||||||
|
- 为什么小模型容易重复和幻觉?
|
||||||
|
- 如何设计一个低成本的 LLM eval?
|
||||||
|
- 如何判断 SFT 是否过拟合?
|
||||||
|
|
||||||
|
## 工业要求
|
||||||
|
|
||||||
|
- 每次训练后必须跑固定评测集。
|
||||||
|
- 评测 prompt 要覆盖:身份、事实、代码、总结、推理、工具调用。
|
||||||
|
- 生成评测要固定随机种子或多次采样。
|
||||||
|
- 不要只展示最好样例,要记录失败样例。
|
||||||
|
- 小模型不要用大模型标准苛责,但要诚实记录边界。
|
||||||
|
|
||||||
|
## 项目中要看的文件
|
||||||
|
|
||||||
|
- `eval_llm.py`
|
||||||
|
- `scripts/eval_toolcall.py`
|
||||||
|
- `scripts/serve_openai_api.py`
|
||||||
|
- `scripts/web_demo.py`
|
||||||
|
|
||||||
|
## 今日固定评测集
|
||||||
|
|
||||||
|
```text
|
||||||
|
1. 你是谁?
|
||||||
|
2. 解释什么是机器学习。
|
||||||
|
3. 为什么天空是蓝色的?
|
||||||
|
4. 请用 Python 写一个斐波那契函数。
|
||||||
|
5. 比较猫和狗作为宠物的优缺点。
|
||||||
|
6. 推荐一些中国美食。
|
||||||
|
7. 把下面这段话总结成三点:...
|
||||||
|
8. 如果明天下雨,我应该如何出门?
|
||||||
|
```
|
||||||
|
|
||||||
|
## 今日注意事项
|
||||||
|
|
||||||
|
- 同一个 prompt 至少用 2 组采样参数测试。
|
||||||
|
- 把失败样例原样保存。
|
||||||
|
- 评估时区分:语言流畅、指令跟随、事实正确、格式正确。
|
||||||
|
|
@ -0,0 +1,84 @@
|
||||||
|
# Day 5 Deep Dive: 生成质量为什么不等于训练 loss
|
||||||
|
|
||||||
|
## 深挖问题 1:eval_llm 如何构造输入?
|
||||||
|
|
||||||
|
打开 `eval_llm.py`。
|
||||||
|
|
||||||
|
关键差异:
|
||||||
|
|
||||||
|
```python
|
||||||
|
if 'pretrain' in args.weight:
|
||||||
|
inputs = tokenizer.bos_token + prompt
|
||||||
|
else:
|
||||||
|
inputs = tokenizer.apply_chat_template(...)
|
||||||
|
```
|
||||||
|
|
||||||
|
说明:
|
||||||
|
|
||||||
|
- pretrain 模型只适合文本续写。
|
||||||
|
- SFT 模型适合 chat template。
|
||||||
|
- 如果推理模板和训练模板不一致,效果会明显变差。
|
||||||
|
|
||||||
|
## 深挖问题 2:采样参数如何影响输出?
|
||||||
|
|
||||||
|
`generate` 中关键参数:
|
||||||
|
|
||||||
|
- `temperature`:logits 除以温度,越高越随机。
|
||||||
|
- `top_p`:只在累计概率前 p 的 token 中采样。
|
||||||
|
- `top_k`:只在概率最高 k 个 token 中采样。
|
||||||
|
- `repetition_penalty`:惩罚已出现 token。
|
||||||
|
|
||||||
|
实验建议:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python eval_llm.py --weight full_sft --temperature 0.3 --top_p 0.9
|
||||||
|
python eval_llm.py --weight full_sft --temperature 0.9 --top_p 0.95
|
||||||
|
```
|
||||||
|
|
||||||
|
观察:
|
||||||
|
|
||||||
|
- 低温更稳定,但可能死板。
|
||||||
|
- 高温更多样,但更容易胡编。
|
||||||
|
|
||||||
|
## 深挖问题 3:为什么 loss 低但生成差?
|
||||||
|
|
||||||
|
常见原因:
|
||||||
|
|
||||||
|
- 数据分布简单,loss 易降但泛化差。
|
||||||
|
- teacher-forcing 训练,推理时错误会累积。
|
||||||
|
- 评测 prompt 超出训练分布。
|
||||||
|
- 模板不一致。
|
||||||
|
- 采样参数过激。
|
||||||
|
- 小模型容量不足。
|
||||||
|
|
||||||
|
## 深挖问题 4:Tool call 如何评测?
|
||||||
|
|
||||||
|
打开 `scripts/eval_toolcall.py`。
|
||||||
|
|
||||||
|
关注:
|
||||||
|
|
||||||
|
- 工具 schema 如何给模型。
|
||||||
|
- 模型如何输出 `<tool_call>`。
|
||||||
|
- `parse_tool_calls` 如何解析。
|
||||||
|
- `execute_tool` 如何执行。
|
||||||
|
- 工具结果如何追加回 messages。
|
||||||
|
|
||||||
|
失败类型分类:
|
||||||
|
|
||||||
|
```text
|
||||||
|
格式失败:JSON 不合法或标签不闭合
|
||||||
|
选择失败:调用了错误工具
|
||||||
|
参数失败:工具对但参数错
|
||||||
|
利用失败:工具结果对,但最终回答错
|
||||||
|
```
|
||||||
|
|
||||||
|
## 今日产出
|
||||||
|
|
||||||
|
写一份 `eval_report.md`,至少包含:
|
||||||
|
|
||||||
|
- 评测命令
|
||||||
|
- prompt 列表
|
||||||
|
- 每个 prompt 的回答
|
||||||
|
- 评分表
|
||||||
|
- 失败样例
|
||||||
|
- 下一步改进建议
|
||||||
|
|
@ -0,0 +1,53 @@
|
||||||
|
# Day 5 Exam
|
||||||
|
|
||||||
|
## A. 代码阅读题
|
||||||
|
|
||||||
|
1. `eval_llm.py` 中 `open_thinking` 如何传给 chat template?
|
||||||
|
2. `temperature` 在 `model.generate` 中如何影响 logits?
|
||||||
|
3. `top_p` 和 `top_k` 可以同时使用吗?有什么效果?
|
||||||
|
4. `eval_toolcall.py` 的多轮工具调用循环何时停止?
|
||||||
|
5. tool call 失败时,应该先看模型输出还是工具执行?
|
||||||
|
|
||||||
|
## B. 评测设计题
|
||||||
|
|
||||||
|
为 MiniMind 设计一个 20 题小评测集,覆盖:
|
||||||
|
|
||||||
|
- 身份/自我认知
|
||||||
|
- 基础事实
|
||||||
|
- 中文表达
|
||||||
|
- 代码
|
||||||
|
- 总结
|
||||||
|
- 简单数学
|
||||||
|
- tool call
|
||||||
|
|
||||||
|
每题给出评分标准:0/1/2 分。
|
||||||
|
|
||||||
|
## C. 失败分析题
|
||||||
|
|
||||||
|
对于以下现象,写出至少 3 个可能原因:
|
||||||
|
|
||||||
|
1. 模型回答流畅但事实错误。
|
||||||
|
2. 模型重复同一句话。
|
||||||
|
3. 模型不按 JSON 输出 tool call。
|
||||||
|
4. 模型回答特别短。
|
||||||
|
5. 模型输出大量 `<think>` 但内容无意义。
|
||||||
|
|
||||||
|
## D. 实验题
|
||||||
|
|
||||||
|
用同一组 prompt,分别测试:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python eval_llm.py --weight full_sft --temperature 0.3 --top_p 0.9
|
||||||
|
python eval_llm.py --weight full_sft --temperature 0.9 --top_p 0.95
|
||||||
|
```
|
||||||
|
|
||||||
|
记录不同采样参数下:
|
||||||
|
|
||||||
|
- 稳定性
|
||||||
|
- 创造性
|
||||||
|
- 幻觉
|
||||||
|
- 重复
|
||||||
|
|
||||||
|
## E. 通过标准
|
||||||
|
|
||||||
|
你能给一个模型输出做结构化诊断,而不是只说“好”或“不好”。
|
||||||
|
|
@ -0,0 +1,63 @@
|
||||||
|
# Day 6 Prereading: 蒸馏与 GRPO/CISPO 后训练
|
||||||
|
|
||||||
|
## 今日目标
|
||||||
|
|
||||||
|
今天选择一个后训练方向深入:白盒蒸馏或 GRPO/CISPO。推荐优先读 GRPO,因为你当前打开的是 `train_grpo.py`,但如果 reward model 没准备好,可以把实跑换成蒸馏代码理解。
|
||||||
|
|
||||||
|
## 必须掌握的基础概念
|
||||||
|
|
||||||
|
### 蒸馏
|
||||||
|
|
||||||
|
- Black-box distillation:学习 teacher 生成的硬标签答案。
|
||||||
|
- White-box distillation:学习 teacher logits/probability distribution。
|
||||||
|
- Temperature:软化概率分布。
|
||||||
|
- KL divergence:衡量 teacher/student 分布差异。
|
||||||
|
|
||||||
|
### GRPO/CISPO
|
||||||
|
|
||||||
|
- On-policy rollout:用当前模型实时采样。
|
||||||
|
- Reward model:给回答打分。
|
||||||
|
- Reference model:限制 policy 偏离原模型。
|
||||||
|
- Advantage:回答比同组其他回答好多少。
|
||||||
|
- KL penalty:防止 reward hacking 和分布漂移。
|
||||||
|
- CISPO:把 clipped ratio 作为权重,保留 logprob 梯度路径。
|
||||||
|
|
||||||
|
## 面试考点
|
||||||
|
|
||||||
|
- DPO、PPO、GRPO 的主要区别是什么?
|
||||||
|
- GRPO 为什么可以不训练 critic?
|
||||||
|
- 为什么每个 prompt 要采样多个 responses?
|
||||||
|
- KL penalty 防什么?
|
||||||
|
- 蒸馏中的 temperature 为什么要乘 `T^2`?
|
||||||
|
- Reward hacking 是什么?
|
||||||
|
|
||||||
|
## 工业要求
|
||||||
|
|
||||||
|
- RL 后训练必须监控 reward、KL、response length、advantage std。
|
||||||
|
- 只看 reward 上升很危险,可能是 reward hacking。
|
||||||
|
- 小模型 RL 容易奖励稀疏,要选择能力边界内任务。
|
||||||
|
- 蒸馏要确认 teacher/student vocab 对齐或正确裁剪。
|
||||||
|
- 后训练前必须保留 SFT baseline,便于回退和对比。
|
||||||
|
|
||||||
|
## 项目中要看的文件
|
||||||
|
|
||||||
|
- `trainer/train_grpo.py`
|
||||||
|
- `trainer/train_agent.py`
|
||||||
|
- `trainer/rollout_engine.py`
|
||||||
|
- `trainer/train_distillation.py`
|
||||||
|
- `dataset/lm_dataset.py::RLAIFDataset`
|
||||||
|
- `trainer/trainer_utils.py::LMForRewardModel`
|
||||||
|
|
||||||
|
## 今日建议
|
||||||
|
|
||||||
|
如果 reward model 已准备好:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd trainer
|
||||||
|
python train_grpo.py --debug_mode --debug_interval 5
|
||||||
|
```
|
||||||
|
|
||||||
|
如果没有 reward model:
|
||||||
|
|
||||||
|
- 先只读代码。
|
||||||
|
- 或选择 `train_distillation.py` 做理解。
|
||||||
|
|
@ -0,0 +1,120 @@
|
||||||
|
# Day 6 Deep Dive: GRPO/CISPO 与蒸馏
|
||||||
|
|
||||||
|
## 深挖问题 1:GRPO 的训练主线是什么?
|
||||||
|
|
||||||
|
代码位置:`trainer/train_grpo.py::grpo_train_epoch`
|
||||||
|
|
||||||
|
流程:
|
||||||
|
|
||||||
|
```text
|
||||||
|
batch prompts
|
||||||
|
-> rollout_engine.rollout 生成多个 completions
|
||||||
|
-> calculate_rewards 得到 reward
|
||||||
|
-> 当前 policy 重新算 per-token logp
|
||||||
|
-> ref_model 算 ref logp
|
||||||
|
-> group rewards -> advantages
|
||||||
|
-> 计算 KL
|
||||||
|
-> GRPO/CISPO policy loss
|
||||||
|
-> optimizer step
|
||||||
|
-> update rollout policy
|
||||||
|
```
|
||||||
|
|
||||||
|
这是 online RL,不是普通离线 SFT。
|
||||||
|
|
||||||
|
## 深挖问题 2:为什么 `num_generations` 是关键?
|
||||||
|
|
||||||
|
GRPO 对同一个 prompt 生成 N 个回答:
|
||||||
|
|
||||||
|
```python
|
||||||
|
grouped_rewards = rewards.view(-1, args.num_generations)
|
||||||
|
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)
|
||||||
|
std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)
|
||||||
|
advantages = (rewards - mean_r) / (std_r + 1e-4)
|
||||||
|
```
|
||||||
|
|
||||||
|
含义:
|
||||||
|
|
||||||
|
- 同题内部比较,降低不同题难度差异影响。
|
||||||
|
- 高于组均值的回答被鼓励。
|
||||||
|
- 低于组均值的回答被压制。
|
||||||
|
- 如果同组 reward 全一样,advantage 接近 0,学习信号消失。
|
||||||
|
|
||||||
|
## 深挖问题 3:reward 是怎么来的?
|
||||||
|
|
||||||
|
代码位置:`calculate_rewards`
|
||||||
|
|
||||||
|
包含:
|
||||||
|
|
||||||
|
- 长度奖励/惩罚。
|
||||||
|
- thinking 格式奖励。
|
||||||
|
- 重复惩罚 `rep_penalty`。
|
||||||
|
- reward model score。
|
||||||
|
|
||||||
|
工业理解:
|
||||||
|
|
||||||
|
- reward 不是客观真理,是训练目标。
|
||||||
|
- 设计不好会鼓励模型钻空子。
|
||||||
|
- 必须同时观察 reward 和真实输出。
|
||||||
|
|
||||||
|
## 深挖问题 4:GRPO 和 CISPO loss 差异
|
||||||
|
|
||||||
|
GRPO:
|
||||||
|
|
||||||
|
```python
|
||||||
|
ratio = exp(new_logp - old_logp)
|
||||||
|
clipped_ratio = clamp(ratio, 1-eps, 1+eps)
|
||||||
|
loss = -min(ratio*A, clipped_ratio*A) + beta*KL
|
||||||
|
```
|
||||||
|
|
||||||
|
CISPO:
|
||||||
|
|
||||||
|
```python
|
||||||
|
clamped_ratio = clamp(ratio, max=epsilon_high).detach()
|
||||||
|
loss = -(clamped_ratio * A * per_token_logps - beta*KL)
|
||||||
|
```
|
||||||
|
|
||||||
|
核心差异:
|
||||||
|
|
||||||
|
- GRPO/PPO-style clip 可能让梯度在 clip 后被截断。
|
||||||
|
- CISPO 把 ratio 当作 detached 权重,梯度仍来自 `log pi`。
|
||||||
|
|
||||||
|
## 深挖问题 5:蒸馏如何训练?
|
||||||
|
|
||||||
|
代码位置:`trainer/train_distillation.py`
|
||||||
|
|
||||||
|
总损失:
|
||||||
|
|
||||||
|
```python
|
||||||
|
loss = alpha * ce_loss + (1 - alpha) * distill_loss
|
||||||
|
```
|
||||||
|
|
||||||
|
其中:
|
||||||
|
|
||||||
|
- CE:学生拟合硬标签。
|
||||||
|
- KL:学生拟合 teacher token 分布。
|
||||||
|
- Temperature:让 teacher 分布更软,暴露 token 间偏好。
|
||||||
|
|
||||||
|
适合理解:
|
||||||
|
|
||||||
|
- teacher/student 双模型加载。
|
||||||
|
- teacher `eval()` 和 `requires_grad_(False)`。
|
||||||
|
- 只在 loss mask 位置做 KL。
|
||||||
|
|
||||||
|
## 今日输出
|
||||||
|
|
||||||
|
任选一个方向写 1 页报告:
|
||||||
|
|
||||||
|
### GRPO 报告
|
||||||
|
|
||||||
|
- rollout 输入输出是什么?
|
||||||
|
- old logp / new logp / ref logp 分别代表什么?
|
||||||
|
- advantage 怎么算?
|
||||||
|
- KL 为什么必要?
|
||||||
|
- 你认为当前 reward 设计有什么风险?
|
||||||
|
|
||||||
|
### 蒸馏报告
|
||||||
|
|
||||||
|
- CE 和 KL 各自训练什么?
|
||||||
|
- temperature 如何影响 teacher distribution?
|
||||||
|
- teacher/student 结构不同时要注意什么?
|
||||||
|
- 为什么 teacher 不需要反向传播?
|
||||||
|
|
@ -0,0 +1,53 @@
|
||||||
|
# Day 6 Exam
|
||||||
|
|
||||||
|
## A. GRPO 代码阅读题
|
||||||
|
|
||||||
|
1. `rollout_engine.rollout` 返回哪些字段?每个字段用于什么?
|
||||||
|
2. `old_per_token_logps` 和 `per_token_logps` 有什么区别?
|
||||||
|
3. `ref_per_token_logps` 为什么用 `torch.no_grad()`?
|
||||||
|
4. `completion_mask` 为什么要截到 EOS?
|
||||||
|
5. `advantages_std` 接近 0 说明什么?
|
||||||
|
|
||||||
|
## B. GRPO 推导题
|
||||||
|
|
||||||
|
给定同一个 prompt 的 4 个 reward:
|
||||||
|
|
||||||
|
```text
|
||||||
|
[1.0, 2.0, 2.0, 5.0]
|
||||||
|
```
|
||||||
|
|
||||||
|
计算:
|
||||||
|
|
||||||
|
- mean
|
||||||
|
- std
|
||||||
|
- 每个 response 的 advantage
|
||||||
|
|
||||||
|
回答:哪个 response 被最强鼓励?哪个被压制?
|
||||||
|
|
||||||
|
## C. CISPO 理解题
|
||||||
|
|
||||||
|
1. CISPO 为什么对 ratio 使用 `.detach()`?
|
||||||
|
2. `epsilon_high` 太大或太小分别有什么风险?
|
||||||
|
3. CISPO 和 GRPO 哪个更像 PPO clip?
|
||||||
|
|
||||||
|
## D. 蒸馏代码阅读题
|
||||||
|
|
||||||
|
1. `distillation_loss` 中 teacher logits 为什么要 `detach()`?
|
||||||
|
2. 为什么 KL 只在 `loss_mask == 1` 的 token 上算?
|
||||||
|
3. `alpha=0.5` 表示什么?
|
||||||
|
4. teacher vocab 大于 student vocab 时,代码如何处理?
|
||||||
|
|
||||||
|
## E. 工业题
|
||||||
|
|
||||||
|
1. RL 训练 reward 上升但人工评测下降,可能是什么问题?
|
||||||
|
2. KL 持续升高说明什么?
|
||||||
|
3. response length 越来越长,可能是 reward 哪部分导致?
|
||||||
|
4. 小模型做数学 RLVR 为什么容易奖励稀疏?
|
||||||
|
|
||||||
|
## F. 通过标准
|
||||||
|
|
||||||
|
你能讲清楚:
|
||||||
|
|
||||||
|
```text
|
||||||
|
SFT 是模仿标准答案;GRPO 是采样多个答案后按相对奖励调整概率;蒸馏是拟合 teacher 的硬答案或软分布。
|
||||||
|
```
|
||||||
|
|
@ -0,0 +1,57 @@
|
||||||
|
# Day 7 Prereading: 总复盘、交付、工程化思维
|
||||||
|
|
||||||
|
## 今日目标
|
||||||
|
|
||||||
|
今天不再堆新内容。你要把一周训练流程整理成可复现报告,并完成最终考核。
|
||||||
|
|
||||||
|
## 必须掌握的基础概念
|
||||||
|
|
||||||
|
- Reproducibility:别人能按你的记录复现结果。
|
||||||
|
- Ablation:只改变一个变量观察影响。
|
||||||
|
- Baseline:每次改动都要和基线比较。
|
||||||
|
- Checkpoint hygiene:权重命名、保存、恢复、版本管理清晰。
|
||||||
|
- Evaluation discipline:固定测试集、记录失败、不挑样例。
|
||||||
|
|
||||||
|
## 面试考点
|
||||||
|
|
||||||
|
- 从零训练一个小 LLM 的完整流程是什么?
|
||||||
|
- Pretrain、SFT、DPO/GRPO、Distillation 的顺序和作用是什么?
|
||||||
|
- 训练中你最关注哪些指标?
|
||||||
|
- 如果效果不好,你如何定位是数据、模型、优化还是评测问题?
|
||||||
|
- 如何把研究代码变成可靠训练 pipeline?
|
||||||
|
|
||||||
|
## 工业要求
|
||||||
|
|
||||||
|
- 命令可复制。
|
||||||
|
- 数据路径明确。
|
||||||
|
- 权重命名明确。
|
||||||
|
- 每个实验有目的、有结果、有结论。
|
||||||
|
- 失败样例和成功样例一样重要。
|
||||||
|
- 不把一次随机生成当作模型能力结论。
|
||||||
|
|
||||||
|
## 今日交付
|
||||||
|
|
||||||
|
1. `training_log.md`
|
||||||
|
2. `eval_report.md`
|
||||||
|
3. `code_understanding.md`
|
||||||
|
4. `final_exam_answers.md`
|
||||||
|
|
||||||
|
建议结构:
|
||||||
|
|
||||||
|
```text
|
||||||
|
环境
|
||||||
|
数据
|
||||||
|
模型配置
|
||||||
|
训练命令
|
||||||
|
loss 记录
|
||||||
|
推理评测
|
||||||
|
失败分析
|
||||||
|
后训练理解
|
||||||
|
下一步计划
|
||||||
|
```
|
||||||
|
|
||||||
|
## 今日注意事项
|
||||||
|
|
||||||
|
- 不要隐藏没跑通的部分。
|
||||||
|
- 不要只写结论,要写证据。
|
||||||
|
- 不要凭感觉说模型变好,要给对比样例。
|
||||||
|
|
@ -0,0 +1,77 @@
|
||||||
|
# Day 7 Deep Dive: 把 MiniMind 训练链路串起来
|
||||||
|
|
||||||
|
## 总链路
|
||||||
|
|
||||||
|
```text
|
||||||
|
Tokenizer
|
||||||
|
-> Pretrain
|
||||||
|
-> SFT
|
||||||
|
-> Evaluation
|
||||||
|
-> Optional: Distillation / DPO / PPO / GRPO / Agent RL
|
||||||
|
-> Convert / Serve / Web Demo
|
||||||
|
```
|
||||||
|
|
||||||
|
## 关键连接 1:模型结构和训练目标
|
||||||
|
|
||||||
|
模型只负责:
|
||||||
|
|
||||||
|
```text
|
||||||
|
input_ids -> logits
|
||||||
|
```
|
||||||
|
|
||||||
|
训练目标由 labels 决定:
|
||||||
|
|
||||||
|
- Pretrain labels:几乎所有文本 token。
|
||||||
|
- SFT labels:assistant token。
|
||||||
|
- DPO labels/masks:chosen/rejected assistant span。
|
||||||
|
- GRPO:生成 completion 后用 reward 调整 logprob。
|
||||||
|
|
||||||
|
真正改变学习行为的,不只是模型,而是目标函数和数据组织。
|
||||||
|
|
||||||
|
## 关键连接 2:SFT 为什么是主线核心
|
||||||
|
|
||||||
|
MiniMind 一周训练目标应以 `pretrain + full_sft` 为核心,因为:
|
||||||
|
|
||||||
|
- pretrain 让模型学会语言建模。
|
||||||
|
- SFT 让模型学会对话模板和助手行为。
|
||||||
|
- 这两步能在个人 GPU 上稳定复现。
|
||||||
|
- RL/蒸馏更适合理解和小规模实验,不适合作为第一周硬交付。
|
||||||
|
|
||||||
|
## 关键连接 3:RL 为什么难
|
||||||
|
|
||||||
|
RL 多了这些不稳定来源:
|
||||||
|
|
||||||
|
- policy 在线生成,数据分布不断变。
|
||||||
|
- reward model 不等于真实人类偏好。
|
||||||
|
- KL 太小会跑偏,太大会学不动。
|
||||||
|
- reward 方差太小没有学习信号。
|
||||||
|
- 小模型能力弱,容易产生无效 rollout。
|
||||||
|
|
||||||
|
所以 RL 训练要看:
|
||||||
|
|
||||||
|
- reward
|
||||||
|
- KL
|
||||||
|
- advantage std
|
||||||
|
- response length
|
||||||
|
- debug samples
|
||||||
|
- 人工评测
|
||||||
|
|
||||||
|
## 关键连接 4:工程上如何扩展
|
||||||
|
|
||||||
|
如果要继续提升:
|
||||||
|
|
||||||
|
1. 先扩大或清洗 SFT 数据。
|
||||||
|
2. 增加固定评测集。
|
||||||
|
3. 调整 `max_seq_len` 和 batch。
|
||||||
|
4. 尝试 MoE,但记录训练成本。
|
||||||
|
5. 尝试 LoRA 做领域适配。
|
||||||
|
6. 再做 GRPO/Agent RL,且必须设计可靠 reward。
|
||||||
|
|
||||||
|
## 最终复盘问题
|
||||||
|
|
||||||
|
你要能回答:
|
||||||
|
|
||||||
|
- 当前模型最强能力是什么?
|
||||||
|
- 最弱能力是什么?
|
||||||
|
- 训练瓶颈是数据、模型容量、训练时长还是评测方式?
|
||||||
|
- 下一周最值得投入的改进是什么?
|
||||||
|
|
@ -0,0 +1,78 @@
|
||||||
|
# Day 7 Final Exam
|
||||||
|
|
||||||
|
## A. 闭卷概念题
|
||||||
|
|
||||||
|
1. 从零训练 MiniMind,按顺序需要哪些阶段?
|
||||||
|
2. Pretrain 和 SFT 的 labels 有什么区别?
|
||||||
|
3. 为什么 `ignore_index=-100` 是训练中的关键设计?
|
||||||
|
4. GRPO 为什么不需要 critic?
|
||||||
|
5. 蒸馏中的 CE 和 KL 分别学什么?
|
||||||
|
6. checkpoint 为什么要保存 optimizer?
|
||||||
|
7. 为什么训练 loss 不能替代生成评测?
|
||||||
|
8. Tool call 能力在 MiniMind 中主要从哪里来?
|
||||||
|
|
||||||
|
## B. 代码定位题
|
||||||
|
|
||||||
|
写出以下功能所在文件和函数:
|
||||||
|
|
||||||
|
1. MiniMind causal LM loss。
|
||||||
|
2. SFT assistant-only labels。
|
||||||
|
3. Pretrain 数据 padding 和 label mask。
|
||||||
|
4. Checkpoint 保存和恢复。
|
||||||
|
5. GRPO reward 计算。
|
||||||
|
6. Rollout engine 抽象。
|
||||||
|
7. 推理时 chat template 构造。
|
||||||
|
8. Tool call 自动测试入口。
|
||||||
|
|
||||||
|
## C. 实战题
|
||||||
|
|
||||||
|
重新从空日志开始,写出你会执行的完整命令:
|
||||||
|
|
||||||
|
1. 环境检查。
|
||||||
|
2. pretrain。
|
||||||
|
3. resume pretrain。
|
||||||
|
4. SFT。
|
||||||
|
5. eval pretrain。
|
||||||
|
6. eval full_sft。
|
||||||
|
7. eval toolcall。
|
||||||
|
|
||||||
|
每条命令后写一句:这一步验证什么。
|
||||||
|
|
||||||
|
## D. Debug 场景题
|
||||||
|
|
||||||
|
给出排查路径:
|
||||||
|
|
||||||
|
1. SFT loss 正常下降,但模型推理输出乱码。
|
||||||
|
2. 模型只会复述用户问题。
|
||||||
|
3. GRPO reward 上升,但人工看回答变差。
|
||||||
|
4. 训练恢复后 step 对不上。
|
||||||
|
5. Tool call JSON 经常非法。
|
||||||
|
|
||||||
|
## E. 最终报告题
|
||||||
|
|
||||||
|
写一份 2 页以内最终报告,必须包含:
|
||||||
|
|
||||||
|
- 训练环境
|
||||||
|
- 数据文件
|
||||||
|
- 模型配置
|
||||||
|
- 训练命令
|
||||||
|
- loss 记录
|
||||||
|
- pretrain vs full_sft 对比
|
||||||
|
- 失败样例
|
||||||
|
- 你对 GRPO 或蒸馏的理解
|
||||||
|
- 下一步最值得做的 3 个实验
|
||||||
|
|
||||||
|
## F. 评分标准
|
||||||
|
|
||||||
|
总分 100:
|
||||||
|
|
||||||
|
- 20:成功训练并推理 `pretrain` / `full_sft`
|
||||||
|
- 20:理解 dataset、template、label mask
|
||||||
|
- 15:理解模型结构和 causal LM loss
|
||||||
|
- 15:理解训练工程和 checkpoint
|
||||||
|
- 15:能做结构化评测和失败分析
|
||||||
|
- 15:理解至少一个后训练算法
|
||||||
|
|
||||||
|
通过线:75。
|
||||||
|
|
||||||
|
优秀线:90,且最终报告里有清晰失败分析和下一步实验设计。
|
||||||
|
|
@ -4,219 +4,405 @@ from transformers.activations import ACT2FN
|
||||||
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
|
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
|
||||||
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
|
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
|
||||||
|
|
||||||
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
|
|
||||||
# MiniMind Config
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
|
# 整体架构协作关系
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
#
|
||||||
|
# MiniMindForCausalLM (最顶层,训练和推理的入口)
|
||||||
|
# ├── MiniMindModel (Transformer 主干)
|
||||||
|
# │ ├── embed_tokens (Embedding层: token_id → 向量)
|
||||||
|
# │ ├── MiniMindBlock × N (N=8 层 Transformer 块,核心计算单元)
|
||||||
|
# │ │ ├── input_layernorm (RMSNorm: 注意力前的归一化)
|
||||||
|
# │ │ ├── Attention (GQA多头注意力 + RoPE位置编码)
|
||||||
|
# │ │ │ ├── q_proj, k_proj, v_proj (线性投影)
|
||||||
|
# │ │ │ ├── q_norm, k_norm (QK归一化,防止注意力分数爆炸)
|
||||||
|
# │ │ │ ├── apply_rotary_pos_emb (RoPE旋转位置编码)
|
||||||
|
# │ │ │ ├── repeat_kv (GQA: 复制KV头以匹配Q头数量)
|
||||||
|
# │ │ │ ├── Flash Attention 或 手动注意力计算
|
||||||
|
# │ │ │ └── o_proj (输出投影)
|
||||||
|
# │ │ ├── post_attention_layernorm (RMSNorm: FFN前的归一化)
|
||||||
|
# │ │ └── mlp (前馈网络,二选一):
|
||||||
|
# │ │ ├── FeedForward (Dense版: SwiGLU激活)
|
||||||
|
# │ │ └── MOEFeedForward (MoE版: Router + 多个Expert)
|
||||||
|
# │ └── norm (最终RMSNorm)
|
||||||
|
# └── lm_head (Linear: 隐藏向量 → 词表logits,与embed_tokens共享权重)
|
||||||
|
#
|
||||||
|
# 数据流 (训练时):
|
||||||
|
# input_ids [B, S] → Embedding → [B, S, 768] → 8层Block → RMSNorm → lm_head
|
||||||
|
# → logits [B, S, 6400] → CrossEntropy(logits[:-1], labels[1:]) → loss
|
||||||
|
#
|
||||||
|
# 数据流 (推理/generate时):
|
||||||
|
# 每步只输入1个新token → 用KV Cache避免重复计算 → 取最后位置logits
|
||||||
|
# → temperature/top_k/top_p采样 → 输出next_token → 循环直到EOS
|
||||||
|
#
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
|
||||||
|
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# MiniMind Config
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 模型的所有超参数配置,对齐 Qwen3 生态
|
||||||
class MiniMindConfig(PretrainedConfig):
|
class MiniMindConfig(PretrainedConfig):
|
||||||
model_type = "minimind"
|
model_type = "minimind"
|
||||||
def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
|
def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
|
||||||
super().__init__(**kwargs)
|
super().__init__(**kwargs)
|
||||||
self.hidden_size = hidden_size
|
# ---- 基础结构参数 ----
|
||||||
self.num_hidden_layers = num_hidden_layers
|
self.hidden_size = hidden_size # 隐藏层维度 d_model,整个模型的"宽度"
|
||||||
self.use_moe = use_moe
|
self.num_hidden_layers = num_hidden_layers # Transformer 层数,模型的"深度"
|
||||||
self.dropout = kwargs.get("dropout", 0.0)
|
self.use_moe = use_moe # 是否使用 MoE (Mixture of Experts)
|
||||||
self.vocab_size = kwargs.get("vocab_size", 6400)
|
self.dropout = kwargs.get("dropout", 0.0) # Dropout率,默认0(不丢弃)
|
||||||
self.bos_token_id = kwargs.get("bos_token_id", 1)
|
|
||||||
self.eos_token_id = kwargs.get("eos_token_id", 2)
|
# ---- 词表与特殊token ----
|
||||||
self.flash_attn = kwargs.get("flash_attn", True)
|
self.vocab_size = kwargs.get("vocab_size", 6400) # 词表大小(极小词表,压缩参数量)
|
||||||
self.num_attention_heads = kwargs.get("num_attention_heads", 8)
|
self.bos_token_id = kwargs.get("bos_token_id", 1) # 句子开始token
|
||||||
self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
|
self.eos_token_id = kwargs.get("eos_token_id", 2) # 句子结束token
|
||||||
self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
|
|
||||||
self.hidden_act = kwargs.get("hidden_act", 'silu')
|
# ---- 注意力参数 ----
|
||||||
|
self.flash_attn = kwargs.get("flash_attn", True) # 是否启用Flash Attention加速
|
||||||
|
self.num_attention_heads = kwargs.get("num_attention_heads", 8) # Q头数量(8个)
|
||||||
|
self.num_key_value_heads = kwargs.get("num_key_value_heads", 4) # KV头数量(4个,GQA模式: 每2个Q头共享1个KV头)
|
||||||
|
self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads) # 每个头的维度 = 768/8 = 96
|
||||||
|
|
||||||
|
# ---- FFN参数 ----
|
||||||
|
self.hidden_act = kwargs.get("hidden_act", 'silu') # 激活函数 (SiLU,用于SwiGLU)
|
||||||
|
# intermediate_size = ceil(768 * π / 64) * 64 = 2432,FFN的中间维度(约hidden_size的3.14倍)
|
||||||
self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
|
self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
|
||||||
self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
|
|
||||||
self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
|
# ---- 位置编码参数 ----
|
||||||
self.rope_theta = kwargs.get("rope_theta", 1e6)
|
self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768) # 最大位置数(32K)
|
||||||
self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)
|
self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6) # RMSNorm的epsilon防除零
|
||||||
|
self.rope_theta = kwargs.get("rope_theta", 1e6) # RoPE基础频率(越大支持越长序列)
|
||||||
|
|
||||||
|
# ---- 权重绑定 ----
|
||||||
|
self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True) # embedding和lm_head共享权重,省一份参数
|
||||||
|
|
||||||
|
# ---- YaRN长文本外推 ----
|
||||||
self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
|
self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
|
||||||
self.rope_scaling = {
|
self.rope_scaling = {
|
||||||
"beta_fast": 32,
|
"beta_fast": 32, # 高频维度边界(这些维度不缩放)
|
||||||
"beta_slow": 1,
|
"beta_slow": 1, # 低频维度边界(这些维度做最大缩放)
|
||||||
"factor": 16,
|
"factor": 16, # 缩放因子(目标长度/原始长度)
|
||||||
"original_max_position_embeddings": 2048,
|
"original_max_position_embeddings": 2048, # 原始训练长度
|
||||||
"attention_factor": 1.0,
|
"attention_factor": 1.0, # 注意力缩放系数
|
||||||
"type": "yarn"
|
"type": "yarn"
|
||||||
} if self.inference_rope_scaling else None
|
} if self.inference_rope_scaling else None
|
||||||
### MoE specific configs (ignored if use_moe = False)
|
|
||||||
self.num_experts = kwargs.get("num_experts", 4)
|
|
||||||
self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
|
|
||||||
self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
|
|
||||||
self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
|
|
||||||
self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)
|
|
||||||
|
|
||||||
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
|
# ---- MoE专用参数 (use_moe=False时忽略) ----
|
||||||
# MiniMind Model
|
self.num_experts = kwargs.get("num_experts", 4) # Expert总数
|
||||||
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
|
self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1) # 每个token激活几个Expert (top-1)
|
||||||
|
self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size) # 每个Expert的FFN中间维度
|
||||||
|
self.norm_topk_prob = kwargs.get("norm_topk_prob", True) # 是否归一化路由概率
|
||||||
|
self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4) # 负载均衡loss系数
|
||||||
|
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 基础组件
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# ---- RMSNorm: 比LayerNorm更简单高效的归一化 ----
|
||||||
|
# LayerNorm: (x - mean) / std * γ + β (需要均值和方差)
|
||||||
|
# RMSNorm: x / RMS(x) * γ (只需均方根,省掉均值和偏置)
|
||||||
class RMSNorm(torch.nn.Module):
|
class RMSNorm(torch.nn.Module):
|
||||||
def __init__(self, dim: int, eps: float = 1e-5):
|
def __init__(self, dim: int, eps: float = 1e-5):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.eps = eps
|
self.eps = eps
|
||||||
self.weight = nn.Parameter(torch.ones(dim))
|
self.weight = nn.Parameter(torch.ones(dim)) # 可学习的缩放参数 γ
|
||||||
|
|
||||||
def norm(self, x):
|
def norm(self, x):
|
||||||
|
# RMS(x) = sqrt(mean(x²)),rsqrt = 1/sqrt
|
||||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
|
# 先转float32算归一化(防精度损失),再转回原dtype
|
||||||
return (self.weight * self.norm(x.float())).type_as(x)
|
return (self.weight * self.norm(x.float())).type_as(x)
|
||||||
|
|
||||||
|
# ---- RoPE: 旋转位置编码 (预计算频率表) ----
|
||||||
|
# 核心思想: 不把位置"加"到embedding上,而是把Q和K向量"旋转"一个与位置相关的角度
|
||||||
|
# 这样 Q_i · K_j 的点积自然编码了相对距离 |i-j|
|
||||||
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
|
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
|
||||||
|
# 每个维度对应一个频率: freq_i = 1 / (θ^(2i/d))
|
||||||
|
# 低维度 → 高频(捕捉局部位置),高维度 → 低频(捕捉远距离位置)
|
||||||
freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
|
freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
|
||||||
if rope_scaling is not None: # YaRN: f'(i) = f(i)((1-γ) + γ/s), where γ∈[0,1] is linear ramp
|
|
||||||
|
# YaRN外推: 当推理长度 > 训练长度时,对不同维度做差异化频率缩放
|
||||||
|
# 高频维度(局部位置信息): 保持不变
|
||||||
|
# 低频维度(远距离位置信息): 频率除以factor缩放
|
||||||
|
# 中间维度: 线性插值过渡
|
||||||
|
if rope_scaling is not None:
|
||||||
orig_max, factor, beta_fast, beta_slow, attn_factor = (
|
orig_max, factor, beta_fast, beta_slow, attn_factor = (
|
||||||
rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
|
rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
|
||||||
rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
|
rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
|
||||||
)
|
)
|
||||||
if end / orig_max > 1.0:
|
if end / orig_max > 1.0: # 只有推理长度超出训练长度时才缩放
|
||||||
inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
|
inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
|
||||||
low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
|
low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
|
||||||
|
# ramp: 0→1的线性斜坡,控制缩放程度
|
||||||
ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
|
ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
|
||||||
|
# 最终频率 = 原频率 * (1-ramp + ramp/factor)
|
||||||
|
# ramp=0 → 保持原频率(高频维度)
|
||||||
|
# ramp=1 → 频率/factor(低频维度)
|
||||||
freqs = freqs * (1 - ramp + ramp / factor)
|
freqs = freqs * (1 - ramp + ramp / factor)
|
||||||
t = torch.arange(end, device=freqs.device)
|
|
||||||
freqs = torch.outer(t, freqs).float()
|
# 生成位置-频率矩阵: [max_pos, dim//2]
|
||||||
|
t = torch.arange(end, device=freqs.device) # 位置序列 [0, 1, 2, ..., end-1]
|
||||||
|
freqs = torch.outer(t, freqs).float() # 外积: 每个位置 × 每个频率
|
||||||
|
# 计算cos和sin,拼接成完整维度 [max_pos, dim]
|
||||||
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
|
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
|
||||||
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
|
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
|
||||||
return freqs_cos, freqs_sin
|
return freqs_cos, freqs_sin
|
||||||
|
|
||||||
|
# ---- 应用RoPE旋转到Q和K ----
|
||||||
|
# 公式: q_rot = q * cos(θ) + rotate_half(q) * sin(θ)
|
||||||
|
# 其中 rotate_half 将向量前后半部分交换并取负: [a,b] → [-b,a] (相当于复数乘法)
|
||||||
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
|
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
|
||||||
def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
|
def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
|
||||||
q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
|
q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
|
||||||
k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
|
k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
|
||||||
return q_embed, k_embed
|
return q_embed, k_embed
|
||||||
|
|
||||||
|
# ---- GQA辅助: 复制KV头以匹配Q头数量 ----
|
||||||
|
# GQA (Grouped Query Attention): Q=8头, KV=4头,每2个Q头共享1对KV
|
||||||
|
# 需要把4个KV头复制一份变成8个,才能和8个Q头做注意力计算
|
||||||
|
# 相比MHA(Q=K=V=8头)节省50%的KV Cache显存
|
||||||
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||||
bs, slen, num_key_value_heads, head_dim = x.shape
|
bs, slen, num_key_value_heads, head_dim = x.shape
|
||||||
if n_rep == 1: return x
|
if n_rep == 1: return x # MHA模式(Q头数=KV头数),无需复制
|
||||||
|
# [B, S, kv_heads, dim] → [B, S, kv_heads, n_rep, dim] → [B, S, q_heads, dim]
|
||||||
return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim))
|
return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim))
|
||||||
|
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 注意力层
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 实现 GQA (Grouped Query Attention) + RoPE + KV Cache + Flash Attention
|
||||||
|
#
|
||||||
|
# 计算流程:
|
||||||
|
# x → Q,K,V投影 → QK归一化 → RoPE旋转 → KV Cache拼接
|
||||||
|
# → GQA复制KV → 注意力计算(Flash或手动) → 输出投影
|
||||||
class Attention(nn.Module):
|
class Attention(nn.Module):
|
||||||
def __init__(self, config: MiniMindConfig):
|
def __init__(self, config: MiniMindConfig):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
|
self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
|
||||||
self.n_local_heads = config.num_attention_heads
|
self.n_local_heads = config.num_attention_heads # Q头数 = 8
|
||||||
self.n_local_kv_heads = self.num_key_value_heads
|
self.n_local_kv_heads = self.num_key_value_heads # KV头数 = 4
|
||||||
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
self.n_rep = self.n_local_heads // self.n_local_kv_heads # 每个KV头对应几个Q头 = 2
|
||||||
self.head_dim = config.head_dim
|
self.head_dim = config.head_dim # 每个头的维度 = 96
|
||||||
self.is_causal = True
|
self.is_causal = True # 因果注意力(只能看到前面的token)
|
||||||
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
|
# Q/K/V/O 线性投影 (无偏置,现代LLM标配)
|
||||||
self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False) # 768 → 8*96=768
|
||||||
self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) # 768 → 4*96=384
|
||||||
self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
|
self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) # 768 → 4*96=384
|
||||||
|
self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) # 768 → 768
|
||||||
|
# QK归一化 (Qwen3的设计,防止注意力分数随训练爆炸)
|
||||||
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
||||||
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
||||||
self.attn_dropout = nn.Dropout(config.dropout)
|
self.attn_dropout = nn.Dropout(config.dropout)
|
||||||
self.resid_dropout = nn.Dropout(config.dropout)
|
self.resid_dropout = nn.Dropout(config.dropout)
|
||||||
self.dropout = config.dropout
|
self.dropout = config.dropout
|
||||||
|
# 检测是否支持Flash Attention (PyTorch 2.0+的SDPA)
|
||||||
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn
|
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn
|
||||||
|
|
||||||
def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
|
def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
|
||||||
bsz, seq_len, _ = x.shape
|
bsz, seq_len, _ = x.shape
|
||||||
|
|
||||||
|
# Step 1: 线性投影得到Q, K, V
|
||||||
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||||
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
|
# reshape成多头格式: [B, S, n_heads, head_dim]
|
||||||
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim) # [B, S, 8, 96]
|
||||||
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) # [B, S, 4, 96]
|
||||||
|
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) # [B, S, 4, 96]
|
||||||
|
|
||||||
|
# Step 2: QK归一化 (防止注意力分数随维度增大而爆炸)
|
||||||
xq, xk = self.q_norm(xq), self.k_norm(xk)
|
xq, xk = self.q_norm(xq), self.k_norm(xk)
|
||||||
|
|
||||||
|
# Step 3: 应用RoPE旋转位置编码
|
||||||
cos, sin = position_embeddings
|
cos, sin = position_embeddings
|
||||||
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
|
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
|
||||||
|
|
||||||
|
# Step 4: KV Cache (推理加速,拼接历史K/V)
|
||||||
|
# 训练时: past_key_value=None,每次处理完整序列
|
||||||
|
# 推理时: 拼接之前步骤缓存的K/V,只需计算新token
|
||||||
if past_key_value is not None:
|
if past_key_value is not None:
|
||||||
xk = torch.cat([past_key_value[0], xk], dim=1)
|
xk = torch.cat([past_key_value[0], xk], dim=1) # 拼接历史K
|
||||||
xv = torch.cat([past_key_value[1], xv], dim=1)
|
xv = torch.cat([past_key_value[1], xv], dim=1) # 拼接历史V
|
||||||
past_kv = (xk, xv) if use_cache else None
|
past_kv = (xk, xv) if use_cache else None
|
||||||
|
|
||||||
|
# Step 5: GQA头复制 + 转置为注意力计算格式
|
||||||
|
# Q: [B, 8, S, 96], K: [B, 8, S, 96](从4头复制到8头), V: 同K
|
||||||
xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
|
xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
|
||||||
|
|
||||||
|
# Step 6: 注意力计算 (两条路径)
|
||||||
if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
|
if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
|
||||||
|
# 路径A: Flash Attention (O(N)显存,更快)
|
||||||
|
# 条件: 非单token推理 + 无自定义mask时可用
|
||||||
output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
|
output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
|
||||||
else:
|
else:
|
||||||
|
# 路径B: 手动注意力 (KV Cache推理 或 有padding mask时)
|
||||||
|
# scores = Q @ K^T / sqrt(d)
|
||||||
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
||||||
|
# 因果mask: 上三角设为-inf,确保位置i只能attend到<=i的位置
|
||||||
if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
|
if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
|
||||||
|
# padding mask: 被padding的位置设为-1e9
|
||||||
if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
|
if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
|
||||||
|
# softmax → dropout → 与V相乘
|
||||||
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
|
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
|
||||||
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
|
|
||||||
|
# Step 7: 合并多头 + 输出投影
|
||||||
|
output = output.transpose(1, 2).reshape(bsz, seq_len, -1) # [B, S, 768]
|
||||||
output = self.resid_dropout(self.o_proj(output))
|
output = self.resid_dropout(self.o_proj(output))
|
||||||
return output, past_kv
|
return output, past_kv
|
||||||
|
|
||||||
|
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 前馈网络 (FFN)
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# ---- SwiGLU FFN (Dense版) ----
|
||||||
|
# 标准FFN: output = down(ReLU(up(x))) 两个矩阵
|
||||||
|
# SwiGLU: output = down(SiLU(gate(x)) * up(x)) 三个矩阵,效果更好
|
||||||
|
# gate_proj 和 up_proj 分别投影到中间维度,gate走激活函数后与up逐元素相乘,再down投影回来
|
||||||
class FeedForward(nn.Module):
|
class FeedForward(nn.Module):
|
||||||
def __init__(self, config: MiniMindConfig, intermediate_size: int = None):
|
def __init__(self, config: MiniMindConfig, intermediate_size: int = None):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
intermediate_size = intermediate_size or config.intermediate_size
|
intermediate_size = intermediate_size or config.intermediate_size # 默认2432
|
||||||
self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
|
self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) # 768 → 2432 (门控)
|
||||||
self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
|
self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False) # 2432 → 768 (降维)
|
||||||
self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
|
self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) # 768 → 2432 (上投影)
|
||||||
self.act_fn = ACT2FN[config.hidden_act]
|
self.act_fn = ACT2FN[config.hidden_act] # SiLU = x * sigmoid(x)
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
|
# SwiGLU: down(act(gate(x)) ⊙ up(x))
|
||||||
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||||
|
# ---- MoE FFN (稀疏版) ----
|
||||||
|
# 每个token被Router分配到top-k个Expert(这里k=1)处理
|
||||||
|
# 相比Dense: 总参数量大(4个Expert),但每个token只激活1个Expert
|
||||||
|
# 效果: 用更多参数量换更强表达力,而计算量(FLOPs)基本不变
|
||||||
|
#
|
||||||
|
# 流程: x → Router(gate) → 选出top-1 Expert → Expert(x) * weight → 输出
|
||||||
class MOEFeedForward(nn.Module):
|
class MOEFeedForward(nn.Module):
|
||||||
def __init__(self, config: MiniMindConfig):
|
def __init__(self, config: MiniMindConfig):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.config = config
|
self.config = config
|
||||||
self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
|
self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) # Router: 768 → 4
|
||||||
self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)])
|
self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)]) # 4个独立的FFN
|
||||||
self.act_fn = ACT2FN[config.hidden_act]
|
self.act_fn = ACT2FN[config.hidden_act]
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
batch_size, seq_len, hidden_dim = x.shape
|
batch_size, seq_len, hidden_dim = x.shape
|
||||||
x_flat = x.view(-1, hidden_dim)
|
x_flat = x.view(-1, hidden_dim) # [B*S, 768] 展平为token序列
|
||||||
scores = F.softmax(self.gate(x_flat), dim=-1)
|
|
||||||
topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
|
# Step 1: Router计算每个token对每个Expert的分数
|
||||||
|
scores = F.softmax(self.gate(x_flat), dim=-1) # [B*S, 4] 概率分布
|
||||||
|
|
||||||
|
# Step 2: 选出每个token的top-k Expert (k=1)
|
||||||
|
topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False) # [B*S, 1]
|
||||||
|
# 归一化权重 (top-1时其实就是1.0,top-k>1时确保权重和为1)
|
||||||
if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
|
if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
|
||||||
|
|
||||||
|
# Step 3: 分发token到各Expert并加权聚合
|
||||||
y = torch.zeros_like(x_flat)
|
y = torch.zeros_like(x_flat)
|
||||||
for i, expert in enumerate(self.experts):
|
for i, expert in enumerate(self.experts):
|
||||||
mask = (topk_idx == i)
|
mask = (topk_idx == i) # 哪些token被分配到Expert i
|
||||||
if mask.any():
|
if mask.any():
|
||||||
token_idx = mask.any(dim=-1).nonzero().flatten()
|
token_idx = mask.any(dim=-1).nonzero().flatten() # 被选中的token索引
|
||||||
weight = topk_weight[mask].view(-1, 1)
|
weight = topk_weight[mask].view(-1, 1) # 对应的路由权重
|
||||||
|
# Expert处理 → 乘权重 → 累加到输出
|
||||||
y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
|
y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
|
||||||
elif self.training:
|
elif self.training:
|
||||||
|
# 训练时: 即使Expert未被选中,也让它参与计算图(0*params)
|
||||||
|
# 否则DDP会报错"某些参数没有梯度"
|
||||||
y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
|
y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
|
||||||
|
|
||||||
|
# Step 4: 负载均衡辅助loss (鼓励token均匀分布到各Expert)
|
||||||
|
# 如果没有这个loss,Router倾向于把所有token都发给同一个Expert("赢者通吃")
|
||||||
if self.training and self.config.router_aux_loss_coef > 0:
|
if self.training and self.config.router_aux_loss_coef > 0:
|
||||||
|
# load: 每个Expert实际接收的token比例 [num_experts]
|
||||||
load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
|
load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
|
||||||
|
# aux_loss = (实际负载分布 · 路由概率分布) * num_experts * coef
|
||||||
|
# 当所有Expert负载均匀时,此loss最小
|
||||||
self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
|
self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
|
||||||
else:
|
else:
|
||||||
self.aux_loss = scores.new_zeros(1).squeeze()
|
self.aux_loss = scores.new_zeros(1).squeeze()
|
||||||
|
|
||||||
return y.view(batch_size, seq_len, hidden_dim)
|
return y.view(batch_size, seq_len, hidden_dim)
|
||||||
|
|
||||||
|
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# Transformer Block (单层)
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 一个标准的 Pre-Norm Transformer Decoder Block:
|
||||||
|
#
|
||||||
|
# x ─────────────────────┐
|
||||||
|
# │ │ (残差连接1)
|
||||||
|
# ├→ RMSNorm → Attention ─┘
|
||||||
|
# │ │
|
||||||
|
# ├──────────────────────┤ (残差连接2)
|
||||||
|
# └→ RMSNorm → FFN/MoE ──┘
|
||||||
|
# │
|
||||||
|
# output
|
||||||
class MiniMindBlock(nn.Module):
|
class MiniMindBlock(nn.Module):
|
||||||
def __init__(self, layer_id: int, config: MiniMindConfig):
|
def __init__(self, layer_id: int, config: MiniMindConfig):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.self_attn = Attention(config)
|
self.self_attn = Attention(config)
|
||||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # 注意力前归一化
|
||||||
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # FFN前归一化
|
||||||
|
# 根据配置选择Dense FFN或MoE FFN
|
||||||
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
|
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
|
||||||
|
|
||||||
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
|
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
|
||||||
|
# 残差连接1: Attention
|
||||||
residual = hidden_states
|
residual = hidden_states
|
||||||
hidden_states, present_key_value = self.self_attn(
|
hidden_states, present_key_value = self.self_attn(
|
||||||
self.input_layernorm(hidden_states), position_embeddings,
|
self.input_layernorm(hidden_states), position_embeddings, # Pre-Norm: 先归一化再进注意力
|
||||||
past_key_value, use_cache, attention_mask
|
past_key_value, use_cache, attention_mask
|
||||||
)
|
)
|
||||||
hidden_states += residual
|
hidden_states += residual # x + Attention(Norm(x))
|
||||||
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
|
|
||||||
|
# 残差连接2: FFN
|
||||||
|
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states)) # x + FFN(Norm(x))
|
||||||
return hidden_states, present_key_value
|
return hidden_states, present_key_value
|
||||||
|
|
||||||
|
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# MiniMindModel (Transformer 主干)
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 职责: token_ids → hidden_states (不含最终的lm_head投影)
|
||||||
|
#
|
||||||
|
# 流程: Embedding → Dropout → N层Block → 最终RMSNorm → hidden_states
|
||||||
class MiniMindModel(nn.Module):
|
class MiniMindModel(nn.Module):
|
||||||
def __init__(self, config: MiniMindConfig):
|
def __init__(self, config: MiniMindConfig):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.config = config
|
self.config = config
|
||||||
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
|
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
|
||||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) # 6400 × 768
|
||||||
self.dropout = nn.Dropout(config.dropout)
|
self.dropout = nn.Dropout(config.dropout)
|
||||||
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
|
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)]) # 8层
|
||||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # 最终归一化
|
||||||
|
|
||||||
|
# 预计算RoPE频率表(注册为buffer,不参与训练但随模型保存/移动)
|
||||||
freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
|
freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
|
||||||
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
|
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
|
||||||
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
|
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
|
||||||
|
|
||||||
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
|
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
|
||||||
batch_size, seq_length = input_ids.shape
|
batch_size, seq_length = input_ids.shape
|
||||||
|
|
||||||
|
# 兼容transformers 5.x的cache格式
|
||||||
if hasattr(past_key_values, 'layers'): past_key_values = None
|
if hasattr(past_key_values, 'layers'): past_key_values = None
|
||||||
past_key_values = past_key_values or [None] * len(self.layers)
|
past_key_values = past_key_values or [None] * len(self.layers)
|
||||||
|
|
||||||
|
# 确定当前序列的起始位置 (用于KV Cache场景下取正确的RoPE)
|
||||||
|
# 训练时: start_pos=0 (处理完整序列)
|
||||||
|
# 推理时: start_pos=已生成token数 (只处理新token)
|
||||||
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
||||||
hidden_states = self.dropout(self.embed_tokens(input_ids))
|
|
||||||
# Recompute RoPE buffers lost during meta-device init (transformers>=5.x)
|
# Embedding: token_id → 向量
|
||||||
|
hidden_states = self.dropout(self.embed_tokens(input_ids)) # [B, S] → [B, S, 768]
|
||||||
|
|
||||||
|
# 获取当前位置段的RoPE (处理transformers>=5.x meta-device初始化导致buffer丢失)
|
||||||
if self.freqs_cos[0, 0] == 0:
|
if self.freqs_cos[0, 0] == 0:
|
||||||
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
|
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
|
||||||
self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
|
self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
|
||||||
position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
|
position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
|
||||||
|
|
||||||
|
# 逐层前向传播
|
||||||
presents = []
|
presents = []
|
||||||
for layer, past_key_value in zip(self.layers, past_key_values):
|
for layer, past_key_value in zip(self.layers, past_key_values):
|
||||||
hidden_states, present = layer(
|
hidden_states, present = layer(
|
||||||
|
|
@ -227,62 +413,121 @@ class MiniMindModel(nn.Module):
|
||||||
attention_mask=attention_mask
|
attention_mask=attention_mask
|
||||||
)
|
)
|
||||||
presents.append(present)
|
presents.append(present)
|
||||||
|
|
||||||
|
# 最终归一化
|
||||||
hidden_states = self.norm(hidden_states)
|
hidden_states = self.norm(hidden_states)
|
||||||
|
|
||||||
|
# 汇总所有MoE层的辅助loss
|
||||||
aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
|
aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
|
||||||
return hidden_states, presents, aux_loss
|
return hidden_states, presents, aux_loss
|
||||||
|
|
||||||
|
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# MiniMindForCausalLM (最顶层: 训练+推理入口)
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 在MiniMindModel基础上加了lm_head(hidden → vocab logits)和loss计算
|
||||||
|
# 继承GenerationMixin以兼容HuggingFace的generate接口
|
||||||
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
|
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
|
||||||
config_class = MiniMindConfig
|
config_class = MiniMindConfig
|
||||||
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
|
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} # 声明权重绑定关系
|
||||||
|
|
||||||
def __init__(self, config: MiniMindConfig = None):
|
def __init__(self, config: MiniMindConfig = None):
|
||||||
self.config = config or MiniMindConfig()
|
self.config = config or MiniMindConfig()
|
||||||
super().__init__(self.config)
|
super().__init__(self.config)
|
||||||
self.model = MiniMindModel(self.config)
|
self.model = MiniMindModel(self.config) # Transformer主干
|
||||||
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
|
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) # 768 → 6400
|
||||||
|
# 权重绑定: lm_head和embedding共享同一份权重矩阵,减少参数量
|
||||||
if self.config.tie_word_embeddings: self.model.embed_tokens.weight = self.lm_head.weight
|
if self.config.tie_word_embeddings: self.model.embed_tokens.weight = self.lm_head.weight
|
||||||
self.post_init()
|
self.post_init() # HuggingFace的初始化后处理
|
||||||
|
|
||||||
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
|
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
|
||||||
|
# Step 1: 通过Transformer主干得到hidden_states
|
||||||
hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
|
hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
|
||||||
|
|
||||||
|
# Step 2: lm_head投影到词表空间
|
||||||
|
# logits_to_keep: 优化——只计算最后N个位置的logits(节省显存)
|
||||||
|
# 推理时只需最后1个位置,RL训练时只需completion部分
|
||||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||||
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
logits = self.lm_head(hidden_states[:, slice_indices, :]) # [B, S, 6400]
|
||||||
|
|
||||||
|
# Step 3: 计算loss (仅训练时有labels)
|
||||||
loss = None
|
loss = None
|
||||||
if labels is not None:
|
if labels is not None:
|
||||||
|
# 自回归: 用位置t的logits预测位置t+1的token
|
||||||
|
# logits[:-1] 预测 labels[1:]
|
||||||
x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
|
x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
|
||||||
loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
|
loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100) # -100位置不算loss
|
||||||
|
|
||||||
return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
|
return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
|
||||||
|
|
||||||
# https://github.com/jingyaogong/minimind/discussions/611
|
# ════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 自回归生成 (推理)
|
||||||
|
# ════════════════════════════════════════════════════════════════════════════
|
||||||
|
# 自定义generate,不用HuggingFace的默认实现
|
||||||
|
# 支持: KV Cache / temperature / top_k / top_p / repetition_penalty / streaming
|
||||||
|
#
|
||||||
|
# 每步循环:
|
||||||
|
# 1. 只输入新token(利用KV Cache)
|
||||||
|
# 2. 取最后位置logits
|
||||||
|
# 3. temperature缩放 → repetition penalty → top_k截断 → top_p截断
|
||||||
|
# 4. 采样得到next_token
|
||||||
|
# 5. 拼接到input_ids,重复直到EOS或达到max长度
|
||||||
@torch.inference_mode()
|
@torch.inference_mode()
|
||||||
def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
|
def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
|
||||||
|
# 支持多条序列并行生成
|
||||||
input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
|
input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
|
||||||
attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
|
attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
|
||||||
past_key_values = kwargs.pop("past_key_values", None)
|
past_key_values = kwargs.pop("past_key_values", None)
|
||||||
finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
|
finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device) # 跟踪每条序列是否已结束
|
||||||
if streamer: streamer.put(input_ids.cpu())
|
if streamer: streamer.put(input_ids.cpu())
|
||||||
|
|
||||||
for _ in range(max_new_tokens):
|
for _ in range(max_new_tokens):
|
||||||
|
# 利用KV Cache: 只输入past_len之后的新token
|
||||||
past_len = past_key_values[0][0].shape[1] if past_key_values else 0
|
past_len = past_key_values[0][0].shape[1] if past_key_values else 0
|
||||||
outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
|
outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
|
||||||
|
# 扩展attention_mask以包含新生成的位置
|
||||||
attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
|
attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
|
||||||
|
|
||||||
|
# 取最后一个位置的logits,除以temperature
|
||||||
|
# temperature < 1 → 更确定(偏向高概率token)
|
||||||
|
# temperature > 1 → 更随机(概率分布更平坦)
|
||||||
logits = outputs.logits[:, -1, :] / temperature
|
logits = outputs.logits[:, -1, :] / temperature
|
||||||
|
|
||||||
|
# Repetition Penalty: 降低已出现token的概率
|
||||||
|
# 正分数的token: score / penalty (概率降低)
|
||||||
|
# 负分数的token: score * penalty (概率更低)
|
||||||
if repetition_penalty != 1.0:
|
if repetition_penalty != 1.0:
|
||||||
for i in range(input_ids.shape[0]):
|
for i in range(input_ids.shape[0]):
|
||||||
seen = torch.unique(input_ids[i]); score = logits[i, seen]; logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
|
seen = torch.unique(input_ids[i]); score = logits[i, seen]; logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
|
||||||
|
|
||||||
|
# Top-K 截断: 只保留概率最高的K个token
|
||||||
if top_k > 0:
|
if top_k > 0:
|
||||||
logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
|
logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
|
||||||
|
|
||||||
|
# Top-P (Nucleus) 截断: 保留累积概率刚好超过P的最小token集合
|
||||||
|
# 动态调整候选数量——高置信时候选少,低置信时候选多
|
||||||
if top_p < 1.0:
|
if top_p < 1.0:
|
||||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
||||||
mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
|
mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
|
||||||
mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
|
mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0 # 保证至少保留1个token
|
||||||
logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
|
logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
|
||||||
|
|
||||||
|
# 采样或贪心选择
|
||||||
next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
|
next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
|
||||||
|
|
||||||
|
# 已结束的序列强制输出EOS (防止padding位置产生垃圾token)
|
||||||
if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
|
if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
|
||||||
|
|
||||||
|
# 拼接新token
|
||||||
input_ids = torch.cat([input_ids, next_token], dim=-1)
|
input_ids = torch.cat([input_ids, next_token], dim=-1)
|
||||||
past_key_values = outputs.past_key_values if use_cache else None
|
past_key_values = outputs.past_key_values if use_cache else None
|
||||||
if streamer: streamer.put(next_token.cpu())
|
if streamer: streamer.put(next_token.cpu())
|
||||||
|
|
||||||
|
# 检查是否所有序列都已生成EOS
|
||||||
if eos_token_id is not None:
|
if eos_token_id is not None:
|
||||||
finished |= next_token.squeeze(-1).eq(eos_token_id)
|
finished |= next_token.squeeze(-1).eq(eos_token_id)
|
||||||
if finished.all(): break
|
if finished.all(): break
|
||||||
|
|
||||||
if streamer: streamer.end()
|
if streamer: streamer.end()
|
||||||
if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
|
if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
|
||||||
return input_ids
|
return input_ids
|
||||||
Loading…
Reference in New Issue