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
QiaohuiGao 2026-06-01 21:33:23 -04:00
parent 3f1a7cc25b
commit b23a05868d
25 changed files with 1927 additions and 91 deletions

1
.gitignore vendored
View File

@ -4,3 +4,4 @@ __pycache__/
out out
website/ website/
docs-minimind/ docs-minimind/
.gstack/

23
learning/README.md Normal file
View File

@ -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 或蒸馏算法理解报告

View File

@ -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,混合精度、性能和部分算子兼容性要单独确认。

View File

@ -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)`,为什么?

48
learning/day01/03_exam.md Normal file
View File

@ -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。

View File

View File

@ -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。

View File

@ -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?

52
learning/day02/03_exam.md Normal file
View File

@ -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。
```

View File

@ -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 是否成功:
```

View File

@ -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()`?

55
learning/day03/03_exam.md Normal file
View File

@ -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`。

View File

@ -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 回答样例:
明显提升:
明显问题:
```

View File

@ -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>` 标签,可能和哪些训练样本有关?

53
learning/day04/03_exam.md Normal file
View File

@ -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。
```

View File

@ -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 组采样参数测试。
- 把失败样例原样保存。
- 评估时区分:语言流畅、指令跟随、事实正确、格式正确。

View File

@ -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 的回答
- 评分表
- 失败样例
- 下一步改进建议

53
learning/day05/03_exam.md Normal file
View File

@ -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. 通过标准
你能给一个模型输出做结构化诊断,而不是只说“好”或“不好”。

View File

@ -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` 做理解。

View File

@ -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 不需要反向传播?

53
learning/day06/03_exam.md Normal file
View File

@ -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 的硬答案或软分布。
```

View File

@ -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 记录
推理评测
失败分析
后训练理解
下一步计划
```
## 今日注意事项
- 不要隐藏没跑通的部分。
- 不要只写结论,要写证据。
- 不要凭感觉说模型变好,要给对比样例。

View File

@ -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。
## 最终复盘问题
你要能回答:
- 当前模型最强能力是什么?
- 最弱能力是什么?
- 训练瓶颈是数据、模型容量、训练时长还是评测方式?
- 下一周最值得投入的改进是什么?

78
learning/day07/03_exam.md Normal file
View File

@ -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,且最终报告里有清晰失败分析和下一步实验设计。

View File

@ -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