feat: add multi-turn tool-use on-policy distillation
HITSZJiahaoLiu
parent
1e6e909f88
commit
f82990b0ed
|
|
@ -750,6 +750,8 @@ python eval_llm.py --weight full_sft
|
|||
|
||||
仓库中提供的 `train_distillation.py` 更适合作为理解白盒蒸馏流程的参考实现：它完整展示了教师/学生双模型加载、`CE + KL` 混合损失、温度缩放、MoE 与 dense 组合蒸馏，以及断点续训和分布式训练等关键细节。
|
||||
|
||||
多轮工具交互的在线蒸馏参考实现见 [`train_agent_opd.py`](trainer/train_agent_opd.py)：学生执行工具并读取结果，冻结教师沿着学生实际生成的交互轨迹提供 token 分布监督。使用方式、单设备限制与验证范围见 [Agent OPD 指南](docs/agent_opd.md)。
|
||||
|
||||
```bash
|
||||
# 方式1
|
||||
torchrun --nproc_per_node 1 train_distillation.py
|
||||
|
|
|
|||
|
|
@ -750,6 +750,8 @@ White-box distillation typically fits the teacher distribution in addition to th
|
|||
|
||||
The `train_distillation.py` script is intended as a reference implementation for understanding the white-box distillation pipeline: it demonstrates teacher/student dual model loading, `CE + KL` mixed loss, temperature scaling, MoE and dense combination distillation, as well as key details such as checkpoint resume and distributed training.
|
||||
|
||||
For on-policy distillation through multi-turn tool interactions, see [`train_agent_opd.py`](trainer/train_agent_opd.py). The student executes tools and reads their results while a frozen teacher supervises token distributions along the student's sampled trajectory. See the [Agent OPD guide](docs/agent_opd.md) for usage, single-device scope and validation limits.
|
||||
|
||||
```bash
|
||||
# Method 1
|
||||
torchrun --nproc_per_node 1 train_distillation.py
|
||||
|
|
|
|||
|
|
@ -0,0 +1,110 @@
|
|||
# Multi-turn tool-use on-policy distillation
|
||||
|
||||
`trainer/train_agent_opd.py` is a single-device reference trainer for distilling a
|
||||
frozen MiniMind teacher along the student's own tool interactions. It uses the
|
||||
local tool implementations in `train_agent.py`; weather, exchange rates and other
|
||||
mock tool results are **not live services**. It does not load a reward model.
|
||||
|
||||
The student generates a turn, its tool calls are executed, and the returned
|
||||
observations condition the next student turn. The teacher sees exactly these
|
||||
student-generated token prefixes and observations. No teacher answer replaces a
|
||||
student rollout. Prompt text, inserted role markers, tool observations and padding
|
||||
are context only. Every generated assistant token, including each turn's first
|
||||
EOS, is a distillation target. Later assistant turns remain targets after earlier
|
||||
EOS tokens.
|
||||
|
||||
## Objective
|
||||
|
||||
Following the on-policy data collection idea in
|
||||
[GKD (Agarwal et al., ICLR 2024)](https://arxiv.org/abs/2306.13649), the default loss is
|
||||
the mean full-vocabulary `T² KL(teacher_T || student_T)` over assistant tokens on
|
||||
the sampled interaction. `--kl_direction reverse_kl` instead uses
|
||||
`T² KL(student_T || teacher_T)` at the same prefixes. This is a GKD-style
|
||||
semi-gradient: sampling and environment execution are detached. It is not a
|
||||
REINFORCE estimator of trajectory-level reverse KL.
|
||||
|
||||
Sampling defaults to temperature 1, top-p 1 and top-k 0. `--temperature` controls
|
||||
student sampling; `--distill_temperature` controls only the loss distributions.
|
||||
Both checkpoints must be native MiniMind weights using the **same tokenizer and
|
||||
token-to-ID mapping**. A teacher with an unrelated vocabulary is not supported.
|
||||
|
||||
## Run
|
||||
|
||||
Run from the repository root, after installing the repository dependencies:
|
||||
|
||||
```bash
|
||||
python trainer/train_agent_opd.py \
|
||||
--student_checkpoint out/full_sft_768.pth \
|
||||
--teacher_checkpoint out/agent_768.pth \
|
||||
--data_path dataset/agent_rl.jsonl \
|
||||
--device cuda:0 --dtype bfloat16 \
|
||||
--max_turns 3 --max_gen_len 256 --max_total_len 2048 \
|
||||
--batch_size 1 --accumulation_steps 4 --max_steps 100
|
||||
```
|
||||
|
||||
Use an independently trained, stronger teacher. The example filenames refer to
|
||||
weights you already have; the trainer does not download them. Architecture flags
|
||||
(`--student_hidden_size`, `--teacher_num_layers`, `--teacher_use_moe`, etc.) must
|
||||
match their checkpoints. Weights are loaded strictly. For a CPU correctness run,
|
||||
use `--device cpu --dtype float32` and small checkpoints. This initial trainer
|
||||
supports one process/device; it does not provide DDP or SGLang training.
|
||||
|
||||
Data may use `conversations` (as in `agent_rl.jsonl`) or `messages`. Tool schemas
|
||||
can appear in the system message's `tools` field or the record's top-level
|
||||
`tools` field. A trailing fixed assistant answer is ignored. `gt` is not required:
|
||||
|
||||
```json
|
||||
{"messages":[{"role":"user","content":"请用工具计算 2+2。"}],"tools":[{"type":"function","function":{"name":"calculate_math","description":"计算数学表达式","parameters":{"type":"object","properties":{"expression":{"type":"string"}},"required":["expression"]}}}]}
|
||||
```
|
||||
|
||||
Only tools declared for the sample are available. Unknown tools and invalid
|
||||
arguments produce an error observation. Context and turn limits stop further
|
||||
interaction while retaining already sampled targets. Context is never silently
|
||||
left-truncated into a different teacher prefix. Oversized initial prompts raise
|
||||
an error; adjust the data or `--max_total_len`.
|
||||
|
||||
The output includes an inference weight file, `out/agent_opd/agent_opd_768.pth`,
|
||||
and a full-precision `agent_opd_768_resume.pth` containing optimizer/scaler state,
|
||||
the next data cursor and random states. Checkpoints are saved after complete
|
||||
optimizer updates, including a final partial accumulation window. Resume with
|
||||
the original arguments plus:
|
||||
|
||||
```bash
|
||||
--resume out/agent_opd/agent_opd_768_resume.pth --max_steps 200
|
||||
```
|
||||
|
||||
`--max_steps` counts total optimizer steps, including resumed steps. Keep the
|
||||
teacher, data, architecture, optimizer and sampling arguments unchanged. The
|
||||
logged `loss`, `mean_turns`, `tool_calls`, `assistant_tokens` and
|
||||
`truncated_fraction` describe training, not task success rates.
|
||||
|
||||
## Validation and limits
|
||||
|
||||
```bash
|
||||
OMP_NUM_THREADS=1 HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false \
|
||||
python -m unittest discover -s tests -p test_agent_opd.py -v
|
||||
```
|
||||
|
||||
Tests use the real repository tokenizer, local tools and tiny CPU MiniMind
|
||||
models. Scripted generation covers deterministic multi-turn boundaries, including
|
||||
tool errors, per-turn EOS and multiple calls; other tests exercise actual model
|
||||
generation, optimization, teacher freezing and exact CPU checkpoint continuation.
|
||||
A controlled two-turn optimization test verifies that teacher divergence falls.
|
||||
No pretrained task-quality improvement or CUDA/MPS performance result is claimed.
|
||||
|
||||
Before claiming better tool use, evaluate the starting student, teacher and OPD
|
||||
checkpoint on held-out prompts with identical decoding and tool limits. Compare
|
||||
against fixed-data distillation as well, and report final-answer accuracy,
|
||||
successful tool execution, truncation, wall time and peak memory. A falling
|
||||
distillation loss alone is not evidence of better answers.
|
||||
|
||||
## 中文说明
|
||||
|
||||
这是面向多轮工具交互的 OPD 参考实现：学生自己生成工具调用，读取工具结果后继续生成，
|
||||
教师沿着学生实际经历的上下文提供分布监督。工具结果参与上下文，但不作为预测标签；每轮
|
||||
学生生成的 EOS 都参与损失，后续轮次不会因第一轮 EOS 而被屏蔽。默认使用前向 KL，支持
|
||||
反向 KL 的 GKD 式半梯度。教师被冻结，数据不需要 `gt` 或标准答案。
|
||||
|
||||
首版仅支持单进程、共享分词器的 MiniMind 教师与学生，复用仓库的本地模拟工具。
|
||||
请根据实际权重指定教师/学生结构。测试验证训练链路和数值正确性；正式模型上的能力提升
|
||||
仍需独立留出集实验。断点文件保存完整精度和随机状态，恢复时沿用原始参数并加上 `--resume`。
|
||||
|
|
@ -0,0 +1,294 @@
|
|||
"""CPU regressions with the real tokenizer, tools and tiny MiniMind models."""
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
||||
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
||||
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from trainer.agent_opd_utils import (
|
||||
AgentTrajectory, collect_trajectory, distillation_loss, pack_trajectories,
|
||||
)
|
||||
from trainer.train_agent import TOOLS
|
||||
from trainer import train_agent_opd as trainer
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
MATH_TOOLS = [TOOLS[0]]
|
||||
MESSAGES = [{"role": "user", "content": "请用工具计算 2+2。"}]
|
||||
CALL = '<tool_call>\n{"name":"calculate_math","arguments":{"expression":"2+2"}}\n</tool_call>'
|
||||
|
||||
|
||||
class ScriptedStudent(torch.nn.Module):
|
||||
"""Deterministic generation at the model boundary, keeping real tool/token IO."""
|
||||
|
||||
def __init__(self, tokenizer, responses):
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(torch.zeros(()))
|
||||
self.responses = responses
|
||||
self.tokenizer = tokenizer
|
||||
self.inputs, self.options = [], []
|
||||
|
||||
def generate(self, input_ids, **kwargs):
|
||||
self.inputs.append(input_ids[0].tolist())
|
||||
self.options.append(kwargs)
|
||||
response = self.responses[len(self.inputs) - 1]
|
||||
ids = self.tokenizer(response, add_special_tokens=False)["input_ids"] if isinstance(response, str) else response
|
||||
ids = ids[:kwargs["max_new_tokens"]]
|
||||
return torch.cat((input_ids, input_ids.new_tensor([ids])), dim=1)
|
||||
|
||||
|
||||
class TrajectoryTests(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
torch.set_num_threads(1)
|
||||
cls.tokenizer = AutoTokenizer.from_pretrained(ROOT / "model")
|
||||
cls.eos = cls.tokenizer.eos_token
|
||||
|
||||
def sample(self, responses, **kwargs):
|
||||
student = ScriptedStudent(self.tokenizer, responses)
|
||||
trajectory = collect_trajectory(student, self.tokenizer, MESSAGES, MATH_TOOLS, **kwargs)
|
||||
return student, trajectory
|
||||
|
||||
def test_tool_result_conditions_next_turn_but_is_not_a_target(self):
|
||||
student, trajectory = self.sample([CALL + self.eos, "结果是 4。" + self.eos])
|
||||
self.assertEqual(trajectory.turns, 2)
|
||||
self.assertEqual(trajectory.tool_calls, 1)
|
||||
self.assertFalse(trajectory.truncated)
|
||||
self.assertIn('"result": "4"', self.tokenizer.decode(student.inputs[1]))
|
||||
assistant_ids = [t for t, keep in zip(trajectory.input_ids, trajectory.assistant_mask) if keep]
|
||||
expected = self.tokenizer(CALL + self.eos, add_special_tokens=False)["input_ids"]
|
||||
expected += self.tokenizer("结果是 4。" + self.eos, add_special_tokens=False)["input_ids"]
|
||||
self.assertEqual(assistant_ids, expected)
|
||||
self.assertEqual(assistant_ids.count(self.tokenizer.eos_token_id), 2)
|
||||
self.assertTrue(student.training)
|
||||
for options in student.options:
|
||||
self.assertEqual((options["top_p"], options["top_k"], options["temperature"]), (1.0, 0, 1.0))
|
||||
self.assertEqual(MESSAGES, [{"role": "user", "content": "请用工具计算 2+2。"}])
|
||||
|
||||
def test_preserves_exact_sampled_ids_and_first_eos_only(self):
|
||||
sampled = self.tokenizer(CALL, add_special_tokens=False)["input_ids"] + [self.tokenizer.eos_token_id] * 3
|
||||
student, trajectory = self.sample([sampled, [self.tokenizer.eos_token_id]])
|
||||
start = len(student.inputs[0])
|
||||
self.assertEqual(trajectory.input_ids[start:start + len(sampled) - 2], sampled[:-2])
|
||||
self.assertEqual(sum(trajectory.assistant_mask), len(sampled) - 1)
|
||||
|
||||
def test_multiple_tool_calls_return_separate_observations(self):
|
||||
second = CALL.replace("2+2", "3*3")
|
||||
student, trajectory = self.sample([CALL + "\n" + second + self.eos, "4 和 9。" + self.eos])
|
||||
context = self.tokenizer.decode(student.inputs[1])
|
||||
self.assertIn('"result": "4"', context)
|
||||
self.assertIn('"result": "9"', context)
|
||||
self.assertEqual(trajectory.tool_calls, 2)
|
||||
|
||||
def test_unknown_tool_and_malformed_arguments_become_observations(self):
|
||||
for call in [CALL.replace("calculate_math", "get_current_weather"),
|
||||
'<tool_call>{"name":[],"arguments":null}</tool_call>',
|
||||
'<tool_call>{"name":"calculate_math","arguments":"invalid"}</tool_call>']:
|
||||
with self.subTest(call=call):
|
||||
student, trajectory = self.sample([call + self.eos, "无法计算。" + self.eos])
|
||||
self.assertIn("unknown tool or invalid arguments", self.tokenizer.decode(student.inputs[1]))
|
||||
self.assertEqual(trajectory.turns, 2)
|
||||
|
||||
def test_non_dictionary_tool_json_does_not_crash(self):
|
||||
_, trajectory = self.sample(['<tool_call>[]</tool_call>' + self.eos])
|
||||
self.assertEqual(trajectory.turns, 1)
|
||||
|
||||
def test_turn_limit_keeps_generated_tokens_without_executing_more_tools(self):
|
||||
with patch("trainer.agent_opd_utils.execute_tool") as execute:
|
||||
_, trajectory = self.sample([CALL + self.eos], max_turns=1)
|
||||
self.assertTrue(trajectory.truncated)
|
||||
execute.assert_not_called()
|
||||
self.assertTrue(trajectory.assistant_mask[-1])
|
||||
|
||||
def test_context_limit_never_slices_an_observation_into_a_training_target(self):
|
||||
probe = ScriptedStudent(self.tokenizer, [CALL + self.eos])
|
||||
prompt = self.tokenizer.apply_chat_template(MESSAGES, tools=MATH_TOOLS, tokenize=False,
|
||||
add_generation_prompt=True, open_thinking=False)
|
||||
limit = len(self.tokenizer(prompt, add_special_tokens=False)["input_ids"]) + 8
|
||||
trajectory = collect_trajectory(probe, self.tokenizer, MESSAGES, MATH_TOOLS, max_total_len=limit)
|
||||
self.assertEqual(len(trajectory.input_ids), limit)
|
||||
self.assertTrue(trajectory.truncated)
|
||||
with self.assertRaisesRegex(ValueError, "Prompt leaves no generation space"):
|
||||
collect_trajectory(probe, self.tokenizer, MESSAGES, MATH_TOOLS, max_total_len=1)
|
||||
|
||||
def test_empty_final_answer_still_trains_eos(self):
|
||||
_, trajectory = self.sample([[self.tokenizer.eos_token_id]])
|
||||
self.assertEqual(sum(trajectory.assistant_mask), 1)
|
||||
self.assertFalse(trajectory.truncated)
|
||||
|
||||
def test_real_model_can_generate_an_on_policy_trajectory(self):
|
||||
torch.manual_seed(3)
|
||||
student = MiniMindForCausalLM(MiniMindConfig(hidden_size=32, num_hidden_layers=1))
|
||||
trajectory = collect_trajectory(student, self.tokenizer, MESSAGES, max_new_tokens=2)
|
||||
self.assertGreaterEqual(sum(trajectory.assistant_mask), 1)
|
||||
self.assertLessEqual(sum(trajectory.assistant_mask), 2)
|
||||
|
||||
|
||||
class LossTests(unittest.TestCase):
|
||||
def test_shifted_masks_include_both_turns_and_exclude_observations_and_padding(self):
|
||||
trajectories = [AgentTrajectory([1, 10, 2, 30, 31, 11, 2], [0, 1, 1, 0, 0, 1, 1], 2, 1, False),
|
||||
AgentTrajectory([1, 12, 2], [0, 1, 1], 1, 0, False)]
|
||||
ids, attention, targets = pack_trajectories(trajectories, 2, "cpu")
|
||||
self.assertEqual(targets.tolist(), [[True, True, False, False, True, True],
|
||||
[True, True, False, False, False, False]])
|
||||
self.assertEqual(attention[1].tolist(), [1, 1, 1, 0, 0, 0, 0])
|
||||
self.assertEqual(ids[1, -1].item(), 2)
|
||||
|
||||
def test_kl_directions_match_independent_probability_calculation(self):
|
||||
student = torch.tensor([[[0.2, -0.4, 1.0]]], requires_grad=True)
|
||||
teacher = torch.tensor([[[0.8, -0.2, 0.1]]], requires_grad=True)
|
||||
for direction in ("forward_kl", "reverse_kl"):
|
||||
with self.subTest(direction=direction):
|
||||
temperature = 1.7
|
||||
ps, pt = (student.detach() / temperature).softmax(-1), (teacher.detach() / temperature).softmax(-1)
|
||||
p, q = (pt, ps) if direction == "forward_kl" else (ps, pt)
|
||||
expected = (p * (p.log() - q.log())).sum() * temperature ** 2
|
||||
loss = distillation_loss(student, teacher, torch.ones(1, 1, dtype=torch.bool), temperature, direction)
|
||||
torch.testing.assert_close(loss, expected)
|
||||
loss.backward()
|
||||
self.assertIsNone(teacher.grad)
|
||||
self.assertGreater(student.grad.abs().sum().item(), 0)
|
||||
student.grad = None
|
||||
|
||||
def test_masked_tokens_have_no_direct_gradient_even_with_invalid_logits(self):
|
||||
student = torch.randn(1, 5, 4, requires_grad=True)
|
||||
teacher = torch.randn(1, 5, 4, requires_grad=True)
|
||||
mask = torch.tensor([[True, False, False, True, False]])
|
||||
with torch.no_grad():
|
||||
student[:, 1] = float("nan")
|
||||
teacher[:, 1] = float("nan")
|
||||
loss = distillation_loss(student, teacher, mask)
|
||||
self.assertTrue(torch.isfinite(loss))
|
||||
loss.backward()
|
||||
self.assertEqual(student.grad[~mask].abs().sum().item(), 0)
|
||||
self.assertIsNone(teacher.grad)
|
||||
|
||||
def test_empty_mask_is_finite_differentiable_zero(self):
|
||||
student = torch.randn(2, 3, 5, requires_grad=True)
|
||||
loss = distillation_loss(student, student.detach(), torch.zeros(2, 3, dtype=torch.bool))
|
||||
loss.backward()
|
||||
self.assertEqual(loss.item(), 0)
|
||||
self.assertEqual(student.grad.abs().sum().item(), 0)
|
||||
|
||||
|
||||
class TrainingTests(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
torch.set_num_threads(1)
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.directory = Path(self.tmp.name)
|
||||
self.student_path, self.teacher_path = self.directory / "student.pth", self.directory / "teacher.pth"
|
||||
config = MiniMindConfig(hidden_size=32, num_hidden_layers=1)
|
||||
torch.manual_seed(1)
|
||||
torch.save(MiniMindForCausalLM(config).state_dict(), self.student_path)
|
||||
torch.manual_seed(2)
|
||||
torch.save(MiniMindForCausalLM(config).state_dict(), self.teacher_path)
|
||||
self.data = self.directory / "prompts.jsonl"
|
||||
self.data.write_text("\n".join(json.dumps({"messages": MESSAGES}) for _ in range(3)) + "\n")
|
||||
|
||||
def args(self, name="run", extra=()):
|
||||
return trainer.parse_args([
|
||||
"--student_checkpoint", str(self.student_path), "--teacher_checkpoint", str(self.teacher_path),
|
||||
"--data_path", str(self.data), "--output_dir", str(self.directory / name),
|
||||
"--student_hidden_size", "32", "--student_num_layers", "1",
|
||||
"--teacher_hidden_size", "32", "--teacher_num_layers", "1",
|
||||
"--device", "cpu", "--max_gen_len", "2", "--learning_rate", "0.001",
|
||||
*extra,
|
||||
])
|
||||
|
||||
def run_train(self, args):
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
return trainer.train(args)
|
||||
|
||||
def checkpoint(self, name):
|
||||
return torch.load(self.directory / name / "agent_opd_32_resume.pth", weights_only=True)
|
||||
|
||||
def test_actual_training_updates_student_only_and_saves_residual_window(self):
|
||||
loaded = []
|
||||
original = trainer.load_model
|
||||
|
||||
def track(*args):
|
||||
model = original(*args)
|
||||
loaded.append(model)
|
||||
return model
|
||||
|
||||
args = self.args(extra=["--accumulation_steps", "2"])
|
||||
with patch.object(trainer, "load_model", side_effect=track):
|
||||
metrics = self.run_train(args)
|
||||
self.assertEqual(len(metrics), 2) # 3 microbatches, including the final partial window.
|
||||
student, teacher = loaded
|
||||
initial = torch.load(self.student_path, weights_only=True)
|
||||
self.assertTrue(any(not torch.equal(initial[k], v) for k, v in student.state_dict().items()))
|
||||
teacher_weights = torch.load(self.teacher_path, weights_only=True)
|
||||
for key, value in teacher.state_dict().items():
|
||||
torch.testing.assert_close(value, teacher_weights[key], rtol=0, atol=0)
|
||||
self.assertTrue(all(not p.requires_grad and p.grad is None for p in teacher.parameters()))
|
||||
saved = self.checkpoint("run")
|
||||
self.assertEqual((saved["epoch"], saved["batch"], saved["step"]), (1, 0, 2))
|
||||
for key, value in student.state_dict().items():
|
||||
torch.testing.assert_close(saved["model"][key], value, rtol=0, atol=0)
|
||||
|
||||
def test_resumed_online_sampling_matches_uninterrupted_training(self):
|
||||
self.run_train(self.args("full", ["--max_steps", "2"]))
|
||||
self.run_train(self.args("resume", ["--max_steps", "1"]))
|
||||
path = self.directory / "resume/agent_opd_32_resume.pth"
|
||||
self.run_train(self.args("resume", ["--max_steps", "2", "--resume", str(path)]))
|
||||
full, resumed = self.checkpoint("full"), self.checkpoint("resume")
|
||||
self.assertEqual((full["epoch"], full["batch"], full["step"]),
|
||||
(resumed["epoch"], resumed["batch"], resumed["step"]))
|
||||
for key in full["model"]:
|
||||
torch.testing.assert_close(full["model"][key], resumed["model"][key], rtol=0, atol=0)
|
||||
|
||||
def test_resume_rejects_a_different_teacher_or_sampling_setup(self):
|
||||
self.run_train(self.args(extra=["--max_steps", "1"]))
|
||||
resume = str(self.directory / "run/agent_opd_32_resume.pth")
|
||||
with self.assertRaisesRegex(ValueError, "same data, teacher, sampling"):
|
||||
self.run_train(self.args(extra=["--resume", resume, "--temperature", "0.8"]))
|
||||
|
||||
def test_agent_dataset_needs_no_ground_truth_answer(self):
|
||||
self.data.write_text(json.dumps({"conversations": [
|
||||
{"role": "system", "content": "Use tools", "tools": json.dumps(MATH_TOOLS)},
|
||||
*MESSAGES, {"role": "assistant", "content": "ignored fixed answer"},
|
||||
]}) + "\n")
|
||||
messages, tools = trainer.load_prompts(self.data)[0]
|
||||
self.assertEqual(messages[-1]["role"], "user")
|
||||
self.assertEqual(tools, MATH_TOOLS)
|
||||
|
||||
def test_multi_turn_training_reduces_teacher_divergence(self):
|
||||
# Control generation only; both models, KL, backprop and optimizer are real.
|
||||
tokenizer = AutoTokenizer.from_pretrained(ROOT / "model")
|
||||
self.data.write_text(json.dumps({"messages": MESSAGES, "tools": MATH_TOOLS}) + "\n")
|
||||
|
||||
def generate(model, input_ids, **kwargs):
|
||||
has_result = '"result": "4"' in tokenizer.decode(input_ids[0])
|
||||
text = ("结果是 4。" if has_result else CALL) + tokenizer.eos_token
|
||||
new_ids = tokenizer(text, add_special_tokens=False)["input_ids"]
|
||||
return torch.cat((input_ids, input_ids.new_tensor([new_ids])), dim=1)
|
||||
|
||||
args = self.args(extra=["--epochs", "12", "--max_gen_len", "128"])
|
||||
with patch.object(MiniMindForCausalLM, "generate", new=generate):
|
||||
metrics = self.run_train(args)
|
||||
self.assertTrue(all(record["mean_turns"] == 2 and record["tool_calls"] == 1 for record in metrics))
|
||||
self.assertLess(metrics[-1]["loss"], metrics[0]["loss"] * 0.5)
|
||||
|
||||
def test_invalid_training_parameters_fail_before_loading_models(self):
|
||||
for flag, value in [("--grad_clip", "0"), ("--learning_rate", "nan"), ("--max_steps", "-1")]:
|
||||
with self.subTest(flag=flag), self.assertRaises(ValueError):
|
||||
self.run_train(self.args(extra=[flag, value]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -0,0 +1,158 @@
|
|||
"""Student-owned tool trajectories and token-level on-policy distillation."""
|
||||
|
||||
import json
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from trainer.train_agent import execute_tool, parse_tool_calls
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentTrajectory:
|
||||
input_ids: list
|
||||
assistant_mask: list
|
||||
turns: int
|
||||
tool_calls: int
|
||||
truncated: bool
|
||||
|
||||
|
||||
def tool_observation(tokenizer, calls, tools, open_thinking=False):
|
||||
"""Render only the environment response and the next assistant prefix."""
|
||||
marker = "<|agent_opd_observation_boundary|>"
|
||||
messages = [{"role": "assistant", "content": marker}]
|
||||
allowed = {tool["function"]["name"] for tool in (tools or [])}
|
||||
for call in calls:
|
||||
name, arguments = call.get("name", ""), call.get("arguments", {})
|
||||
if isinstance(arguments, str):
|
||||
try:
|
||||
arguments = json.loads(arguments)
|
||||
except json.JSONDecodeError:
|
||||
arguments = None
|
||||
result = None
|
||||
if isinstance(name, str) and name in allowed and isinstance(arguments, dict):
|
||||
result = execute_tool(name, arguments)
|
||||
if result is None:
|
||||
result = {"error": "unknown tool or invalid arguments"}
|
||||
messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)[:2048]})
|
||||
rendered = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True,
|
||||
tools=tools, open_thinking=open_thinking,
|
||||
)
|
||||
_, found, observation = rendered.partition(marker)
|
||||
if not found:
|
||||
raise ValueError("The chat template must preserve assistant content")
|
||||
return tokenizer(observation, add_special_tokens=False)["input_ids"]
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def collect_trajectory(student, tokenizer, messages, tools=None, *, max_turns=3,
|
||||
max_new_tokens=256, max_total_len=2048, temperature=1.0,
|
||||
open_thinking=False, autocast_ctx=None):
|
||||
"""Sample each assistant turn, execute local tools, and preserve sampled IDs.
|
||||
|
||||
Only student-generated tokens (including each turn's first EOS) are targets.
|
||||
Environment observations remain visible context with no direct loss.
|
||||
"""
|
||||
if min(max_turns, max_new_tokens, max_total_len) < 1 or temperature <= 0:
|
||||
raise ValueError("Generation limits and temperature must be positive")
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True,
|
||||
tools=tools, open_thinking=open_thinking,
|
||||
)
|
||||
ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
|
||||
if not ids or len(ids) >= max_total_len:
|
||||
raise ValueError("Prompt leaves no generation space; increase --max_total_len")
|
||||
mask = [False] * len(ids)
|
||||
model = getattr(student, "module", student)
|
||||
device = next(model.parameters()).device
|
||||
was_training = model.training
|
||||
model.eval()
|
||||
turns, call_count, truncated = 0, 0, False
|
||||
try:
|
||||
for turn in range(max_turns):
|
||||
inputs = torch.tensor([ids], device=device)
|
||||
with autocast_ctx if autocast_ctx is not None else nullcontext():
|
||||
output = model.generate(
|
||||
input_ids=inputs, attention_mask=torch.ones_like(inputs),
|
||||
max_new_tokens=min(max_new_tokens, max_total_len - len(ids)),
|
||||
do_sample=True, temperature=temperature, top_p=1.0, top_k=0,
|
||||
repetition_penalty=1.0, num_beams=1, use_cache=True,
|
||||
eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
new_ids = output[0, len(ids):].tolist()
|
||||
if tokenizer.eos_token_id in new_ids:
|
||||
new_ids = new_ids[:new_ids.index(tokenizer.eos_token_id) + 1]
|
||||
if not new_ids:
|
||||
truncated = True
|
||||
break
|
||||
ids.extend(new_ids)
|
||||
mask.extend([True] * len(new_ids))
|
||||
turns += 1
|
||||
if new_ids[-1] != tokenizer.eos_token_id:
|
||||
truncated = True
|
||||
break
|
||||
text = tokenizer.decode(new_ids, skip_special_tokens=True)
|
||||
calls = [call for call in parse_tool_calls(text) if isinstance(call, dict)]
|
||||
call_count += len(calls)
|
||||
if not calls:
|
||||
break
|
||||
if turn + 1 == max_turns:
|
||||
truncated = True
|
||||
break
|
||||
observation = tool_observation(tokenizer, calls, tools, open_thinking)
|
||||
if observation[:1] == [tokenizer.eos_token_id]:
|
||||
observation = observation[1:] # EOS is already in the sampled IDs.
|
||||
if len(ids) + len(observation) >= max_total_len:
|
||||
truncated = True
|
||||
break
|
||||
ids.extend(observation)
|
||||
mask.extend([False] * len(observation))
|
||||
finally:
|
||||
model.train(was_training)
|
||||
return AgentTrajectory(ids, mask, turns, call_count, truncated)
|
||||
|
||||
|
||||
def pack_trajectories(trajectories, pad_token_id, device):
|
||||
if not trajectories:
|
||||
raise ValueError("At least one trajectory is required")
|
||||
if any(len(t.input_ids) != len(t.assistant_mask) for t in trajectories):
|
||||
raise ValueError("Trajectory tokens and masks must have the same length")
|
||||
length = max(len(t.input_ids) for t in trajectories)
|
||||
ids, attention, targets = [], [], []
|
||||
for trajectory in trajectories:
|
||||
size = len(trajectory.input_ids)
|
||||
ids.append(trajectory.input_ids + [pad_token_id] * (length - size))
|
||||
attention.append([1] * size + [0] * (length - size))
|
||||
targets.append(trajectory.assistant_mask + [False] * (length - size))
|
||||
return (
|
||||
torch.tensor(ids, dtype=torch.long, device=device),
|
||||
torch.tensor(attention, dtype=torch.long, device=device),
|
||||
torch.tensor(targets, dtype=torch.bool, device=device)[:, 1:],
|
||||
)
|
||||
|
||||
|
||||
def distillation_loss(student_logits, teacher_logits, target_mask,
|
||||
temperature=1.0, direction="forward_kl"):
|
||||
"""Full-vocabulary KL averaged over assistant tokens on sampled prefixes.
|
||||
|
||||
Sampling is detached (the GKD semi-gradient); this is not a REINFORCE
|
||||
estimator of trajectory-level reverse KL. Teacher gradients are blocked.
|
||||
"""
|
||||
if student_logits.shape != teacher_logits.shape or student_logits.ndim != 3:
|
||||
raise ValueError("Student and teacher logits must share [batch, time, vocab]")
|
||||
if target_mask.shape != student_logits.shape[:2]:
|
||||
raise ValueError("Target mask must align with the shifted logits")
|
||||
if temperature <= 0 or direction not in {"forward_kl", "reverse_kl"}:
|
||||
raise ValueError("Invalid distillation temperature or KL direction")
|
||||
if not target_mask.any():
|
||||
return student_logits[:, :0, :].sum()
|
||||
student_logps = F.log_softmax(student_logits[target_mask].float() / temperature, dim=-1)
|
||||
teacher_logps = F.log_softmax(teacher_logits.detach()[target_mask].float() / temperature, dim=-1)
|
||||
if direction == "forward_kl":
|
||||
loss = F.kl_div(student_logps, teacher_logps, log_target=True, reduction="batchmean")
|
||||
else:
|
||||
loss = F.kl_div(teacher_logps, student_logps, log_target=True, reduction="batchmean")
|
||||
return loss * temperature ** 2
|
||||
|
|
@ -0,0 +1,235 @@
|
|||
"""Single-device on-policy distillation through MiniMind's local tool environment."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
import datasets # noqa: F401 # Keep pyarrow before torch on Windows (issue #771).
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from trainer.agent_opd_utils import collect_trajectory, pack_trajectories, distillation_loss
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def load_prompts(path):
|
||||
samples = []
|
||||
with open(path, encoding="utf-8") as stream:
|
||||
for line in stream:
|
||||
if not line.strip():
|
||||
continue
|
||||
sample = json.loads(line)
|
||||
messages = sample.get("conversations", sample.get("messages", []))
|
||||
if messages and messages[-1].get("role") == "assistant":
|
||||
messages = messages[:-1]
|
||||
if not messages or messages[-1].get("role") not in {"user", "tool"}:
|
||||
raise ValueError("Each sample must end in a user/tool prompt or a removable assistant answer")
|
||||
tools = sample.get("tools")
|
||||
for message in messages:
|
||||
if message.get("role") == "system" and message.get("tools"):
|
||||
tools = message["tools"]
|
||||
if isinstance(tools, str):
|
||||
tools = json.loads(tools)
|
||||
samples.append((messages, tools))
|
||||
if not samples:
|
||||
raise ValueError("The prompt dataset is empty")
|
||||
return samples
|
||||
|
||||
|
||||
def load_model(path, config, device):
|
||||
weights = torch.load(path, map_location="cpu", weights_only=True)
|
||||
weights = weights.get("model", weights)
|
||||
model = MiniMindForCausalLM(config)
|
||||
model.load_state_dict(weights, strict=True)
|
||||
return model.to(device)
|
||||
|
||||
|
||||
def training_signature(args):
|
||||
# A resumed run must retain its teacher, data order, and optimization setup.
|
||||
adjustable = {"resume", "epochs", "max_steps", "save_interval", "output_dir"}
|
||||
return {key: value for key, value in vars(args).items() if key not in adjustable}
|
||||
|
||||
|
||||
def save_checkpoint(args, model, optimizer, scaler, epoch, batch, step):
|
||||
directory = Path(args.output_dir)
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
suffix = "_moe" if args.student_use_moe else ""
|
||||
prefix = directory / f"agent_opd_{args.student_hidden_size}{suffix}"
|
||||
weights = {name: tensor.detach().cpu() for name, tensor in model.state_dict().items()}
|
||||
state = {
|
||||
"model": weights, "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict(),
|
||||
"epoch": epoch, "batch": batch, "step": step, "signature": training_signature(args),
|
||||
"torch_rng": torch.get_rng_state(), "python_rng": random.getstate(),
|
||||
"cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [],
|
||||
"mps_rng": torch.mps.get_rng_state() if next(model.parameters()).device.type == "mps" else None,
|
||||
}
|
||||
for filename, content in [(f"{prefix}_resume.pth", state),
|
||||
(f"{prefix}.pth", {key: value.half() for key, value in weights.items()})]:
|
||||
torch.save(content, filename + ".tmp")
|
||||
os.replace(filename + ".tmp", filename)
|
||||
|
||||
|
||||
def train(args):
|
||||
if int(os.environ.get("WORLD_SIZE", "1")) > 1:
|
||||
raise ValueError("Agent OPD currently supports one training process/device")
|
||||
for key in ("epochs", "batch_size", "accumulation_steps", "max_turns", "max_gen_len",
|
||||
"max_total_len", "save_interval"):
|
||||
if getattr(args, key) < 1:
|
||||
raise ValueError(f"--{key} must be positive")
|
||||
if args.max_steps < 0 or any(not math.isfinite(value) or value <= 0 for value in
|
||||
(args.learning_rate, args.temperature, args.distill_temperature, args.grad_clip)):
|
||||
raise ValueError("Steps must be nonnegative; learning rate, temperatures and grad clip must be positive")
|
||||
device = torch.device(args.device)
|
||||
if device.type == "cpu" and args.dtype != "float32":
|
||||
raise ValueError("Use --dtype float32 for the CPU reference run")
|
||||
torch.manual_seed(args.seed)
|
||||
random.seed(args.seed)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
||||
if tokenizer.pad_token_id is None or tokenizer.eos_token_id is None:
|
||||
raise ValueError("The shared tokenizer must define PAD and EOS")
|
||||
configs = {
|
||||
who: MiniMindConfig(hidden_size=getattr(args, f"{who}_hidden_size"),
|
||||
num_hidden_layers=getattr(args, f"{who}_num_layers"),
|
||||
use_moe=bool(getattr(args, f"{who}_use_moe")))
|
||||
for who in ("student", "teacher")
|
||||
}
|
||||
if len(tokenizer) != configs["student"].vocab_size:
|
||||
raise ValueError("Both MiniMind checkpoints must use the supplied shared tokenizer/vocabulary")
|
||||
if args.max_total_len > min(config.max_position_embeddings for config in configs.values()):
|
||||
raise ValueError("--max_total_len exceeds the model context length")
|
||||
student = load_model(args.student_checkpoint, configs["student"], device)
|
||||
teacher = load_model(args.teacher_checkpoint, configs["teacher"], device).eval().requires_grad_(False)
|
||||
optimizer = torch.optim.AdamW(student.parameters(), lr=args.learning_rate)
|
||||
scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda" and args.dtype == "float16")
|
||||
dtype = getattr(torch, args.dtype)
|
||||
autocast_ctx = nullcontext() if dtype == torch.float32 else torch.autocast(device.type, dtype=dtype)
|
||||
samples = load_prompts(args.data_path)
|
||||
batches_per_epoch = math.ceil(len(samples) / args.batch_size)
|
||||
start_epoch, start_batch, step = 0, 0, 0
|
||||
if args.resume:
|
||||
state = torch.load(args.resume, map_location="cpu", weights_only=True)
|
||||
if state["signature"] != training_signature(args):
|
||||
raise ValueError("Resume requires the same data, teacher, sampling and optimizer arguments")
|
||||
student.load_state_dict(state["model"], strict=True)
|
||||
optimizer.load_state_dict(state["optimizer"])
|
||||
scaler.load_state_dict(state["scaler"])
|
||||
start_epoch, start_batch, step = state["epoch"], state["batch"], state["step"]
|
||||
torch.set_rng_state(state["torch_rng"])
|
||||
random.setstate(state["python_rng"])
|
||||
if state["cuda_rng"]:
|
||||
torch.cuda.set_rng_state_all(state["cuda_rng"])
|
||||
if state["mps_rng"] is not None:
|
||||
torch.mps.set_rng_state(state["mps_rng"])
|
||||
if args.max_steps and step >= args.max_steps:
|
||||
return []
|
||||
student.train()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
metrics = []
|
||||
cursor_epoch, cursor_batch = start_epoch, start_batch
|
||||
for epoch in range(start_epoch, args.epochs):
|
||||
order = list(range(len(samples)))
|
||||
random.Random(args.seed + epoch).shuffle(order)
|
||||
first_batch = start_batch if epoch == start_epoch else 0
|
||||
for window in range(first_batch, batches_per_epoch, args.accumulation_steps):
|
||||
count = min(args.accumulation_steps, batches_per_epoch - window)
|
||||
total_loss, turns, calls, tokens, truncated, trajectories_seen = 0.0, 0, 0, 0, 0, 0
|
||||
for batch in range(window, window + count):
|
||||
indices = order[batch * args.batch_size:(batch + 1) * args.batch_size]
|
||||
trajectories = [collect_trajectory(
|
||||
student, tokenizer, *samples[index], max_turns=args.max_turns,
|
||||
max_new_tokens=args.max_gen_len, max_total_len=args.max_total_len,
|
||||
temperature=args.temperature, open_thinking=args.open_thinking,
|
||||
autocast_ctx=autocast_ctx,
|
||||
) for index in indices]
|
||||
input_ids, attention_mask, target_mask = pack_trajectories(
|
||||
trajectories, tokenizer.pad_token_id, device,
|
||||
)
|
||||
if not target_mask.any():
|
||||
raise ValueError("No assistant tokens were generated for this batch")
|
||||
with torch.no_grad(), autocast_ctx:
|
||||
teacher_logits = teacher(input_ids, attention_mask=attention_mask).logits[:, :-1]
|
||||
with autocast_ctx:
|
||||
output = student(input_ids, attention_mask=attention_mask)
|
||||
kd_loss = distillation_loss(output.logits[:, :-1], teacher_logits, target_mask,
|
||||
args.distill_temperature, args.kl_direction)
|
||||
loss = kd_loss + output.aux_loss if configs["student"].use_moe else kd_loss
|
||||
if not torch.isfinite(loss):
|
||||
raise FloatingPointError("Non-finite OPD loss")
|
||||
scaler.scale(loss / count).backward()
|
||||
total_loss += loss.item() / count
|
||||
tokens += int(target_mask.sum())
|
||||
turns += sum(t.turns for t in trajectories)
|
||||
calls += sum(t.tool_calls for t in trajectories)
|
||||
truncated += sum(t.truncated for t in trajectories)
|
||||
trajectories_seen += len(trajectories)
|
||||
del teacher_logits, output, loss, kd_loss
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(student.parameters(), args.grad_clip, error_if_nonfinite=True)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
step += 1
|
||||
cursor_batch = window + count
|
||||
cursor_epoch = epoch
|
||||
if cursor_batch == batches_per_epoch:
|
||||
cursor_epoch, cursor_batch = epoch + 1, 0
|
||||
record = {"step": step, "loss": total_loss, "assistant_tokens": tokens,
|
||||
"mean_turns": turns / trajectories_seen, "tool_calls": calls,
|
||||
"truncated_fraction": truncated / trajectories_seen}
|
||||
metrics.append(record)
|
||||
print(json.dumps(record), flush=True)
|
||||
if step % args.save_interval == 0:
|
||||
save_checkpoint(args, student, optimizer, scaler, cursor_epoch, cursor_batch, step)
|
||||
if args.max_steps and step >= args.max_steps:
|
||||
break
|
||||
if args.max_steps and step >= args.max_steps:
|
||||
break
|
||||
save_checkpoint(args, student, optimizer, scaler, cursor_epoch, cursor_batch, step)
|
||||
return metrics
|
||||
|
||||
|
||||
def parse_args(argv=None):
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--student_checkpoint", required=True, help="Native MiniMind .pth weights")
|
||||
parser.add_argument("--teacher_checkpoint", required=True, help="Frozen teacher using the same tokenizer")
|
||||
parser.add_argument("--data_path", default=str(ROOT / "dataset/agent_rl.jsonl"))
|
||||
parser.add_argument("--tokenizer_path", default=str(ROOT / "model"))
|
||||
parser.add_argument("--output_dir", default=str(ROOT / "out/agent_opd"))
|
||||
parser.add_argument("--resume", default="", help="Full-precision *_resume.pth checkpoint")
|
||||
for who in ("student", "teacher"):
|
||||
parser.add_argument(f"--{who}_hidden_size", type=int, default=768)
|
||||
parser.add_argument(f"--{who}_num_layers", type=int, default=8)
|
||||
parser.add_argument(f"--{who}_use_moe", type=int, choices=[0, 1], default=0)
|
||||
parser.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", choices=["float32", "bfloat16", "float16"], default="float32")
|
||||
parser.add_argument("--epochs", type=int, default=1)
|
||||
parser.add_argument("--batch_size", type=int, default=1)
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-6)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--max_steps", type=int, default=0, help="Optimizer steps; 0 runs all epochs")
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument("--max_turns", type=int, default=3)
|
||||
parser.add_argument("--max_gen_len", type=int, default=256)
|
||||
parser.add_argument("--max_total_len", type=int, default=2048)
|
||||
parser.add_argument("--temperature", type=float, default=1.0)
|
||||
parser.add_argument("--distill_temperature", type=float, default=1.0)
|
||||
parser.add_argument("--kl_direction", choices=["forward_kl", "reverse_kl"], default="forward_kl")
|
||||
parser.add_argument("--open_thinking", action="store_true")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
train(parse_args())
|
||||
Loading…
Reference in New Issue