On-Policy Self-Distillation

Code source and cost

• The complete implementation comes from KMnO4-zx/llm-agent-rl-lab/04-opsd. This page was checked against commit 52c2f1c.
• The measured PyTRIO cost was ¥41.81 for the 100-step training run. Five complete AIME25 evaluations—Base Model and Steps 25, 50, 75, and 100—sampled 18.05M tokens and cost ¥83.61, for a combined total of ¥125.42.
• The accumulated step time for 100 training steps was about 2 hours 6 minutes, excluding data preparation and full evaluation. See the SwanLab run for training curves.
• These are measurements from one reference run, not fixed prices. This is a multi-file data, training, and evaluation project, so this page explains the main code and complete logic while the source repository remains the runnable reference.

Introduction

OPSD stands for On-Policy Self-Distillation. It assigns two roles to the same initial model:

• the Student sees only the problem and samples from its current policy;
• the Teacher additionally sees a reference solution but does not generate a second answer;
• the Teacher computes per-token logprobs along the exact Student completion;
• the Student is updated from the token-level discrepancy.

An intuitive description is:

Student: solve the problem closed-book
Teacher: read the reference solution and grade the Student's original answer token by token

The most important boundary is:

The Teacher never samples a replacement "expert answer." The Student generates the only training trajectory, and the Teacher performs a logprob forward pass over that trajectory.

Training therefore remains on the states actually visited by the current Student policy.

How OPSD differs from other methods

OPSD combines on-policy trajectories, dense token-level feedback, and self-distillation without deploying a larger Teacher model.

Objective implemented in this example

The OPSD paper discusses both full-vocabulary logit distillation and sampled-token distillation. This PyTRIO example implements the latter: it compares Student and Teacher logprobs only for tokens actually sampled by the Student.

For a Student token \hat{y}_t:

$\operatorname{reverse\_kl}_t = \log p_S(\hat{y}_t \mid x, \hat{y}_{<t}) - \log p_T(\hat{y}_t \mid x, y^*, \hat{y}_{<t})$

$A_t = -\beta \cdot \operatorname{reverse\_kl}_t = \beta(\log p_T - \log p_S)$

If the Teacher assigns the token a higher logprob than the Student, its advantage is positive; if the Teacher prefers it less, its probability is pushed down.

This is a sampled-token reverse-KL implementation, not the full-vocabulary JSD used in the paper's main experiments. Keep this implementation boundary in mind when interpreting results.

Configuration

Project layout

The OPSD example consists of data preparation, training, evaluation, and analysis files:

04-opsd/
├── 00-datasets.py       # Download and validate OPSD / AIME25 data
├── 00-eval-aime25.py    # Shared base/checkpoint evaluation
├── 01-opsd-async.py     # Recommended asynchronous OPSD training
├── 01-opsd-sync.py      # Synchronous version for step-by-step reading
└── analysis.py          # Aggregate AIME25 checkpoint results

Use the complete source directory to run the example. The following sections use the asynchronous implementation to explain one training step.

Environment and data

Only a networked CPU machine is needed locally. Student sampling, Teacher logprobs, LoRA forward/backward, and optimizer updates run on the remote PyTRIO service.

git clone https://github.com/KMnO4-zx/llm-agent-rl-lab.git
cd llm-agent-rl-lab

uv sync
trio login
swanlab login

Download pinned revisions of siyanzhao/Openthoughts_math_30k_opsd and yentinglin/aime_2025:

uv run python 04-opsd/00-datasets.py

The script saves and validates:

04-opsd/datasets/
├── openthoughts_math_30k_opsd/   # 29,434 training examples
└── aime_2025/                     # 30 evaluation problems

Each training record contains at least:

problem   # Visible to Student and Teacher
solution  # Visible only to the Teacher

solution is not an SFT label. It only changes the Teacher's conditional distribution as privileged information; the Student's targets still come from its own completion.

Core logic

One training step can be summarized as:

for step in range(total_steps):
    student_sampler = refresh_latest_student_weights()

    rollouts = await asyncio.gather(
        *[
            student_sample_then_teacher_score(problem)
            for problem in batch
        ]
    )

    datums = build_importance_sampling_datums(rollouts)
    await forward_backward(datums)
    await optim_step()

The essential details are the following alignment constraints.

1. Use two prompts for the same model

The Student sees only the problem:

def build_student_prompt_ids(tokenizer, problem, enable_thinking):
    user_message = (
        f"Problem: {problem.strip()}\n\n"
        "Please reason step by step, and put your final answer within \\boxed{}."
    )
    return render_chat_prompt(tokenizer, user_message, enable_thinking)

The Teacher additionally sees the reference solution:

def build_teacher_prompt_ids(tokenizer, problem, solution, enable_thinking):
    user_message = (
        f"Problem: {problem.strip()}\n\n"
        "Here is a reference solution to this problem:\n"
        "=== Reference Solution Begin ===\n"
        f"{solution.strip()}\n"
        "=== Reference Solution End ===\n\n\n"
        f"{TEACHER_TRANSITION}\n\n"
        f"{STUDENT_INSTRUCTION}"
    )
    return render_chat_prompt(tokenizer, user_message, enable_thinking)

Both roles share the same tokenizer and initial model. The Teacher's additional capability comes from its privileged context rather than additional parameters.

2. Create a trainable Student and a fixed Teacher

01-opsd-async.py creates a LoRA TrainingClient and a fixed SamplingClient without a model_path:

service_client = trio.ServiceClient()

training_client = await service_client.create_lora_training_client_async(
    base_model=args.base_model,
    rank=args.lora_rank,
    seed=args.seed,
    train_attn=True,
    train_mlp=True,
    train_unembed=args.train_unembed,
)

teacher_client = await service_client.create_sampling_client_async(
    base_model=args.base_model,
)

Only the Student LoRA is updated. The Teacher remains the step-0 base policy and has no optimizer.

3. Let the Student generate the on-policy trajectory

At the start of each step, refresh the Student sampler:

student_sampler = (
    await training_client.save_weights_and_get_sampling_client_async()
)

The Student then samples while seeing only the problem:

sample_result = await student_sampler.sample_async(
    prompt=trio.ModelInput.from_ints(student_prompt_ids),
    num_samples=args.group_size,
    sampling_params=sampling_params,
    return_text=False,
)

The reference configuration uses group_size=1 and sampler_refresh_steps=1. Each problem yields one completion, and every new step uses Student weights from the previous optimizer update.

4. Compute Teacher logprobs for that same completion

After Student sampling, the Teacher receives:

teacher_prompt_ids + student_completion_ids

The Teacher calls compute_logprobs_async(), not sample_async():

all_ids = teacher_prompt_ids + completion_ids
all_logprobs = await teacher_client.compute_logprobs_async(
    trio.ModelInput.from_ints(all_ids)
)
teacher_logprobs = all_logprobs[len(teacher_prompt_ids):]

These three arrays must have exactly the same length:

Student completion tokens
Student rollout logprobs
Teacher completion logprobs

The source validates lengths and None values instead of silently truncating. Token-level reverse KL is undefined if they are misaligned.

5. Convert the logprob difference into token advantages

For each valid Student completion:

student_lps = [float(value) for value in sequence.logprobs]
reverse_kl = np.asarray(student_lps) - np.asarray(teacher_lps)
advantages = -args.kl_penalty_coef * reverse_kl

GRPO derives advantage from final rewards across several complete trajectories. OPSD derives advantage from the Teacher–Student logprob difference at every Student token. A wrong final answer can therefore still carry training signal when the Teacher disagrees with intermediate tokens.

6. Use the prompt as context and train only the completion

build_opd_datum() applies the same autoregressive shift to every field and fills the prompt interval with zeros:

prompt_loss_len = len(student_prompt_ids) - 1
input_ids = student_prompt_ids + completion_ids[:-1]

target_ids = [0] * prompt_loss_len + completion_ids
padded_logprobs = [0.0] * prompt_loss_len + old_logprobs
padded_advantages = [0.0] * prompt_loss_len + advantages.tolist()

datum = trio.Datum(
    model_input=trio.ModelInput.from_ints(input_ids),
    loss_fn_inputs={
        "target_tokens": np.asarray(target_ids, dtype=np.int64),
        "logprobs": np.asarray(padded_logprobs, dtype=np.float32),
        "advantages": np.asarray(padded_advantages, dtype=np.float32),
    },
)

Prompt tokens remain context but do not participate in policy optimization. Only tokens generated by the Student completion enter the importance_sampling loss.

7. Run different problems concurrently but preserve per-problem order

Within one problem, Student sampling must finish before the Teacher scores the resulting completion. Different problems can run concurrently:

rollouts = await asyncio.gather(
    *(rollout_and_track(row) for row in batch)
)

The async implementation shares an asyncio.Semaphore(32) across Student sampling and Teacher scoring. All rollouts in a batch come from one Student checkpoint and finish before the optimizer update, so concurrency does not violate the on-policy boundary.

8. Update the Student

After flattening all Datums from the current step:

fwd_bwd_future = await training_client.forward_backward_async(
    datums,
    loss_fn="importance_sampling",
)
optim_future = await training_client.optim_step_async(adam)

fwd_bwd_result = await fwd_bwd_future
await optim_future

PyTRIO handles remote forward/backward, the LoRA optimizer, and checkpoints. Local code controls the Student and Teacher prompts, rollouts, advantages, batch boundary, and update timing.

Run training

Run the 100-step asynchronous OPSD experiment:

uv run python 04-opsd/01-opsd-async.py \
    --steps 100 \
    --batch-size 32 \
    --group-size 1 \
    --max-tokens 1024 \
    --sample-size 0 \
    --save-every-steps 25 \
    --max-concurrency 32 \
    --swanlab-mode online

Every 25 steps, it saves:

*-state            # Optimizer state for resuming training
*-sampler_weights  # Weights for sampling and AIME25 evaluation

Selected metrics from the first 100 steps of the reference run were:

trainer/loss_mean: 0.0577 → 0.0425
reverse_kl_mean:   0.0644 → 0.0471
reverse_kl_std:    0.4378 → 0.3687

These changes show that the Student moved closer to the privileged Teacher, but they do not by themselves prove improved mathematical ability. An independent benchmark is still required.

AIME25 evaluation

Evaluate the base model:

uv run python 04-opsd/00-eval-aime25.py \
    --val-n 12 \
    --max-tokens 38912 \
    --temperature 1.0 \
    --enable-thinking false \
    --output 04-opsd/eval-results/aime25-base.jsonl

Evaluate the Step 100 checkpoint:

uv run python 04-opsd/00-eval-aime25.py \
    --val-n 12 \
    --max-tokens 38912 \
    --temperature 1.0 \
    --enable-thinking false \
    --model-path trio://<your_sampler_weights_path> \
    --output 04-opsd/eval-results/aime25-sampler-steps100.jsonl

Each model state produces 30 problems × 12 samples = 360 completions. max_tokens=38,912 is a per-completion upper bound, not actual usage; each reference evaluation consumed about 3.40M–3.86M sample tokens.

Reference results:

Step 100 improved Average@12 by 1.11 percentage points and Pass@12 by 6.67 points over the base model. Intermediate checkpoints were not monotonic, and this was one 30-problem evaluation without multiple training seeds. The evidence supports a working training loop with a small positive change, not a claim of stable reproduction of the paper's gains.

Metrics to monitor

Lower loss and reverse KL only show that the Student is approaching the Teacher. Whether the Teacher uses the reference solution reliably—and whether an independent benchmark improves—must still be evaluated separately.

For complete code, cost screenshots, experiment figures, and a longer discussion, see the source repository's OPSD README.