feat(opd): [Task.39] EOPD cross-architecture distillation (arXiv:2603.07079) - #274
Open
buaacoder wants to merge 6 commits into
Open
feat(opd): [Task.39] EOPD cross-architecture distillation (arXiv:2603.07079)#274buaacoder wants to merge 6 commits into
buaacoder wants to merge 6 commits into
Conversation
buaacoder
requested review from
Aurelius84,
NINGBENZHE and
Yangruipis
as code owners
August 14, 2026 02:38
Implement EOPD from arXiv:2603.07079 — entropy-gated forward KL loss that supplements standard OPD reverse KL at high-entropy teacher tokens. Key changes: - Add `--use-eopd`, `--eopd-entropy-threshold`, `--eopd-fkl-coef` CLI args - Implement `compute_eopd_fkl_loss()` with entropy mask gating - Plumb `teacher_entropy` through megatron forward and transfer pipeline - Fix megatron OPD path: compute teacher top-K log-probs during rollout, map `teacher_topk_*` → `opd_topk_*` in actor, add megatron batch keys - Add 10 unit tests (loss function + argument validation) - Add smoke test script Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Enable EOPD (entropy-gated forward KL) for --opd-type=sglang by computing per-token teacher entropy server-side and wiring it through the OPD data pipeline. This allows cross-architecture distillation (e.g. Qwen3-8B → Qwen3-0.6B) where the teacher has genuine high-entropy tokens. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…y-gated FKL loss Implement Entropy-aware On-Policy Distillation (EOPD) from arXiv:2603.07079 for cross-architecture teacher-student setups (e.g. Qwen3-8B → Qwen3-0.6B-Base). Key changes: - Add entropy-gated forward KL loss scaling via teacher_advantage in Megatron loss - Add SGLang entropy monkey-patch for teacher prefill entropy extraction - Add batch student prefill for colocate OPD rollout - Add EOS token replacement (--opd-eos-replace) for teacher prefill alignment - Add teacher advantage computation (--opd-teacher-advantage) with loss/additive modes - Add OPD colocate GPU allocation in controller and utils Tested with Qwen3-8B teacher → Qwen3-0.6B-Base student on MATH-500. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
buaacoder
force-pushed
the
feat/eopd-loss
branch
from
August 14, 2026 04:38
8bdcc2e to
d29cc3d
Compare
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- test_eopd_requires_loss_mode: test both-nonzero case (invalid) instead of advantage-mode case (valid) - test_eopd_loss.py: apply ruff formatting Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
buaacoder
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implement Entropy-aware On-Policy Distillation (EOPD) from arXiv:2603.07079 for cross-architecture teacher→student setups in Relax.
--opd-teacher-advantageflagLogitsProcessorto compute and return per-token entropy during teacher prefill--opd-eos-replacereplaces student EOS tokens with teacher-generated continuations before teacher prefillloss(multiply FKL by entropy mask) andadditive(add entropy-weighted FKL to pg loss)Experiment Results (Qwen3-8B → Qwen3-0.6B-Base, MATH-500 mean@8)
Note: Our experiments use lr=1e-6 constant schedule (paper uses 3e-6 cosine), so absolute numbers and method ordering differ. See Issue #244 for detailed analysis.
Changed Files
relax/backends/megatron/loss.py— EOPD FKL loss scaling with teacher_advantagerelax/backends/sglang/sglang_engine.py— Entropy patch integration in scheduler subprocessrelax/core/controller.py— OPD colocate GPU allocationrelax/engine/rollout/on_policy_distillation.py— Teacher entropy/EOS replace/teacher_log_probs handlingrelax/engine/rollout/sglang_rollout.py— Batch student prefill for colocate OPDrelax/utils/opd/opd_main_worker.py— Legacy topk parsing, entropy_from_topkrelax/utils/opd/opd_sglang_entropy_patch.py— SGLang LogitsProcessor entropy monkey-patchrelax/utils/opd/opd_utils.py— Teacher advantage computation, EOS replace, validationrelax/utils/utils.py— OPD colocate dp_size computationTests
tests/engine/rollout/test_on_policy_distillation_eos_replace.py— EOS replacement unit teststests/integration/test_eopd_smoke.py— End-to-end EOPD FKL loss + gradient smoke testtests/utils/test_eopd_reference_parity.py— Element-wise parity vs reference FKL formulastests/utils/test_opd_teacher_advantage.py— Teacher advantage computation testsRelated
🤖 Generated with Claude Code