Skip to main content
On-policy distillation (OPD) trains a student model on its own rollouts while using a teacher model’s token-level probabilities as the distillation signal. In Miles, the teacher signal is converted into a per-token reverse-KL penalty and applied after the selected RL advantage estimator has produced token advantages. This lets the same OPD recipe compose with GRPO, PPO, REINFORCE++, GSPO, and other estimators.

​
Key Arguments

​
How It Works

OPD modifies the advantage computation by subtracting a KL penalty term that encourages the student to match the teacher’s output distribution: A^t=At−λopd⋅DKL(Pstudent∥Pteacher)t\hat{A}_t = A_t - \lambda_{\text{opd}} \cdot D_{\text{KL}}(P_{\text{student}} \| P_{\text{teacher}})_t Where AtA_t is the original advantage from the base estimator (e.g., GRPO), λopd\lambda_{\text{opd}} is --opd-kl-coef, and DKLD_{\text{KL}} is the token-level reverse KL divergence. The implementation follows the additive OPD training recipe described in the Thinking Machines OPD blog, with an additional SGLang top-k reward mode from Rethinking On-Policy Distillation.

​
Rethinking OPD Top-K Reward

SGLang OPD supports the top-k token reward recipe from Rethinking On-Policy Distillation. Set --opd-log-prob-top-k above zero to request student rollout top-logprobs, score the same sequence with the teacher, and aggregate a weighted reverse-KL estimate over a selected token set at each response position. The token set is controlled by --opd-top-k-strategy: --opd-reward-weight-mode controls whether each selected token is weighted by student probability, teacher probability, or uniformly. For compatibility, --opd-log-prob-top-k=0 keeps the original sampled-token OPD path.

​
Two Teacher Modes

​
SGLang Mode (--opd-type sglang)

The teacher runs on an external SGLang server. Teacher log-probs are obtained during the rollout phase. When to use: The teacher has a different architecture from the student, or the teacher is too large to load alongside the training model. How it works:
  1. An external SGLang server runs the teacher model.
  2. During rollout, the custom reward function (miles.rollout.on_policy_distillation.reward_func) sends each sample to the teacher server to obtain token-level log-probs.
  3. With --opd-log-prob-top-k=0, the custom post-processing function trims sampled-token teacher log-probs to the response span and stores them in sample.teacher_log_probs.
  4. With --opd-log-prob-top-k>0, it computes the Rethinking OPD weighted top-k reverse-KL estimate and stores it in sample.opd_reverse_kl.
  5. During training, the stored OPD penalty is subtracted from the selected estimator’s advantages.
Configuration:

​
Multi-Teacher Routing (SGLang mode only)

--opd-teacher-urls routes each sample to a task-specific teacher, e.g. a math specialist for math prompts and a code specialist for code prompts. Each sample is still scored by exactly one teacher, so scoring cost is identical to single-teacher OPD and the loss is unchanged — teacher_log_probs / opd_reverse_kl are per-sample and do not care which teacher produced them. How it works:
  1. Tag each prompt with a teacher name in its dataset metadata column (read via --metadata-key, default metadata):
  2. Map names to teacher endpoints with --opd-teacher-urls NAME=URL .... The reserved name default is the fallback for samples whose name is missing or unknown; without a default entry such samples raise an error (failing loudly beats silently distilling from the wrong teacher).
  3. reward_func resolves the URL per sample; everything downstream is unchanged.
Configuration (on top of the SGLang-mode flags above; --rm-url is ignored when the routing map is set):
Notes: All teachers must share the student’s tokenizer — scoring sends input_ids and gathers per-token-id log-probs. Works with both the sampled-token path and the top-k path (student-side scoring still goes to the student router). For throughput, point multiple names (or one name backed by an sglang router) at replicas; --opd-teacher-urls is for different teachers, not load balancing.

​
Megatron Mode (--opd-type megatron)

The teacher model is loaded directly into Megatron via --opd-teacher-load. Teacher log-probs are computed during the training forward pass. When to use: The teacher has the same architecture as the student/reference model and fits in GPU memory. How it works:
  1. The teacher model is loaded as an additional Megatron model during initialization.
  2. During the training forward pass, the teacher model computes log-probs for each sample.
  3. The KL penalty is computed inline and applied to advantages.
Configuration:
Note: The teacher checkpoint must be in Megatron format (torch_dist or torch). You can convert from HuggingFace format using tools/convert_hf_to_torch_dist.py.

​
Running the Examples

Complete example scripts are provided in examples/on_policy_distillation/:

​
SGLang Teacher

​
Megatron Teacher

​
Preliminary Results

Using Qwen3-8B-Base model SFT-ed on part of the OpenThoughts3-1.2M dataset, on-policy distillation with a Qwen3-32B teacher on the remaining data yields:

​
References