On-Policy Distillation on a Single Free-Tier GPU

A from-scratch implementation of on-policy distillation (OPD), distilling Qwen2.5-1.5B-Instruct into Qwen2.5-0.5B-Instruct on GSM8K, trained end-to-end on one free-tier Colab T4.

The loss is written directly in PyTorch rather than configured through a trainer, because the point of the project was to understand the objective, its practical constraints, and its failure modes rather than to optimize for a benchmark number.

Headline result: the student's pass@1 stayed roughly flat while its any-of-4 success rate fell sharply. The model did not demonstrate a reliable capability improvement --- instead, its sampled behavior became substantially narrower. That signal is easy to miss if evaluation reports only pass@1.

What on-policy distillation is

Standard distillation is off-policy: the student trains on text generated by the teacher, then gets evaluated on text it generates. The distribution seen during training therefore differs from the distribution encountered at inference.

On-policy distillation closes that gap by having the student generate its own trajectory and having the teacher evaluate that same trajectory:

                    ┌────────────────────┐
                    │ Current Student    │
                    │ Qwen2.5-0.5B       │
                    └─────────┬──────────┘
                              │
                           sample
                              │
                              ▼
                    ┌────────────────────┐
                    │ Student-generated  │
                    │ trajectory         │
                    └─────────┬──────────┘
                              │
                        same tokens
                              │
                              ▼
                    ┌────────────────────┐
                    │ Teacher            │
                    │ Qwen2.5-1.5B       │
                    └─────────┬──────────┘
                              │
                       top-K logits
                              │
                              ▼
                    ┌────────────────────┐
                    │ Truncated          │
                    │ reverse KL         │
                    └─────────┬──────────┘
                              │
                              ▼
                    ┌────────────────────┐
                    │ Update Student     │
                    └────────────────────┘

The training loop is:

1. Sample a completion from the current student policy (no gradient).
2. Run one teacher forward pass over the student's sampled tokens (no gradient).
3. Compute token-level reverse KL between student and teacher distributions.
4. Backpropagate the loss through the student.

With γ=0, supervision is local to the current token rather than propagating a future reward signal backward. There is no REINFORCE term or explicit outcome reward. The teacher supplies a dense token-level training signal while only performing forward passes.

This makes the setup substantially simpler than outcome-reward RL, although it does not eliminate the optimization and distribution-shift problems that can arise during on-policy training.

Results

The reported checkpoint was selected at step 20, before the training collapse described below.

Evaluation used a single seed and 100 problems for greedy evaluation. The any-of-4 experiment used 50 problems with four sampled completions per problem.

Note: These are empirical results from this single experimental run, not benchmark claims.

Evaluation results

Interpreting the results

The pass@1 differences are not interpretable at this evaluation size.

With only 100 problems, the uncertainty around a 0.11 proportion is large. The +2 point GSM8K change and −5 point SVAMP change are therefore not evidence of a reliable improvement or degradation.

The any-of-4 result is the most interesting diagnostic signal.

The fraction of problems for which at least one of four sampled completions was correct fell from 18% to 4%.

This is a much larger change than the pass@1 difference and moves in the opposite direction. It is consistent with the hypothesis that the current reverse-KL objective narrowed the student's useful sampling distribution.

However, this is still a small, single-seed experiment. It should be treated as a strong diagnostic signal rather than a statistically established claim that OPD universally causes diversity collapse.

A larger evaluation, additional seeds, and explicit diversity metrics such as unique completions per problem would be required to establish that conclusion.

The key lesson from this experiment is therefore not that the distilled model is better. It is that pass@1 alone can hide substantial changes in sampled behavior.

Training instability

The training run entered a degenerate state:

• KL decreased toward zero
• completion length approached the generation cap
• accuracy fell toward zero
• outputs eventually degraded into repeated tokens or symbol-like sequences

Two learning-rate experiments showed the same qualitative behavior:

Learning rate Approximate collapse onset

1e-5 ~step 25 1e-6 ~step 40

Lowering the learning rate delayed the degradation but did not prevent it.

This suggests that, within this small-model configuration, the problem was not simply one excessively large update. The reported checkpoint was therefore selected at step 20, before the observed drift.

Grad norms near collapse were also relatively calm (around 11.4 in the observed run and well below the clipping threshold), providing no evidence of a single exploding-gradient event. The degradation appeared gradual over multiple steps.

Failure modes hit and fixed

These are documented because they consumed a substantial part of the project and several of them did not immediately look like training failures.

FP16 overflow on Turing

The Colab T4 does not support BF16. Running the student in FP16 produced non-finite logits, which eventually caused torch.multinomial to fail with a device-side CUDA assertion.

Greedy decoding initially hid the problem because argmax can still return an index even when the underlying logits are corrupted.

Fix: run the student in FP32 because it is the model being sampled from and backpropagated through, while keeping the frozen teacher in FP16 for inference.

Destroyed tied embeddings

Qwen2.5-0.5B ties its input embedding and output projection weights. Gradients through the LM head therefore also update the token embedding matrix.

An early experiment corrupted this tied parameter block and produced severely degraded generations, including vocabulary-like output.

Fix: freeze the tied embedding. This also removes a large parameter block from the optimizer while leaving the transformer body trainable.

Reentrant gradient checkpointing

The default reentrant checkpointing implementation caused problems with the training setup.

Fix: use:

use_reentrant=False

and enable input gradients where required.

Incorrect completion-length measurement

Using len(row) on batched generation output does not necessarily represent the actual completion length of an individual example because generated tensors can include padding.

Fix: count non-padding completion tokens when measuring generation length.

Miscalibrated gradient guard

An initial gradient-norm skip threshold of 10.0 rejected essentially every training step. The observed normal range for this setup was closer to 20--32.

Fix: log the gradient distribution before selecting a safety threshold.

Implementation notes

Truncated reverse KL

Full-vocabulary KL over Qwen's approximately 151K-token vocabulary is memory-intensive on a 16GB T4.

The implementation therefore restricts the objective to the teacher's top-64 tokens and renormalizes both student and teacher distributions over that support.

This is a teacher-top-K renormalized reverse-KL approximation, not the exact full-vocabulary reverse KL.

That distinction matters because truncating the support changes the objective and can introduce optimization behavior that differs from full-vocabulary KL.

Logit alignment

Causal LM logits at position i predict the token at position i+1.

The student and teacher logits therefore need to be shifted relative to the generated-token mask.

This is an easy bug to introduce because a misaligned implementation can still produce a smoothly decreasing loss while training the wrong positions.

A self-KL sanity check is run before training. Feeding identical logits into the objective should produce approximately zero KL.

Sampling temperature

Student rollouts are sampled at T=1.0.

Lowering the sampling temperature reduces exploration of the student's current policy and can substantially alter the on-policy distribution. Temperature is therefore treated as part of the training configuration rather than merely an inference parameter.

Memory

The target hardware is a 16GB NVIDIA T4.

The setup uses:

• Student: FP32
• Teacher: FP16
• 8-bit Adam
• Gradient checkpointing
• Batch size: 2
• Frozen tied embedding

8-bit optimizer states substantially reduce optimizer memory, while gradient checkpointing reduces activation memory.

The resulting configuration fits within the constraints of a free-tier T4 without requiring multiple GPUs.

Demo: Load and Run the Model

You can load the model directly with 🤗 Transformers:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "YOUR_USERNAME/qwen2.5-0.5b-opd-gsm8k"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    device_map="auto"
)

print("Successfully loaded:", MODEL_ID)

question = """A store has 240 apples. It sells 35% of them in the morning.
In the afternoon, it sells 48 more apples. How many apples are left?"""

messages = [
    {
        "role": "system",
        "content": "Solve the math problem. Reason step by step, then give the final answer as: #### <number>"
    },
    {
        "role": "user",
        "content": question
    }
]

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

inputs = tokenizer(text, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=False,
    pad_token_id=tokenizer.eos_token_id
)

answer = tokenizer.decode(
    outputs[0][inputs["input_ids"].shape[1]:],
    skip_special_tokens=True
)

print(answer)

Expected usage

The model is intended for small-scale mathematical reasoning experiments. The example above uses greedy decoding (do_sample=False) and the same system prompt used during evaluation.

Replace YOUR_USERNAME with your Hugging Face username after uploading the model.

Important: This checkpoint is an experimental OPD model. The reported evaluation results are small-scale and should not be interpreted as evidence that this model consistently outperforms the base Qwen2.5-0.5B-Instruct model.

Evaluation methodology

The evaluation intentionally measures more than a single accuracy number.

1. In-domain pass@1

GSM8K is used as the primary in-domain mathematical reasoning benchmark.

2. Out-of-domain pass@1

SVAMP is used as a separate arithmetic reasoning dataset with different questions and phrasing.

The purpose is to check whether behavior transfers beyond the training distribution.

3. Any-of-4 success

Four sampled completions are generated for each problem.

A problem counts as successful if at least one of the four completions produces the correct numerical answer.

This is reported as any-of-4 success rather than as a formal diversity metric. A future version should additionally report unique completion and unique-answer rates to directly measure output diversity.

4. Format compliance

The requested output format is:

#### <number>

Format compliance is tracked separately from mathematical correctness.

5. Completion length

Mean completion length is tracked because generation-length inflation or generation truncation can mask changes in model behavior.

A future evaluation should additionally report the fraction of generations that hit the maximum token limit.

Running it

pip install transformers datasets accelerate matplotlib bitsandbytes
python on_policy_distillation.py
python evals.py

The training run takes approximately 40 minutes on a T4, depending on the environment.

Restart the Colab runtime between training and evaluation. The scripts load their own models and the available 16GB VRAM is insufficient to keep the full training setup and evaluation models resident simultaneously.

Limitations

This is a small-scale implementation and not a research result demonstrating that OPD improves mathematical reasoning.

• Single seed; no repeated-run error bars.
• Only 20 training steps are reported for the selected checkpoint.
• Batch size is 2, corresponding to approximately 40 sampled training examples.
• No compute-matched comparison against SFT, GRPO, or other distillation baselines.
• The teacher itself scores only 0.17 pass@1 on the evaluated GSM8K subset, limiting the available capability ceiling in this experiment.
• The evaluation uses a small subset rather than the full benchmark.
• Format compliance is close to zero across the models, so the reported correctness numbers primarily measure whether the correct number appeared somewhere in the generated text rather than strict adherence to the #### <number> protocol.
• The any-of-4 experiment is based on only 50 problems and four samples per problem.
• No explicit unique-completion metric was included in this first version.
• The current implementation uses a teacher-top-K approximation rather than full-vocabulary reverse KL.
• Same-tokenizer student/teacher setup only. Cross-tokenizer OPD introduces additional alignment problems.

The experiment should therefore be interpreted as an investigation of one OPD configuration on a small model and constrained hardware, not as evidence about the general effectiveness of OPD.

What this experiment taught me

The main result was not a benchmark improvement.

It was learning how easy it is for an on-policy training loop to look healthy while the policy is moving in an undesirable direction.

A decreasing KL loss does not by itself imply improved reasoning.

A stable gradient norm does not imply stable behavior.

A small increase in pass@1 does not necessarily imply increased capability.

And a model that can still solve individual problems can simultaneously lose useful diversity across sampled trajectories.

Those were the failure modes this implementation was built to expose.

References

• On-Policy Distillation --- Thinking Machines Lab, 2025
• GKD: On-Policy Distillation of Language Models --- Agarwal et al., 2023
• MiniLLM: On-Policy Distillation of Large Language Models --- Gu et al., 2023
• awesome-on-policy-distillation --- literature index