Consider Computing Teacher Logits Using Training Engine for On-Policy Distillation

@ahxt https://github.com/THUDM/slime/pull/673

Problem Statement

In the current implementation of on-policy distillation (introduced in commit 12dd6b2), teacher_log_probs are computed by querying the inference engine (SGLang server) via the /generate endpoint with return_logprob=True. However, there is already a non-negligible precision gap between the training engine (Megatron/FSDP) and the inference engine (SGLang), which may introduce inconsistencies in the distillation process.
Like mentioned in https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda

Current Implementation

Currently, teacher_log_probs are obtained as follows:

1. During rollout: The teacher model runs on SGLang inference servers, and log probabilities are extracted from the inference engine's response (reward["meta_info"]["input_token_logprobs"]).

2. During training: The student model computes log probabilities using the training engine (Megatron/FSDP) via compute_log_prob().

This creates a mismatch where:

• Teacher logits: Computed by inference engine (SGLang)
• Student logits: Computed by training engine (Megatron/FSDP)

Proposed Solution

Since the training engine and inference engine may have different numerical precisions, quantization strategies, or implementation details, it would be more consistent to compute both teacher and student logits using the same training engine. This would:

1. Eliminate precision mismatches: Both teacher and student logits would be computed under identical numerical conditions.
2. Improve distillation consistency: The advantage signal (teacher_log_prob - student_log_prob) would be more accurate.
3. Better align with training dynamics: The teacher logits would reflect the same computational path as the student during training.

Proposed Implementation

Instead of querying the inference engine for teacher logits, we could:

1. Option A: Load the teacher model weights into the training engine and compute logits using compute_log_prob() with model_tag="teacher" (similar to how reference model logits are computed).

2. Option B: Add a flag to allow users to choose between inference-engine logits (current) and training-engine logits (proposed).

Questions

1. Is there a specific reason why teacher logits are currently computed from the inference engine? (e.g., performance, memory constraints, or architectural limitations)

2. Are there any plans to support computing teacher logits using the training engine?

3. Would this change require significant modifications to the current architecture, or could it be implemented as an optional feature?

Related Code

• Teacher logits extraction: examples/on_policy_distillation/on_policy_distillation.py:25-39
• Advantage calculation: slime/backends/megatron_utils/loss.py:289-302
• Student logits computation: slime/backends/megatron_utils/actor.py:compute_log_prob()

Additional Context

The precision gap between training and inference engines is a known issue in many RL frameworks. Using the training engine for both teacher and student logits would ensure numerical consistency and potentially improve distillation quality.

Start with teacher-logit extraction in examples/on_policy_distillation/on_policy_distillation.py:25-39, then trace student computation through slime/backends/megatron_utils/actor.py:compute_log_prob() and advantage handling in loss.py:289-302. Determine whether training-engine teacher logits are feasible and whether an optional path is needed; done requires an agreed implementation and validation of consistent teacher and student log probabilities.