[DO NOT REVIEW][NOT FOR LAND] on-policy distillation example

[casteryh]

Summary

This PR adds a new online distillation app that demonstrates how a smaller student model (Qwen3-1.7B) can learn from a larger teacher model (Qwen3-32B) via KL divergence on generated completions.

This is an example of online distillation where:

• Student model generates rollouts from prompts
• Teacher model (frozen) provides target logprobs for the same completions
• Training objective: Minimize KL divergence between student and teacher distributions
• No rewards or advantages - pure distillation objective

Key Components

Architecture:

• StudentGenerator (vLLM): Generates completions from prompts using student model
• StudentTrainer: Trains student model to match teacher distributions
• TeacherModel (frozen): Provides target logprobs for distillation
• DatasetActor: Provides prompts (GSM8K)
• ReplayBuffer: Batches episodes for training

Loss Function:

def distillation_loss(student_logprobs, teacher_logprobs, padding_mask):
    # Forward KL: KL(teacher || student)
    kl = teacher_logprobs.exp() * (teacher_logprobs - student_logprobs)
    return kl.mean()

Implementation Details

Based on apps/grpo/main.py with key differences:

• ✅ Removed RewardActor and ComputeAdvantages (not needed for distillation)
• ✅ Replaced simple_grpo_loss with distillation_loss (pure KL divergence)
• ✅ Simplified Episode dataclass (removed reward and advantage fields)
• ✅ Renamed policy → student_generator, ref_model → teacher_model for clarity
• ✅ Updated collate function to work without advantages

Files Added

• apps/distillation/main.py: Main training loop
• apps/distillation/qwen3_distillation.yaml: Config for Qwen3-1.7B → Qwen3-32B distillation

Usage

python -m apps.distillation.main --config apps/distillation/qwen3_distillation.yaml

Test Plan

• ✅ Verified Python syntax with py_compile
• ✅ Config follows same pattern as GRPO configs
• ✅ All pre-commit hooks passed

cc @wukaixingxp