composer_replication/recipes/prime_rl/composer_loss.py

"""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"]