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

> Explore on-policy distillation with SGLang in Miles. Learn how token-level knowledge transfer from teacher to student models enhances performance. A complete technical guide.

- Repository: [RadixArk/miles](https://github.com/radixark/miles)
- Tags: deep-dive
- Published: 2026-09-06

---

**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`](https://github.com/radixark/miles/blob/main/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`](https://github.com/radixark/miles/blob/main/miles/utils/arguments.py) ([lines 1655-1700](https://github.com/radixark/miles/blob/main/miles/utils/arguments.py#L1655-L1700)):

| 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

```bash
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`](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py) ([lines 26-42](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L26-L42)) transforms CLI strings into a routing dictionary:

```python

# 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](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L45-L74)) 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

```python
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](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L95-L119)) builds the JSON payload sent to SGLang teachers:

```python
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](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L138-L140)) constructs this endpoint:

```python
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](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L81-L147)).

### Token Set Strategies

Controlled by `--opd-top-k-strategy` ([lines 81-85](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L81-L85)):

| 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

```python

# 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`](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py) ([lines 52-148](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L52-L148)) coordinates the full workflow:

```python
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](https://github.com/radixark/miles/blob/main/miles/rollout/on_policy_distillation.py#L303-L336)) attaches computed tensors to samples:

```python
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.

```bash
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:

```json
{"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:

```python
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`](https://github.com/radixark/miles/blob/main/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`](https://github.com/radixark/miles/blob/main/tests/e2e/short/test_qwen2.5_0.5B_opd_sglang.py) | End-to-end validation with live SGLang server |

Run fast tests:

```bash
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.