ADAPT

ADAPT is a framework for amortizing distillation across both size and variant axes of a model family. It consists of two variants:

• Two-phase distillation — distill on a pretraining corpus, then switch to instruction data partway through to allow for effective size interpolation of post-trained models on generation and reasoning tasks.
• Weight-delta transfer — carry a distillation delta learned against a base teacher onto a student pruned from that teacher's post-trained counterpart. Approximates the effect of two-phase distillation and allows for zero-shot transfer of distillation effects across post-trained model variants.

This repo contains code for ADAPT from our EMNLP 2026 Findings paper Thinking at the Right Size: Amortized Distillation Across Post-Trained LLMs.

The models used in this paper are accessible via Huggingface through this collection.

Table of contents

• Installation
• Distillation training
• Weight-delta transfer
• Patching and evaluation
• Repository layout
• Citation

Installation

conda create -n adapt python==3.12
conda activate adapt
pip3 install -r requirements.txt

To reproduce the environment used for the paper experiments, use requirements_dev.txt. The requirements have only been tested on Linux.

LiveCodeBench is not on PyPI. It is only needed for the livecodebench task:

pip install "git+https://github.com/LiveCodeBench/LiveCodeBench.git"

Distillation training

train/train.py prunes a student from the teacher and distills it with a combination of cross-entropy, KL divergence on the logits, and cosine embedding loss between each student layer's hidden states and its corresponding teacher block.

TEACHER="Qwen/Qwen3-4B-Instruct-2507"

torchrun --standalone --nnodes=1 --nproc-per-node=4 -m train.train \
    --teacher_model_name_or_path $TEACHER \
    --save_directory "/path/to/save/directory" \
    --dataset "the_pile_and_nemotron_thinking_off" \
    --fsdp_config $TEACHER \
    --teacher_fsdp True

Key options:

• teacher_model_name_or_path: Hugging Face reference or local path for the teacher. Defines the layer mapping and initializes the student.

• save_directory: where checkpoints are written. A subdirectory encoding the run configuration is created inside it.

• dataset: one of EleutherAI/the_pile_deduplicated, nemotron_thinking, nemotron_thinking_off, the_pile_and_nemotron_thinking, or the_pile_and_nemotron_thinking_off. The Nemotron variants stream from the hub; see scripts/README.md for which file each reads and how to point them at a local dump.

• fsdp_config: set to the teacher model name so FSDP wraps the right modules (see train/training_args.py).

• first_layers_to_keep / last_layers_to_keep / alternate_every_n_layers: which teacher layers survive pruning. The two end counts default to the teacher's family and normally need not be set:

  Teacher family	first_layers_to_keep	last_layers_to_keep
  llama	2	1
  qwen	1	2
  olmo	1	2
  anything else	1	2

  Passing either flag explicitly overrides that side only.

Two-phase training

The the_pile_and_nemotron_* datasets train on The Pile for phase_1_steps and then switch to Nemotron for phase_2_steps. Batch size, gradient accumulation, and loss weights can differ between phases:

    --phase_1_steps 240 --phase_2_steps 293 \
    --phase_1_batch_size 2 --phase_1_gradient_accumulation_steps 256 \
    --phase_2_batch_size 4 --phase_2_gradient_accumulation_steps 256 \
    --phase_1_kl_loss_weight 0.1  --phase_1_cosine_loss_weight 10.0 \
    --phase_2_kl_loss_weight 0.0  --phase_2_cosine_loss_weight 0.0 \
    --save_phase_transition_checkpoint True

The dataset decides where to switch by counting examples (batch_size x gradient_accumulation_steps x num_devices), so the per-phase values passed to the trainer must match what the dataset sees, or the data switch and the trainer's transition will land on different steps.

Loss weights left unset for a phase fall back to the global --cross_entropy_loss_weight, --kl_loss_weight, and --cosine_loss_weight. All arguments are documented in train/training_args.py, train/model_args.py, and train/data_args.py.

Weight-delta transfer

Given a base teacher and its post-trained counterpart, distilling once against the base teacher is enough: the delta that distillation learned transfers onto a student pruned from the post-trained teacher.

With $\theta_{ub}$ and $\theta_{uf}$ the untrained students pruned from the base and finetuned teachers, and $\theta_{sb}$ the student distilled against the base teacher:

$$\theta_{sf} \approx \theta_{uf} + \alpha(\theta_{sb} - \theta_{ub})$$

First build the two untrained students. This is cheap — no training happens:

for TEACHER in "Qwen/Qwen3-4B-Base" "Qwen/Qwen3-4B-Instruct-2507"; do
    python3 -m train.save_untrained_student \
        --teacher_model_name_or_path "$TEACHER" \
        --save_directory "/path/to/untrained_student" \
        --alternate_every_n_layers 2
done

Then distill against the base teacher, and transfer the delta at evaluation time. No gradient steps are taken in this step:

python3 -m evaluate.evaluate \
    --teacher_model_name_or_path "Qwen/Qwen3-4B-Instruct-2507" \
    --student_model_name_or_path "/path/to/base-distilled-student" \
    --weight_delta_student \
    --untrained_student_base "/path/to/untrained_student/Qwen3-4B-Base_..." \
    --untrained_student_finetuned "/path/to/untrained_student/Qwen3-4B-Instruct-2507_..." \
    --weight_delta_alpha 1.0 \
    --num_layers_to_patch 4 \
    --tasks mmlu_redux

The teacher here is the finetuned teacher: the transferred student is patched with post-trained teacher blocks, producing an instruct-family intermediate model without ever distilling against the instruct teacher. The pruning arguments must match across all three students, since the delta is applied parameter by parameter. Leaving the layer-count flags unset keeps them consistent, since all three default from their teacher's family.

--weight_delta_teacher applies the delta on the teacher side instead: the student layers are shifted by the teacher's own post-training delta, and --match_layers {first,last,average} selects whether a student layer takes its block's first teacher layer, its last, or the block average.

The same transfer can initialize a training run rather than an evaluation, by passing --weight_delta_student to train/train.py along with the same three paths.

Patching and evaluation

patching/patch.py:build_intermediate_model constructs intermediate models. evaluate/evaluate.py builds one and evaluates it with vLLM:

python3 -m evaluate.evaluate \
    --teacher_model_name_or_path "Qwen/Qwen3-4B-Instruct-2507" \
    --student_model_name_or_path "/path/to/student" \
    --save_directory "/path/to/results" \
    --num_layers_to_patch 4 \
    --tasks "mmlu_redux"

Options:

• num_layers_to_patch: number of student layers to replace with their teacher blocks. Patching every patchable layer recovers the teacher; patching none leaves the student.
• patch_first_k_layers: patch from the first layers rather than the last. Llama and OLMo teachers use first-k by default; pass --override_llama_patching or --override_olmo_patching to disable that and fall back to this flag.
• tasks: comma-separated. One or more of aime, gsm8k, ifeval, math500, mmlu_redux, humaneval_plus, mbpp_plus, livecodebench.
• dtype: weight dtype, bfloat16 by default.
• max_new_tokens, temperature, top_p, top_k, min_p, presence_penalty, seed: sampling configuration.
• log_n_results: print this many generations per task, for inspection.

Repository layout

patching/    building intermediate models; weight-delta operations
train/       distillation training, dataset pipelines, untrained-student export
evaluate/    the evaluation entrypoint, per-task evaluators, grading utilities
scripts/     shell wrappers for the common workflows

scripts/README.md covers running these under a job scheduler and pointing the Nemotron pipelines at a local data dump.

evaluate/utils/instructions*.py are vendored verbatim from Google Research's IFEval and carry their original Apache-2.0 headers.

Citation

@misc{zhou2026thinkingrightsizeamortized,
      title={Thinking at the Right Size: Amortized Distillation Across Post-Trained LLMs}, 
      author={Yan Zhou and Sara Kangaslahti and Jonathan Geuter and Nihal V. Nayak and Marco Fumero and Francesco Locatello and David Alvarez-Melis},
      year={2026},
      eprint={2608.22854},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2608.22854}, 
}