[Feature] Mopd (Multi-Teacher On-Policy distillation) supported

[leoyuppieqnew]

feat: Multi-Teacher On-Policy Distillation (MOPD)

Summary

Add Multi-Teacher On-Policy Distillation (MOPD) support to slime, enabling a single student model to distill knowledge from multiple domain-specific teachers simultaneously with importance sampling (IS) correction for stable off-policy training.

34 files changed, +9040 / -50 lines

Motivation

Standard on-policy distillation (OPD) supports only a single teacher. In practice, different domains (math, code, reasoning) benefit from specialized teacher models. MOPD extends OPD to aggregate knowledge from multiple domain-specific teachers into a single student, using per-teacher reverse KL advantages averaged across domains with IS-weight correction for stable training when the student policy diverges from the sampling policy.

Algorithm

MOPD supports three distillation strategies, controlled by --mopd-distill-type:

Token-Level Mode (token_level, default)

Uses the sampled token's log-prob difference as a point estimate of the reverse KL divergence:

reverse_kl_d(y_t) = sg[log π_d(y_t) - log π_θ(y_t)]

Training loss:

w_t = sg[π_θ(y_t) / μ_θ(y_t)]  clipped to [ε_low, ε_high]   # IS weight
Â_MOPD,t = (1/D) Σ_d (reverse_kl_d + α · Â_ORM)              # avg across D teachers
L = -E[1/|y| Σ_t w_t · Â_MOPD,t · log π_θ(y_t)]              # proxy policy loss

• Memory: Negligible (1 scalar per token per teacher)
• Accuracy: Underestimates the true KL (only evaluates at sampled positions)
• Teacher modes: SGLang and Megatron

Top-K Mode (top_k, recommended)

A memory-efficient approximation of the full-vocab KL. Only stores the teacher's top-k logits and indices, with an analytical tail correction:

D_KL(π_θ ∥ π_d) ≈ KL_topk + KL_tail

Top-K part (exact over the teacher's top-k tokens):

KL_topk = Σ_{y ∈ top-k} π_s(y) [log π_s(y) - log π_t(y)]

Tail correction (approximates the remaining vocabulary):

KL_tail ≈ π_s_tail · log(π_s_tail / π_t_tail)

where:

• π_s_tail = 1 - Σ_{y ∈ top-k} π_s(y) — student's exact tail mass (computed via all-reduce across TP ranks)
• Teacher tail mass estimation:
  • SGLang mode: π_t_tail = 1 - Σ exp(log_prob_t(y)) — exact (SGLang returns full-vocab log-probs)
  • Megatron mode: π_t_tail ≈ (V - V_eff) / V — uniform assumption over non-top-k tokens

Full loss:

w_t = sg[π_θ(y_t) / μ_θ(y_t)]  clipped to [ε_low, ε_high]    # IS weight
L_topk_kl = (1/D) Σ_d (1/|y| Σ_t w_t · KL_topk+d(π_θ ∥ π_d))  # IS-corrected approx KL
L = L_topk_kl + α · L_pg                                         # combined with PG loss

• Memory: B × R × k × 2 × 4B / TP (~97% reduction vs full_vocab for k=1024, V=152K)
• Accuracy: Very close to full_vocab (top-k tokens capture >99% probability mass)
• Teacher modes: SGLang and Megatron
• TP-aware: SGLang returns global token IDs → converted to per-TP-shard local indices with -inf padding for out-of-shard entries; Megatron provides local indices directly

Full-Vocabulary Mode (full_vocab)

Computes the exact reverse KL over the entire vocabulary:

D_KL(π_θ ∥ π_d) = Σ_y π_θ(y) [log π_θ(y) - log π_d(y)]

Training loss:

w_t = sg[π_θ(y_t) / μ_θ(y_t)]  clipped to [ε_low, ε_high]   # IS weight
L_fv_kl = (1/D) Σ_d (1/|y| Σ_t w_t · D_KL(π_θ ∥ π_d))        # IS-corrected KL loss
L = L_fv_kl + α · L_pg                                         # combined with PG loss

• Memory: Very high — B × R × V × 4B / TP (stores full teacher logits)
• Accuracy: Exact KL (gold standard)
• Teacher modes: Megatron only (SGLang cannot efficiently return full-vocab logits)

Comparison

	token_level	top_k	full_vocab
KL accuracy	Point estimate	Approximate (~99%+)	Exact
Teacher data per token	1 scalar	k×2 values	V values
Memory overhead	≈ 0	O(k)	O(V)
Teacher mode	SGLang / Megatron	SGLang / Megatron	Megatron only
Gradient	Through policy loss only	Through full student softmax	Through full student softmax

Key Features

• Multi-teacher distillation: Aggregate knowledge from multiple domain-specific teachers into a single student
• Importance sampling (IS) correction: Clipped IS weights w_t = sg[π_θ/μ_θ] ensure stable training
• ORM combination: Coefficient α blends reverse KL with standard ORM advantages
• Two teacher modes:
  • sglang — teachers on external SGLang servers (can have different architectures)
  • megatron — teachers loaded into Megatron (same architecture required)
• Per-sample domain routing: Route samples to specific teachers via mopd_domains metadata field
• Non-colocate mode: Separate actor training GPUs from SGLang rollout GPUs
• Qwen3.5 VL MoE support: Bridge-mode weight sync for vision encoders via HfWeightIteratorBridge

Architecture

┌─────────────────────────────────────────────────────────────────┐
│                         Rollout Phase                            │
├─────────────────────────────────────────────────────────────────┤
│  mopd.py::reward_func()                                         │
│    → Query teacher SGLang servers (concurrent, per-domain)      │
│    → Request top_logprobs_num=k per position                    │
│    → Extract top-k logprobs + token indices                     │
│    → Store in sample.mopd_teacher_topk_{logits,indices}         │
│                                                                  │
│  rollout.py::_split_train_data_by_dp()                          │
│    → Partition mopd dict data by DP rank                        │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│                        Training Phase                            │
├─────────────────────────────────────────────────────────────────┤
│  actor.py::_get_rollout_data()                                  │
│    → Pop mopd_teacher_topk_{logits,indices} dict                │
│    → Convert global token IDs → per-TP-shard local indices      │
│    → Pad out-of-shard entries with -inf logit / 0 index         │
│    → Store as rollout_data["mopd_teacher_{domain}_topk_*"]      │
│                                                                  │
│  model.py::forward_step()                                       │
│    → Dynamic batch_keys (base + per-domain topk keys)           │
│    → get_batch() → DataIterator                                 │
│                                                                  │
│  loss.py::policy_loss_function()                                │
│    → get_responses(apply_temperature=False) → student logits    │
│    → apply_mopd_topk_to_loss()                                  │
│      → vocab_parallel_topk_reverse_kl() [TP-aware]             │
│      → IS-weighted KL loss + optional PG loss                   │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│                      Weight Sync Phase                           │
├─────────────────────────────────────────────────────────────────┤
│  update_weight_from_distributed.py                              │
│    → if megatron_to_hf_mode == "bridge":                        │
│        _update_weights_bridge() via HfWeightIteratorBridge      │
│    → else:                                                       │
│        _send_weights() via template method pattern              │
└─────────────────────────────────────────────────────────────────┘

New Files

File	Description
slime/rollout/mopd.py	SGLang-mode MOPD reward function & post-processing (780 lines)
slime/utils/ppo_utils.py (+369)	vocab_parallel_topk_reverse_kl, vocab_parallel_reverse_kl
slime_plugins/megatron_bridge/qwen35_vl_moe.py	Qwen3.5 VL MoE bridge plugin (1263 lines)
scripts/models/qwen3.5-397B-A17B.sh	Model config for 397B-A17B MoE
tools/convert_torch_dist_to_hf_parallel.py	Parallel distributed checkpoint conversion
tools/merge_missing_keys.py	Merge missing VL encoder keys into checkpoint
tools/patch_attention_gate_on_cluster.py	Attention gate patching utility
tests/test_mopd.py	Unit tests for MOPD token-level pipeline (555 lines)
tests/test_mopd_full_vocab.py	Unit tests for full-vocab distillation (1041 lines)
tests/test_mopd_sglang_topk_pipeline.py	Unit tests for SGLang top-k pipeline (747 lines)

Modified Files (key changes)

File	Changes
slime/backends/megatron_utils/actor.py (+317)	SGLang top-k data preparation: global→local TP index conversion
slime/backends/megatron_utils/loss.py (+829)	TopK KL loss, full-vocab KL loss, IS-weighted loss integration
slime/backends/megatron_utils/model.py (+59)	Dynamic batch_keys for per-domain MOPD teacher data
slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py (+49)	Bridge mode via HfWeightIteratorBridge for VL MoE
slime/utils/arguments.py (+376)	MOPD argument definitions, validation, auto-configuration
slime/ray/rollout.py (+86)	MOPD dict-typed data partitioning across DP ranks
slime/utils/types.py (+15)	messages field on Sample dataclass
slime/utils/data.py (+10)	filter_long_prompt multimodal input handling

Example Scripts

Script	Setup
run-qwen35-397B-A17B-mopd-topk-sglang.sh	288 GPUs (256 train + 32 rollout), non-colocate, SGLang top-k
run-qwen35-35B-A3B-mopd-topk-sglang.sh	48 GPUs (32 train + 16 rollout), non-colocate, SGLang top-k

Usage

# 1. Start teacher SGLang server
python3 -m sglang.launch_server \
    --model-path /path/to/teacher \
    --host 0.0.0.0 --port 13141 \
    --tp 8 --ep-size 16 \
    --mem-fraction-static 0.7

# 2. Convert student to Megatron format
source scripts/models/qwen3.5-397B-A17B.sh
PYTHONPATH=/root/Megatron-LM torchrun --nproc_per_node=8 \
    tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /path/to/student \
    --save /path/to/student_torch_dist

# 3. Run MOPD training
export MOPD_TEACHER_URLS='{"math":"http://teacher-host:8300/generate"}'
export MOPD_TEACHERS_JSON='[{"name":"math_teacher","domain":"math"}]'

ray job submit --address="http://127.0.0.1:8265" \
    -- python3 train.py \
    --use-mopd \
    --mopd-teacher-mode sglang \
    --mopd-distill-type top_k \
    --mopd-topk-k 128 \
    --mopd-alpha 0.0 \
    --mopd-eps-low 0.2 \
    --mopd-eps-high 5.0 \
    ${MODEL_ARGS[@]} \
    ...

Key Arguments

Argument	Description	Default
--use-mopd	Enable MOPD (mutually exclusive with --use-opd)	—
--mopd-teachers	JSON list of teacher configs (or env MOPD_TEACHERS_JSON)	required
--mopd-teacher-mode	sglang or megatron	auto-detected
--mopd-distill-type	token_level, top_k, or full_vocab	token_level
--mopd-topk-k	Number of top-k tokens (for top_k mode)	1024
--mopd-alpha	ORM combination coefficient (0 = pure distillation)	0.0
--mopd-eps-low	IS weight lower clipping bound	0.2
--mopd-eps-high	IS weight upper clipping bound	5.0
--mopd-teacher-loads	Megatron teacher checkpoint paths (megatron mode)	—
--mopd-sampling-logprobs-key	Key for sampling log-probs in IS computation	rollout_log_probs

Training Results

The following figure shows key training metrics from an MOPD top-k distillation run, confirming stable and effective convergence:

Key observations:

• mopd_topk_kl / mopd_topk_kl/origin / mopd_topk_kl/enhanced: KL divergence between student and teacher decreases steadily, indicating the student successfully absorbs teacher knowledge across domains.
• mopd_is_weight_mean: Importance sampling weights remain tightly centered around 1.0, validating that IS clipping (eps_low=0.2, eps_high=5.0) effectively prevents variance explosion.
• train_rollout_logprob_abs_diff: The absolute difference between training and rollout log-probs stays small and stable, confirming on-policy consistency.
• grad_norm / entropy_loss / loss: Gradient norm stays bounded, entropy decreases gradually, and the overall loss converges — all signs of healthy training dynamics.
• ppo_kl: Remains well-controlled, indicating the student policy is not drifting excessively from the reference policy.

Testing

python -m pytest tests/test_mopd.py tests/test_mopd_full_vocab.py tests/test_mopd_sglang_topk_pipeline.py -v

Compatibility

• Requires SGLang with top_logprobs_num support (recent versions)
• Teacher vocabulary size must match student's padded_vocab_size (for SGLang top-k mode)
• Bridge mode (--megatron-to-hf-mode bridge) required for Qwen3.5 VL MoE weight sync
• Non-colocate mode requires separate GPU allocation for actor and rollout

[leoyuppieqnew]

@zhuzilin Hello！Would you mind helping review it when you’re available? I’d really appreciate it. Thank you!