Teacher-Guidance Distillation Kit

Train a small retrieval agent ("student") to internalise the step-by-step feedback of a strong "teacher" model, and evaluate it against the teacher itself — on four multi-hop QA datasets, with a leakage-audited split, resumable scripts, and an LLM judge.

The kit ships the data, the harness, and one command per stage:

Every long-running script writes to a fixed directory, appends results as they land, skips finished work on re-run, and keeps a status.json you can watch.

What is in the box

data/questions/<ds>/     2,000 questions each for HotpotQA, 2WikiMultihopQA and MuSiQue,
                         1,999 for StrategyQA, with their per-question document sets
                         (gzipped, with provenance manifests)
data/episodes/           7,999 teacher-guided episodes (student granite-4.1-3b,
                         teacher DeepSeek-V4-Flash), gzipped, plus an index
data/benchmarks/         MMLU, GSM8K and HellaSwag eval items for the forgetting check
data/splits/             test question files, pool assignment, stats, leakage report.
                         The SFT train/dev files are NOT shipped: build them once with
                         `make data` (~2 min, byte-identical on every machine)
agentsim/                the simulation harness (prompts, tools, teacher critic, metrics)
tgd/                     library code shared by the scripts
scripts/                 the stage commands above
slurm/                   sbatch templates + a one-command pipeline for HPC clusters
tests/                   unit tests, an offline end-to-end smoke test, a GPU smoke test
docs/                    OVERVIEW, EXPERIMENTS, MODELS, DATASET, DATA, TRAINING,
                         EVALUATION, FORGETTING, STABILITY, TROUBLESHOOTING, PROVIDERS,
                         REPRODUCE, RESULTS

The four evaluation arms

All arms answer the same questions with the same corpus, tools, step budget and metrics; only the agent differs.

Two split groups, one held-out set

• uniform — train on the trainable 90 % of every dataset; test on the held-out 10 % of every dataset (747 questions). Does internalised guidance work?
• lodo (leave-one-dataset-out) — four folds; each trains on three datasets and is tested on the entire fourth dataset (never seen in any form) plus the held-out 10 % of the three training datasets. Does it transfer to an unseen dataset?

The 10 % held-out pool is a salted hash of the question id, so it is identical on every machine and is excluded from every split. scripts/check_leakage.py re-derives the proof from the files (id disjointness, question-text containment, placeholder hygiene, hash consistency). See docs/DATA.md.

Quick start

bash setup_env.sh                      # .venv (CPU), .venv_train (GPU), .venv_vllm (GPU)
make data                              # build the SFT train/dev files (~2 min; required before training)
cp .env.example .env                   # add your teacher / judge endpoint + key
.venv/bin/python -m pytest tests -q    # unit tests, seconds
bash tests/smoke_offline.sh            # whole pipeline with mock models, no GPU/API, ~1 min

# reproduce the four-arm comparison on the uniform split (Slurm):
TEACHER=oai-teacher/<model> JUDGE=oai-judge/<model> bash slurm/run_pipeline.sh -p <gpu-partition>

Without Slurm, run the same stages by hand — docs/REPRODUCE.md lists every command. When something breaks, or finishes but looks wrong, docs/TROUBLESHOOTING.md is organised by symptom and leads with the failures that produce no error at all. To run your own experiments — a different student, teacher, dataset or training objective — start from docs/EXPERIMENTS.md. For swapping the student specifically, docs/MODELS.md lists what the kit adapts automatically and what to check first.

Bring your own models

• Student: any Hugging Face causal LM with a chat template (--model); served by vLLM or in-process by transformers.
• Teacher / judge / any API model: a provider-prefixed id. oai-<name>/<model> talks to any OpenAI-compatible endpoint configured as OAI_<NAME>_BASE_URL / OAI_<NAME>_API_KEY; EdenAI, NVIDIA NIM, OpenRouter-style, Ollama and the native OpenAI/Anthropic/Google SDKs are also wired. Comma-separated ids form a fallback chain. See docs/PROVIDERS.md.
• Mock: mock/<anything> is an offline stand-in used by the smoke tests.

Reference results

With granite-4.1-3b as student, DeepSeek-V4-Flash as teacher and Kimi-K2.6 as judge, on the 747 held-out questions (judge-correct): base 29.8 %, guided 60.5 %, trained 65.5 %, teacher 72.3 %. The trained student pays 1.4 points of general ability for that gain (docs/FORGETTING.md) — and, trained with plain SFT, it can only be decoded greedily: sampling it at temperature 0.3 produces nothing usable. That failure, how to detect it and two ways to prevent it are in docs/STABILITY.md; it is the single most surprising result in this kit, and the reason the loss-path guard exists. The trained student needs no teacher at inference and 0.43× the tokens of the guided student. Full tables, per-dataset numbers, the leave-one-dataset-out transfer results and confidence intervals are in docs/RESULTS.md.

Requirements

Python 3.11–3.13 (PyTorch publishes no wheels for 3.14 yet, and setup_env.sh will fail opaquely if python3 on your PATH is newer). One GPU with ≥ 40 GB for the 3B student at full precision (training with LoRA and 8k-token sequences peaks around 26 GB; smaller students need less). The CPU environment is enough for data building, the teacher-alone arm, judging and results.

On a consumer GPU. --load-4bit loads the base weights as NF4 and trains the LoRA adapter on top, which brings a 3B student from 7.3 GB of weights to ~2.3 GB and fits an 8 GB card. Pass it to scripts/train_sft.py and to scripts/eval.py --student hf: an adapter trained on quantized weights must be evaluated on the same quantized weights, or the arm is not the model that was trained. Expect absolute accuracy below the same recipe at bf16; comparisons between arms that share the setting are unaffected.

Platform. Linux, macOS and Windows. CI runs the unit tests on all three (Python 3.11–3.13 on Linux, 3.12 on macOS and Windows), the whole offline pipeline on Linux and Windows, and make data on Windows under the legacy cp1252 code page — because the portability bugs worth catching are the ones that produce no error on the machine they were written on. Every text file is opened as UTF-8 explicitly (the corpora are Wikipedia-derived: the first 4,000 lines of the HotpotQA corpus carry 6,888 characters cp1252 cannot represent), the virtualenv interpreter is located rather than assumed, and console output degrades instead of raising where it cannot be encoded.

The shell entry points (setup_env.sh, tests/smoke_offline.sh, make) need a POSIX shell; on Windows that is Git Bash, which ships with Git. Two things stay Linux-only by nature: slurm/ and scripts/serve_vllm.sh, since vLLM publishes no Windows wheels — evaluate with --student hf there instead.

Disk: about 30 GB. The three virtual environments take ~15 GB (the CUDA and vLLM wheels dominate), the model cache ~7 GB for a 3B student, the repository ~230 MB, and the built SFT files ~1.5 GB. Set HF_HOME if your home directory is small or quota'd.

License and attribution

The code in this repository is licensed under the Apache License 2.0 (LICENSE), which includes the agentsim/ harness vendored here.

The data keeps its own terms. Question sets and corpora derive from HotpotQA (CC BY-SA 4.0), 2WikiMultihopQA (Apache-2.0), MuSiQue (CC BY 4.0) and StrategyQA (MIT); each dataset's manifest next to its question file records source, homepage and license. The episodes were generated with the models named in every record's student_model and teacher_models_used fields. NOTICE summarises all of it.