feat(opd): [Task.39] EOPD cross-architecture distillation (arXiv:2603.07079)

relax/backends/megatron/actor.py

@@ -534,6 +534,7 @@ def compute_log_prob(
         num_microbatches: list[int],
         store_prefix: str = "",
         collect_topk: bool = False,
+        with_entropy: bool | None = None,
     ) -> dict[str, list[torch.Tensor]]:
         with timer(f"{store_prefix}log_probs"):
             log_prob_func = get_log_probs_and_entropy
@@ -546,6 +547,7 @@ def compute_log_prob(
                 data_iterator,
                 num_microbatches,
                 store_prefix=store_prefix,
+                with_entropy=with_entropy,
             )
 
     def _run_step_evaluation(self, rollout_id: int, *, end_update_weight: bool = False) -> None:
@@ -847,8 +849,14 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
                             num_microbatches_logprobs,
                             store_prefix="teacher_",
                             collect_topk=self.args.use_opd and self.args.opd_log_prob_top_k > 0,
+                            with_entropy=getattr(self.args, "use_eopd", False) or None,
                         )
                     )
+                    if self.args.use_opd and getattr(self.args, "opd_token_selection", "") == "teacher_topk":
+                        if "teacher_topk_token_ids" in rollout_data:
+                            rollout_data["opd_topk_token_ids"] = rollout_data["teacher_topk_token_ids"]
+                        if "teacher_topk_log_probs" in rollout_data:
+                            rollout_data["opd_topk_teacher_log_probs"] = rollout_data["teacher_topk_log_probs"]
 
                 self._switch_model("old_actor" if self.args.keep_old_actor else "actor")
                 if not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics:
@@ -1116,8 +1124,19 @@ def _hybrid_forward_subbatch(self, sub_batch: RolloutBatch) -> None:
                     os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough"
                 self._switch_model("teacher")
                 sub_batch.update(
-                    self.compute_log_prob(data_iterator_logprobs, num_microbatches_logprobs, store_prefix="teacher_")
+                    self.compute_log_prob(
+                        data_iterator_logprobs,
+                        num_microbatches_logprobs,
+                        store_prefix="teacher_",
+                        collect_topk=self.args.use_opd and self.args.opd_log_prob_top_k > 0,
+                        with_entropy=getattr(self.args, "use_eopd", False) or None,
+                    )
                 )
+                if self.args.use_opd and getattr(self.args, "opd_token_selection", "") == "teacher_topk":
+                    if "teacher_topk_token_ids" in sub_batch:
+                        sub_batch["opd_topk_token_ids"] = sub_batch["teacher_topk_token_ids"]
+                    if "teacher_topk_log_probs" in sub_batch:
+                        sub_batch["opd_topk_teacher_log_probs"] = sub_batch["teacher_topk_log_probs"]
 
             # Actor forward
             self._switch_model("old_actor" if self.args.keep_old_actor else "actor")
@@ -2175,6 +2194,7 @@ def _gather_cp_output_for_transfer_queue(self, output_dict, rollout_data):
             "ref_log_probs",
             "rollout_log_probs",
             "teacher_log_probs",
+            "teacher_entropy",
             "values",
             "advantages",
             "returns",

relax/backends/megatron/loss.py

@@ -13,6 +13,7 @@
 from relax.utils.misc import load_function
 from relax.utils.opd.opd_utils import (
     apply_opd_to_advantages,
+    compute_eopd_fkl_loss,
     compute_opd_topk_log_probs,
     compute_policy_opd_loss,
     resolve_opd_gather_topk_token_ids,
@@ -423,7 +424,9 @@ def get_log_probs_and_entropy(
 
         if with_topk:
             k = min(max(int(resolved_topk_k), 1), int(logits_chunk.size(-1)))
-            topk_token_ids_list.append(torch.topk(logits_chunk, k=k, dim=-1).indices)
+            topk_ids = torch.topk(logits_chunk, k=k, dim=-1).indices
+            topk_token_ids_list.append(topk_ids)
+            topk_log_probs_list.append(compute_opd_topk_log_probs(logits_chunk, [topk_ids], 0))
 
         if gather_topk_token_ids is not None:
             topk_log_probs_list.append(compute_opd_topk_log_probs(logits_chunk, gather_topk_token_ids, sample_idx))
@@ -435,7 +438,7 @@ def get_log_probs_and_entropy(
         res["entropy"] = entropy_list
     if with_topk:
         res["topk_token_ids"] = topk_token_ids_list
-    if gather_topk_token_ids is not None:
+    if topk_log_probs_list:
         res["topk_log_probs"] = topk_log_probs_list
 
     # we need to turn the all gather kv into zigzag ring attn kv
@@ -1090,6 +1093,17 @@ def policy_loss_function(
     if opd_loss is not None:
         loss = loss + opd_loss
 
+    eopd_fkl_loss, eopd_reported_loss = compute_eopd_fkl_loss(
+        args=args,
+        batch=batch,
+        log_probs_and_entropy=log_probs_and_entropy,
+    )
+    if eopd_fkl_loss is not None:
+        if getattr(args, "opd_teacher_advantage", False):
+            num_samples = len(batch["loss_masks"])
+            eopd_fkl_loss = eopd_fkl_loss * num_samples
+        loss = loss + eopd_fkl_loss
+
     if log_probs.numel() == 0:
         loss += 0 * logits.sum()
 
@@ -1130,6 +1144,7 @@ def policy_loss_function(
         reported_loss["kl_loss"] = kl_loss.clone().detach()
 
     reported_loss.update(opd_reported_loss)
+    reported_loss.update(eopd_reported_loss)
 
     if args.get_mismatch_metrics or args.use_tis:
         # Aggregate mismatch/TIS/RS related metrics with the *pre-RS* masks.

relax/backends/megatron/model.py

@@ -35,7 +35,7 @@
 from relax.utils.megatron_bridge_utils import patch_megatron_model
 from relax.utils.megatron_peft_utils import is_lora_enabled
 from relax.utils.memory_utils import clear_memory
-from relax.utils.opd.opd_utils import consume_opd_train_data
+from relax.utils.opd.opd_utils import consume_opd_train_data, get_megatron_opd_batch_keys
 from relax.utils.timer import timer
 from relax.utils.training.ppo_utils import (
     install_critic_value_head_runtime_check,
@@ -705,6 +705,7 @@ def forward_only(
     num_microbatches: Sequence[int],
     store_prefix: str = "",
     per_sample_output: bool = True,
+    with_entropy: bool | None = None,
 ) -> dict[str, list[torch.Tensor]]:
     """Run forward passes only and collect non-loss outputs (e.g., logprobs).
 
@@ -842,7 +843,7 @@ def forward_step(
             unconcat_tokens=unconcat_tokens,
             total_lengths=total_lengths,
             response_lengths=response_lengths,
-            with_entropy=args.use_rollout_entropy,
+            with_entropy=with_entropy if with_entropy is not None else args.use_rollout_entropy,
             max_seq_lens=batch.get("max_seq_lens", None),
             padded_total_lengths=batch.get("padded_total_lengths", None),
             loss_masks=batch.get("loss_masks", None),
@@ -1013,6 +1014,7 @@ def forward_step(
             _opd_keys: list[str] = []
             if args.use_opd:
                 consume_opd_train_data(_opd_keys, args)
+                _opd_keys.extend(get_megatron_opd_batch_keys(args))
             batch = get_batch(
                 data_iterator,
                 [

relax/backends/sglang/sglang_engine.py

@@ -162,31 +162,43 @@ def _to_local_gpu_id(physical_gpu_id: int) -> int:
 
 
 def _patched_run_scheduler_process(*args, **kwargs):
-    """Scheduler-subprocess entry used for the routing-replay path.
+    """Scheduler-subprocess entry that applies env-gated patches.
 
-    This wrapper is only installed when ``--optimize-routing-replay`` is
-    enabled (see ``_launch_server_with_patches``), so the routing-replay async
-    D→H patch is applied **unconditionally** here, preserving the original
-    behavior.
+    Patches that target ``LogitsProcessor`` or other scheduler-resident objects
+    must be applied here — the scheduler is a separate ``mp.Process`` and does
+    NOT inherit monkey-patches from the http-server process.
     """
-    from relax.backends.sglang.routing_replay_patch import apply_patch
+    if os.environ.get("RELAX_OPTIMIZE_ROUTING_REPLAY", "0") == "1":
+        from relax.backends.sglang.routing_replay_patch import apply_patch
 
-    apply_patch()
+        apply_patch()
+
+    if os.environ.get("RELAX_OPD_ENTROPY_PATCH", "0") == "1":
+        from relax.utils.opd.opd_sglang_entropy_patch import apply_opd_entropy_patch
+
+        apply_opd_entropy_patch()
 
     from sglang.srt.managers.scheduler import run_scheduler_process
 
     return run_scheduler_process(*args, **kwargs)
 
 
+def _needs_scheduler_patches() -> bool:
+    return (
+        os.environ.get("RELAX_OPTIMIZE_ROUTING_REPLAY", "0") == "1"
+        or os.environ.get("RELAX_OPD_ENTROPY_PATCH", "0") == "1"
+    )
+
+
 def _launch_server_with_patches(server_args: ServerArgs):
     """Top-level picklable ``multiprocessing.Process`` target that applies the
     SGLang patches, each gated by its own env flag so any combination is valid:
 
     - main process: OPD pre-expanded multimodal patch
       (``RELAX_OPD_PREEXPANDED_PATCH=1``).
     - scheduler subprocess: routing-replay (``RELAX_OPTIMIZE_ROUTING_REPLAY=1``)
-      installs ``_patched_run_scheduler_process``, which applies the
-      routing-replay patch unconditionally.
+      and entropy patch (``RELAX_OPD_ENTROPY_PATCH=1``) are applied via
+      ``_patched_run_scheduler_process``.
     """
     from sglang.srt.entrypoints.http_server import launch_server
 
@@ -195,7 +207,7 @@ def _launch_server_with_patches(server_args: ServerArgs):
 
         apply_opd_preexpanded_patch()
 
-    if Envs.RELAX_OPTIMIZE_ROUTING_REPLAY:
+    if _needs_scheduler_patches():
         launch_server(server_args, run_scheduler_process_func=_patched_run_scheduler_process)
     else:
         launch_server(server_args)
@@ -243,9 +255,11 @@ def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
     optimize = Envs.RELAX_OPTIMIZE_ROUTING_REPLAY
     opd_patch = Envs.RELAX_OPD_PREEXPANDED_PATCH
     per_pos = Envs.RELAX_OPD_PER_POS_TOKEN_IDS
+    entropy_patch = os.environ.get("RELAX_OPD_ENTROPY_PATCH", "0") == "1"
     logger.info(
         "Launching SGLang server with independently-gated patches: "
-        f"routing_replay={optimize}, opd_preexpanded={opd_patch}, per_pos_token_ids={per_pos}"
+        f"routing_replay={optimize}, opd_preexpanded={opd_patch}, "
+        f"per_pos_token_ids={per_pos}, entropy={entropy_patch}"
     )
 
     p = multiprocessing.Process(target=_launch_server_with_patches, args=(server_args,))

relax/core/controller.py

@@ -43,6 +43,7 @@
 from relax.utils.logging_utils import get_logger
 from relax.utils.misc import load_function
 from relax.utils.opd.opd_utils import (
+    is_managed_opd_teacher_colocate,
     maybe_start_managed_opd_teacher,
     set_managed_opd_teacher_on_actor_service,
     shutdown_managed_opd_teacher,
@@ -518,6 +519,9 @@ def register_all_serve(self):
                 continue
             num_serves, num_gpus = self.config.resource.get(role)
             assert num_serves == 1, f"Currently only support num_serves=1 for {role}, but received {num_serves=}"
+            if str(role) == "actor" and is_managed_opd_teacher_colocate(self.config):
+                _, rollout_gpus = self.config.resource.get("rollout", (1, num_gpus))
+                num_gpus = rollout_gpus
             self._health_manager.mark_healthy(role)
             logger.info(f"Service {role} start creating.")
 

relax/distributed/ray/teacher_manager.py

@@ -1,6 +1,8 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
 
+import os
+
 import ray
 from ray.util.placement_group import remove_placement_group
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
@@ -40,6 +42,11 @@ def _build_teacher_engine_env(args) -> dict[str, str]:
         # passed through from the driver. RELAX_OPD_PREEXPANDED_PATCH affects the
         # teacher engine only; RELAX_OPD_PER_POS_TOKEN_IDS affects teacher + student.
         "RELAX_OPD_PREEXPANDED_PATCH": str(int(Envs.RELAX_OPD_PREEXPANDED_PATCH)),
+        "RELAX_OPD_ENTROPY_PATCH": (
+            "1"
+            if getattr(args, "use_eopd", False) and getattr(args, "opd_type", "") == "sglang"
+            else os.environ.get("RELAX_OPD_ENTROPY_PATCH", "0")
+        ),
         "RELAX_OPD_PER_POS_TOKEN_IDS": str(int(Envs.RELAX_OPD_PER_POS_TOKEN_IDS)),
         "RELAX_OPD_TOKEN_IDS_LOGPROB_K": Envs.RELAX_OPD_TOKEN_IDS_LOGPROB_K,
         "SGL_JIT_DEEPGEMM_PRECOMPILE": "false",