"""PRIME-RL composer loss adapter. Per ADR-006, PRIME-RL exposes a ``CustomLossConfig`` that takes an importable function. This module supplies that function: a thin adapter that maps PRIME-RL's ``LossInputs`` struct onto the framework's 3-channel loss composition. Channel status (v0): 1. **DPPO + KL on the importance-sampling ratio** — implemented to match PRIME-RL's upstream ``default_loss_fn`` byte-for-byte. 2. **SDPO / OPSD** — deferred (raises ``NotImplementedError`` when enabled). PRIME-RL v0.5 exposes log-probs, not full logits, and SDPO requires the full vocabulary distribution. 3. **Trace-replay DPO** — out of scope for this recipe; emits a warning if ``beta_dpo != 0``. LossInputs shape (verified against PrimeIntellect-ai/prime-rl ``src/prime_rl/trainer/rl/loss.py`` lines 13-22): .. code-block:: python @dataclass class LossInputs: trainer_logprobs: Float[Tensor, ' seq'] # current policy log-probs inference_logprobs: Float[Tensor, ' seq'] # rollout-time policy log-probs teacher_logprobs: Float[Tensor, ' seq'] | None advantages: Float[Tensor, ' seq'] # per-token advantage loss_mask: Bool[Tensor, ' seq'] # which tokens count PRIME-RL calls the loss function once per sample, not on a batched ``(B, T)`` tensor. PRIME-RL's ``default_loss_fn`` (upstream) ----------------------------------------- Verbatim from ``prime_rl/trainer/rl/loss.py`` lines 116-165 and the ``DefaultLossConfig`` defaults at ``packages/prime-rl-configs/src/prime_rl/configs/trainer.py`` lines 412-425:: def default_loss_fn(inputs, loss_config): # line 133-135 log_importance_ratio = trainer_logprobs - inference_logprobs importance_ratio = exp(log_importance_ratio) mismatch_kl = importance_ratio - log_importance_ratio - 1 # line 137: NOTE — probability-space diff, not log-ratio probs_diff = exp(trainer_logprobs) - exp(inference_logprobs) # lines 138-139 dppo_invalid_mask_high = probs_diff > loss_config.dppo_mask_high dppo_invalid_mask_low = probs_diff < -loss_config.dppo_mask_low # lines 140-142: sign-of-advantage gate positive_advantages = advantages > 0 dppo_invalid_mask = where(positive_advantages, dppo_invalid_mask_high, dppo_invalid_mask_low) # lines 147-148 drop_mask = loss_mask & dppo_invalid_mask keep_mask = loss_mask & ~dppo_invalid_mask # lines 150-153 advantages = loss_config.adv_tau * advantages pg_loss = keep_mask * advantages * importance_ratio kl_loss = loss_mask * log_importance_ratio**2 loss = (-pg_loss + loss_config.kl_tau * kl_loss).sum() Defaults: ``dppo_mask_low=0.2``, ``dppo_mask_high=0.2``, ``adv_tau=1.0``, ``kl_tau=1e-3`` — all ``Field(..., ge=0)``. Three things this differs from a textbook PPO-clip: 1. The mask gate is on **probability-space** ``probs_diff``, not on the log-ratio. ``-loss_config.dppo_mask_low`` flips the sign so ``dppo_mask_low`` is itself non-negative. 2. The policy-gradient term is multiplied by ``importance_ratio`` (= ``exp(trainer_lp - inference_lp)``), giving a proper IS-corrected gradient — not a plain REINFORCE on ``trainer_lp``. 3. The mask is **conditioned on advantage sign**: a positive-advantage token is dropped when ``probs_diff`` exceeds ``dppo_mask_high`` (we'd be upweighting it too aggressively); a negative-advantage token is dropped when ``probs_diff`` falls below ``-dppo_mask_low`` (we'd be downweighting it too aggressively). Zero-advantage tokens are never DPPO-masked. The reduction is a plain ``sum()`` (PRIME-RL's outer ``compute_loss`` divides by ``loss_scale``); we mirror that. License: MIT (matches the rest of the framework). PRIME-RL is Apache-2; we reference its algorithm and convention but vendor no code. Upstream parity: VERIFIED byte-for-byte (max abs diff 0.00e+00) against PrimeIntellect-ai/prime-rl @ f510ef6 across 24 cases. See ``PARITY_VERIFIED.md`` and reproduce with ``verify_parity.sh`` (isolated venv, no vLLM/pydantic deps). Upstream has since refactored the importance-ratio into ``compute_importance_ratio_and_mismatch_kl`` — the line-references above predate that extraction but the math is unchanged; re-run verify_parity.sh after any upstream bump. """ from __future__ import annotations from collections import namedtuple from typing import Any # PRIME-RL's setup_loss_fns expects loss functions to return a LossOutputs # struct with `.loss` (scalar Tensor) and `.metrics` (dict). When PRIME-RL is # installed we use the upstream dataclass directly so isinstance() checks in # any downstream code keep working; otherwise we fall back to a structurally # equivalent NamedTuple that exposes the same attribute access. # # Upstream definition (prime_rl/trainer/rl/loss.py lines 24-29): # @dataclass # class LossOutputs: # loss: Float[Tensor, ""] # metrics: dict[str, Tensor] try: # pragma: no cover - exercised only when prime-rl is installed from prime_rl.trainer.rl.loss import ( # type: ignore[import-not-found] LossOutputs, ) except Exception: # noqa: BLE001 - missing module, version skew, or jaxtyping LossOutputs = namedtuple("LossOutputs", ["loss", "metrics"]) # type: ignore[misc,assignment] def loss_fn( inputs: Any, # PRIME-RL's LossInputs — typed as Any to avoid hard import *, alpha_sdpo: float = 0.0, beta_dpo: float = 0.0, dppo_mask_high: float = 0.2, dppo_mask_low: float = 0.2, adv_tau: float = 1.0, kl_tau: float = 1e-3, ) -> Any: # Returns a torch.Tensor (scalar) matching PRIME-RL's contract """Composer 3-channel loss adapted to PRIME-RL's ``LossInputs`` struct. Channel 1 mirrors PRIME-RL's ``default_loss_fn`` exactly so configs from PRIME-RL's own examples translate. Channels 2 and 3 are deferred — see module docstring. Args: inputs: PRIME-RL ``LossInputs`` (duck-typed). All tensor fields are expected to be 1-D with shape ``(seq,)``. alpha_sdpo: weight on channel 2 (SDPO). Must be 0 in v0; >0 raises :class:`NotImplementedError`. beta_dpo: weight on channel 3 (DPO). Non-zero emits a warning; channel 3 is not yet wired in this recipe. dppo_mask_high: upper DPPO masking threshold on ``exp(trainer_lp) - exp(inference_lp)``. Tokens with **positive advantage** whose ``probs_diff`` exceeds this value are dropped. PRIME-RL default: ``0.2``. Must be >= 0. dppo_mask_low: magnitude of the lower DPPO masking threshold. Tokens with **negative advantage** whose ``probs_diff`` is below ``-dppo_mask_low`` are dropped. PRIME-RL default: ``0.2``. Must be >= 0 (note: PRIME-RL stores the magnitude; the sign flip is internal to the comparison). adv_tau: temperature on the advantage term. PRIME-RL default ``1.0``. Must be >= 0. kl_tau: temperature on the KL term ``log_importance_ratio**2``. PRIME-RL default ``1e-3``. Must be >= 0. Returns: :class:`LossOutputs` with ``loss`` (scalar ``torch.Tensor``) and ``metrics`` (``dict[str, Tensor | float]``). PRIME-RL's outer ``compute_loss`` reads ``out.loss``, divides by ``loss_scale``, and calls ``.backward()``; the ``metrics`` dict is forwarded to the logger. When PRIME-RL is installed this is upstream's ``LossOutputs`` dataclass; otherwise it is a structurally equivalent ``namedtuple`` defined at the top of this module. Raises: ValueError: if any of ``trainer_logprobs``, ``inference_logprobs``, ``advantages``, ``loss_mask`` is not 1-D, or any of ``dppo_mask_high``, ``dppo_mask_low``, ``adv_tau``, ``kl_tau`` is negative. NotImplementedError: if ``alpha_sdpo > 0`` (channel 2 is deferred). """ import torch # lazy — keep module importable without torch installed # PRIME-RL enforces these via Pydantic Field(..., ge=0); we mirror it. for name, val in ( ("dppo_mask_high", dppo_mask_high), ("dppo_mask_low", dppo_mask_low), ("adv_tau", adv_tau), ("kl_tau", kl_tau), ): if val < 0: raise ValueError( f"{name} must be >= 0 (PRIME-RL config contract); got {val}" ) advantages = inputs.advantages trainer_lp = inputs.trainer_logprobs inference_lp = inputs.inference_logprobs loss_mask = inputs.loss_mask # --- Shape validation ------------------------------------------------- # PRIME-RL passes per-sample (seq,) tensors. Reject (B, T) explicitly so # callers don't silently get the wrong reduction. for name, t in ( ("trainer_logprobs", trainer_lp), ("inference_logprobs", inference_lp), ("advantages", advantages), ("loss_mask", loss_mask), ): if t.dim() != 1: raise ValueError( f"PRIME-RL loss_fn expects 1-D (seq,) tensors per " f"PRIME-RL's LossInputs contract; got {name} with shape " f"{tuple(t.shape)} (dim={t.dim()}). PRIME-RL calls the loss " f"function once per sample, not on a batched (B, T) tensor." ) # --- Channel 1: DPPO + KL on the importance ratio -------------------- # Mirrors prime_rl/trainer/rl/loss.py default_loss_fn lines 133-153. log_importance_ratio = trainer_lp - inference_lp importance_ratio = torch.exp(log_importance_ratio) # NOTE: probability-space diff, NOT log-ratio. This is the key # divergence from a naive PPO-clip implementation. probs_diff = torch.exp(trainer_lp) - torch.exp(inference_lp) dppo_invalid_mask_high = probs_diff > dppo_mask_high dppo_invalid_mask_low = probs_diff < -dppo_mask_low positive_advantages = advantages > 0 # Sign-of-advantage gate: positive-advantage tokens use the "high" # threshold; negative-advantage tokens use the "low" threshold. # Zero-advantage tokens fall through ``positive_advantages == False``, # so they are gated by the (negative-advantage) low check; in practice # zero-advantage tokens contribute zero to ``pg_loss`` regardless. dppo_invalid_mask = torch.where( positive_advantages, dppo_invalid_mask_high, dppo_invalid_mask_low ) # loss_mask may be bool; combine via boolean ops to match upstream # exactly, then cast to the working dtype for the multiply. if loss_mask.dtype != torch.bool: loss_mask_bool = loss_mask.to(torch.bool) else: loss_mask_bool = loss_mask keep_mask_bool = loss_mask_bool & ~dppo_invalid_mask keep_mask = keep_mask_bool.to(trainer_lp.dtype) loss_mask_f = loss_mask_bool.to(trainer_lp.dtype) scaled_advantages = adv_tau * advantages pg_loss = keep_mask * scaled_advantages * importance_ratio kl_loss = loss_mask_f * log_importance_ratio**2 total = (-pg_loss + kl_tau * kl_loss).sum() # --- Channel 2: SDPO/OPSD — DEFERRED in PRIME-RL recipe v0 ----------- # # Wave 13 cross-model review caught that an earlier draft applied # `unsqueeze(-1)` to log-prob tensors before generalized_jsd_loss, # which calls log_softmax(dim=-1). Softmax of a 1-element vector is # exactly 1.0; its log is 0. The SDPO term was mathematically # degenerate (always 0), silently disabling channel 2 while reporting # alpha_sdpo>0 in the config. Until PRIME-RL exposes full logits we # refuse to fake the channel: teacher_lp = getattr(inputs, "teacher_logprobs", None) if alpha_sdpo > 0: raise NotImplementedError( "SDPO channel in the PRIME-RL recipe is deferred. PRIME-RL " "v0.5 exposes (seq,) log-probs through LossInputs but not " "full vocabulary logits, and SDPO/OPSD requires the full " "distribution. Set alpha_sdpo=0.0 to silence this and use " "channel 1 (DPPO+KL) only. teacher_logprobs is " f"{'present' if teacher_lp is not None else 'absent'} in this " "call but unused. For the SDPO channel, use the TRL host " "(composer_replication.trainer.ComposerReplicationTrainer with " "alpha_sdpo>0), which has full logits — see ADR-008. " "See docs/research/WAVE_13_FINAL_REVIEW.md Finding 1." ) # --- Channel 3: not supported in PRIME-RL recipe v0 ------------------- if beta_dpo != 0.0: import warnings warnings.warn( "PRIME-RL recipe v0 does not support DPO channel; " "set beta_dpo=0.0 to silence this warning.", stacklevel=2, ) return LossOutputs(loss=total, metrics={"channel_1_pg_loss": float(total.detach())}) __all__ = ["loss_fn", "LossOutputs"]