111 lines
5.9 KiB
Markdown
111 lines
5.9 KiB
Markdown
# 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`。
|