On-Policy Distillation with SGLang in Miles: A Complete Technical Guide

Miles implements on-policy distillation (OPD) as a KL-penalty layer that queries external SGLang teacher servers during rollout, enabling token-level knowledge transfer from powerful teachers to student models.

On-policy distillation with SGLang in the Miles framework provides a flexible, production-ready mechanism for distilling knowledge from large teacher models without loading them into GPU memory. This integration treats SGLang as an external inference service, allowing per-sample teacher routing, top-k token-level KL estimation, and seamless composition with any advantage estimator. According to the radixark/miles source code, the implementation spans argument parsing, dynamic URL resolution, and specialized reverse-KL computation.

Architecture Overview

The SGLang OPD system in Miles follows a teacher-as-a-service pattern. The student model generates rollouts, then queries external SGLang endpoints for teacher log-probabilities. These are compared against student log-probs to compute a KL penalty that drives distillation.

Key design principles from miles/rollout/on_policy_distillation.py:

• Decoupled inference: Teachers run on separate SGLang servers, not the training GPU
• Per-sample routing: Different samples can hit different teachers via metadata annotations
• Token-level granularity: Optional top-k scoring enables fine-grained distillation objectives

Configuration and CLI Arguments

OPD behavior is controlled through flags added in miles/utils/arguments.py (lines 1655-1700):

Flag	Purpose
--use-opd	Master switch enabling on-policy distillation
--opd-type=sglang	Selects SGLang backend (vs. local or other backends)
--opd-teacher-urls	Routing map in NAME=URL format for multi-teacher setups
--opd-teacher-key	Metadata key to lookup teacher name (default: opd_teacher)
--sglang-router-ip / --sglang-router-port	Student-side scoring server address
--opd-log-prob-top-k	Enables top-k token scoring (0 = full vocab)
--opd-top-k-strategy	Token selection strategy: only-student, only-teacher, intersection, union, xor
--opd-reward-weight-mode	Weight normalization: student_p, teacher_p, or none

Basic SGLang OPD Launch

python -m miles.train \
  --use-opd \
  --opd-type sglang \
  --sglang-router-ip 10.0.0.1 \
  --sglang-router-port 3000 \
  --opd-teacher-urls default=http://10.0.0.1:3001/generate

Teacher URL Parsing and Routing

The parse_teacher_urls function in miles/rollout/on_policy_distillation.py (lines 26-42) transforms CLI strings into a routing dictionary:

# Input: ["math=http://10.0.0.2:3002/generate", "default=http://10.0.0.1:3001/generate"]

# Output: {"math": "http://10.0.0.2:3002/generate", "default": "http://10.0.0.1:3001/generate"}

def parse_teacher_urls(url_list: List[str]) -> Dict[str, str]:
    result = {}
    for entry in url_list:
        if "=" not in entry:
            raise ValueError(f"Invalid teacher URL format: {entry}")
        name, url = entry.split("=", 1)
        result[name.strip()] = url.strip()
    return result

Per-Sample Teacher Resolution

The _teacher_url_for_sample function (lines 45-74) implements fallback logic:

1. Extract teacher name from sample metadata using --opd-teacher-key
2. If name exists in routing map → return that URL
3. If no routing map configured → fall back to --rm-url
4. If name unknown → use reserved default teacher, or raise an error if absent

from miles.rollout.on_policy_distillation import parse_teacher_urls, _teacher_url_for_sample

args = Namespace(
    opd_teacher_urls=parse_teacher_urls([
        "math=http://10.0.0.2:3002/generate",
        "code=http://10.0.0.3:3003/generate",
        "default=http://10.0.0.1:3001/generate"
    ]),
    opd_teacher_key="subject",  # metadata key to read

    rm_url="http://10.0.0.1:3001/generate",
)

sample = Sample(
    tokens=[...],
    metadata={"subject": "math"}  # routes to math teacher

)

url = _teacher_url_for_sample(args, sample)  # → http://10.0.0.2:3002/generate

SGLang Payload Construction and Scoring

Teacher Scoring Request

The _score_payload function (lines 95-119) builds the JSON payload sent to SGLang teachers:

def _score_payload(input_ids: List[int], top_k: int = 0, token_positions: Optional[List[int]] = None) -> dict:
    payload = {
        "input_ids": input_ids,           # prompt + generated tokens

        "sampling_params": {
            "temperature": 0,             # deterministic scoring

            "max_new_tokens": 0,          # no generation, only scoring

        },
        "return_logprob": True,           # required for distillation

    }
    if top_k > 0:
        payload["sampling_params"]["top_logprobs_num"] = top_k
    if token_positions:
        payload["token_ids_logprob_positions"] = token_positions
    return payload

Student-Side Scoring Endpoint

When --opd-log-prob-top-k > 0, Miles needs student log-probs on teacher top-k positions. The _student_score_url helper (lines 138-140) constructs this endpoint:

def _student_score_url(args) -> str:
    return f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"

Top-K Reverse KL Computation

The core distillation objective is implemented in _compute_topk_reverse_kl (lines 81-147).

Token Set Strategies

Controlled by --opd-top-k-strategy (lines 81-85):

Strategy	Description	Use Case
only-student	Student's top-k tokens only	Prevent teacher bias
only-teacher	Teacher's top-k tokens only	Aggressive distillation
intersection	Tokens in both top-k sets	Safe, high-confidence transfer
union	Tokens in either top-k set	Comprehensive coverage
xor	Tokens in exactly one top-k set	Focus on disagreement

KL Penalty Calculation

For each position, the function:

1. Builds log-prob dictionaries from teacher and student responses
2. Selects token IDs per the chosen strategy
3. Normalizes weights if --opd-reward-weight-mode is set
4. Aggregates w * (s_logp - t_logp) across selected tokens

# Conceptual structure of reverse-KL computation

for position in sequence:
    teacher_logps = get_teacher_logprobs(position, selected_tokens)
    student_logps = get_student_logprobs(position, selected_tokens)
    
    # Weight determination

    if weight_mode == "student_p":
        weights = softmax(student_logps)
    elif weight_mode == "teacher_p":
        weights = softmax(teacher_logps)
    else:
        weights = uniform  # none mode

    
    # Reverse KL: E_teacher[log p_teacher - log p_student]

    reverse_kl = sum(w * (s - t) for w, s, t in zip(weights, student_logps, teacher_logps))

Reward Function Orchestration

The reward_func in miles/rollout/on_policy_distillation.py (lines 52-148) coordinates the full workflow:

async def reward_func(args, sample: Sample) -> Dict[str, Any]:
    # 1. Resolve teacher URL

    teacher_url = _teacher_url_for_sample(args, sample)
    
    # 2. Get teacher log-probs (always required)

    teacher_response = await _post_json(teacher_url, _score_payload(sample.tokens, args.opd_log_prob_top_k))
    
    # 3. Optional: student-on-teacher pass for top-k alignment

    if args.opd_log_prob_top_k > 0:
        student_positions = extract_topk_positions(teacher_response)
        student_response = await _post_json(
            _student_score_url(args),
            _score_payload(sample.tokens, token_positions=student_positions)
        )
    
    # 4. Compute reverse KL

    reverse_kl = _compute_topk_reverse_kl(teacher_response, student_response, args)
    
    return {"opd_reverse_kl": reverse_kl, "reward": 0.0}  # pure OPD mode

All HTTP requests respect --sglang-router-request-timeout-secs via the _post_json helper.

Post-Processing and Training Integration

The post_process_rewards function (lines 303-336) attaches computed tensors to samples:

def post_process_rewards(samples: List[Sample], reward_outputs: List[Dict]):
    for sample, output in zip(samples, reward_outputs):
        sample.opd_reverse_kl = output["opd_reverse_kl"]  # tensor stored for loss computation

        # scalar reward is 0.0 in pure OPD mode; mix with task rewards if desired

This design allows the KL penalty to drive learning while remaining compatible with:

• Pure distillation: reward = 0.0, loss = reverse KL only
• Mixed objectives: add task-specific rewards to the scalar
• Arbitrary advantage estimators: GRPO, PPO, RLOO, etc.

Advanced Configuration: Multi-Teacher Self-Distillation

The examples/on_policy_distillation/qwen3_5_35b_selfdistill/ directory demonstrates a production setup where a 35B model distills into itself across specialized domains.

python -m miles.train \
  --use-opd \
  --opd-type sglang \
  --opd-teacher-urls \
    reasoning=http://reasoning-cluster:3000/generate \
    coding=http://code-cluster:3001/generate \
    default=http://general-cluster:3002/generate \
  --opd-teacher-key skill_domain \
  --opd-log-prob-top-k 10 \
  --opd-top-k-strategy intersection \
  --opd-reward-weight-mode teacher_p

With corresponding dataset metadata:

{"text": "Solve for x: 2x + 5 = 13", "skill_domain": "reasoning"}
{"text": "Write a Python function to...", "skill_domain": "coding"}
{"text": "Translate to French...", "skill_domain": "default"}

Debugging and Manual Invocation

For development, call reward_func directly:

import asyncio
from argparse import Namespace
from miles.rollout.on_policy_distillation import reward_func
from miles.utils.types import Sample

args = Namespace(
    sglang_router_ip="127.0.0.1",
    sglang_router_port=3000,
    opd_type="sglang",
    opd_log_prob_top_k=5,
    opd_top_k_strategy="union",
    opd_reward_weight_mode="student_p",
    opd_teacher_urls=None,  # uses rm_url fallback

    rm_url="http://127.0.0.1:3001/generate",
    opd_teacher_key="opd_teacher",
    sglang_router_request_timeout_secs=60,
)

sample = Sample(
    tokens=[151644, 872, 198, 198, 14958, 13],  # Qwen chat format tokens

    response_length=3,
    metadata={"opd_teacher": "default"}
)

result = asyncio.run(reward_func(args, sample))
print(f"Reverse KL shape: {result['opd_reverse_kl'].shape}")
print(f"Sample tokens: {result.get('teacher_tokens_scored', [])}")

Testing and Validation

Test File	Coverage
tests/fast/rollout/test_on_policy_distillation.py	Unit tests for URL parsing, routing logic, top-k strategies, reverse-KL math
tests/e2e/short/test_qwen2.5_0.5B_opd_sglang.py	End-to-end validation with live SGLang server

Run fast tests:

pytest tests/fast/rollout/test_on_policy_distillation.py -v

Summary

• On-policy distillation with SGLang in Miles enables scalable knowledge transfer without loading teachers on training GPUs
• Per-sample routing via --opd-teacher-urls and metadata keys supports complex multi-teacher deployments
• Token-level top-k scoring with configurable strategies (intersection, union, etc.) provides fine-grained control over distillation objectives
• Reverse KL computation incorporates optional weight normalization and integrates with any advantage estimator
• Pure OPD mode returns zero scalar reward, using only the KL penalty for gradient updates

Frequently Asked Questions

What hardware requirements does the SGLang teacher need?

The teacher SGLang server requires sufficient GPU memory to host the teacher model—typically 2-4× the student model size for inference. The training nodes only need GPUs for the student model. Network latency between training and teacher clusters directly impacts rollout throughput; co-locate in the same data center when possible.

Can I mix SGLang teachers with local reward models?

Yes. The --opd-type sglang flag only affects how teacher log-probabilities are obtained. You can still configure --rm-type and related flags for separate reward models. The scalar reward from OPD is 0.0 by default, so add task rewards through standard Miles mechanisms if needed.

How do I debug why a sample routes to the wrong teacher?

Set --opd-teacher-key explicitly and inspect sample metadata. The _teacher_url_for_sample function raises clear errors for unknown teacher names when no default entry exists. Enable verbose logging to see resolved URLs per sample, or temporarily patch reward_func to log all routing decisions.

What is the performance impact of --opd-log-prob-top-k versus top-k=0?

Top-k=0 (full vocabulary) requires one teacher request per sample but yields more accurate KL estimates. Top-k>0 adds a second request to the student scoring server for token alignment, increasing latency by 30-50% typically but reducing computational load for both teacher and student during the forward pass. Choose based on your latency budget and distillation quality requirements.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}