diff --git a/README.md b/README.md index dadc82c..85f65e3 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/README_en.md b/README_en.md index 95fb03a..8b91899 100644 --- a/README_en.md +++ b/README_en.md @@ -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 diff --git a/docs/agent_opd.md b/docs/agent_opd.md new file mode 100644 index 0000000..731339c --- /dev/null +++ b/docs/agent_opd.md @@ -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`。 diff --git a/tests/test_agent_opd.py b/tests/test_agent_opd.py new file mode 100644 index 0000000..11aefc1 --- /dev/null +++ b/tests/test_agent_opd.py @@ -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 = '\n{"name":"calculate_math","arguments":{"expression":"2+2"}}\n' + + +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"), + '{"name":[],"arguments":null}', + '{"name":"calculate_math","arguments":"invalid"}']: + 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(['[]' + 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() diff --git a/trainer/agent_opd_utils.py b/trainer/agent_opd_utils.py new file mode 100644 index 0000000..ef3b755 --- /dev/null +++ b/trainer/agent_opd_utils.py @@ -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 diff --git a/trainer/train_agent_opd.py b/trainer/train_agent_opd.py new file mode 100644 index 0000000..6063931 --- /dev/null +++ b/trainer/train_agent_opd.py @@ -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())