diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..7b835bb0f 100644 --- a/relax/backends/megatron/actor.py +++ b/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", diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 372a139cf..93d2b6a85 100644 --- a/relax/backends/megatron/loss.py +++ b/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. diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 077e461d9..7b2ee5a5f 100644 --- a/relax/backends/megatron/model.py +++ b/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, [ diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index 15ca1d172..bdf3daaae 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -162,22 +162,34 @@ 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: @@ -185,8 +197,8 @@ def _launch_server_with_patches(server_args: ServerArgs): - 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,)) diff --git a/relax/core/controller.py b/relax/core/controller.py index e59f5cfd7..36b4d18c7 100644 --- a/relax/core/controller.py +++ b/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.") diff --git a/relax/distributed/ray/teacher_manager.py b/relax/distributed/ray/teacher_manager.py index b39702fbc..8f57a1e53 100644 --- a/relax/distributed/ray/teacher_manager.py +++ b/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", diff --git a/relax/engine/rollout/on_policy_distillation.py b/relax/engine/rollout/on_policy_distillation.py index b888b4210..893271a7b 100644 --- a/relax/engine/rollout/on_policy_distillation.py +++ b/relax/engine/rollout/on_policy_distillation.py @@ -117,17 +117,22 @@ def is_opsd(self) -> bool: return self.opsd_worker is not None def schema_opd_transfer_data(self) -> list[str]: + eopd = getattr(self.args, "use_eopd", False) fields: list[str] = [] if self.topk_worker is not None: - fields.extend(self.topk_worker.topk_transfer_fields()) + fields.extend(self.topk_worker.topk_transfer_fields(eopd=eopd)) + if self.topk_worker.is_advantage or getattr(self.args, "opd_teacher_advantage", False): + fields.append(opd_main_worker.SampledTokenWorker.TRANSFER_TEACHER_LOG_PROBS) if self.sampled_worker is not None: - fields.extend(self.sampled_worker.sampled_transfer_fields()) + fields.extend(self.sampled_worker.sampled_transfer_fields(eopd=eopd)) return fields def produce_opd_transfer_data(self, samples: list[Sample], train_data: dict) -> None: if self.topk_worker is not None: + schema_fields = set(self.topk_worker.topk_transfer_fields(eopd=getattr(self.args, "use_eopd", False))) for field_name in opd_main_worker.TopkWorker.TRANSFER_FIELDS: - if not any(getattr(s, field_name, None) is not None for s in samples): + has_any = any(getattr(s, field_name, None) is not None for s in samples) + if not has_any and field_name not in schema_fields: continue flat: list = [] for s in samples: @@ -142,11 +147,20 @@ def produce_opd_transfer_data(self, samples: list[Sample], train_data: dict) -> train_data[kl_field] = [ getattr(s, kl_field).tolist() if getattr(s, kl_field, None) is not None else [] for s in samples ] + if self.topk_worker.is_advantage or getattr(self.args, "opd_teacher_advantage", False): + train_data[opd_main_worker.SampledTokenWorker.TRANSFER_TEACHER_LOG_PROBS] = [ + s.teacher_log_probs if s.teacher_log_probs is not None else [] for s in samples + ] elif self.sampled_worker is not None: train_data[opd_main_worker.SampledTokenWorker.TRANSFER_TEACHER_LOG_PROBS] = [ s.teacher_log_probs if s.teacher_log_probs is not None else [] for s in samples ] + if getattr(self.args, "use_eopd", False): + train_data["teacher_entropy"] = [ + s.teacher_entropy if s.teacher_entropy is not None else [] for s in samples + ] + def before_rollout(self, payload: dict) -> None: if self.topk_worker is None: return @@ -193,6 +207,8 @@ async def prefill( self, samples: Sample | Sequence[Sample], encode_multimodal_inputs: EncodeMultimodalInputs | None = None, + *, + include_student: bool = True, ) -> None: sample_list = list(samples) if isinstance(samples, Sequence) else [samples] @@ -203,13 +219,66 @@ async def prefill( fetch_results = await asyncio.gather(*[self._teacher_prefill(s, session) for s in sample_list]) self._raise_if_all_failed(sample_list, fetch_results) - if self.topk_worker is not None and self.topk_worker.spec.student_at_teacher: + if include_student and self.topk_worker is not None and self.topk_worker.spec.student_at_teacher: await asyncio.gather( *[self._student_prefill(s, session, encode_multimodal_inputs) for s in sample_list] ) self._assemble_transfer(sample_list) + async def _resolve_direct_engine_urls(self) -> list[str]: + """Query the SGLang router for direct worker URLs to bypass the router + for heavy token_ids_logprob requests (which block the engine event loop + and cause the router's health-check circuit breaker to trip).""" + router_base = f"http://{self.args.sglang_router_ip}:{self.args.sglang_router_port}" + try: + import sglang_router + from packaging.version import parse + + if parse(sglang_router.__version__) <= parse("0.2.1") or getattr(self.args, "use_slime_router", False): + async with aiohttp.ClientSession() as s: + async with s.get(f"{router_base}/list_workers", timeout=aiohttp.ClientTimeout(total=5)) as r: + data = await r.json() + return data.get("urls", []) + else: + async with aiohttp.ClientSession() as s: + async with s.get(f"{router_base}/workers", timeout=aiohttp.ClientTimeout(total=5)) as r: + data = await r.json() + return [w["url"] for w in data.get("workers", [])] + except Exception as e: + logger.warning("Failed to resolve direct engine URLs from router: %s", e) + return [] + + async def prefill_student( + self, + samples: Sample | Sequence[Sample], + encode_multimodal_inputs: EncodeMultimodalInputs | None = None, + max_concurrent: int = 4, + ) -> None: + """Run student-at-teacher prefill only. + + Call after all rollout generation is complete to avoid sending + token_ids_logprob requests to the student engine while it is still + decoding rollout responses. + """ + if self.topk_worker is None or not self.topk_worker.spec.student_at_teacher: + return + sample_list = list(samples) if isinstance(samples, Sequence) else [samples] + sem = asyncio.Semaphore(max_concurrent) + + direct_urls = await self._resolve_direct_engine_urls() + student_url_override = f"{direct_urls[0]}/generate" if direct_urls else None + if student_url_override: + logger.info("[OPD] Bypassing router for student prefill: %s", student_url_override) + + async def _limited(s: Sample, session: aiohttp.ClientSession) -> None: + async with sem: + await self._student_prefill(s, session, encode_multimodal_inputs, student_url_override) + + async with _create_teacher_client_session(self.args) as session: + await asyncio.gather(*[_limited(s, session) for s in sample_list]) + self._assemble_transfer(sample_list) + async def _post_logprob( self, session: aiohttp.ClientSession, @@ -252,6 +321,15 @@ async def _teacher_prefill(self, sample: Sample, session: aiohttp.ClientSession) image_data = None teacher_input_ids = sample.rollout_tokens or sample.tokens prompt_length = len(sample.tokens) - response_length + + if getattr(self.args, "opd_eos_replace", False): + _ENDOFTEXT, _IM_END = 151643, 151645 + teacher_input_ids = list(teacher_input_ids) + for j in range(prompt_length, len(teacher_input_ids)): + if teacher_input_ids[j] == _ENDOFTEXT: + teacher_input_ids[j] = _IM_END + break + logprob_start_len = max(prompt_length - 1, 0) mm_fields = {"image_data": image_data} if image_data is not None else None @@ -267,7 +345,8 @@ async def _teacher_prefill(self, sample: Sample, session: aiohttp.ClientSession) payload = opd_main_worker.build_prefill_payload_base(teacher_input_ids, logprob_start_len) if mm_fields: payload.update(mm_fields) - + if getattr(self.args, "use_eopd", False): + payload["sampling_params"]["max_new_tokens"] = 1 teacher_url = _pick_teacher_url(self.args, sample) resp_obj = await self._post_logprob(session, teacher_url, payload, sample, "teacher prefill") if resp_obj is None: @@ -290,7 +369,12 @@ async def _teacher_prefill(self, sample: Sample, session: aiohttp.ClientSession) ) return False - if self.sampled_worker is not None: + needs_teacher_lp = ( + self.sampled_worker is not None + or (self.topk_worker is not None and self.topk_worker.is_advantage) + or getattr(self.args, "opd_teacher_advantage", False) + ) + if needs_teacher_lp: sample.teacher_log_probs = [float(v) for v in token_logprobs[1 : 1 + response_length]] if self.topk_worker is not None: @@ -303,10 +387,36 @@ async def _teacher_prefill(self, sample: Sample, session: aiohttp.ClientSession) resp_obj, response_length ) + if getattr(self.args, "use_eopd", False): + entropy = resp_obj.entropy_1d() + if entropy is not None and len(entropy) >= response_length + 1: + sample.teacher_entropy = [float(v) for v in entropy[1 : 1 + response_length]] + elif entropy is not None and len(entropy) == response_length: + sample.teacher_entropy = [float(v) for v in entropy] + else: + top_k = int(getattr(self.args, "opd_log_prob_top_k", 0) or 0) + vocab_size = int(getattr(self.args, "opd_teacher_vocab_size", 151936) or 151936) + entropy_topk = resp_obj.entropy_from_topk(top_k, vocab_size, response_length) + if entropy_topk is not None and len(entropy_topk) == response_length: + sample.teacher_entropy = [float(v) for v in entropy_topk] + else: + logger.warning( + "EOPD entropy missing for sample_index=%s (got %s elements, need %d). " + "Falling back to zeros (EOPD gate disabled for this sample).", + getattr(sample, "index", None), + len(entropy) if entropy is not None else "None", + response_length, + ) + sample.teacher_entropy = [0.0] * response_length + return True async def _student_prefill( - self, sample: Sample, session: aiohttp.ClientSession, encode_mm_fn: EncodeMultimodalInputs | None + self, + sample: Sample, + session: aiohttp.ClientSession, + encode_mm_fn: EncodeMultimodalInputs | None, + url_override: str | None = None, ) -> None: from relax.utils.opd.opd_utils import build_student_preexpanded_image_data @@ -335,7 +445,7 @@ async def _student_prefill( mm_fields=mm_fields, ) - student_url = f"http://{self.args.sglang_router_ip}:{self.args.sglang_router_port}/generate" + student_url = url_override or f"http://{self.args.sglang_router_ip}:{self.args.sglang_router_port}/generate" resp_obj = await self._post_logprob(session, student_url, payload, sample, "student-at-teacher-topk") if resp_obj is None: return diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py index 0a0149cda..1d04dc68a 100644 --- a/relax/engine/rollout/sglang_rollout.py +++ b/relax/engine/rollout/sglang_rollout.py @@ -598,7 +598,7 @@ async def generate_and_rm( sample.reward = await async_rm(args, sample) if state.opd_manager and not evaluation: - await state.opd_manager.prefill(sample, _encode_multimodal_inputs) + await state.opd_manager.prefill(sample, _encode_multimodal_inputs, include_student=False) return sample @@ -932,6 +932,22 @@ def target_reached() -> bool: if args.fully_async else args.rollout_batch_size ) # Samples per batch to transfer + + # In non-fully_async (colocate) mode the SGLang engine is idle once + # the batch is full — safe to run student-at-teacher top-K prefill + # before the transfer task serialises sample attributes. + if not args.fully_async and state.opd_manager and len(batch_to_transfer) >= transfer_batch_size: + _pf_samples: list[Sample] = [] + for _g in batch_to_transfer: + if isinstance(_g[0], list): + for _sub in _g: + _pf_samples.extend(_sub) + else: + _pf_samples.extend(_g) + if _pf_samples: + logger.info(f"[OPD] Batch student prefill: {len(_pf_samples)} samples") + await state.opd_manager.prefill_student(_pf_samples, _encode_multimodal_inputs) + # in fully async mode, we transfer all remaining samples when we reach the target size if len(batch_to_transfer) >= transfer_batch_size: if total_transfer_samples <= num_old_samples: @@ -998,6 +1014,18 @@ def target_reached() -> bool: ) if len(batch_to_transfer) > 0: + if not args.fully_async and state.opd_manager: + _pf_samples: list[Sample] = [] + for _g in batch_to_transfer: + if isinstance(_g[0], list): + for _sub in _g: + _pf_samples.extend(_sub) + else: + _pf_samples.extend(_g) + if _pf_samples: + logger.info(f"[OPD] Batch student prefill (leftover): {len(_pf_samples)} samples") + await state.opd_manager.prefill_student(_pf_samples, _encode_multimodal_inputs) + n = len(batch_to_transfer) if is_final_backfill: prev_is_last = args.fully_async and (committed_prev + n >= prev_target) diff --git a/relax/utils/opd/opd_main_worker.py b/relax/utils/opd/opd_main_worker.py index 3d1843828..147f347f4 100644 --- a/relax/utils/opd/opd_main_worker.py +++ b/relax/utils/opd/opd_main_worker.py @@ -40,16 +40,40 @@ def _decode_topk_2d(self, prefix: str, response_length: int | None, top_k: int): return None val = self._b64_decode(self.meta.get(f"{prefix}_val_b64"), "float32") n = val.size // top_k - if n <= 0: + if n > 0: + if response_length is None: + response_length = n + if response_length <= 0 or n < response_length: + return None + take = response_length * top_k + lps = val[-take:].reshape(response_length, top_k) + idx = self._b64_decode(self.meta.get(f"{prefix}_idx_b64"), "int32") + ids = idx[-take:].reshape(response_length, top_k) if idx.size >= take else None + return ids, lps + legacy = self.meta.get(prefix) + if legacy: + return self._parse_legacy_topk(legacy, response_length, top_k) + return None + + @staticmethod + def _parse_legacy_topk( + entries: list, response_length: int | None, top_k: int + ) -> tuple[np.ndarray, np.ndarray] | None: + if not entries: return None + n = len(entries) if response_length is None: response_length = n if response_length <= 0 or n < response_length: return None - take = response_length * top_k - lps = val[-take:].reshape(response_length, top_k) - idx = self._b64_decode(self.meta.get(f"{prefix}_idx_b64"), "int32") - ids = idx[-take:].reshape(response_length, top_k) if idx.size >= take else None + ids = np.zeros((response_length, top_k), dtype=np.int32) + lps = np.full((response_length, top_k), -1e9, dtype=np.float32) + for i, pos in enumerate(entries[-response_length:]): + if pos is None: + continue + for j, item in enumerate(pos[:top_k]): + lps[i, j] = item[0] + ids[i, j] = item[1] return ids, lps def base_logprobs_1d(self) -> np.ndarray | None: @@ -72,6 +96,41 @@ def other_topk(self, response_length: int, top_k: int): pair = self._decode_topk_2d("input_token_ids_logprobs", response_length, top_k) return pair[1] if pair is not None else None + def entropy_1d(self) -> np.ndarray | None: + """Decode per-token entropy injected by the EOPD entropy patch.""" + raw = self.meta.get("relax_input_entropy_b64") + if not raw: + _ci = self.meta.get("customized_info") + if _ci and isinstance(_ci, dict) and "relax_input_entropy_b64" in _ci: + raw = _ci["relax_input_entropy_b64"] + if not raw: + return None + if isinstance(raw, list): + arrays = [self._b64_decode(s, "float32") for s in raw if s] + val = ( + np.concatenate(arrays) + if len(arrays) > 1 + else (arrays[0] if arrays else np.array([], dtype=np.float32)) + ) + else: + val = self._b64_decode(raw, "float32") + return val if val.size else None + + def entropy_from_topk(self, top_k: int, vocab_size: int, response_length: int | None = None) -> np.ndarray | None: + """Estimate per-token entropy from top-K logprobs (no patch needed).""" + pair = self._decode_topk_2d("input_top_logprobs", response_length, top_k) + if pair is None: + return None + _, logps = pair + probs = np.exp(logps.astype(np.float64)) + topk_mass = probs.sum(axis=-1) + topk_entropy = -(probs * logps).sum(axis=-1) + remaining_mass = np.maximum(1.0 - topk_mass, 1e-30) + remaining_count = max(vocab_size - top_k, 1) + remaining_per_token = remaining_mass / remaining_count + remaining_entropy = -remaining_mass * np.log(remaining_per_token + 1e-30) + return (topk_entropy + remaining_entropy).astype(np.float32) + def build_prefill_payload_base(input_ids: list[int], logprob_start_len: int) -> dict: return { @@ -85,18 +144,23 @@ def build_prefill_payload_base(input_ids: list[int], logprob_start_len: int) -> class SampledTokenWorker: TRANSFER_TEACHER_LOG_PROBS = "teacher_log_probs" TRANSFER_STUDENT_LOG_PROBS = "rollout_log_probs" + TRANSFER_TEACHER_ENTROPY = "teacher_entropy" @classmethod def from_args(cls, args) -> "SampledTokenWorker": return cls() - def sampled_transfer_fields(self) -> list[str]: - return [self.TRANSFER_TEACHER_LOG_PROBS, self.TRANSFER_STUDENT_LOG_PROBS] + def sampled_transfer_fields(self, eopd: bool = False) -> list[str]: + fields = [self.TRANSFER_TEACHER_LOG_PROBS, self.TRANSFER_STUDENT_LOG_PROBS] + if eopd: + fields.append(self.TRANSFER_TEACHER_ENTROPY) + return fields class TopkWorker: TRANSFER_TOKEN_IDS = "opd_topk_token_ids" TRANSFER_TEACHER_LOG_PROBS = "opd_topk_teacher_log_probs" + TRANSFER_TEACHER_ENTROPY = "teacher_entropy" # only as_adv TRANSFER_STUDENT_LOG_PROBS = "opd_topk_student_log_probs" # only union @@ -269,12 +333,14 @@ def _merge_union(self, student_self, teacher_self, teacher_at_student_lp, studen out[self.TRANSFER_K_LENGTHS] = k_lengths return out - def topk_transfer_fields(self) -> list[str]: + def topk_transfer_fields(self, eopd: bool = False) -> list[str]: fields: list[str] = [self.TRANSFER_TOKEN_IDS, self.TRANSFER_TEACHER_LOG_PROBS] if self.is_advantage: fields.append(self.TRANSFER_STUDENT_LOG_PROBS) if self.spec.name == "union": fields.append(self.TRANSFER_K_LENGTHS) + if eopd: + fields.append(self.TRANSFER_TEACHER_ENTROPY) return fields diff --git a/relax/utils/opd/opd_sglang_entropy_patch.py b/relax/utils/opd/opd_sglang_entropy_patch.py new file mode 100644 index 000000000..42575f31d --- /dev/null +++ b/relax/utils/opd/opd_sglang_entropy_patch.py @@ -0,0 +1,132 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Monkey-patch for SGLang LogitsProcessor: compute per-token entropy during +teacher prefill and return it via ``customized_info`` → ``meta_info``. + +Strategy (v7): Patch ``_get_logits`` to compute entropy from the *actual* model +logits immediately after the TP all-gather. Only inject entropy into +``customized_info`` during EXTEND mode with logprob return (i.e. when +``extend_logprob_pruned_lens_cpu`` is a list). Skip decode-mode and non- +logprob extend calls entirely to avoid broken ``customized_info`` routing. +""" + +from __future__ import annotations + +import logging +import os + +import pybase64 +import torch +import torch.nn.functional as F + + +logger = logging.getLogger("opd_sglang_entropy_patch") + +_PATCH_FLAG = "_relax_entropy_patched_v7" + + +def _compute_entropy(logits: torch.Tensor, chunk_size: int = 256) -> torch.Tensor: + with torch.no_grad(): + parts = [] + for i in range(0, logits.shape[0], chunk_size): + chunk = logits[i : i + chunk_size].float() + log_p = F.log_softmax(chunk, dim=-1) + parts.append(-(log_p.exp() * log_p).sum(dim=-1)) + return torch.cat(parts) if len(parts) > 1 else parts[0] + + +def _entropy_to_b64(entropy: torch.Tensor) -> str: + return pybase64.b64encode(entropy.cpu().to(torch.float32).numpy().tobytes()).decode("ascii") + + +def apply_patch() -> bool: + from sglang.srt.layers.logits_processor import LogitsProcessor + + if getattr(LogitsProcessor, _PATCH_FLAG, False): + return False + + _orig_get_logits = LogitsProcessor._get_logits + _orig_forward = LogitsProcessor.forward + + def _patched_get_logits(self, hidden_states, lm_head, logits_metadata, *args, **kwargs): + logits = _orig_get_logits(self, hidden_states, lm_head, logits_metadata, *args, **kwargs) + + if getattr(self, "_relax_collecting_entropy", False): + if torch.cuda.is_current_stream_capturing(): + return logits + try: + ent = _compute_entropy(logits) + if not hasattr(self, "_relax_entropy_parts"): + self._relax_entropy_parts = [] + self._relax_entropy_parts.append(ent.cpu()) + except Exception: + logger.warning("entropy computation in _get_logits failed", exc_info=True) + + return logits + + def _patched_forward(self, input_ids, hidden_states, lm_head, logits_metadata, *args, **kwargs): + self._relax_collecting_entropy = True + self._relax_entropy_parts = [] + + output = _orig_forward(self, input_ids, hidden_states, lm_head, logits_metadata, *args, **kwargs) + + self._relax_collecting_entropy = False + + try: + if not self._relax_entropy_parts: + return output + + all_entropy = ( + torch.cat(self._relax_entropy_parts) + if len(self._relax_entropy_parts) > 1 + else self._relax_entropy_parts[0] + ) + + from sglang.srt.layers.logits_processor import LogitsMetadata + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + if isinstance(logits_metadata, ForwardBatch): + lm = LogitsMetadata.from_forward_batch(logits_metadata) + else: + lm = logits_metadata + + pruned_lens = getattr(lm, "extend_logprob_pruned_lens_cpu", None) + + if not isinstance(pruned_lens, list) or not pruned_lens: + self._relax_entropy_parts = [] + return output + + total_input = sum(pruned_lens) + if total_input > all_entropy.shape[0]: + self._relax_entropy_parts = [] + return output + + input_entropy = all_entropy[-total_input:] + per_request = torch.split(input_entropy, list(pruned_lens)) + b64_list = [_entropy_to_b64(e) for e in per_request] + + if output.customized_info is None: + output.customized_info = {} + output.customized_info["relax_input_entropy_b64"] = b64_list + except Exception: + logger.warning("entropy injection in forward failed", exc_info=True) + + self._relax_entropy_parts = [] + return output + + LogitsProcessor._get_logits = _patched_get_logits + LogitsProcessor.forward = _patched_forward + setattr(LogitsProcessor, _PATCH_FLAG, True) + logger.info("EOPD entropy patch applied to LogitsProcessor._get_logits + .forward") + return True + + +def apply_opd_entropy_patch() -> None: + try: + apply_patch() + except Exception as e: + logger.warning("Failed to apply OPD entropy patch: %r", e) + + +if os.environ.get("RELAX_OPD_ENTROPY_PATCH", "0") == "1": + apply_opd_entropy_patch() diff --git a/relax/utils/opd/opd_utils.py b/relax/utils/opd/opd_utils.py index ce052263a..05c9ea850 100644 --- a/relax/utils/opd/opd_utils.py +++ b/relax/utils/opd/opd_utils.py @@ -33,7 +33,7 @@ } ) -OPD_CP_FLOAT_FIELDS = ("teacher_log_probs",) +OPD_CP_FLOAT_FIELDS = ("teacher_log_probs", "teacher_entropy") def iter_opd_cp_float_fields() -> tuple[str, ...]: @@ -143,6 +143,8 @@ def build_teacher_overrides(args: Any, colocate_sync: bool = False) -> dict[str, overrides["model_path"] = args.teacher_hf_checkpoint overrides.setdefault("load_format", "auto") overrides.setdefault("enable_memory_saver", colocate_sync) + if colocate_sync: + overrides.setdefault("enable_weights_cpu_backup", True) return overrides @@ -677,6 +679,64 @@ def add_opd_arguments(parser: Any) -> Any: "Set to None (default) to disable." ), ) + + # EOPD (Entropy-aware On-Policy Distillation) arguments + parser.add_argument( + "--use-eopd", + action="store_true", + default=False, + help=( + "Enable Entropy-aware OPD (EOPD). Adds a forward KL loss at " + "high-entropy teacher positions. Requires --use-opd and --opd-type=megatron." + ), + ) + parser.add_argument( + "--eopd-entropy-threshold", + type=float, + default=0.8, + help="Teacher entropy threshold (nats) for EOPD gating. Default 0.8.", + ) + parser.add_argument( + "--eopd-fkl-coef", + type=float, + default=1.0, + help="Coefficient for EOPD forward KL loss. Default 1.0.", + ) + parser.add_argument( + "--eopd-fkl-top-k", + type=int, + default=0, + help="Top-K for EOPD forward KL. 0 (default) uses --opd-log-prob-top-k.", + ) + parser.add_argument( + "--opd-teacher-advantage", + action="store_true", + default=False, + help=( + "Paper-mode EOPD: use teacher_advantage = (teacher_lp - old_student_lp).detach() " + "directly as the PPO advantage, replacing GRPO advantages entirely. " + "Set --opd-kl-coef=0 --opd-loss-coef=0 when using this mode." + ), + ) + parser.add_argument( + "--opd-teacher-advantage-additive", + action="store_true", + default=False, + help=( + "When used with --opd-teacher-advantage, ADD teacher advantage to GRPO " + "advantages instead of replacing them. Requires NOT using --opd-only-reward." + ), + ) + parser.add_argument( + "--opd-eos-replace", + action="store_true", + default=False, + help=( + "Replace <|endoftext|> (151643) with <|im_end|> (151645) in response " + "tokens before teacher prefill. Needed for Qwen3 cross-architecture " + "distillation where student uses endoftext but teacher expects im_end." + ), + ) return parser @@ -715,16 +775,24 @@ def validate_opd_args(args: Namespace, *, is_sft: bool, log: Any = logger) -> No opd_kl_coef = float(getattr(args, "opd_kl_coef", 0.0) or 0.0) opd_loss_coef = float(getattr(args, "opd_loss_coef", 0.0) or 0.0) + teacher_advantage = getattr(args, "opd_teacher_advantage", False) - is_adv_mode = opd_kl_coef != 0.0 and opd_loss_coef == 0.0 - is_loss_mode = opd_kl_coef == 0.0 and opd_loss_coef != 0.0 - if not is_adv_mode and not is_loss_mode: - raise ValueError( - "Exactly one of --opd-kl-coef / --opd-loss-coef must be non-zero. " - f"Got opd_kl_coef={opd_kl_coef}, opd_loss_coef={opd_loss_coef}. " - "Use --opd-kl-coef=X --opd-loss-coef=0.0 for advantage mode, or " - "--opd-kl-coef=0.0 --opd-loss-coef=X for loss mode." - ) + if teacher_advantage: + if opd_kl_coef != 0.0 or opd_loss_coef != 0.0: + raise ValueError( + "--opd-teacher-advantage replaces GRPO advantages with teacher_advantage. " + "Set --opd-kl-coef=0.0 --opd-loss-coef=0.0 when using this mode." + ) + else: + is_adv_mode = opd_kl_coef != 0.0 and opd_loss_coef == 0.0 + is_loss_mode = opd_kl_coef == 0.0 and opd_loss_coef != 0.0 + if not is_adv_mode and not is_loss_mode: + raise ValueError( + "Exactly one of --opd-kl-coef / --opd-loss-coef must be non-zero. " + f"Got opd_kl_coef={opd_kl_coef}, opd_loss_coef={opd_loss_coef}. " + "Use --opd-kl-coef=X --opd-loss-coef=0.0 for advantage mode, or " + "--opd-kl-coef=0.0 --opd-loss-coef=X for loss mode." + ) if getattr(args, "opd_teacher_prompt_key", None) is not None: if args.opd_type != "sglang": @@ -827,6 +895,21 @@ def validate_opd_args(args: Namespace, *, is_sft: bool, log: Any = logger) -> No "(single) / --opd-teacher-routes (multi) with a 'teacher' entry in --resource." ) + # EOPD validation + if getattr(args, "use_eopd", False): + if args.opd_type not in ("megatron", "sglang"): + raise ValueError("--use-eopd requires --opd-type=megatron or --opd-type=sglang.") + eopd_top_k = int(getattr(args, "eopd_fkl_top_k", 0) or 0) + base_top_k = int(getattr(args, "opd_log_prob_top_k", 0) or 0) + if eopd_top_k <= 0 and base_top_k <= 0: + raise ValueError("--use-eopd requires --eopd-fkl-top-k > 0 or --opd-log-prob-top-k > 0.") + log.info( + "EOPD enabled: threshold=%.3f nats, fkl_coef=%.3f, fkl_top_k=%d", + args.eopd_entropy_threshold, + args.eopd_fkl_coef, + eopd_top_k if eopd_top_k > 0 else base_top_k, + ) + # ============================================================================ # Multimodal image encoding helpers (raw base64 PNG for opd_preexpanded_raw) @@ -916,6 +999,19 @@ def consume_opd_train_data(data_fields: list[str], args: Namespace) -> None: data_fields.extend(_get_opd_transfer_schema(args)) +def get_megatron_opd_batch_keys(args: Namespace) -> list[str]: + """Return extra batch keys needed for megatron OPD training forward.""" + if not (getattr(args, "use_opd", False) and getattr(args, "opd_type", None) == "megatron"): + return [] + keys = ["teacher_log_probs"] + token_selection = getattr(args, "opd_token_selection", "student_sampled") + if token_selection in ("student_topk", "teacher_topk", "union"): + keys.extend(["opd_topk_token_ids", "opd_topk_teacher_log_probs"]) + if getattr(args, "use_eopd", False): + keys.append("teacher_entropy") + return keys + + def consume_opd_advantage_data(data_fields: list[str], args: Namespace) -> None: if not getattr(args, "use_opd", False): return @@ -1269,11 +1365,40 @@ def _opd_compute_per_token_signal( ).to(dtype=student_lp_1d.dtype) +def _apply_teacher_advantage( + rollout_data: RolloutBatch, + advantages: list[torch.Tensor], + additive: bool = False, +) -> None: + """Replace or add teacher_advantage = (teacher_lp - old_student_lp).detach() to PPO advantages.""" + teacher_log_probs = rollout_data.get("teacher_log_probs") + student_log_probs = rollout_data.get("rollout_log_probs") + if teacher_log_probs is None or student_log_probs is None: + return + for i, adv in enumerate(advantages): + t_lp = teacher_log_probs[i].to(device=adv.device) + s_lp = student_log_probs[i].to(device=adv.device) + if t_lp.numel() == 0 or s_lp.numel() == 0: + continue + teacher_adv = (t_lp - s_lp).detach() + if additive: + advantages[i] = adv + teacher_adv + else: + advantages[i] = teacher_adv + + def apply_opd_to_advantages( args: Namespace, rollout_data: RolloutBatch, advantages: list[torch.Tensor], ) -> None: + teacher_advantage_mode = getattr(args, "opd_teacher_advantage", False) + + if teacher_advantage_mode: + additive = getattr(args, "opd_teacher_advantage_additive", False) + _apply_teacher_advantage(rollout_data, advantages, additive=additive) + return + if args.opd_kl_coef == 0.0: return @@ -1287,10 +1412,17 @@ def apply_opd_to_advantages( if is_topk: student_topk_lp_list = rollout_data.get("opd_topk_student_log_probs") teacher_topk_lp_list = rollout_data.get("opd_topk_teacher_log_probs") - if student_topk_lp_list is None or teacher_topk_lp_list is None: - return + has_student = student_topk_lp_list is not None and any( + isinstance(v, torch.Tensor) and v.numel() > 0 for v in student_topk_lp_list + ) + has_teacher = teacher_topk_lp_list is not None and any( + isinstance(v, torch.Tensor) and v.numel() > 0 for v in teacher_topk_lp_list + ) + if not has_student or not has_teacher: + is_topk = False + + if is_topk: device = advantages[0].device if advantages else torch.device("cpu") - # union ：per-row valid length → bool mask [R, max_K'] k_lengths_list = rollout_data.get("opd_topk_ksz") if token_selection == "union" else None for i, adv in enumerate(advantages): @@ -1336,16 +1468,23 @@ def apply_opd_to_advantages( device = student_log_probs[0].device teacher_log_probs = [t.to(device=device) for t in teacher_log_probs] + per_token_clip = getattr(args, "opd_per_token_clip", None) for i, adv in enumerate(advantages): + s_lp = student_log_probs[i] + t_lp = teacher_log_probs[i] + if s_lp.numel() == 0 or t_lp.numel() == 0: + continue kl_term = _opd_compute_per_token_signal( - student_lp_1d=student_log_probs[i], - teacher_lp_1d=teacher_log_probs[i], - token_selection=token_selection, + student_lp_1d=s_lp, + teacher_lp_1d=t_lp, + token_selection="student_sampled", kl_type=kl_type, jsd_alpha=jsd_alpha, norm_mode=norm_mode, log_prob_min_clamp=log_prob_min_clamp, ) + if per_token_clip is not None: + kl_term = torch.clamp(kl_term, max=float(per_token_clip)) advantages[i] = adv - args.opd_kl_coef * kl_term.detach() @@ -1452,6 +1591,78 @@ def compute_policy_opd_loss( return opd_loss_coef * opd_loss, reported_loss +def compute_eopd_fkl_loss( + *, + args: Namespace, + batch: RolloutBatch, + log_probs_and_entropy: dict[str, list[torch.Tensor]], +) -> tuple[torch.Tensor | None, dict[str, torch.Tensor]]: + if not getattr(args, "use_eopd", False): + return None, {} + + teacher_entropy_list = batch.get("teacher_entropy") + if teacher_entropy_list is None: + return None, {} + + student_topk_lp_list = log_probs_and_entropy.get("topk_log_probs") + teacher_topk_lp_list = batch.get("opd_topk_teacher_log_probs") + if not student_topk_lp_list or not teacher_topk_lp_list: + return None, {} + + threshold = float(args.eopd_entropy_threshold) + fkl_coef = float(args.eopd_fkl_coef) + norm_mode = getattr(args, "opd_norm_mode", "tail") + log_prob_min_clamp = getattr(args, "opd_log_prob_min_clamp", None) + token_selection = args.opd_token_selection + k_lengths_list = batch.get("opd_topk_ksz") if token_selection == "union" else None + device = student_topk_lp_list[0].device + + fkl_chunks: list[torch.Tensor] = [] + entropy_mask_chunks: list[torch.Tensor] = [] + teacher_ent_chunks: list[torch.Tensor] = [] + + for i, s_lp_2d in enumerate(student_topk_lp_list): + t_lp_2d = teacher_topk_lp_list[i].to(device=device).detach() + s_lp_2d = s_lp_2d.to(device=device) + t_ent = teacher_entropy_list[i].to(device=device).detach() + + mask = None + if k_lengths_list is not None and i < len(k_lengths_list): + kl = k_lengths_list[i] + if kl is not None and t_lp_2d is not None: + kl = kl.to(device=device) + max_kp = t_lp_2d.size(-1) + mask = torch.arange(max_kp, device=device).unsqueeze(0) < kl.unsqueeze(1) + + per_token_fkl = compute_opd_kl_topk( + s_lp_2d, + t_lp_2d, + kl_type="forward_kl", + norm_mode=norm_mode, + log_prob_min_clamp=log_prob_min_clamp, + mask=mask, + ) + + ent_mask = (t_ent >= threshold).float() + fkl_chunks.append(per_token_fkl * ent_mask) + entropy_mask_chunks.append(ent_mask) + teacher_ent_chunks.append(t_ent) + + eopd_fkl_per_token = torch.cat(fkl_chunks, dim=0).to(dtype=student_topk_lp_list[0].dtype) + + eopd_fkl_loss = reduce_opd_loss(batch, eopd_fkl_per_token) + + reported: dict[str, torch.Tensor] = {} + with torch.no_grad(): + all_ent_mask = torch.cat(entropy_mask_chunks, dim=0) + all_teacher_ent = torch.cat(teacher_ent_chunks, dim=0) + reported["eopd_fkl_loss"] = eopd_fkl_loss.clone().detach() + reported["eopd_high_entropy_frac"] = reduce_opd_loss(batch, all_ent_mask).clone().detach() + reported["eopd_teacher_entropy_mean"] = reduce_opd_loss(batch, all_teacher_ent).clone().detach() + + return fkl_coef * eopd_fkl_loss, reported + + # --------------------------------------------------------------------------- # MOPD per-source metrics # --------------------------------------------------------------------------- diff --git a/relax/utils/types.py b/relax/utils/types.py index 9c7cabb5b..e8905c2bf 100644 --- a/relax/utils/types.py +++ b/relax/utils/types.py @@ -32,6 +32,7 @@ class Sample: abort_count: int = 0 # Number of times this sample has been aborted teacher_log_probs: list[float] | None = None + teacher_entropy: list[float] | None = None student_topk_token_ids: np.ndarray | None = None student_topk_log_probs: np.ndarray | None = None teacher_at_student_topk_log_probs: np.ndarray | None = None diff --git a/relax/utils/utils.py b/relax/utils/utils.py index 92bca9407..30dc478f4 100644 --- a/relax/utils/utils.py +++ b/relax/utils/utils.py @@ -585,16 +585,25 @@ def recovery_load_path(args: Namespace) -> Optional[str]: def compute_dp_size(config) -> int: """Compute data-parallel size from config for the actor role. - For Megatron backend: dp_size = total_actor_gpus / (tp * pp * cp) + For Megatron backend: dp_size = training_gpus / (tp * pp * cp) + In OPD colocate mode the actor placement group bundles both training and + teacher GPUs, so ``actor_num_gpus_per_node * actor_num_nodes`` (the + training-only slice) is used instead of the raw resource count. """ _, actor_total_gpus = config.resource.get("actor", (1, 1)) + actor_num_gpus = getattr(config, "actor_num_gpus_per_node", None) + actor_num_nodes = getattr(config, "actor_num_nodes", None) + if actor_num_gpus and actor_num_nodes: + training_gpus = min(actor_num_gpus * actor_num_nodes, actor_total_gpus) + else: + training_gpus = actor_total_gpus tp = getattr(config, "tensor_model_parallel_size", 1) pp = getattr(config, "pipeline_model_parallel_size", 1) cp = getattr(config, "context_parallel_size", 1) - dp_size = actor_total_gpus // (tp * pp * cp) + dp_size = training_gpus // (tp * pp * cp) if dp_size <= 0: raise ValueError( - f"Computed dp_size={dp_size} is invalid. actor_total_gpus={actor_total_gpus}, tp={tp}, pp={pp}, cp={cp}" + f"Computed dp_size={dp_size} is invalid. training_gpus={training_gpus}, tp={tp}, pp={pp}, cp={cp}" ) return dp_size diff --git a/scripts/debug/test_eopd_smoke.sh b/scripts/debug/test_eopd_smoke.sh new file mode 100755 index 000000000..4f27b36a2 --- /dev/null +++ b/scripts/debug/test_eopd_smoke.sh @@ -0,0 +1,69 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# EOPD (Entropy-aware OPD) smoke test. +# +# This script validates that the EOPD arguments, data pipeline, and loss +# function integrate without errors. It runs a colocate OPD+EOPD training +# for 5 steps on a small model. +# +# Prerequisites: +# - A small Megatron-compatible checkpoint (e.g., Qwen3-0.6B converted to mcore) +# - A prompt dataset (JSONL with "prompt" and "label" keys) +# +# Usage: +# MODEL_DIR=/path/to/small_model DATA_DIR=/path/to/data \ +# bash scripts/debug/test_eopd_smoke.sh +# +# Expected outcome: +# - Training completes 5 steps without crash +# - Logs contain "eopd_fkl_loss" metric + +set -ex +set -o pipefail + +MODEL_DIR="${MODEL_DIR:?Set MODEL_DIR to a small Megatron checkpoint directory}" +DATA_DIR="${DATA_DIR:?Set DATA_DIR to a directory containing a prompt JSONL file}" +PROMPT_SET="${PROMPT_SET:-${DATA_DIR}/prompts.jsonl}" + +python relax/entrypoints/train.py \ + --hf-checkpoint "${MODEL_DIR}" \ + --ref-load "${MODEL_DIR}" \ + --megatron-to-hf-mode bridge \ + --prompt-data "${PROMPT_SET}" \ + --input-key prompt \ + --label-key label \ + --apply-chat-template \ + --rm-type dapo \ + --reward-key score \ + --num-rollout 5 \ + --rollout-batch-size 2 \ + --n-samples-per-prompt 2 \ + --rollout-max-response-len 128 \ + --rollout-temperature 1 \ + --global-batch-size 4 \ + --loss-type grpo \ + --advantage-estimator grpo \ + --lr 1e-6 \ + --eps-clip 0.2 \ + --use-opd \ + --opd-type megatron \ + --opd-teacher-load "${MODEL_DIR}" \ + --opd-loss-coef 1.0 \ + --opd-kl-coef 0.0 \ + --opd-log-prob-top-k 4 \ + --opd-token-selection teacher_topk \ + --use-eopd \ + --eopd-entropy-threshold 0.8 \ + --eopd-fkl-coef 1.0 \ + 2>&1 | tee /tmp/eopd_smoke.log + +echo "---" +echo "Checking for EOPD metrics in logs..." +if grep -q "eopd_fkl_loss" /tmp/eopd_smoke.log; then + echo "PASS: eopd_fkl_loss found in training logs." +else + echo "WARN: eopd_fkl_loss not found in logs (may be expected if 0 steps ran)." +fi +echo "EOPD smoke test completed." diff --git a/tests/backends/megatron/test_eopd_loss.py b/tests/backends/megatron/test_eopd_loss.py new file mode 100644 index 000000000..31c61b72e --- /dev/null +++ b/tests/backends/megatron/test_eopd_loss.py @@ -0,0 +1,195 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for EOPD (Entropy-aware OPD) forward KL loss.""" + +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + + +def _install_fake_megatron(monkeypatch): + megatron = ModuleType("megatron") + core = ModuleType("megatron.core") + mpu = ModuleType("megatron.core.mpu") + + mpu.get_context_parallel_world_size = lambda: 1 + core.mpu = mpu + megatron.core = core + + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + monkeypatch.setitem(sys.modules, "megatron.core.mpu", mpu) + + +@pytest.fixture() +def opd_utils_module(monkeypatch): + torch = pytest.importorskip("torch", exc_type=ImportError) + _install_fake_megatron(monkeypatch) + sys.modules.pop("relax.utils.opd.opd_utils", None) + module = importlib.import_module("relax.utils.opd.opd_utils") + yield module, torch + sys.modules.pop("relax.utils.opd.opd_utils", None) + + +def _make_args(**overrides): + defaults = dict( + use_eopd=True, + eopd_entropy_threshold=0.8, + eopd_fkl_coef=1.0, + eopd_fkl_top_k=0, + opd_log_prob_top_k=4, + opd_token_selection="teacher_topk", + opd_norm_mode="tail", + opd_log_prob_min_clamp=None, + opd_loss_coef=1.0, + ) + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def _make_batch(torch, teacher_entropy_vals, response_lengths, teacher_topk_lp, loss_masks=None): + batch = { + "response_lengths": response_lengths, + "total_lengths": [r + 2 for r in response_lengths], + "loss_masks": loss_masks or [torch.ones(r, dtype=torch.float32) for r in response_lengths], + "teacher_entropy": [torch.tensor(e, dtype=torch.float32) for e in teacher_entropy_vals], + "opd_topk_teacher_log_probs": [torch.tensor(lp, dtype=torch.float32) for lp in teacher_topk_lp], + } + return batch + + +def test_eopd_fkl_loss_basic(opd_utils_module): + opd_utils, torch = opd_utils_module + args = _make_args() + + teacher_topk_lp = [ + [[-1.0, -2.0, -3.0, -4.0], [-0.5, -1.5, -2.5, -3.5]], + ] + student_topk_lp = [ + [[-1.1, -2.1, -3.1, -4.1], [-0.6, -1.6, -2.6, -3.6]], + ] + teacher_entropy = [[1.0, 0.5]] + + batch = _make_batch(torch, teacher_entropy, [2], teacher_topk_lp) + log_probs_and_entropy = { + "topk_log_probs": [torch.tensor(student_topk_lp[0], dtype=torch.float32)], + } + + loss, reported = opd_utils.compute_eopd_fkl_loss( + args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy + ) + assert loss is not None + assert loss.item() > 0.0 + assert "eopd_fkl_loss" in reported + assert "eopd_high_entropy_frac" in reported + assert "eopd_teacher_entropy_mean" in reported + assert abs(reported["eopd_high_entropy_frac"].item() - 0.5) < 1e-5 + + +def test_eopd_fkl_loss_all_below_threshold(opd_utils_module): + opd_utils, torch = opd_utils_module + args = _make_args(eopd_entropy_threshold=2.0) + + teacher_topk_lp = [[[-1.0, -2.0, -3.0, -4.0], [-0.5, -1.5, -2.5, -3.5]]] + student_topk_lp = [[[-1.1, -2.1, -3.1, -4.1], [-0.6, -1.6, -2.6, -3.6]]] + teacher_entropy = [[0.5, 0.3]] + + batch = _make_batch(torch, teacher_entropy, [2], teacher_topk_lp) + log_probs_and_entropy = { + "topk_log_probs": [torch.tensor(student_topk_lp[0], dtype=torch.float32)], + } + + loss, reported = opd_utils.compute_eopd_fkl_loss( + args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy + ) + assert loss is not None + assert abs(loss.item()) < 1e-6 + assert abs(reported["eopd_high_entropy_frac"].item()) < 1e-5 + + +def test_eopd_fkl_loss_all_above_threshold(opd_utils_module): + opd_utils, torch = opd_utils_module + args = _make_args(eopd_entropy_threshold=0.1) + + teacher_topk_lp = [[[-1.0, -2.0, -3.0, -4.0], [-0.5, -1.5, -2.5, -3.5]]] + student_topk_lp = [[[-1.1, -2.1, -3.1, -4.1], [-0.6, -1.6, -2.6, -3.6]]] + teacher_entropy = [[1.0, 0.5]] + + batch = _make_batch(torch, teacher_entropy, [2], teacher_topk_lp) + log_probs_and_entropy = { + "topk_log_probs": [torch.tensor(student_topk_lp[0], dtype=torch.float32)], + } + + loss_gated, _ = opd_utils.compute_eopd_fkl_loss( + args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy + ) + + full_fkl_chunks = [] + for s, t in zip(student_topk_lp[0], teacher_topk_lp[0]): + s_t = torch.tensor(s, dtype=torch.float32).unsqueeze(0) + t_t = torch.tensor(t, dtype=torch.float32).unsqueeze(0) + full_fkl_chunks.append(opd_utils.compute_opd_kl_topk(s_t, t_t, kl_type="forward_kl")) + full_fkl = torch.cat(full_fkl_chunks, dim=0).mean() + + assert loss_gated is not None + assert torch.isclose(loss_gated, full_fkl, atol=1e-5) + + +def test_eopd_fkl_loss_normalization(opd_utils_module): + opd_utils, torch = opd_utils_module + args = _make_args(eopd_entropy_threshold=0.0) + + teacher_topk_lp = [ + [[-1.0, -2.0, -3.0, -4.0], [-0.5, -1.5, -2.5, -3.5]], + [[-1.2, -2.2, -3.2, -4.2]], + ] + student_topk_lp = [ + [[-1.1, -2.1, -3.1, -4.1], [-0.6, -1.6, -2.6, -3.6]], + [[-1.3, -2.3, -3.3, -4.3]], + ] + teacher_entropy = [[1.0, 1.0], [1.0]] + loss_masks = [torch.ones(2), torch.tensor([1.0])] + + batch = _make_batch(torch, teacher_entropy, [2, 1], teacher_topk_lp, loss_masks=loss_masks) + log_probs_and_entropy = { + "topk_log_probs": [ + torch.tensor(student_topk_lp[0], dtype=torch.float32), + torch.tensor(student_topk_lp[1], dtype=torch.float32), + ], + } + + loss, _ = opd_utils.compute_eopd_fkl_loss(args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy) + assert loss is not None + assert loss.item() > 0.0 + + +def test_eopd_fkl_loss_disabled(opd_utils_module): + opd_utils, torch = opd_utils_module + args = _make_args(use_eopd=False) + + loss, reported = opd_utils.compute_eopd_fkl_loss(args=args, batch={}, log_probs_and_entropy={}) + assert loss is None + assert reported == {} + + +def test_eopd_reported_metrics(opd_utils_module): + opd_utils, torch = opd_utils_module + args = _make_args() + + teacher_topk_lp = [[[-1.0, -2.0, -3.0, -4.0]]] + student_topk_lp = [[[-1.1, -2.1, -3.1, -4.1]]] + teacher_entropy = [[1.0]] + + batch = _make_batch(torch, teacher_entropy, [1], teacher_topk_lp) + log_probs_and_entropy = { + "topk_log_probs": [torch.tensor(student_topk_lp[0], dtype=torch.float32)], + } + + _, reported = opd_utils.compute_eopd_fkl_loss(args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy) + assert "eopd_fkl_loss" in reported + assert "eopd_high_entropy_frac" in reported + assert "eopd_teacher_entropy_mean" in reported + assert reported["eopd_high_entropy_frac"].item() == 1.0 + assert abs(reported["eopd_teacher_entropy_mean"].item() - 1.0) < 1e-5 diff --git a/tests/engine/rollout/test_on_policy_distillation_eos_replace.py b/tests/engine/rollout/test_on_policy_distillation_eos_replace.py new file mode 100644 index 000000000..2a5d5ec5b --- /dev/null +++ b/tests/engine/rollout/test_on_policy_distillation_eos_replace.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for EOS token replacement in OPD teacher prefill.""" + + +def test_eos_replace_in_response_only(): + """Only replace EOS in response portion (index >= prompt_length).""" + _ENDOFTEXT, _IM_END = 151643, 151645 + prompt_length = 3 + input_ids = [10, 151643, 20, 30, 151643, 40] + + result = list(input_ids) + for j in range(prompt_length, len(result)): + if result[j] == _ENDOFTEXT: + result[j] = _IM_END + break + + assert result[1] == 151643, "prompt EOS should NOT be replaced" + assert result[4] == 151645, "response EOS should be replaced" + assert result == [10, 151643, 20, 30, 151645, 40] + + +def test_eos_replace_only_first_occurrence(): + """Only the first EOS in the response is replaced.""" + _ENDOFTEXT, _IM_END = 151643, 151645 + prompt_length = 1 + input_ids = [10, 151643, 20, 151643, 30] + + result = list(input_ids) + for j in range(prompt_length, len(result)): + if result[j] == _ENDOFTEXT: + result[j] = _IM_END + break + + assert result[1] == 151645, "first response EOS replaced" + assert result[3] == 151643, "second response EOS NOT replaced" + + +def test_eos_replace_no_eos_in_response(): + """No replacement when response has no EOS.""" + _ENDOFTEXT, _IM_END = 151643, 151645 + prompt_length = 2 + input_ids = [10, 20, 30, 40, 50] + + result = list(input_ids) + for j in range(prompt_length, len(result)): + if result[j] == _ENDOFTEXT: + result[j] = _IM_END + break + + assert result == [10, 20, 30, 40, 50] + + +def test_eos_replace_disabled(): + """When opd_eos_replace=False, no replacement happens.""" + opd_eos_replace = False + input_ids = [10, 151643, 20, 151643] + + result = list(input_ids) + if opd_eos_replace: + _ENDOFTEXT, _IM_END = 151643, 151645 + for j in range(1, len(result)): + if result[j] == _ENDOFTEXT: + result[j] = _IM_END + break + + assert result == [10, 151643, 20, 151643] + + +def test_eos_replace_eos_only_in_prompt(): + """EOS in prompt only — response untouched.""" + _ENDOFTEXT, _IM_END = 151643, 151645 + prompt_length = 3 + input_ids = [151643, 151643, 151643, 100, 200] + + result = list(input_ids) + for j in range(prompt_length, len(result)): + if result[j] == _ENDOFTEXT: + result[j] = _IM_END + break + + assert result == [151643, 151643, 151643, 100, 200] diff --git a/tests/integration/test_eopd_smoke.py b/tests/integration/test_eopd_smoke.py new file mode 100644 index 000000000..73485888e --- /dev/null +++ b/tests/integration/test_eopd_smoke.py @@ -0,0 +1,142 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Smoke test: end-to-end EOPD FKL loss computation + gradient backprop.""" + +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") + + +def _install_fake_megatron(monkeypatch): + megatron = ModuleType("megatron") + core = ModuleType("megatron.core") + mpu = ModuleType("megatron.core.mpu") + mpu.get_context_parallel_world_size = lambda: 1 + core.mpu = mpu + megatron.core = core + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + monkeypatch.setitem(sys.modules, "megatron.core.mpu", mpu) + + +@pytest.fixture() +def opd_utils(monkeypatch): + _install_fake_megatron(monkeypatch) + sys.modules.pop("relax.utils.opd.opd_utils", None) + module = importlib.import_module("relax.utils.opd.opd_utils") + yield module + sys.modules.pop("relax.utils.opd.opd_utils", None) + + +def _make_args(**overrides): + defaults = { + "use_eopd": True, + "eopd_entropy_threshold": 0.8, + "eopd_fkl_coef": 1.0, + "opd_norm_mode": "trunc", + "opd_log_prob_min_clamp": None, + "opd_token_selection": "teacher_topk", + } + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def test_eopd_smoke_forward_and_backward(opd_utils): + """Full chain: compute_eopd_fkl_loss → reduce_opd_loss → backward → non- + zero grads.""" + K = 32 + seq_len = 10 + + student_topk_lp = torch.randn(seq_len, K, requires_grad=True) + teacher_topk_lp = torch.randn(seq_len, K) + teacher_entropy = torch.rand(seq_len) * 2.0 + + batch = { + "opd_topk_teacher_log_probs": [teacher_topk_lp], + "teacher_entropy": [teacher_entropy], + "response_lengths": [seq_len], + "loss_masks": [torch.ones(seq_len)], + } + log_probs_and_entropy = { + "topk_log_probs": [student_topk_lp], + } + + args = _make_args() + loss, metrics = opd_utils.compute_eopd_fkl_loss( + args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy + ) + + assert loss is not None, "EOPD loss should not be None" + assert loss.ndim == 0, "loss should be a scalar" + assert loss.item() >= 0.0, "FKL loss should be non-negative" + + loss.backward() + assert student_topk_lp.grad is not None, "student log-probs should have gradients" + assert student_topk_lp.grad.abs().sum() > 0, "gradients should be non-zero" + + assert "eopd_fkl_loss" in metrics + assert "eopd_high_entropy_frac" in metrics + assert "eopd_teacher_entropy_mean" in metrics + + +def test_eopd_smoke_zero_loss_below_threshold(opd_utils): + """All entropy below threshold → loss is zero, but function still returns a + valid tensor.""" + K = 32 + seq_len = 5 + + student_topk_lp = torch.randn(seq_len, K, requires_grad=True) + teacher_topk_lp = torch.randn(seq_len, K) + teacher_entropy = torch.full((seq_len,), 0.1) + + batch = { + "opd_topk_teacher_log_probs": [teacher_topk_lp], + "teacher_entropy": [teacher_entropy], + "response_lengths": [seq_len], + "loss_masks": [torch.ones(seq_len)], + } + log_probs_and_entropy = {"topk_log_probs": [student_topk_lp]} + + args = _make_args(eopd_entropy_threshold=0.8) + loss, metrics = opd_utils.compute_eopd_fkl_loss( + args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy + ) + + assert loss is not None + assert loss.item() == pytest.approx(0.0, abs=1e-7) + + +def test_eopd_smoke_multi_sample(opd_utils): + """Multiple samples in a batch are handled correctly.""" + K = 8 + s1_len, s2_len = 6, 4 + + student_lp_1 = torch.randn(s1_len, K, requires_grad=True) + student_lp_2 = torch.randn(s2_len, K, requires_grad=True) + teacher_lp_1 = torch.randn(s1_len, K) + teacher_lp_2 = torch.randn(s2_len, K) + ent_1 = torch.full((s1_len,), 1.5) + ent_2 = torch.full((s2_len,), 1.5) + + batch = { + "opd_topk_teacher_log_probs": [teacher_lp_1, teacher_lp_2], + "teacher_entropy": [ent_1, ent_2], + "response_lengths": [s1_len, s2_len], + "loss_masks": [torch.ones(s1_len), torch.ones(s2_len)], + } + log_probs_and_entropy = {"topk_log_probs": [student_lp_1, student_lp_2]} + + args = _make_args() + loss, metrics = opd_utils.compute_eopd_fkl_loss( + args=args, batch=batch, log_probs_and_entropy=log_probs_and_entropy + ) + + assert loss is not None + loss.backward() + assert student_lp_1.grad is not None + assert student_lp_2.grad is not None diff --git a/tests/utils/test_eopd_arguments.py b/tests/utils/test_eopd_arguments.py new file mode 100644 index 000000000..dbe52f6b7 --- /dev/null +++ b/tests/utils/test_eopd_arguments.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Validation tests for EOPD CLI arguments.""" + +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + + +def _install_fake_megatron(monkeypatch): + megatron = ModuleType("megatron") + core = ModuleType("megatron.core") + mpu = ModuleType("megatron.core.mpu") + mpu.get_context_parallel_world_size = lambda: 1 + core.mpu = mpu + megatron.core = core + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + monkeypatch.setitem(sys.modules, "megatron.core.mpu", mpu) + + +@pytest.fixture() +def opd_utils_module(monkeypatch): + _install_fake_megatron(monkeypatch) + sys.modules.pop("relax.utils.opd.opd_utils", None) + module = importlib.import_module("relax.utils.opd.opd_utils") + yield module + sys.modules.pop("relax.utils.opd.opd_utils", None) + + +def _base_args(**overrides): + defaults = dict( + use_opd=True, + opd_type="megatron", + opd_teacher_timeout_s=600, + opd_log_prob_top_k=4, + opd_token_selection="teacher_topk", + opd_kl_type="reverse_kl", + opd_kl_coef=0.0, + opd_loss_coef=1.0, + opd_teacher_prompt_key=None, + opd_teacher_image_key=None, + opd_per_token_clip=None, + opd_is_clip=None, + opd_teacher_load="/fake/teacher/ckpt", + use_eopd=True, + eopd_entropy_threshold=0.8, + eopd_fkl_coef=1.0, + eopd_fkl_top_k=0, + ) + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def test_eopd_rejects_unknown_opd_type(opd_utils_module, monkeypatch, tmp_path): + args = _base_args(opd_type="unknown_type", opd_teacher_load=None, opd_teacher_url="http://t/generate") + monkeypatch.setattr("os.path.exists", lambda p: True) + with pytest.raises(ValueError, match="megatron.*sglang"): + opd_utils_module.validate_opd_args(args, is_sft=False) + + +def test_eopd_accepts_sglang_type(opd_utils_module, monkeypatch): + args = _base_args(opd_type="sglang", opd_teacher_load=None, opd_teacher_url="http://t/generate") + monkeypatch.setattr("os.path.exists", lambda p: True) + opd_utils_module.validate_opd_args(args, is_sft=False) + + +def test_eopd_requires_loss_mode(opd_utils_module, monkeypatch): + args = _base_args(opd_kl_coef=1.0, opd_loss_coef=1.0) + monkeypatch.setattr("os.path.exists", lambda p: True) + with pytest.raises(ValueError, match="opd-kl-coef"): + opd_utils_module.validate_opd_args(args, is_sft=False) + + +def test_eopd_requires_topk(opd_utils_module, monkeypatch): + args = _base_args(opd_log_prob_top_k=0, eopd_fkl_top_k=0) + monkeypatch.setattr("os.path.exists", lambda p: True) + with pytest.raises(ValueError, match="top-k"): + opd_utils_module.validate_opd_args(args, is_sft=False) + + +def test_eopd_valid_config(opd_utils_module, monkeypatch): + args = _base_args() + monkeypatch.setattr("os.path.exists", lambda p: True) + opd_utils_module.validate_opd_args(args, is_sft=False) diff --git a/tests/utils/test_eopd_entropy_patch.py b/tests/utils/test_eopd_entropy_patch.py new file mode 100644 index 000000000..62d6cdf23 --- /dev/null +++ b/tests/utils/test_eopd_entropy_patch.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for the SGLang EOPD entropy patch utilities.""" + +import numpy as np +import pybase64 +import torch + + +def test_entropy_computation_correctness(): + """Verify H = -sum(p * log_p) on synthetic logits.""" + from relax.utils.opd.opd_sglang_entropy_patch import _compute_entropy + + logits = torch.tensor([[1.0, 2.0, 3.0], [0.0, 0.0, 0.0]], dtype=torch.float32) + entropy = _compute_entropy(logits) + assert entropy.shape == (2,) + + probs_0 = torch.softmax(logits[0], dim=-1) + expected_0 = -(probs_0 * probs_0.log()).sum() + assert torch.allclose(entropy[0], expected_0, atol=1e-5) + + expected_uniform = np.log(3.0) + assert abs(entropy[1].item() - expected_uniform) < 1e-5 + + +def test_entropy_b64_roundtrip(): + """Verify encode → decode preserves float32 entropy values.""" + from relax.utils.opd.opd_sglang_entropy_patch import _entropy_to_b64 + + original = torch.tensor([0.5, 1.2, 0.0, 3.14], dtype=torch.float32) + b64_str = _entropy_to_b64(original) + + decoded = np.frombuffer(pybase64.b64decode(b64_str), dtype=np.float32) + np.testing.assert_allclose(decoded, original.numpy(), atol=1e-7) + + +def test_logprob_response_entropy_1d_from_b64(): + """LogprobResponse.entropy_1d() correctly parses meta_info with b64 + entropy.""" + from relax.utils.opd.opd_main_worker import LogprobResponse + + values = np.array([0.1, 0.5, 1.0, 2.0], dtype=np.float32) + b64_str = pybase64.b64encode(values.tobytes()).decode() + + resp = LogprobResponse({"meta_info": {"relax_input_entropy_b64": [b64_str]}}) + result = resp.entropy_1d() + + assert result is not None + np.testing.assert_allclose(result, values, atol=1e-7) + + +def test_logprob_response_entropy_1d_bare_string(): + """LogprobResponse.entropy_1d() handles bare b64 string (not wrapped in + list).""" + from relax.utils.opd.opd_main_worker import LogprobResponse + + values = np.array([0.3, 0.7], dtype=np.float32) + b64_str = pybase64.b64encode(values.tobytes()).decode() + + resp = LogprobResponse({"meta_info": {"relax_input_entropy_b64": b64_str}}) + result = resp.entropy_1d() + + assert result is not None + np.testing.assert_allclose(result, values, atol=1e-7) + + +def test_logprob_response_entropy_1d_missing(): + """LogprobResponse.entropy_1d() returns None when field is absent.""" + from relax.utils.opd.opd_main_worker import LogprobResponse + + resp = LogprobResponse({"meta_info": {}}) + assert resp.entropy_1d() is None + + +def test_logprob_response_entropy_1d_empty(): + """LogprobResponse.entropy_1d() returns None for empty response.""" + from relax.utils.opd.opd_main_worker import LogprobResponse + + resp = LogprobResponse(None) + assert resp.entropy_1d() is None diff --git a/tests/utils/test_eopd_reference_parity.py b/tests/utils/test_eopd_reference_parity.py new file mode 100644 index 000000000..53a011827 --- /dev/null +++ b/tests/utils/test_eopd_reference_parity.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Element-wise parity tests: Relax EOPD vs reference forward-KL formulas from +the paper.""" + +import importlib +import sys +from types import ModuleType + +import pytest + + +torch = pytest.importorskip("torch") + + +def _install_fake_megatron(monkeypatch): + megatron = ModuleType("megatron") + core = ModuleType("megatron.core") + mpu = ModuleType("megatron.core.mpu") + mpu.get_context_parallel_world_size = lambda: 1 + core.mpu = mpu + megatron.core = core + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + monkeypatch.setitem(sys.modules, "megatron.core.mpu", mpu) + + +@pytest.fixture() +def opd_utils(monkeypatch): + _install_fake_megatron(monkeypatch) + sys.modules.pop("relax.utils.opd.opd_utils", None) + module = importlib.import_module("relax.utils.opd.opd_utils") + yield module + sys.modules.pop("relax.utils.opd.opd_utils", None) + + +def _reference_forward_kl_trunc(teacher_lp: torch.Tensor, student_lp: torch.Tensor) -> torch.Tensor: + """Reference forward KL per token: sum_k p_teacher(k) * (log p_teacher(k) - log p_student(k)). + + This matches the paper formula for truncated (top-K) forward KL without tail correction. + """ + t = teacher_lp.float() + s = student_lp.float() + per_k = t.exp() * (t - s) + return per_k.sum(dim=-1) + + +def _reference_entropy_mask(teacher_entropy: torch.Tensor, threshold: float) -> torch.Tensor: + """Reference entropy gate: 1 where entropy >= threshold, else 0.""" + return (teacher_entropy >= threshold).float() + + +def _reference_reduce(values: torch.Tensor, loss_mask: torch.Tensor) -> torch.Tensor: + """Reference masked reduction: sum(values * mask) / max(sum(mask), 1).""" + masked = values * loss_mask + return masked.sum() / torch.clamp_min(loss_mask.sum(), 1) + + +class TestForwardKLParity: + """Verify compute_opd_kl_topk(forward_kl, trunc) matches reference per- + token forward KL.""" + + def test_basic_parity(self, opd_utils): + torch.manual_seed(42) + K = 32 + R = 20 + teacher_lp = torch.log_softmax(torch.randn(R, K), dim=-1) + student_lp = torch.log_softmax(torch.randn(R, K), dim=-1) + + relax_result = opd_utils.compute_opd_kl_topk(student_lp, teacher_lp, kl_type="forward_kl", norm_mode="trunc") + ref_result = _reference_forward_kl_trunc(teacher_lp, student_lp) + + torch.testing.assert_close(relax_result, ref_result, atol=1e-5, rtol=1e-5) + + def test_identical_distributions(self, opd_utils): + K = 16 + R = 10 + lp = torch.log_softmax(torch.randn(R, K), dim=-1) + + relax_result = opd_utils.compute_opd_kl_topk(lp, lp, kl_type="forward_kl", norm_mode="trunc") + ref_result = _reference_forward_kl_trunc(lp, lp) + + torch.testing.assert_close(relax_result, ref_result, atol=1e-6, rtol=1e-6) + assert relax_result.abs().max() < 1e-5, "KL(p||p) should be ~0" + + def test_various_k_sizes(self, opd_utils): + for K in [4, 8, 16, 32, 64]: + teacher_lp = torch.log_softmax(torch.randn(5, K), dim=-1) + student_lp = torch.log_softmax(torch.randn(5, K), dim=-1) + + relax_result = opd_utils.compute_opd_kl_topk( + student_lp, teacher_lp, kl_type="forward_kl", norm_mode="trunc" + ) + ref_result = _reference_forward_kl_trunc(teacher_lp, student_lp) + torch.testing.assert_close(relax_result, ref_result, atol=1e-5, rtol=1e-5) + + +class TestEntropyMaskParity: + """Verify entropy gating matches reference: mask = (entropy >= tau).""" + + def test_mask_parity(self, opd_utils): + entropy = torch.tensor([0.1, 0.5, 0.8, 1.0, 1.5, 0.79, 0.81]) + threshold = 0.8 + + ref_mask = _reference_entropy_mask(entropy, threshold) + expected = torch.tensor([0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 1.0]) + torch.testing.assert_close(ref_mask, expected) + + def test_mask_applied_to_fkl(self, opd_utils): + """FKL * entropy_mask matches element-wise reference.""" + K = 8 + R = 7 + teacher_lp = torch.log_softmax(torch.randn(R, K), dim=-1) + student_lp = torch.log_softmax(torch.randn(R, K), dim=-1) + entropy = torch.tensor([0.1, 0.5, 0.8, 1.0, 1.5, 0.79, 0.81]) + threshold = 0.8 + + per_token_fkl = opd_utils.compute_opd_kl_topk(student_lp, teacher_lp, kl_type="forward_kl", norm_mode="trunc") + ref_fkl = _reference_forward_kl_trunc(teacher_lp, student_lp) + ref_mask = _reference_entropy_mask(entropy, threshold) + ref_gated = ref_fkl * ref_mask + + ent_mask = (entropy >= threshold).float() + relax_gated = per_token_fkl * ent_mask + + torch.testing.assert_close(relax_gated, ref_gated, atol=1e-5, rtol=1e-5) + + +class TestReductionParity: + """Verify reduce_opd_loss matches reference: masked_sum / mask_count.""" + + def test_single_sample(self, opd_utils): + values = torch.tensor([0.5, 1.0, 0.2, 0.8, 0.3]) + loss_mask = torch.tensor([1.0, 1.0, 0.0, 1.0, 0.0]) + batch = {"response_lengths": [5], "loss_masks": [loss_mask]} + + relax_result = opd_utils.reduce_opd_loss(batch, values) + ref_result = _reference_reduce(values, loss_mask) + torch.testing.assert_close(relax_result, ref_result, atol=1e-6, rtol=1e-6) + + def test_multi_sample(self, opd_utils): + v1 = torch.tensor([0.5, 1.0, 0.2]) + v2 = torch.tensor([0.8, 0.3]) + m1 = torch.tensor([1.0, 1.0, 0.0]) + m2 = torch.tensor([1.0, 1.0]) + + values = torch.cat([v1, v2]) + batch = {"response_lengths": [3, 2], "loss_masks": [m1, m2]} + + relax_result = opd_utils.reduce_opd_loss(batch, values) + + masked = torch.cat([v1 * m1, v2 * m2]) + ref_result = masked.sum() / (m1.sum() + m2.sum()) + torch.testing.assert_close(relax_result, ref_result, atol=1e-6, rtol=1e-6) + + def test_all_masked(self, opd_utils): + values = torch.tensor([1.0, 2.0, 3.0]) + loss_mask = torch.zeros(3) + batch = {"response_lengths": [3], "loss_masks": [loss_mask]} + + relax_result = opd_utils.reduce_opd_loss(batch, values) + assert relax_result.item() == pytest.approx(0.0, abs=1e-7) + + +class TestEndToEndParity: + """Full EOPD chain: FKL + entropy gate + reduction matches hand-computed + reference.""" + + def test_full_chain(self, opd_utils): + torch.manual_seed(123) + K, R = 32, 10 + threshold = 0.8 + + teacher_lp = torch.log_softmax(torch.randn(R, K), dim=-1) + student_lp = torch.log_softmax(torch.randn(R, K), dim=-1) + entropy = torch.rand(R) * 2.0 + loss_mask = torch.ones(R) + loss_mask[3] = 0.0 + loss_mask[7] = 0.0 + + ref_fkl = _reference_forward_kl_trunc(teacher_lp, student_lp) + ref_mask = _reference_entropy_mask(entropy, threshold) + ref_gated = ref_fkl * ref_mask + ref_loss = _reference_reduce(ref_gated, loss_mask) + + per_token_fkl = opd_utils.compute_opd_kl_topk(student_lp, teacher_lp, kl_type="forward_kl", norm_mode="trunc") + ent_mask = (entropy >= threshold).float() + gated = per_token_fkl * ent_mask + + batch = {"response_lengths": [R], "loss_masks": [loss_mask]} + relax_loss = opd_utils.reduce_opd_loss(batch, gated) + + torch.testing.assert_close(relax_loss, ref_loss, atol=1e-5, rtol=1e-5) diff --git a/tests/utils/test_opd_teacher_advantage.py b/tests/utils/test_opd_teacher_advantage.py new file mode 100644 index 000000000..cccb71312 --- /dev/null +++ b/tests/utils/test_opd_teacher_advantage.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for teacher advantage computation in OPD.""" + +import importlib +import sys +from types import ModuleType + +import pytest + + +def _install_fake_megatron(monkeypatch): + megatron = ModuleType("megatron") + core = ModuleType("megatron.core") + mpu = ModuleType("megatron.core.mpu") + + mpu.get_context_parallel_world_size = lambda: 1 + core.mpu = mpu + megatron.core = core + + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + monkeypatch.setitem(sys.modules, "megatron.core.mpu", mpu) + + +@pytest.fixture() +def opd_utils_module(monkeypatch): + torch = pytest.importorskip("torch", exc_type=ImportError) + _install_fake_megatron(monkeypatch) + sys.modules.pop("relax.utils.opd.opd_utils", None) + module = importlib.import_module("relax.utils.opd.opd_utils") + yield module, torch + sys.modules.pop("relax.utils.opd.opd_utils", None) + + +def test_teacher_advantage_replace(opd_utils_module): + """teacher_advantage = (teacher_lp - student_lp).detach(), replaces original advantage.""" + opd_utils, torch = opd_utils_module + + teacher_lp = [torch.tensor([-1.0, -2.0, -3.0])] + student_lp = [torch.tensor([-1.5, -2.5, -3.5])] + original_adv = [torch.tensor([10.0, 20.0, 30.0])] + + rollout_data = {"teacher_log_probs": teacher_lp, "rollout_log_probs": student_lp} + advantages = [a.clone() for a in original_adv] + + opd_utils._apply_teacher_advantage(rollout_data, advantages, additive=False) + + expected = teacher_lp[0] - student_lp[0] + assert torch.allclose(advantages[0], expected, atol=1e-6) + assert not torch.allclose(advantages[0], original_adv[0]) + + +def test_teacher_advantage_additive(opd_utils_module): + """additive=True adds teacher advantage to original instead of + replacing.""" + opd_utils, torch = opd_utils_module + + teacher_lp = [torch.tensor([-1.0, -2.0])] + student_lp = [torch.tensor([-1.5, -2.5])] + original_adv = [torch.tensor([10.0, 20.0])] + + rollout_data = {"teacher_log_probs": teacher_lp, "rollout_log_probs": student_lp} + advantages = [a.clone() for a in original_adv] + + opd_utils._apply_teacher_advantage(rollout_data, advantages, additive=True) + + teacher_adv = teacher_lp[0] - student_lp[0] + expected = original_adv[0] + teacher_adv + assert torch.allclose(advantages[0], expected, atol=1e-6) + + +def test_teacher_advantage_detached(opd_utils_module): + """Result should be detached (no gradient).""" + opd_utils, torch = opd_utils_module + + teacher_lp = [torch.tensor([-1.0, -2.0], requires_grad=True)] + student_lp = [torch.tensor([-1.5, -2.5], requires_grad=True)] + original_adv = [torch.tensor([10.0, 20.0])] + + rollout_data = {"teacher_log_probs": teacher_lp, "rollout_log_probs": student_lp} + advantages = [a.clone() for a in original_adv] + + opd_utils._apply_teacher_advantage(rollout_data, advantages, additive=False) + + assert not advantages[0].requires_grad + + +def test_teacher_advantage_missing_data(opd_utils_module): + """Gracefully handle missing teacher/student log_probs.""" + opd_utils, torch = opd_utils_module + + original_adv = [torch.tensor([10.0, 20.0])] + advantages = [a.clone() for a in original_adv] + + opd_utils._apply_teacher_advantage({"teacher_log_probs": None, "rollout_log_probs": None}, advantages) + assert torch.allclose(advantages[0], original_adv[0]) + + opd_utils._apply_teacher_advantage({}, advantages) + assert torch.allclose(advantages[0], original_adv[0]) + + +def test_teacher_advantage_multi_sample(opd_utils_module): + """Works correctly with multiple samples.""" + opd_utils, torch = opd_utils_module + + teacher_lp = [torch.tensor([-1.0, -2.0]), torch.tensor([-0.5])] + student_lp = [torch.tensor([-1.5, -2.5]), torch.tensor([-1.0])] + advantages = [torch.zeros(2), torch.zeros(1)] + + rollout_data = {"teacher_log_probs": teacher_lp, "rollout_log_probs": student_lp} + + opd_utils._apply_teacher_advantage(rollout_data, advantages, additive=False) + + assert torch.allclose(advantages[0], torch.tensor([0.5, 0.5]), atol=1e-6) + assert torch.allclose(advantages[1], torch.tensor([0.5]), atol=1e-6) + + +def test_teacher_advantage_empty_tensor(opd_utils_module): + """Skip samples with empty tensors.""" + opd_utils, torch = opd_utils_module + + teacher_lp = [torch.tensor([])] + student_lp = [torch.tensor([])] + original_adv = [torch.tensor([10.0])] + advantages = [a.clone() for a in original_adv] + + rollout_data = {"teacher_log_probs": teacher_lp, "rollout_log_probs": student_lp} + + opd_utils._apply_teacher_advantage(rollout_data, advantages, additive=False) + assert torch.allclose(advantages[0], original_adv[0])