从零开始的LLM 8.LLM中的蒸馏（知识内化）

发表于 更新于

前言

本文旨在学习知识蒸馏（Distillation）的基本机制、核心原理及其在模型训练中的作用，并进一步深入理解 Context Distillation 的原理、训练流程以及代码实现。

经过前面的学习，我们已经系统地完成了大语言模型（LLM）的基础部分，主要参考了以下资料：

  • 《动手学深度学习》
  • 《蘑菇书》
  • 《Happy-LLM》
  • 《Hello-Agents》
  • 《all-in-rag》

到这里，可以认为已经具备了继续深入学习 Agent 的基础能力。

接下来的内容将不再像之前那样按照固定课程循序渐进，而是更多围绕当前 Agent 领域的研究方向展开。学习内容既会涉及工程实践（如各种 Agent 框架、工作流、记忆系统等），也会涉及算法与模型层面的探索（如多模态、持续学习等）。

  1. Generative Agents
  2. Letta
  3. A-MEM
  4. Voyager

这条路线主要围绕 Agent 的持续学习（Continual Learning） 展开，也是我目前最感兴趣的方向：如何让一个 Agent 在长期运行过程中不断积累经验、提升能力，而不是每次都从零开始。

目前来看，大致有两种思路：

  1. 不修改模型本身。 仅通过 Prompt、记忆系统（Memory）、RAG、工具调用等机制，使 Agent 能够不断积累知识和经验，实现能力的持续增长。
  2. 修改模型本身。 基于 Agent 长期运行过程中积累的对话、笔记和总结等数据，在合适的时机对模型进行持续训练（Continual Fine-tuning），使模型真正将这些经验内化，而不仅仅依赖外部记忆。

由于这一阶段更多是阅读论文、分析开源项目、验证实验以及整理个人理解，而不是单纯跟着教程学习，因此资料收集、实验和整理都会比之前花费更多时间。

原计划是学习 A-MEM，但某天晚上突然产生了一个关于“实现知识内部参数化”的想法，因此暂时调整方向，优先对该想法进行了实验验证。几天后查阅相关资料时发现，这一思路实际上与 Context Distillation 存在较高的关联。

后续复盘发现，最初的方案并不完全属于传统的 Distillation 方法，而更接近于 Distillation + 判别器（Discriminator） 的混合框架。然而，经过进一步分析后发现，该方向存在一定局限性：判别器只能对结果进行评估，难以有效引导 Agent 生成正确答案，导致探索过程过于依赖试错，效率较低。

因此，近期开始结合 DPO（Direct Preference Optimization）等方法，对原有思路进行重新设计和优化。

本章节将先对 Distillation 的基础概念、核心机制以及相关方法进行总结，为后续 Context Distillation 和改进方案的研究提供基础。

1. Response Distillation

1.1 蒸馏的基本概念

知识蒸馏（Knowledge Distillation, KD）是一种模型压缩（Model Compression）和知识迁移（Knowledge Transfer）方法，其核心思想是利用一个性能更强、规模更大的模型（Teacher Model，教师模型）指导一个规模较小的模型（Student Model，学生模型）进行学习，使学生模型在减少参数量和计算成本的同时，尽可能保留教师模型中的有效能力。

在大语言模型（Large Language Model, LLM）中，Response Distillation（响应蒸馏）是一种常见的蒸馏方式，其目标是让学生模型学习教师模型面对相同输入时的输出行为。

具体来说，教师模型首先根据输入问题生成高质量回答：Teacher(question)=Teacher_answer

然后利用这些输入输出数据训练学生模型：Student(question)=Student_answer

通过大量样本训练，使学生模型逐渐学习教师模型的知识表达方式、回答结构以及语言模式。

需要注意的是，在 Response Distillation 中，Teacher 和 Student 接收相同的问题，区别仅在于模型能力不同。而后续介绍的 Context Distillation 则会进一步改变这一点，即 Teacher 会额外拥有 Student 不具备的上下文信息。

1.2 LLM 中 Response Distillation 的基本流程

在大语言模型中，Response Distillation 通常采用教师模型生成数据，再训练学生模型的方式。

首先，使用教师模型对输入问题进行回答：Teacher(question)=Teacher_answer

例如：

1
2
Question:
牛顿第二定律是什么？

教师模型生成：

1
2
Teacher_answer:
牛顿第二定律是 F=ma，表示物体受到的合外力等于质量乘以加速度。

随后，将问题和教师模型生成的答案组合成训练样本：Prompt=Question+Teacher_answer

形成：

1
2
3
4
5
Input:
牛顿第二定律是什么？

Target:
牛顿第二定律是 F=ma，表示物体受到的合外力等于质量乘以加速度。

之后，使用这些样本训练学生模型。

学生模型接收相同的问题。

训练过程中，采用 Teacher Forcing（教师强制） 的方式，将教师模型生成的答案作为学生模型的训练目标。

例如教师模型生成：

1
2
Teacher_answer:
牛顿第二定律是 F=ma。

经过 tokenizer 后：

1
[牛顿, 第二, 定律, 是, F, =, ma, 。]

训练时，学生模型根据问题以及前面已经给出的正确 token，逐位置预测下一个 token：

1
2
3
4
5
输入：
问题 + 已知前面的 token

预测目标：
当前对应位置的 Teacher token

例如：

输入 学生模型预测 Target（教师答案对应 token）
牛顿第二定律是什么？ 预测第1个 token 牛顿
牛顿第二定律是什么？+牛顿 预测第2个 token 第二
牛顿第二定律是什么？+牛顿+第二 预测第3个 token 定律

在每一个位置，学生模型都会输出整个词表的概率分布，然后与该位置真实的 Teacher token 计算交叉熵损失。

例如第一个位置：

学生模型预测：

1
2
3
牛顿：0.8
苹果：0.1
汽车：0.1

对应 Target：

1
牛顿

该位置损失：Loss_1=-log(0.8)

随后，对所有 token 位置分别计算损失：`Loss_1,Loss_2,Loss_3,...,Loss_T

最后将所有位置的损失求和（实际训练中通常取平均）：\(Loss=\frac{1}{T}\sum_{t=1}^{T}Loss_t\)

其中：

  • (T)：教师答案包含的 token 数量；
  • (\(Loss_t\))：第 (t) 个 token 位置的交叉熵损失。

通过最小化这个 Loss，学生模型参数不断更新，使其在每个 token 位置上更倾向于生成教师答案中的 token。

因此，LLM 中的 Response Distillation 本质上仍然是语言模型的 next-token prediction（下一个 token 预测）训练过程，区别在于训练目标不是人工标注答案，而是由教师模型生成的答案。

1.3 Response Distillation 的作用

Response Distillation 的核心作用，是将教师模型（Teacher Model）的能力迁移到规模更小的学生模型（Student Model）中，使学生模型在降低参数规模、计算成本和部署成本的情况下，仍然能够获得接近教师模型的回答效果。

对于大语言模型而言，大规模模型通常具有更丰富的知识、更强的语言理解能力以及更复杂的推理能力，但同时需要大量计算资源。例如，一个数百亿参数规模的模型虽然能够生成高质量回答，但在本地部署、边缘设备或者实时应用中可能受到显存和计算成本限制。

Response Distillation 通过让教师模型生成大量高质量回答，然后利用这些数据训练学生模型，使学生模型学习教师模型在不同问题下的回答模式。

例如：

教师模型：

1
2
3
4
5
6
Question:
解释牛顿第二定律

Teacher:
牛顿第二定律表示物体受到的合外力等于质量乘以加速度，
数学表达式为 F=ma，其中 F 表示合外力，m 表示质量，a 表示加速度。

学生模型初始状态：

1
2
3
4
5
Question:
解释牛顿第二定律

Student:
牛顿第二定律是能量守恒。

经过蒸馏训练后：

1
2
3
4
Student:
牛顿第二定律是 F=ma，
表示物体的加速度与受到的合外力成正比，
与物体质量成反比。

可以看到，学生模型并没有直接复制教师模型的参数，而是通过大量输入输出样本学习教师模型的行为，使自身参数逐渐具备类似的回答能力。

因此，Response Distillation 主要具有以下作用：

  1. 降低模型部署成本

大型模型通常需要大量 GPU 显存和计算资源，而蒸馏后的学生模型参数规模更小，可以运行在资源有限的设备上。

  1. 提升小模型能力

小模型由于参数量有限，通常知识储备和推理能力不足。通过学习教师模型生成的数据，可以获得更强的回答能力以及更符合任务要求的输出方式。

  1. 迁移教师模型能力

虽然学生模型规模更小，但通过蒸馏可以学习教师模型中的部分核心能力，例如：

  • 知识表达能力；
  • 指令遵循能力；
  • 推理模式；
  • 输出格式。

因此，Response Distillation 可以理解为：

使用一个能力更强的大模型作为教师，通过生成高质量回答，将其输出行为和部分能力迁移到更小的模型中，使小模型在较低计算成本下获得接近大模型的回答效果。

需要注意的是，“公司将员工文档全部拿来蒸馏知识”这种说法容易产生混淆。

严格来说，传统意义上的知识蒸馏（Knowledge Distillation, KD）需要存在明确的 Teacher Model → Student Model 结构，即通过教师模型的输出、行为或中间表示，将已有模型中的知识迁移到学生模型中。

而如果只是直接利用员工文档、技术文档等原始文本进行训练，本质上是从语料中学习知识，更接近 知识注入（Knowledge Injection） 或 领域持续预训练（Continued Pretraining），并不属于严格意义上的知识蒸馏。

真正符合蒸馏定义的方式是：先利用一个能力更强的教师模型读取文档并生成高质量训练数据（例如 Question/Answer 对），再使用这些数据训练学生模型，使学生模型学习教师模型的输出行为。

2. Context Distillation

2.1 Context Distillation 的基本概念

Context Distillation（上下文蒸馏）是一种特殊形式的知识蒸馏方法，其核心思想是利用教师模型（Teacher Model）在额外上下文（Context）条件下生成的行为，训练学生模型（Student Model）在缺少该上下文的情况下完成相同任务，从而将上下文中的知识或能力迁移到学生模型参数中。

与传统 Response Distillation 不同，Context Distillation 的主要区别在于 Teacher 和 Student 接收到的信息不同。

在 Context Distillation 中：

Teacher 接收额外上下文：Teacher( context + question )=answer

Student 不再提供 Context：Student( question )=answer

训练目标是让：Student(question) ≈ Teacher(context + question)

通过这种方式，原本需要依赖外部 Context 的知识，可以逐渐被模型吸收并存储在参数中。

2.2 Context Distillation 的基本流程

Context Distillation 通常包含两个步骤。

第一步：Teacher 根据 Context 生成答案

首先，将额外信息作为 Context 提供给教师模型。

例如：

1
2
3
4
5
Context:
巴黎是法国的首都，也是法国最大的城市之一。

Question:
法国的首都是哪里？

Teacher 根据 Context 生成：

1
2
Answer:
法国的首都是巴黎。

即：Teacher(Context + Question) = Answer

此时 Teacher 的回答依赖于额外提供的 Context。

第二步：Student 学习去除 Context 后的回答能力

随后，将 Context 移除，只保留 Question：

1
2
Question:
法国的首都是哪里？

使用 Teacher 生成的答案作为训练目标：

1
2
3
4
5
Input:
法国的首都是哪里？

Target:
法国的首都是巴黎。

训练 Student：Student(Question) = Answer

训练过程与普通语言模型训练类似：

  • 输入 Question；
  • 预测 Answer 中的每个 token；
  • 对 Answer 部分计算交叉熵损失；
  • 通过反向传播更新模型参数。

最终，使 Student 在没有 Context 的情况下，也能够完成原本依赖 Context 的任务。

2.3 Context Distillation 的作用

Context Distillation 的主要作用是实现知识和能力的内部化。

在传统场景中，模型通常需要依赖外部信息：

1
2
3
4
5
Context + Question
        ↓
       LLM
        ↓
      Answer

例如：

  • 企业知识库；
  • 专业领域文档；
  • 特定任务规则；
  • 长期记忆信息。

如果没有 Context，模型可能无法正确回答。

而通过 Context Distillation：

1
2
3
4
5
Question
    ↓
Student Model
    ↓
Answer

模型可以直接根据自身参数生成答案。

因此，Context Distillation 可以用于：

  • 将外部知识注入模型参数；
  • 减少推理阶段对上下文输入的依赖；
  • 提升模型在特定领域任务中的表现；
  • 构建更加轻量化的专用模型。

Context Distillation 可以理解为：

将原本存在于外部 Context 中的信息，通过蒸馏训练转化为模型参数中的内部知识，使模型能够在缺少显式上下文的情况下完成相同任务。

3. 思考（可以跳过，直接看 4. Context Distillation 实验）

我最初的想法与 Context Distillation 有一定相似之处。

整体思路是：通过一个更强大的 Teacher 模型生成问题，然后构造两个 Student 模型：

  • 一个 Student 不进行训练，推理时输入 knowledge + question，利用外部知识生成答案；
  • 另一个 Student 进行训练，但推理时只输入 question，希望通过训练使模型逐渐学习并内化这些知识。

随后，通过一个判别器 D 对两个 Student 的回答进行比较。当训练不断进行时，如果判别器无法有效区分两个模型的输出质量，则认为训练后的 Student 已经具备了与外部知识增强模型相近的能力，即完成了知识内化。

在实际实现过程中，最开始尝试的方向偏向 Context Distillation。由于最初这个想法只是一个概念，并没有明确具体的优化方法，因此实现过程中逐渐偏向了 Context Distillation 的方案。

但是完成实验后发现效果并不理想。重新分析后，我开始重新思考判别器机制本身的问题。

说人话就是写成了context distillation的实现。后面重新整理思绪回去想原来的方案。

根据 GAIL（Generative Adversarial Imitation Learning） 的思想，单纯的判别器只能判断：

当前生成结果是否更接近目标结果。

但是它无法直接提供：

模型应该如何改进。

因此，如果希望利用判别器有效指导 Student 优化，需要额外构造价值函数 V（类似 RLHF 中的 Value Model），通过计算生成结果的优势函数（Advantage）提供更加明确的优化方向。

然而，价值函数本身的构建较为复杂，同时强化学习训练过程通常具有较大的不稳定性。因此，在探索其他方法后，我开始从 DPO（Direct Preference Optimization） 的角度重新思考知识内化问题。

基于原本 Context Distillation 的思路，我进一步探索不同优化方式在知识内化任务中的作用，希望找到一种更加稳定的方法，使 Student 模型能够在不依赖外部知识输入的情况下学习并保留目标知识。

3.1 自回归训练（Next-token Prediction）

无论是：

  • 预训练（Pre-training）
  • SFT（Supervised Fine-tuning）
  • Context Distillation

本质上都属于自回归语言建模。

目标都是：

根据前面的 token，预测下一个 token： \[ P(x_t|x_1,...,x_{t-1}) \] 这种训练方式主要解决：

模型如何生成合理的文本。

例如：

经过训练后，模型能够学习到一定的知识表示。例如：

1
牛顿第二定律是 F=ma

因此，当面对问题：

1
解释牛顿第二定律

模型能够生成相关回答，而不是：

1
我不知道

但是，仅依靠 next-token prediction 存在一定限制。

首先，它优化的是文本概率分布，而不是答案质量。因此，模型虽然可能掌握相关知识，但生成结果仍然可能存在：

  • 错误事实；
  • 幻觉（Hallucination）；
  • 逻辑不完整；
  • 不符合用户需求的回答。

其次，next-token prediction 无法直接告诉模型：

多个可能答案中，哪个答案更加符合人类偏好。

例如，模型可能生成多个语义合理的回答，但是无法判断：

  • 哪个更加准确；
  • 哪个更加简洁；
  • 哪个更加符合用户需求。

因此，next-token prediction 主要解决：

模型如何生成合理文本。

而不是：

模型应该选择什么样的答案。

3.2 基于生成反馈的优化（Generation-based Optimization）

这类方法的核心思想是：

让模型主动生成候选答案，再通过额外的反馈信号指导模型优化。

整体流程：

1
2
3
4
5
6
7
Question
    ↓
模型生成答案
    ↓
Reward / Critique / 人工评价
    ↓
优化模型参数

反馈信号可以来自：

  • Reward Model；
  • 人工偏好评价；
  • 规则函数；
  • 验证器（Verifier）等。

代表方法包括：

  • RLHF；
  • PPO；
  • GRPO；
  • Rejection Sampling。

例如 GRPO：

模型针对同一个问题生成多个答案：

1
2
3
4
5
6
7
8
9
10
11
Question:
解释牛顿第二定律

Answer A:
F=ma，表示物体受到的合外力等于质量乘以加速度。

Answer B:
牛顿第二定律描述了物体运动变化的规律。

Answer C:
牛顿第二定律说明物体速度保持不变。

然后通过 reward 函数对不同答案进行评价：

1
Reward(A) > Reward(B) > Reward(C)

提高高 reward 答案的生成概率。

这种方法的优势在于：

模型可以通过探索生成空间，逐渐找到更加符合目标的答案。

但是，这类方法也存在一些限制。

首先，模型优化高度依赖外部评价机制的质量。

无论是 Reward Model、人工评价还是规则函数，本质上都是告诉模型：

当前生成结果是否更符合目标。

如果评价函数本身存在偏差，模型可能会优化错误目标，甚至产生 reward hacking（奖励投机）现象。

其次，生成式优化依赖模型自身的探索能力，而探索本身是一件非常低效的事情。

模型需要不断尝试生成不同答案，然后通过 reward 判断哪些结果更好。

但是，对于一个具体问题，我们无法确定模型需要探索多少次才能生成满足要求的答案。

例如：

对于问题：

1
9+8是多少？

如果模型当前能力不足，它可能连续生成：

1
2
3
4
5
6
7
9+8=16

9+8=18

9+8=15

9+8=20

直到某一次随机生成：

1
9+8=17

之后 reward 才能告诉模型：

这个答案更好。

但是我们无法提前知道：

  • 模型需要生成 5 次？
  • 100 次？
  • 还是永远无法生成正确答案？

因此，生成式优化存在明显的探索效率问题。

此外，模型只能在当前生成分布（generation distribution）附近进行优化。

如果正确答案根本不在模型当前能够生成的范围内，那么即使 reward 能够判断正确答案，模型也无法直接获得该答案。

例如：

模型始终认为：

1
牛顿第二定律是 E=mc²。

而正确答案：

1
牛顿第二定律是 F=ma。

并不在模型当前生成空间中。

此时 reward 只能降低错误答案的概率，却无法告诉模型如何产生新的正确知识。

模型可能会陷入：

  • 持续生成错误答案；
  • 在错误答案附近优化；
  • 无法探索到正确答案区域。

因此，生成反馈优化通常要求：

  • 模型本身已经具备一定基础能力；
  • 正确答案位于模型可探索的生成空间内；
  • reward 或评价函数能够准确反映优化目标。

换句话说：

生成反馈优化不仅依赖评价函数是否准确，也依赖模型是否能够探索到具有优化价值的答案空间。

这也是为什么实际 LLM 训练通常遵循：

1
2
3
4
5
Pre-training
↓
SFT / Distillation
↓
RLHF / GRPO

先让模型具备基本能力，再利用生成反馈方法进一步优化行为。

3.3 基于偏好的优化（Preference Optimization）

与基于生成反馈的方法不同，偏好优化不要求模型主动探索答案，而是直接利用已有的偏好数据进行训练。

其核心思想：

如果一个答案比另一个答案更好，则提高模型生成该答案的概率。

训练数据形式： \[ (x,y_w,y_l) \] 其中：

  • x：prompt；
  • \(y_w\)：chosen（更优答案）；
  • \(y_l\)：rejected（较差答案）。

例如：

1
2
3
4
5
6
7
8
Question:
9+8是多少？

Chosen:
9+8=17

Rejected:
9+8=18

优化目标： \[ P(chosen|x)>P(rejected|x) \] 代表方法包括：

  • DPO
  • IPO
  • KTO
  • ORPO

与生成式优化相比：

RLHF / GRPO：

1
2
3
4
5
模型生成答案
↓
Reward评价
↓
优化模型

Preference Optimization：

1
2
3
4
5
已有好答案和坏答案
↓
比较两者差异
↓
调整模型概率分布

因此，它避免了模型自主探索答案空间的问题，而是直接利用已有偏好信号进行优化。

但是，这类方法也存在限制：

它只能知道：

1
chosen > rejected

而不知道：

1
为什么 chosen 更好

因此 preference pair 通常需要满足：

  • chosen 和 rejected 内容相似；
  • 长度接近；
  • 主要体现一个核心差异。

否则模型可能无法判断真正的优化目标。

3.4 DPO（Direct Preference Optimization）

DPO 对我来说最核心的思想是：

不需要额外训练 Reward Model，而是直接通过偏好数据对模型的输出概率分布进行修正。

DPO训练数据： \[ (x,y_w,y_l) \] 其中：

  • x：prompt；
  • \(y_w\)：chosen；
  • \(y_l\)：rejected。

DPO 的核心 loss： \[ L_{DPO}(\theta) = -\log\sigma \left( \beta \left( \log \frac{\pi_\theta(y_w|x)} {\pi_{ref}(y_w|x)} - \log \frac{\pi_\theta(y_l|x)} {\pi_{ref}(y_l|x)} \right) \right) \] 其中：

  • \(\pi_\theta\)：当前正在训练的模型；
  • \(\pi_{ref}\)：参考模型（通常是训练前冻结的模型）；
  • \(y_w\)：chosen；
  • \(y_l\)：rejected；
  • \(\beta\)：控制偏好优化强度；
  • \(\sigma\)：sigmoid函数。

公式核心可以理解为： \[ \log\pi_\theta(y_w|x) - \log\pi_\theta(y_l|x) \] 也就是：

让模型生成 chosen 的概率越来越高，而 rejected 的概率越来越低。

例如：

Prompt：

1
9+8是多少？

两个回答：

Chosen:

1
9+8=17

Rejected:

1
9+8=18

DPO希望： \[ \log\pi(y_w|x) > \log\pi(y_l|x) \]

Token级理解

chosen：

1
9+8=17

概率： \[ \log P(9)+\log P(+)+\log P(8)+\log P(=)+\log P(17) \] rejected：

1
9+8=18

概率： \[ \log P(9)+\log P(+)+\log P(8)+\log P(=)+\log P(18) \] 前面的 token：

1
9+8=

完全相同。

因此在： \[ \log\pi(y_w)-\log\pi(y_l) \] 中会抵消。

真正产生主要梯度的位置来自：

1
17

和：

1
18

之间的概率差异。

模型通过这种方式调整内部概率分布：

提高正确答案 token 的概率，降低错误答案 token 的概率。

DPO中的长度问题

假设：

Chosen：

1
8+9=17

Rejected：

1
8+9经过计算等于18

Rejected 包含更多 token。

如果直接计算： \[ \sum_i logP(token_i) \] 长回答会累积更多概率项，因此可能产生长度偏置。

因此实际训练中通常需要考虑：

  • token平均 log probability；
  • 长度归一化；
  • 其他长度校正方法。

DPO为什么要求chosen和rejected接近？

DPO只知道：

1
A > B

但是不知道：

1
为什么A更好。

例如：

较好的数据：

1
2
3
4
5
6
7
8
Chosen:

北京是中国首都。


Rejected:

北京是中国最大的城市。

模型容易学习：

1
首都 ≠ 最大城市

而：

1
2
3
4
5
6
7
8
Chosen:

北京是中国首都。


Rejected:

从秦朝开始介绍中国历史……

差异包括：

  • 长度；
  • 信息量；
  • 内容；
  • 风格。

模型无法判断真正优化目标。

3.5 Critique 与 DPO 结合

DPO 的训练数据本质上只有：(prompt, chosen, rejected)

其中：

  • chosen：更优答案；
  • rejected：较差答案。

因此，DPO 本身并不知道：

为什么 chosen 比 rejected 更好。

它只能学习：P(chosen |prompt) > P(rejected |prompt)

也就是说，DPO 获得的是结果层面的偏好信号，而不是明确的错误原因或改进方向。

因此，可以引入 Critique（评价/纠错信息）增强 preference 数据质量。

方法1：Critique 作为 answer 的一部分

这种方式最接近普通 DPO。

核心思想：

将错误分析和修正过程直接作为模型回答的一部分，让 DPO 同时学习答案和解释过程。

例如：

1
2
3
4
5
6
7
8
9
{
  "prompt": "9+5=?",

  "chosen":
  "答案是14。\n\n原因：9+5=14，所以15是错误的。",

  "rejected":
  "答案是15。\n\n原因：9+5等于15。"
}

此时：

  • critique 属于 chosen/rejected 的内容；
  • DPO loss 会直接作用于整个回答序列；
  • 模型不仅学习最终答案，还学习如何解释错误和进行纠正。

这种方式适用于：

  • 数学推理；
  • 代码修正；
  • 事实纠错；
  • 需要解释过程的任务。

但是，它也存在一个问题：

DPO 并不知道具体是哪部分导致 chosen 更优。

例如：

chosen：

1
2
答案是14。
原因：9+5=14，所以15错误。

相比 rejected：

1
2
答案是15。
原因：9+5等于15。

模型可能学习到：

  • 正确答案；
  • 更详细的解释；
  • 特定回答格式；

而无法明确区分：

是因为答案正确，还是因为解释更加完整。

方法2：Critique 用于生成 preference pair (推荐)

这种方式在实际研究中更加常见。

核心思想：

Critique 不直接参与 DPO 训练，而是用于帮助构造更加准确的 chosen/rejected 数据。

流程如下：

Step 1：生成候选回答

例如：

1
2
3
4
5
Question:
9+5=?

Answer:
9+5=15。

Step 2：生成 Critique

由更强模型或 Critic 判断：

1
2
3
Critique:
该回答错误。
9+5应该等于14，而不是15。

Step 3：根据 Critique 修正答案

得到：

1
2
3
4
5
chosen:
9+5=14。

rejected:
9+5=15。

最终用于 DPO：

1
2
3
4
5
6
{
  "prompt": "9+5=?",

  "chosen": "9+5=14。",
  "rejected": "9+5=15。"
}

此时：

  • Critique 不参与 DPO loss；
  • Critique 的作用是提高 preference pair 的质量；
  • DPO 只负责学习 chosen 和 rejected 之间的偏好关系。

两种方式的区别

方法 Critique作用 是否进入DPO Loss 学习目标
Critique作为answer一部分 直接增强回答内容 是 学习答案 + 解释过程
Critique生成preference pair 辅助构造训练数据 否 学习答案之间的偏好关系

对于知识内化任务，我认为第二种方式更符合原本的目标。

原因是：

如果目标是让 Student 学习外部知识，那么核心问题是：

如何让模型从错误答案转向正确答案。

此时更重要的是构造：

1
2
3
4
5
错误回答
        ↓
修正方向
        ↓
正确回答

而不是让模型学习额外的解释格式。

因此，一个更合理的流程是：

1
2
3
4
5
6
7
8
9
Teacher生成问题
↓
Student生成候选答案
↓
Critique模型分析错误
↓
生成chosen/rejected pair
↓
DPO修正Student模型概率分布

这样，Critique 负责提供纠错信息，而 DPO 负责完成最终的模型行为修正。

3.6 SFT与DPO的关系

SFT解决：

模型应该学习什么内容。

训练数据：(prompt,answer)

目标：\(-\log\pi_\theta(y|x)\)

直接提高正确答案概率。

DPO解决：

多个答案中哪个更加符合偏好。

数据：\((prompt,chosen,rejected)\)

优化：\(P(chosen)>P(rejected)\)

因此：

SFT更偏向：

1
能力学习

DPO更偏向：

1
行为调整和偏好对齐

通常训练流程：

1
2
3
4
5
Pre-training
↓
SFT
↓
DPO / RLHF

对于知识内化任务，更合理的方式可能是：

1
2
3
4
5
Teacher生成知识
↓
SFT / Distillation学习知识
↓
DPO优化回答质量

而不是单独依靠 DPO 完成知识注入。

4. Context Distillation 实验

在前面的 Context Distillation 部分中，已经介绍了其基本原理。

简单来说，Context Distillation 的目标是：

通过包含额外知识的模型输出，训练一个不依赖外部知识输入的模型，使其将知识逐渐内化到模型参数中。

本实验希望验证：

在推理阶段去除外部 knowledge 后，Student 模型是否能够仅通过 question 完成原本需要依赖知识才能完成的任务。

4.1 知识定义

首先需要确定需要让模型内化的目标知识。

本实验选择构造一个简单规则世界：

1
2
3
4
5
6
在 Zeta 世界中：

• 玩家初始拥有 10 点能量；
• 红色按钮：增加 5 点能量；
• 蓝色按钮：减少 3 点能量；
• 紫色按钮：将当前能量翻倍。

该规则作为模型需要学习的外部知识。

4.2 知识预训练

在正式进行 Context Distillation 前，首先对目标模型进行一次简单预训练。

目的：

提前调整模型的概率分布，使模型具备一定的规则文本表示能力，类似于学生在正式学习前进行预习。

需要注意：

该步骤并不是 Context Distillation 必需步骤，只是实验中的一种尝试。

训练代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskType
from tqdm.auto import tqdm

from pathlib import Path
class RuleMemorizer:
    def __init__(
        self,
        rules,
        torch_dtype=torch.bfloat16,
        model_id="Qwen/Qwen3-0.6B",
        local_root="./LLM_models",
        device_map="auto",
    ):
        self.rules = rules
        self.text = rules  # 训练文本就是 rules 本身，不额外套模板


        # 用模型名作为本地子目录,把斜杠替换成下划线避免多级目录歧义
        self.local_dir = Path(local_root) / model_id.replace("/", "_")  # LLM_models\Qwen_Qwen3-0.6B


        self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.local_dir,
            torch_dtype=torch_dtype,
            device_map=device_map,
        )
        print(next(self.model.parameters()).device)
        # 每隔 1 层训练一层，即只对 0, 2, 4, ... 层挂 LoRA，其余层跳过不训练
        num_layers = self.model.config.num_hidden_layers
        layers_to_transform = list(range(0, num_layers, 2))

        lora_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=8,
            lora_alpha=16,
            lora_dropout=0.0,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
            layers_to_transform=layers_to_transform,
        )
        self.model = get_peft_model(self.model, lora_config)
        self.model.print_trainable_parameters()

        # 只有一条样本：labels 与 input_ids 完全一致，即对全部 token 计算 loss，
        # 是纯 next-token prediction（预训练式），不是只对答案部分计算 loss 的 SFT 方式。
        ids = self.tokenizer(self.text)["input_ids"]
        # 手动补 eos，让模型学到"这段文本读完就该结束"，而不是学到无限接续下去
        if self.tokenizer.eos_token_id is not None and ids[-1] != self.tokenizer.eos_token_id:
            ids.append(self.tokenizer.eos_token_id)
        device = self.model.device
        self.input_ids = torch.tensor([ids], device=device)
        self.labels = self.input_ids.clone()

    def train(self, num_episodes=30, learning_rate=1e-5):
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=learning_rate)

        self.model.train()
        for episode in tqdm(range(num_episodes), desc="memorizing rules"):
            optimizer.zero_grad()
            # 不需要显式传 attention_mask，原因是你的输入只有一条样本，而且没有 padding。
            out = self.model(input_ids=self.input_ids, labels=self.labels)
            out.loss.backward()
            optimizer.step()

    # def save(self, output_dir="./student_b_memorized"):
    #     self.model.save_pretrained(output_dir)
    #     self.tokenizer.save_pretrained(output_dir)

    def save(
            self,
            output_root="./trained_models",
            suffix="rulemem",
            merge=True,
    ):
        root = Path(output_root)
        root.mkdir(parents=True, exist_ok=True)

        base_name = f"{self.local_dir.name}_{suffix}"

        idx = 1
        while True:
            output_dir = root / f"{base_name}_{idx:03d}"
            if not output_dir.exists():
                break
            idx += 1

        if merge:
            # 合并 LoRA 到原模型
            model_to_save = self.model.merge_and_unload()
        else:
            # 仅保存 LoRA Adapter
            model_to_save = self.model

        model_to_save.save_pretrained(output_dir)
        self.tokenizer.save_pretrained(output_dir)

        print(f"Saved to: {output_dir}")

if __name__ == "__main__":
    rules='''
    在 Zeta 世界中：
• 红色按钮增加 5 点能量
• 蓝色按钮减少 3 点能量
• 紫色按钮把当前能量翻倍
• 玩家初始能量为 10
    '''

    model_id="Qwen/Qwen2.5-3B-Instruct"
    # model_id="Qwen/Qwen3-0.6B"
    # model_name_or_path 换成 student B 的实际模型；rules 换成 self.rules 的实际内容
    memorizer = RuleMemorizer(model_id=model_id, rules=rules)
    memorizer.train(num_episodes=100, learning_rate=1e-5)
    memorizer.save()

4.3 问题数据构造

接下来需要构造用于 Context Distillation 的问题数据。

最初尝试使用 teacher.py 自动生成问题：

但是实验过程中发现，直接使用 ChatGPT 生成的问题质量更高，因此最终采用人工设计 + ChatGPT 辅助生成的方式构造测试问题。

例如：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
{
  "question": "请给出 Zeta 世界的全部规则。"
},
{
  "question": "红色按钮有什么作用？"
},
{
  "question": "蓝色按钮有什么作用？"
},
{
  "question": "紫色按钮有什么作用？"
},
{
  "question": "当前能量为12点，按一次红色按钮后是多少？"
},
{
  "question": "当前能量为12点，按一次蓝色按钮后是多少？"
}

4.4 带知识上下文的 Student 生成答案

准备好问题数据后，需要生成用于训练 Student 的参考答案。

这里使用目标模型，并在 prompt 中加入 Zeta 世界规则，使模型能够基于外部 knowledge 进行回答。

该模型可以看作：Student A（Context-aware Student）

其输入包含：prompt = knowledge + question

例如：

1
2
3
4
5
6
7
8
9
10
11
Knowledge:

在 Zeta 世界中：
• 玩家初始拥有10点能量；
• 红色按钮增加5点能量；
• 蓝色按钮减少3点能量；
• 紫色按钮将当前能量翻倍。


Question:
当前能量为12点，按一次红色按钮后是多少？

模型生成：

1
2
3
4
5
Answer:
红色按钮增加5点能量。
当前能量：
12 + 5 = 17
因此最终能量为17点。

这些生成结果作为后续 Student B 训练时的目标答案。

生成代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
"""
student_better.py

"更好的学生"（Student A / student_better）：本地加载一个 HF 模型（不接 LoRA，
就是普通 from_pretrained 加载），把 rules 放进 context 里去回答教师生成的问题。

流程：
1. 从 generated/ 目录下读取教师生成的最新 question_pool_<N>.json（只有 question，没有 answer）
2. 对每个还没回答过的 question，结合 rules 生成 answer
3. 把 {question, answer} 重新组合，增量保存到 generated/qa_pool_<当前条数>.json
   （文件名同样以条数结尾，旧的计数文件会被清理，只保留最新一份）
   如果某条生成的答案长度达到了 max_new_tokens（说明很可能是被截断的不完整回答），
   则跳过这条，不写入 qa_pool，但仍标记为"已处理"，避免下次重复生成同一条。

不需要配置 LoRA，直接用本地已经下载好的底座模型。
"""

import os
import re
import json
import glob
from pathlib import Path

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM


STUDENT_A_PROMPT_TEMPLATE = """以下是 Zeta 世界的规则：

{rules}

问题：
{question}

你是一名遵守规则的学生。请严格依据上述规则回答问题。

要求：
1. 不要假设规则中不存在的信息。
2. 如果需要计算，请逐步展示关键计算过程。
3. 每一步操作后的状态必须正确。
4. 最终答案必须满足问题要求。
5. 如果无法根据规则得到答案，请明确说明。

请按照以下格式输出：

推理：
（简要说明使用了哪些规则，以及关键计算过程）

回答：
（最终答案）
"""


class StudentBetter:
    def __init__(
        self,
        rules_path: str,
        model_id: str = "Qwen/Qwen3-0.6B",
        local_root: str = "LLM_models",
        input_dir: str = "generated",
        output_dir: str = "generated",
        torch_dtype=torch.bfloat16,
        device_map="auto",
        max_new_tokens: int = 512,
    ):
        self.rules = self._load_rules(rules_path)
        self.input_dir = input_dir
        self.output_dir = output_dir
        self.max_new_tokens = max_new_tokens

        os.makedirs(self.output_dir, exist_ok=True)

        # 用模型名作为本地子目录，把斜杠替换成下划线避免多级目录歧义
        self.local_dir = Path(local_root) / model_id.replace("/", "_")

        self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.local_dir,
            torch_dtype=torch_dtype,
            device_map=device_map,
        )
        self.model.eval()

        # 已经回答过的 qa 对（支持断点续跑）
        self.qa_pairs = self._load_existing_qa_pool()
        self.answered_questions = {item["question"] for item in self.qa_pairs}
        # 曾经处理过但因为长度触顶被跳过的问题，同样不再重复生成
        self.skipped_questions = set()

    # ------------------------------------------------------------------ #
    # 基础工具
    # ------------------------------------------------------------------ #
    @staticmethod
    def _load_rules(rules_path: str) -> str:
        with open(rules_path, "r", encoding="utf-8") as f:
            return f.read()

    def _load_latest_questions(self):
        """读取 input_dir 下条数最多的 question_pool_*.json，返回 question 字符串列表。"""
        pattern = os.path.join(self.input_dir, "question_pool_*.json")
        files = glob.glob(pattern)
        if not files:
            raise FileNotFoundError(
                f"在 {self.input_dir} 下没有找到 question_pool_*.json，请先运行 teacher.py 生成问题。"
            )

        def count_of(path):
            m = re.search(r"question_pool_(\d+)\.json$", path)
            return int(m.group(1)) if m else -1

        latest = max(files, key=count_of)
        with open(latest, "r", encoding="utf-8") as f:
            items = json.load(f)
        return [item["question"] for item in items]

    def _existing_qa_files(self):
        pattern = os.path.join(self.output_dir, "qa_pool_*.json")
        return glob.glob(pattern)

    def _load_existing_qa_pool(self):
        files = self._existing_qa_files()
        if not files:
            return []

        def count_of(path):
            m = re.search(r"qa_pool_(\d+)\.json$", path)
            return int(m.group(1)) if m else -1

        latest = max(files, key=count_of)
        with open(latest, "r", encoding="utf-8") as f:
            return json.load(f)

    def _save_qa_pool(self):
        """按当前条数命名保存，并清理旧的计数文件，只保留最新一份。"""
        count = len(self.qa_pairs)
        new_path = os.path.join(self.output_dir, f"qa_pool_{count}.json")

        with open(new_path, "w", encoding="utf-8") as f:
            json.dump(self.qa_pairs, f, ensure_ascii=False, indent=2)

        for old_path in self._existing_qa_files():
            if old_path != new_path:
                os.remove(old_path)

        return new_path

    # ------------------------------------------------------------------ #
    # 生成
    # ------------------------------------------------------------------ #
    def _call_model(self, prompt: str):
        """返回 (解码后的文本, 新生成的 token 数)。"""
        messages = [
            # {
            #     "role": "system",
            #     "content": "/no_think"
            # },
            {
                "role": "user",
                "content": prompt
            }
        ]

        inputs = self.tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
            enable_thinking=False,
        ).to(self.model.device)

        with torch.no_grad():
            output_ids = self.model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                do_sample=True,
                temperature=0.7,
                pad_token_id=self.tokenizer.eos_token_id,
            )

        print(next(self.model.parameters()).device)
        new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
        num_new_tokens = len(new_tokens)
        print(num_new_tokens)
        # print(output_ids[0][-10:])
        # print(self.tokenizer.eos_token_id)

        text = self.tokenizer.decode(
            new_tokens,
            skip_special_tokens=True,
        ).strip()

        return text, num_new_tokens

    @staticmethod
    def _parse_answer(raw_text: str):
        """从模型输出中解析出"回答："后面的内容，解析失败则原样返回整段文本。"""
        a_match = re.search(r"回答[：:]\s*(.+)", raw_text, re.DOTALL)
        if a_match:
            answer = a_match.group(1).strip()
            if answer:
                return answer
        # 没匹配到格式的话，直接把整段输出当作回答，避免丢数据
        return raw_text.strip()

    def answer_one(self, question: str):
        """
        生成一条回答。

        返回 (answer, hit_limit)：
        - hit_limit=True 表示生成的新 token 数达到了 max_new_tokens，
          很可能是被截断的不完整回答，调用方应跳过这条，不写入 qa_pool。
        """
        prompt = STUDENT_A_PROMPT_TEMPLATE.format(rules=self.rules, question=question)
        raw, num_new_tokens = self._call_model(prompt)
        hit_limit = num_new_tokens >= self.max_new_tokens
        return self._parse_answer(raw), hit_limit

    # ------------------------------------------------------------------ #
    # 主流程
    # ------------------------------------------------------------------ #
    def run(self):
        """
        读取教师生成的全部问题，对其中还没回答过的问题逐一生成答案，
        每完成一条就立即落盘一次（qa_pool_<当前条数>.json）。
        若某条答案长度达到 max_new_tokens 上限，视为疑似截断，跳过不保存，
        但同样标记为已处理，避免重复生成。
        """
        questions = self._load_latest_questions()
        pending = [
            q for q in questions
            if q not in self.answered_questions and q not in self.skipped_questions
        ]

        if not pending:
            print("[StudentBetter] 所有问题都已经回答过，无需再生成。")
            return self.qa_pairs

        print(f"[StudentBetter] 共 {len(questions)} 条问题，其中 {len(pending)} 条待回答。")

        for question in pending:
            answer, hit_limit = self.answer_one(question)

            if hit_limit:
                self.skipped_questions.add(question)
                print(
                    f"[StudentBetter] 答案长度达到 max_new_tokens={self.max_new_tokens}，"
                    f"疑似截断，跳过该条，不写入 qa_pool -> {question[:50]}..."
                )
                continue

            self.qa_pairs.append({"question": question, "answer": answer})
            self.answered_questions.add(question)

            saved_path = self._save_qa_pool()
            print(f"[StudentBetter] 第 {len(self.qa_pairs)} 条已生成并保存 -> {saved_path}")

        print(f"[StudentBetter] 完成，共 {len(self.qa_pairs)} 条 QA 数据。")
        return self.qa_pairs


# ------------------------------------------------------------------ #
# 直接在 IDE 里跑：改这里的变量就行，不用传命令行参数
# ------------------------------------------------------------------ #
RULES_PATH = "./zeta_world.md"
MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"  # "Qwen/Qwen3-0.6B"
# LOCAL_ROOT = "LLM_models"
LOCAL_ROOT = r"F:\code\python\teacher_student_train\pre_train\LLM_models"
INPUT_DIR = "generated"
OUTPUT_DIR = "generated"
MAX_NEW_TOKENS = 512


def main():
    student = StudentBetter(
        rules_path=RULES_PATH,
        model_id=MODEL_ID,
        local_root=LOCAL_ROOT,
        input_dir=INPUT_DIR,
        output_dir=OUTPUT_DIR,
        max_new_tokens=MAX_NEW_TOKENS,
    )
    student.run()


if __name__ == "__main__":
    main()

4.5 Student模型训练

完成参考答案生成后，开始对 Student 模型进行训练。

训练目标是：

去除外部 rule / knowledge 输入，使 Student 仅根据 question 生成对应 answer，从而验证知识是否能够被内化到模型参数中。

训练数据形式：

1
question → answer

其中：

  • question：原始问题；
  • answer：由带 knowledge 的 Student 生成的参考答案。

训练阶段主要流程如下：

首先对模型进行 LoRA 配置。为了在训练速度和效果之间进行平衡，实验中选择每间隔一层 Transformer layer 注入 LoRA adapter，而不是对所有层进行训练。

该方式可以减少需要更新的参数数量，提高训练速度，同时保持较好的训练效果。需要注意的是，由于基础模型参数仍然需要加载，因此该方法对显存占用降低有限。

由于实验数据量较小，训练过程中采用单样本训练方式，每次输入一个样本进行参数更新。

训练过程本质仍然属于 Next-token Prediction（下一词预测）。

模型输入：

1
question

注意这里不再包含 rule / knowledge。

模型根据 question 逐 token 预测目标 answer：

1
2
3
4
5
question
↓
Student
↓
answer

训练过程中，将模型生成的 token 概率与真实 answer 中对应 token 进行比较，并计算交叉熵损失。

通过不断优化 loss，使模型逐渐提高生成目标 answer 的概率。

训练代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
"""
student.py

Student B（真正要被"蒸馏"的学生）：
- 本地加载底座模型（和 student_better 一样的加载方式），但套上 LoRA
- LoRA 只作用在"每隔一层"的 transformer 层上（层间隔可配置），不是全部层
- prompt 中不包含 rules，只有 question
- 用 student_better 生成的 {question, answer} 作为监督信号，训练时把 prompt 部分的
  label 设成 -100，只对 answer 部分计算 loss（把知识/规则能力蒸馏进参数里）
- 只取 qa_pool 中的前 train_ratio（默认 90%）条记录参与训练，其余部分不用于本次训练
"""

import os
import re
import json
import glob
import random
import shutil
from pathlib import Path

import torch
from torch.optim import AdamW
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import LoraConfig, get_peft_model, PeftModel, TaskType
from tqdm import tqdm


# STUDENT_B_PROMPT_TEMPLATE = """
# {question}
#
# 你是一名学生，请根据提供的内容，对老师的问题进行解答。
# 除了必要的推理，尽量简洁。
#
# 请按照以下格式输出：
# 回答：xxx
# """
STUDENT_B_PROMPT_TEMPLATE = """
问题：
{question}

你是一名遵守规则的学生。请严格依据上述规则回答问题。

要求：
1. 不要假设规则中不存在的信息。
2. 如果需要计算，请逐步展示关键计算过程。
3. 每一步操作后的状态必须正确。
4. 最终答案必须满足问题要求。
5. 如果无法根据规则得到答案，请明确说明。

请按照以下格式输出：

推理：
（简要说明使用了哪些规则，以及关键计算过程）

回答：
（最终答案）
"""


class Student:
    def __init__(
        self,
        model_id: str = "Qwen/Qwen3-0.6B",
        local_root: str = "LLM_models",
        input_dir: str = "generated",
        output_dir: str = "checkpoints/student_lora",
        qa_pool_path: str = None,  # 手动指定要读取的 qa_pool json 路径，None 则自动从 input_dir 下挑选
        torch_dtype=torch.bfloat16,
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
        # LoRA 相关
        lora_r: int = 8,
        lora_alpha: int = 16,
        lora_dropout: float = 0.05,
        layer_stride: int = 2,   # 每隔几层用一次 LoRA（2 = 隔一层）
        layer_start: int = 0,   # 从第几层开始算起
        # 训练相关
        learning_rate: float = 1e-4,
        num_epochs: int = 3,
        train_ratio: float = 0.9,   # 只取 qa_pool 中的这个比例用于训练
        shuffle_before_split: bool = False,  # 切分前是否先打乱顺序，默认不打乱，直接取前 train_ratio 条
        seed: int = 42,
        save_every_n_epochs: int = 5,   # 每隔多少个 epoch 自动保存一次
        max_checkpoints: int = 5,       # 最多保留几份自动保存的 checkpoint（滚动删除最旧的）
        resume_from_checkpoint: bool = True,  # 是否自动从 output_dir 下最新的 checkpoint 继续训练
    ):
        self.input_dir = input_dir
        self.output_dir = output_dir
        self.qa_pool_path = qa_pool_path
        self.device = device
        self.learning_rate = learning_rate
        self.num_epochs = num_epochs
        self.train_ratio = train_ratio
        self.shuffle_before_split = shuffle_before_split
        self.seed = seed
        self.save_every_n_epochs = save_every_n_epochs
        self.max_checkpoints = max_checkpoints
        self.resume_from_checkpoint = resume_from_checkpoint
        self.start_epoch = 0  # 已经训练过的 epoch 数，resume 时会被更新
        self._saved_checkpoint_dirs = []  # 记录已保存的 checkpoint 目录，按保存顺序排列（旧 -> 新）

        os.makedirs(self.output_dir, exist_ok=True)

        # 用模型名作为本地子目录，把斜杠替换成下划线避免多级目录歧义
        self.local_dir = Path(local_root) / model_id.replace("/", "_")

        self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
        base_model = AutoModelForCausalLM.from_pretrained(
            self.local_dir,
            torch_dtype=torch_dtype,
        ).to(self.device)

        # 找 output_dir 下是否已经有滚动保存的 checkpoint，决定是新建 LoRA 还是继续训练
        latest_ckpt_dir, latest_ckpt_epoch = self._find_latest_checkpoint()

        if self.resume_from_checkpoint and latest_ckpt_dir is not None:
            print(f"[Student] 检测到已有 checkpoint -> {latest_ckpt_dir}，将从第 {latest_ckpt_epoch} 轮之后继续训练。")
            self.model = PeftModel.from_pretrained(base_model, latest_ckpt_dir, is_trainable=True)
            self.start_epoch = latest_ckpt_epoch
            # 把已存在的 checkpoint 目录按 epoch 从旧到新记录下来，保持滚动删除窗口的连续性
            self._saved_checkpoint_dirs = self._list_existing_checkpoints()
        else:
            if self.resume_from_checkpoint and latest_ckpt_dir is None:
                print("[Student] resume_from_checkpoint=True，但没有找到已有 checkpoint，将从头开始训练。")

            target_modules = self._build_stride_target_modules(
                base_model, layer_stride=layer_stride, layer_start=layer_start
            )
            print(f"[Student] LoRA 作用的 target_modules 数量: {len(target_modules)}")
            print(f"[Student] 示例: {target_modules[:7]}")

            lora_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM,
                r=lora_r,
                lora_alpha=lora_alpha,
                lora_dropout=lora_dropout,
                target_modules=target_modules,
                bias="none",
            )
            self.model = get_peft_model(base_model, lora_config)

        self.model.print_trainable_parameters()

    # ------------------------------------------------------------------ #
    # 查找 output_dir 下已有的 checkpoint
    # ------------------------------------------------------------------ #
    @staticmethod
    def _epoch_of_checkpoint_dir(path: str) -> int:
        m = re.search(r"checkpoint-(\d+)$", path)
        return int(m.group(1)) if m else -1

    def _list_existing_checkpoints(self):
        """返回 output_dir 下所有 checkpoint-<epoch> 目录，按 epoch 从旧到新排序。"""
        pattern = os.path.join(self.output_dir, "checkpoint-*")
        dirs = [d for d in glob.glob(pattern) if os.path.isdir(d)]
        dirs = [d for d in dirs if self._epoch_of_checkpoint_dir(d) >= 0]
        dirs.sort(key=self._epoch_of_checkpoint_dir)
        return dirs

    def _find_latest_checkpoint(self):
        """返回 (最新 checkpoint 目录, 对应的 epoch 数)，没有则返回 (None, 0)。"""
        dirs = self._list_existing_checkpoints()
        if not dirs:
            return None, 0
        latest = dirs[-1]
        return latest, self._epoch_of_checkpoint_dir(latest)

    # ------------------------------------------------------------------ #
    # 按"每隔一层"构造 LoRA 的 target_modules
    # ------------------------------------------------------------------ #
    @staticmethod
    def _build_stride_target_modules(base_model, layer_stride: int, layer_start: int):
        """
        只挑选 layer_start, layer_start+layer_stride, layer_start+2*layer_stride, ...
        这些层里的 attention/mlp 投影层，作为 LoRA 的 target_modules。
        target_modules 用"层内后缀名"的形式给出（如 "layers.0.self_attn.q_proj"），
        peft 内部用 endswith 匹配，所以不需要写完整的模块路径。
        """
        num_layers = base_model.config.num_hidden_layers
        selected_layers = list(range(layer_start, num_layers, layer_stride))
        print(f"[Student] 模型共 {num_layers} 层，选中做 LoRA 的层: {selected_layers}")

        # 常见 Qwen/Llama 系架构里，attention 和 mlp 的投影层名字
        proj_suffixes = [
            "self_attn.q_proj",
            "self_attn.k_proj",
            "self_attn.v_proj",
            "self_attn.o_proj",
            "mlp.gate_proj",
            "mlp.up_proj",
            "mlp.down_proj",
        ]

        target_modules = [
            f"layers.{layer_idx}.{suffix}"
            for layer_idx in selected_layers
            for suffix in proj_suffixes
        ]
        return target_modules

    # ------------------------------------------------------------------ #
    # 数据加载
    # ------------------------------------------------------------------ #
    def _load_latest_qa_pool(self):
        if self.qa_pool_path is not None:
            if not os.path.isfile(self.qa_pool_path):
                raise FileNotFoundError(f"指定的 qa_pool_path 不存在: {self.qa_pool_path}")
            with open(self.qa_pool_path, "r", encoding="utf-8") as f:
                data = json.load(f)
            print(f"[Student] 使用手动指定的数据源 {self.qa_pool_path}（{len(data)} 条）作为训练数据。")
            return data

        pattern = os.path.join(self.input_dir, "qa_pool_*.json")
        files = glob.glob(pattern)
        if not files:
            raise FileNotFoundError(
                f"在 {self.input_dir} 下没有找到 qa_pool_*.json，请先运行 student_better.py 生成数据，"
                f"或通过 qa_pool_path 手动指定要读取的文件。"
            )

        loaded = []
        for path in files:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            loaded.append((len(data), path, data))

        # 按实际读取到的条数（而不是文件名里的数字）挑出最多的那份
        loaded.sort(key=lambda item: item[0])
        count, path, data = loaded[-1]
        print(f"[Student] 使用 {path}（实际 {count} 条）作为训练数据。")
        return data

    def _split_train_pairs(self, qa_pairs):
        """
        只取 qa_pairs 中的 train_ratio（默认 90%）条记录用于训练，
        剩余部分本次不参与训练（可用于后续验证/评估）。
        """
        total = len(qa_pairs)
        train_size = int(total * self.train_ratio)
        # 边界保护：至少留 1 条用于训练（只要 qa_pairs 本身非空）
        if total > 0:
            train_size = max(1, min(train_size, total))

        indices = list(range(total))
        if self.shuffle_before_split:
            random.Random(self.seed).shuffle(indices)

        train_indices = indices[:train_size]
        train_pairs = [qa_pairs[i] for i in train_indices]

        print(
            f"[Student] qa_pool 共 {total} 条，按 train_ratio={self.train_ratio} "
            f"取 {len(train_pairs)} 条用于本次训练。"
        )
        return train_pairs

    # ------------------------------------------------------------------ #
    # 构造训练样本：prompt 部分 label 设为 -100，只学 answer
    # ------------------------------------------------------------------ #
    def build_training_example(self, question: str, answer: str):
        prompt = STUDENT_B_PROMPT_TEMPLATE.format(question=question)
        full_text = prompt + answer

        prompt_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
        full_ids = self.tokenizer(full_text, return_tensors="pt").input_ids.to(self.device)
        prompt_len = prompt_ids.shape[1]

        labels = full_ids.clone()
        labels[:, :prompt_len] = -100

        attention_mask = torch.ones_like(full_ids)
        # 因为是逐条样本训练（batch_size=1，没有做 padding），full_ids 里每一个位置都是真实 token，没有 padding，
        # 所以 torch.ones_like(full_ids) 就是生成一个和 full_ids 形状一样、全是 1 的 mask——本质上等于"全部位置都参与计算"。
        return full_ids, attention_mask, labels

    # ------------------------------------------------------------------ #
    # 训练
    # ------------------------------------------------------------------ #
    def _save_checkpoint(self, epoch_num: int):
        """
        保存一份 checkpoint 到 output_dir/checkpoint-<epoch_num>，
        并维护一个最多 max_checkpoints 份的滚动窗口：
        超出上限时删除最旧的一份，只保留最新的 max_checkpoints 份。
        """
        ckpt_dir = os.path.join(self.output_dir, f"checkpoint-{epoch_num}")
        self.model.save_pretrained(ckpt_dir)
        self.tokenizer.save_pretrained(ckpt_dir)
        print(f"[Student] 第 {epoch_num} 轮 checkpoint 已保存 -> {ckpt_dir}")

        self._saved_checkpoint_dirs.append(ckpt_dir)
        while len(self._saved_checkpoint_dirs) > self.max_checkpoints:
            oldest_dir = self._saved_checkpoint_dirs.pop(0)
            if os.path.isdir(oldest_dir):
                shutil.rmtree(oldest_dir)
                print(f"[Student] 超出最多保留 {self.max_checkpoints} 份，删除最旧 checkpoint -> {oldest_dir}")

    def train(self):
        qa_pairs_all = self._load_latest_qa_pool()
        qa_pairs = self._split_train_pairs(qa_pairs_all)

        if self.start_epoch >= self.num_epochs:
            print(
                f"[Student] 已训练到第 {self.start_epoch} 轮，"
                f"达到/超过目标 num_epochs={self.num_epochs}，无需继续训练。"
            )
            return

        print(
            f"[Student] 读取到 {len(qa_pairs)} 条训练数据，"
            f"从第 {self.start_epoch + 1} 轮训练到第 {self.num_epochs} 轮。"
        )

        trainable_params = [p for p in self.model.parameters() if p.requires_grad]
        optimizer = AdamW(trainable_params, lr=self.learning_rate)

        self.model.train()
        for epoch_num in range(self.start_epoch + 1, self.num_epochs + 1):
            total_loss = 0.0
            pbar = tqdm(
                enumerate(qa_pairs),
                total=len(qa_pairs),
                desc=f"[Student] epoch {epoch_num}/{self.num_epochs}",
            )
            for step, item in pbar:
                input_ids, attention_mask, labels = self.build_training_example(
                    item["question"], item["answer"]
                )

                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    labels=labels,
                )
                loss = outputs.loss

                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

                total_loss += loss.item()
                pbar.set_postfix(loss=f"{loss.item():.4f}")

            avg_loss = total_loss / len(qa_pairs)
            print(f"[Student] epoch {epoch_num} 完成，平均 loss={avg_loss:.4f}")

            if self.save_every_n_epochs > 0 and epoch_num % self.save_every_n_epochs == 0:
                self._save_checkpoint(epoch_num)

        self.model.save_pretrained(self.output_dir)
        self.tokenizer.save_pretrained(self.output_dir)
        print(f"[Student] LoRA 适配器已保存到 -> {self.output_dir}")


# ------------------------------------------------------------------ #
# 直接在 IDE 里跑：改这里的变量就行，不用传命令行参数
# ------------------------------------------------------------------ #
# MODEL_ID = "Qwen/Qwen3-0.6B"
# LOCAL_ROOT = "LLM_models"
INPUT_DIR = "generated"
OUTPUT_DIR = "checkpoints/student_lora"
QA_POOL_PATH = "generated/qa_pool_27.json"  # 手动指定要读取的 qa_pool json 路径，如 "generated/qa_pool_200.json"；None 则自动挑选
MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"  # "Qwen/Qwen3-0.6B"
LOCAL_ROOT = r"F:\code\python\teacher_student_train\pre_train\trained_models"


LORA_R = 8
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
LAYER_STRIDE = 2      # 每隔一层用 LoRA
LAYER_START = 0

LEARNING_RATE = 1e-4
NUM_EPOCHS = 30
TRAIN_RATIO = 0.9      # 只用 qa_pool 中 90% 的记录训练
SAVE_EVERY_N_EPOCHS = 5   # 每 5 轮自动保存一次
MAX_CHECKPOINTS = 5       # 最多保留 5 份自动保存的 checkpoint（滚动删除最旧的）
RESUME_FROM_CHECKPOINT = True  # 自动从 output_dir 下最新的 checkpoint 继续训练


def main():
    student = Student(
        model_id=MODEL_ID,
        local_root=LOCAL_ROOT,
        input_dir=INPUT_DIR,
        output_dir=OUTPUT_DIR,
        qa_pool_path=QA_POOL_PATH,
        lora_r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        layer_stride=LAYER_STRIDE,
        layer_start=LAYER_START,
        learning_rate=LEARNING_RATE,
        num_epochs=NUM_EPOCHS,
        train_ratio=TRAIN_RATIO,
        save_every_n_epochs=SAVE_EVERY_N_EPOCHS,
        max_checkpoints=MAX_CHECKPOINTS,
        resume_from_checkpoint=RESUME_FROM_CHECKPOINT,
    )
    student.train()


if __name__ == "__main__":
    main()

4.6 实验划分与结果验证

为了验证 Student 模型是否真正完成知识内化，而非仅依赖外部 Prompt 提供的规则信息，将实验数据划分为：

  • 90%：用于 Student 模型训练；
  • 10%：用于测试模型知识内化效果。

测试集中的问题不会出现在训练数据中，用于评估模型对未见问题的泛化能力。测试阶段，通过对比模型训练前后，以及不同知识注入方式下的回答结果，分析 Context Distillation 对知识迁移和推理能力的影响。

测试问题示例：

1
最终能量为12点，最后一步使用蓝色按钮，那么之前是多少？

该问题属于测试集（10%），未参与 Student 模型训练。

Zeta 世界规则如下：

1
2
3
4
5
在 Zeta 世界中：
• 玩家初始为 10 点能量
• 红色按钮增加 5 点能量
• 蓝色按钮减少 3 点能量
• 紫色按钮把当前能量翻倍

（1）训练前模型

在训练前，模型未接触 Zeta 世界规则，因此只能根据已有语言模式进行回答：

1
2
3
4
5
6
7
你: 最终能量为12点，最后一步使用蓝色按钮，那么之前是多少？

模型: 根据您提供的信息，最终能量为12点，并且知道最后一步是使用了蓝色按钮。在大多数情况下，蓝色按钮可能会增加或提升当前的能量值。因此，在最后一步之前，能量值应该是12点减去蓝色按钮的增加量。

不过，没有具体的蓝色按钮增加的能量数值，我无法给出确切的数字。如果蓝色按钮每次增加固定的能量值，例如每次增加3点，那么之前的能量值就是12点减去蓝色按钮增加的能量点数。

请提供更多关于蓝色按钮的具体信息，这样我才能给出更准确的答案。

此时模型由于缺少任务相关知识，无法完成有效推理。

（2）预训练后的 Student 模型

经过规则数据预训练后，模型已经接触了 Zeta 世界相关知识，因此能够给出一个合理形式的回答：

1
2
3
你: 最终能量为12点，最后一步使用蓝色按钮，那么之前是多少？

模型: 根据您提供的信息，如果最终的能量值是12点，并且最后一跳（即最后一步）使用了蓝色按钮，那么在使用蓝色按钮前的能量值应该是7点。这是因为蓝色按钮通常赋予一定的能量值，使得玩家的总能量达到一个特定的目标值。在这个例子中，假设蓝色按钮提供5点能量（这是一个常见的设定），那么从7点增加到12点就只需要使用这个按钮一次。

相比训练前，模型已经能够围绕该任务生成一个形式上合理的回答，不再直接拒绝回答或要求补充规则信息。

但是，该答案仍然是错误的。模型虽然学习到了部分任务相关模式，但并未稳定掌握规则之间的对应关系，因此在需要进行反向推理时出现错误。

（3）Prompt 注入 Knowledge 后的预训练后的 Student 模型

为了进一步分析模型在显式规则输入下的表现，将相同问题输入带规则 Prompt 的 Student 模型：

1
2
3
4
5
question:
最终能量为12点，最后一步使用蓝色按钮，那么之前是多少？

answer:
7

可以发现，即使模型能够直接访问规则信息，仍然可能产生错误答案。

这说明：

  • Prompt 中包含知识，并不代表模型能够正确理解和执行该知识；
  • 对于反向推理、组合推理等任务，小规模模型仍可能无法稳定利用上下文中的规则；
  • 模型可能学习到了局部模式，而没有形成稳定的规则执行能力。

（4）Context Distillation 后 Student 模型

经过 30 轮 Context Distillation 训练后，在不提供额外规则 Prompt 的情况下，模型能够直接根据问题完成正确推理：

1
2
3
你: 最终能量为12点，最后一步使用蓝色按钮，那么之前是多少？

模型: 根据你提供的信息，最终能量为12点，且最后一步使用了蓝色按钮。假设每按一次蓝色按钮会减少3点能量，那么在最后一步使用蓝色按钮之前，能量应该是 \(12 + 3 = 15\) 点。

实验结果分析

通过上述实验结果可以观察到：

  1. 预训练能够提升模型对于目标任务的适应能力。

    在预训练前，由于模型未接触 Zeta 世界相关规则，面对该类问题时只能根据已有语言模式进行猜测，或者要求提供额外信息。经过预训练后，模型已经能够围绕目标任务生成形式上合理的回答，说明部分任务相关知识已经被模型学习，并能够在生成过程中被调用。

    但是，此阶段模型主要表现为对任务模式的拟合，并不代表其已经具备稳定的规则推理能力。在涉及反向推导、组合计算等问题时，模型仍可能出现规则应用错误。

  1. 仅依赖 Prompt 注入 Knowledge 存在一定局限性。

    实验中发现，即使将完整规则信息显式加入 Prompt，Student 模型仍可能在部分复杂推理任务中产生错误答案。

    这表明：

    • 外部 Knowledge Context 能够为模型提供必要的信息；
    • 但模型并不一定能够稳定地理解和利用这些信息；
    • 对于需要多步推理或反向推导的问题，小规模模型仍可能出现规则调用失败。

    因此，将知识放入上下文并不等价于模型真正掌握该知识，模型还需要通过训练过程学习如何有效利用这些规则。

  1. Context Distillation 能够促进规则知识向模型参数迁移。

    经过 Context Distillation 训练后，Student 模型无需额外输入规则 Prompt，即可恢复相关知识并完成正确推理。这说明 Context Distillation 能够使模型逐渐学习带 Knowledge Context 条件下的输出模式。

    从生成分布角度来看，Context Distillation 可以理解为：

通过大量带 Knowledge Context 的正确示范，使 Student 模型逐渐拟合 Teacher 在规则条件下的生成分布，从而将原本依赖外部 Context 的知识迁移到模型参数中。

与直接 Prompt 注入相比，Context Distillation 不仅使模型记忆规则内容，还提升了模型在规则约束下进行推理的稳定性。

  1. 实验局限性与后续改进方向。

    需要注意的是，当前实验中的部分错误回答可能受到基础模型能力的影响。例如，即使提供完整规则信息，较小规模模型仍可能无法正确完成复杂推理任务。

    因此，后续实验可以考虑：

    • 使用能力更强的基础模型作为 Student 模型，以降低模型自身推理能力不足带来的影响；
    • 使用更强的 Teacher 模型生成 Knowledge Context 和监督数据，提高训练样本质量；
    • 扩大测试集规模，并增加更多复杂规则组合任务，以进一步验证 Context Distillation 的知识内化效果。

    通过提升模型能力和监督数据质量，可以进一步区分模型推理能力限制与知识内化方法本身带来的性能差异。

总体而言，实验结果表明，Context Distillation 能够有效促进模型从依赖外部 Knowledge Context向利用自身参数完成知识调用与推理的转变，实现特定领域知识的参数化内化。

最后贴一下项目文件结构