"""Unit tests for the PRIME-RL composer-loss adapter. Verifies parity with PRIME-RL's upstream ``default_loss_fn`` (``src/prime_rl/trainer/rl/loss.py`` lines 116-165). Hand-computed expected values use the upstream formula; the parity test at the bottom imports PRIME-RL itself (skip-marked when not installed) and compares outputs end-to-end. License: MIT. """ from __future__ import annotations import math from dataclasses import dataclass from typing import Optional import pytest import torch import warnings from composer_replication.recipes.prime_rl.composer_loss import LossOutputs, loss_fn def _loss_value(result) -> torch.Tensor: """Return the scalar loss tensor from either a LossOutputs struct or a bare Tensor. The recipe wraps its return in LossOutputs to satisfy PRIME-RL's setup_loss_fns contract; tests written against the older bare-Tensor return path keep working through this helper. """ if isinstance(result, torch.Tensor): return result # LossOutputs: dataclass (upstream) or namedtuple (fallback). return result.loss # Try to import PRIME-RL upstream for the parity test; skip-mark if # unavailable. PRIME-RL pulls in heavy deps (jaxtyping, beartype) and # is not part of the framework's own test environment. # # Visibility: when the import fails we emit a UserWarning at module load # so the skip is *visible* in pytest output ("PytestUnhandledThreadExceptionWarning" # is too noisy; UserWarning is captured by pytest's default filterwarnings # and printed in the run summary). Without this, CI without prime-rl # silently never runs the parity test and a real divergence could go # undetected for releases at a time. try: from prime_rl.trainer.rl.loss import ( # type: ignore[import-not-found] LossInputs as PrimeRLLossInputs, default_loss_fn as prime_rl_default_loss_fn, ) from prime_rl.configs.trainer import ( # type: ignore[import-not-found] DefaultLossConfig as PrimeRLDefaultLossConfig, ) _HAS_PRIME_RL = True except Exception: # noqa: BLE001 — broad: missing module, version skew, etc. _HAS_PRIME_RL = False warnings.warn( "prime-rl is not importable in this environment; the upstream " "parity test (test_parity_with_prime_rl_default_loss_fn) will be " "skipped. The shadow-parity test below still runs against an " "in-file reference reimplementation.", UserWarning, stacklevel=2, ) # --------------------------------------------------------------------- # Test double — duck-typed stand-in for PRIME-RL's LossInputs # --------------------------------------------------------------------- @dataclass class FakeLossInputs: trainer_logprobs: torch.Tensor inference_logprobs: torch.Tensor advantages: torch.Tensor loss_mask: torch.Tensor teacher_logprobs: Optional[torch.Tensor] = None def _make_inputs( seq: int = 8, *, same_logprobs: bool = True, teacher: bool = False, seed: int = 0, ) -> FakeLossInputs: """Build a realistic (seq,) LossInputs stand-in. Uses ``requires_grad`` on ``trainer_logprobs`` so callers can also sanity-check that the loss is differentiable end-to-end. Default log-probs are clamped to a moderate negative range so ``exp(trainer_lp) - exp(inference_lp)`` stays inside the 0.2 PRIME-RL default DPPO band — i.e. tokens are not all DPPO-masked by chance. """ g = torch.Generator().manual_seed(seed) # Negative log-probs in [-2, -0.5] keep exp() in roughly [0.13, 0.6] # so probs_diff differences stay tiny under small perturbation. trainer = -(0.5 + 1.5 * torch.rand(seq, generator=g)) trainer = trainer.detach().clone().requires_grad_(True) if same_logprobs: # Tiny perturbation -> probs_diff ~ 0, no DPPO masking. inference = trainer.detach().clone() + 0.001 * torch.randn( seq, generator=g ) else: inference = -(0.5 + 1.5 * torch.rand(seq, generator=g)) advantages = torch.randn(seq, generator=g) loss_mask = torch.ones(seq, dtype=torch.bool) teacher_lp = torch.randn(seq, generator=g) if teacher else None return FakeLossInputs( trainer_logprobs=trainer, inference_logprobs=inference, advantages=advantages, loss_mask=loss_mask, teacher_logprobs=teacher_lp, ) # --------------------------------------------------------------------- # Reference re-implementation (independent restatement of upstream). # Used by hand-computed expected-value tests so we don't accidentally # encode our own bugs as ground truth. # # SHADOW-PARITY MAPPING # --------------------- # The body below is structurally identical to PRIME-RL's # ``default_loss_fn`` at ``src/prime_rl/trainer/rl/loss.py`` lines # 116-153 (commit pinned by /tmp/prime-rl-clone clone). The mapping, # line-by-line, is: # # upstream line 133-135 -> ``log_ir = ...``, # ``ir = torch.exp(log_ir)`` # (we elide the unused ``mismatch_kl`` # term — upstream returns it as a metric # only; we drop metrics in the reference # because our channel-1 loss is a scalar # and we compare ``.loss`` only.) # upstream line 137 -> ``probs_diff = exp(trainer_lp) - exp(inference_lp)`` # upstream line 138 -> ``invalid_high = probs_diff > dppo_mask_high`` # upstream line 139 -> ``invalid_low = probs_diff < -dppo_mask_low`` # upstream line 140 -> ``pos_adv = advantages > 0`` # upstream line 142 -> ``invalid = where(pos_adv, invalid_high, invalid_low)`` # upstream line 148 -> ``keep = loss_mask & ~invalid`` # (upstream uses ``& is_masked``; we # pre-cast ``loss_mask`` via ``to(bool)``) # upstream line 150 -> ``adv_tau * advantages`` (inlined) # upstream line 151 -> ``pg = keep_f * (adv_tau * advantages) * ir`` # upstream line 152 -> ``kl = lm_f * log_ir**2`` # upstream line 153 -> ``return (-pg + kl_tau * kl).sum()`` # # Differences (intentional, do not affect ``.loss``): # * upstream returns ``LossOutputs(loss=..., metrics={...})``; we # return only the loss scalar because the seven metric entries # (lines 155-163) don't influence backward and are validated # separately in ``test_parity_with_prime_rl_default_loss_fn``. # * upstream casts via ``loss_mask & is_masked`` (Bool & Bool); our # ``keep_f.to(trainer_lp.dtype)`` matches exactly because both # ``keep_mask`` and ``loss_mask`` are bool tensors broadcast to # ``trainer_lp.dtype`` for the float multiply. # --------------------------------------------------------------------- def _reference_default_loss( trainer_lp: torch.Tensor, inference_lp: torch.Tensor, advantages: torch.Tensor, loss_mask: torch.Tensor, *, dppo_mask_high: float, dppo_mask_low: float, adv_tau: float, kl_tau: float, ) -> torch.Tensor: log_ir = trainer_lp - inference_lp ir = torch.exp(log_ir) probs_diff = torch.exp(trainer_lp) - torch.exp(inference_lp) invalid_high = probs_diff > dppo_mask_high invalid_low = probs_diff < -dppo_mask_low pos_adv = advantages > 0 invalid = torch.where(pos_adv, invalid_high, invalid_low) keep = loss_mask.to(torch.bool) & ~invalid keep_f = keep.to(trainer_lp.dtype) lm_f = loss_mask.to(trainer_lp.dtype) pg = keep_f * (adv_tau * advantages) * ir kl = lm_f * log_ir**2 return (-pg + kl_tau * kl).sum() # --------------------------------------------------------------------- # Test 1 — finite scalar on realistic (seq,) tensors # --------------------------------------------------------------------- def test_returns_finite_scalar(): inputs = _make_inputs(seq=16) result = loss_fn(inputs, alpha_sdpo=0.0, beta_dpo=0.0) # Must be a LossOutputs (dataclass when prime-rl is installed, # NamedTuple fallback otherwise). PRIME-RL's setup_loss_fns reads # ``.loss`` and ``.metrics`` from this struct. assert hasattr(result, "loss") and hasattr(result, "metrics"), ( f"loss_fn must return a LossOutputs-shaped struct; got {type(result)}" ) assert isinstance(result.metrics, dict) assert "channel_1_pg_loss" in result.metrics out = result.loss assert isinstance(out, torch.Tensor) assert out.shape == (), f"expected scalar, got shape {tuple(out.shape)}" assert torch.isfinite(out).item() # Differentiable: gradient flows to trainer_logprobs. out.backward() assert inputs.trainer_logprobs.grad is not None assert torch.isfinite(inputs.trainer_logprobs.grad).all().item() # --------------------------------------------------------------------- # Test 2 — DPPO mask drops tokens whose probs_diff exceeds dppo_mask_high # (advantage-conditioned: positive advantages use the high gate) # --------------------------------------------------------------------- def test_dppo_mask_high_drops_positive_advantage_outliers(): """Token with positive advantage and probs_diff > dppo_mask_high is dropped. Build a 4-token sample where token 0 has ``probs_diff`` huge and positive (trainer prob ~ 1, inference prob ~ 0) AND positive advantage. Tokens 1..3 have tiny probs_diff. With the upstream sign-conditioned gate, only token 0 should be dropped. """ # trainer_lp ~ 0 -> exp ~ 1; inference_lp = -10 -> exp ~ 4.5e-5. # probs_diff[0] ~ 1.0 >> dppo_mask_high (0.2). trainer_lp = torch.tensor( [0.0, math.log(0.30), math.log(0.40), math.log(0.50)], requires_grad=True, ) inference_lp = torch.tensor( [-10.0, math.log(0.31), math.log(0.39), math.log(0.51)] ) advantages = torch.tensor([+5.0, +1.0, -1.0, +1.0]) mask = torch.ones(4, dtype=torch.bool) inputs = FakeLossInputs( trainer_logprobs=trainer_lp, inference_logprobs=inference_lp, advantages=advantages, loss_mask=mask, ) out = _loss_value(loss_fn( inputs, alpha_sdpo=0.0, beta_dpo=0.0, dppo_mask_high=0.2, dppo_mask_low=0.2, adv_tau=1.0, kl_tau=1e-3, )) expected = _reference_default_loss( trainer_lp.detach(), inference_lp, advantages, mask, dppo_mask_high=0.2, dppo_mask_low=0.2, adv_tau=1.0, kl_tau=1e-3, ) assert torch.isclose(out, expected, atol=1e-5), ( f"got {out.item()}, expected {expected.item()}" ) # Token 0 was DPPO-dropped from pg_loss but still contributes to kl_loss # (loss_mask gates KL, not the DPPO mask). The pg gradient on token 0 # should be zero; KL contributes a small grad. We assert the pg path # is masked by checking the gradient magnitude is dominated by the # tiny kl_tau * 2 * log_ir term, not by the +5 advantage. out.backward() g0 = inputs.trainer_logprobs.grad[0].item() # If pg weren't masked, |g0| would be on the order of # advantage * importance_ratio * 1 ~ 5 * exp(10) ~ 1e5. # With pg masked, |g0| is on the order of # 2 * kl_tau * log_ir ~ 2 * 1e-3 * 10 = 0.02. assert abs(g0) < 1.0, ( f"DPPO mask should suppress the pg gradient on token 0; got |g0|={abs(g0)}" ) # --------------------------------------------------------------------- # Test 3 — DPPO mask catches the lower bound on negative-advantage tokens # --------------------------------------------------------------------- def test_dppo_mask_low_drops_negative_advantage_outliers(): """Symmetric coverage: probs_diff < -dppo_mask_low drops a NEGATIVE-adv token.""" # Token 0: trainer prob ~ 0, inference prob ~ 1, so probs_diff ~ -1. # Negative advantage -> the low gate applies -> dropped. trainer_lp = torch.tensor( [-10.0, math.log(0.30), math.log(0.40)], requires_grad=True ) inference_lp = torch.tensor( [0.0, math.log(0.31), math.log(0.39)] ) advantages = torch.tensor([-5.0, +1.0, -1.0]) mask = torch.ones(3, dtype=torch.bool) inputs = FakeLossInputs( trainer_logprobs=trainer_lp, inference_logprobs=inference_lp, advantages=advantages, loss_mask=mask, ) out = _loss_value(loss_fn(inputs, alpha_sdpo=0.0, beta_dpo=0.0)) expected = _reference_default_loss( trainer_lp.detach(), inference_lp, advantages, mask, dppo_mask_high=0.2, dppo_mask_low=0.2, adv_tau=1.0, kl_tau=1e-3, ) assert torch.isclose(out, expected, atol=1e-5) # --------------------------------------------------------------------- # Test 4 — sign-conditioning: a positive-advantage token whose probs_diff # is *negative* (and large in magnitude) is NOT dropped, because the # high gate doesn't fire on a negative probs_diff. # --------------------------------------------------------------------- def test_dppo_mask_sign_conditioned_on_advantage(): """A positive-advantage token with probs_diff < -dppo_mask_low survives. PRIME-RL's gate is ``where(positive_advantages, invalid_high, invalid_low)``. For positive advantages it only checks the upper bound, so ``probs_diff = -0.9`` with a positive advantage is KEPT; with a negative advantage it would be DROPPED. """ # Token 0: probs_diff = exp(-10) - exp(0) ~ -1. Massively negative. trainer_lp_pos = torch.tensor([-10.0], requires_grad=True) inference_lp_pos = torch.tensor([0.0]) adv_pos = torch.tensor([+1.0]) mask = torch.ones(1, dtype=torch.bool) inputs_pos = FakeLossInputs( trainer_logprobs=trainer_lp_pos, inference_logprobs=inference_lp_pos, advantages=adv_pos, loss_mask=mask, ) out_pos = _loss_value(loss_fn(inputs_pos, alpha_sdpo=0.0, beta_dpo=0.0)) # With positive advantage the LOW bound is not checked; the token is # KEPT. pg = +1 * exp(-10 - 0) = ~4.5e-5; kl = (-10)^2 = 100. # loss = -pg + 1e-3 * 100 ~ 0.1. expected_pos = _reference_default_loss( trainer_lp_pos.detach(), inference_lp_pos, adv_pos, mask, dppo_mask_high=0.2, dppo_mask_low=0.2, adv_tau=1.0, kl_tau=1e-3, ) assert torch.isclose(out_pos, expected_pos, atol=1e-5) # Sanity: token wasn't masked, so kl_tau alone shouldn't dominate to # zero — loss should be ~0.1, definitely not zero. assert out_pos.item() > 0.05 # Same probs_diff but negative advantage -> DROPPED from pg. trainer_lp_neg = torch.tensor([-10.0], requires_grad=True) inputs_neg = FakeLossInputs( trainer_logprobs=trainer_lp_neg, inference_logprobs=inference_lp_pos, advantages=torch.tensor([-1.0]), loss_mask=mask, ) out_neg = _loss_value(loss_fn(inputs_neg, alpha_sdpo=0.0, beta_dpo=0.0)) expected_neg = _reference_default_loss( trainer_lp_neg.detach(), inference_lp_pos, torch.tensor([-1.0]), mask, dppo_mask_high=0.2, dppo_mask_low=0.2, adv_tau=1.0, kl_tau=1e-3, ) assert torch.isclose(out_neg, expected_neg, atol=1e-5) # --------------------------------------------------------------------- # Test 5 — alpha_sdpo=0 must not raise (channel 2 disabled) # --------------------------------------------------------------------- def test_alpha_sdpo_zero_does_not_raise(): inputs = _make_inputs(seq=6, teacher=True) out = _loss_value(loss_fn(inputs, alpha_sdpo=0.0, beta_dpo=0.0)) assert torch.isfinite(out).item() # --------------------------------------------------------------------- # Test 6 — alpha_sdpo>0 still raises NotImplementedError # --------------------------------------------------------------------- def test_alpha_sdpo_nonzero_raises_not_implemented(): inputs = _make_inputs(seq=6, teacher=True) with pytest.raises(NotImplementedError, match="SDPO"): loss_fn(inputs, alpha_sdpo=0.5, beta_dpo=0.0) def test_alpha_sdpo_nonzero_no_teacher_also_raises(): """Defensive: even without teacher_logprobs, alpha_sdpo>0 must fail rather than silently no-op.""" inputs = _make_inputs(seq=6, teacher=False) with pytest.raises(NotImplementedError): loss_fn(inputs, alpha_sdpo=0.5, beta_dpo=0.0) # --------------------------------------------------------------------- # Test 7 — shape validation: (seq,) accepted, (B, T) rejected # --------------------------------------------------------------------- def test_advantages_shape_validates_seq_accepted(): inputs = _make_inputs(seq=12) out = _loss_value(loss_fn(inputs, alpha_sdpo=0.0, beta_dpo=0.0)) assert out.shape == () def test_advantages_shape_validates_bt_rejected(): B, T = 2, 4 bad = FakeLossInputs( trainer_logprobs=torch.zeros(B, T, requires_grad=True), inference_logprobs=torch.zeros(B, T), advantages=torch.zeros(B, T), loss_mask=torch.ones(B, T, dtype=torch.bool), ) with pytest.raises(ValueError, match="1-D"): loss_fn(bad, alpha_sdpo=0.0, beta_dpo=0.0) # --------------------------------------------------------------------- # Test 8 — beta_dpo != 0 emits a warning but does not raise # --------------------------------------------------------------------- def test_beta_dpo_nonzero_warns(): inputs = _make_inputs(seq=8) with pytest.warns(UserWarning, match="DPO channel"): out = _loss_value(loss_fn(inputs, alpha_sdpo=0.0, beta_dpo=0.3)) assert torch.isfinite(out).item() # --------------------------------------------------------------------- # Test 9 — config-validation knobs match PRIME-RL Field(..., ge=0) # --------------------------------------------------------------------- @pytest.mark.parametrize( "kw", [ {"dppo_mask_high": -0.1}, {"dppo_mask_low": -0.1}, {"adv_tau": -0.1}, {"kl_tau": -0.1}, ], ) def test_negative_knobs_rejected(kw): inputs = _make_inputs(seq=4) with pytest.raises(ValueError, match=">= 0"): loss_fn(inputs, alpha_sdpo=0.0, beta_dpo=0.0, **kw) # --------------------------------------------------------------------- # Test 10 — disabling masking via wide bounds gives plain DPPO+KL on all # tokens. This pins the "pure IS-corrected REINFORCE + KL" baseline. # --------------------------------------------------------------------- def test_dppo_bounds_can_be_disabled(): """Setting bounds to a huge value disables DPPO masking. At dppo_mask_high=dppo_mask_low=1e6, ``probs_diff`` never exceeds the threshold so ``keep_mask == loss_mask`` and the loss reduces to the plain DPPO+KL on the whole sequence. """ seq = 4 trainer_lp = torch.tensor( [math.log(0.10), math.log(0.30), math.log(0.20), math.log(0.40)], requires_grad=True, ) inference_lp = torch.tensor( [math.log(0.11), math.log(0.31), math.log(0.21), math.log(0.39)] ) advantages = torch.tensor([+1.0, -1.0, +0.5, -0.5]) mask = torch.ones(seq, dtype=torch.bool) inputs = FakeLossInputs( trainer_logprobs=trainer_lp, inference_logprobs=inference_lp, advantages=advantages, loss_mask=mask, ) out = _loss_value(loss_fn( inputs, alpha_sdpo=0.0, beta_dpo=0.0, dppo_mask_high=1e6, dppo_mask_low=1e6, adv_tau=1.0, kl_tau=1e-3, )) expected = _reference_default_loss( trainer_lp.detach(), inference_lp, advantages, mask, dppo_mask_high=1e6, dppo_mask_low=1e6, adv_tau=1.0, kl_tau=1e-3, ) assert torch.isclose(out, expected, atol=1e-6) # --------------------------------------------------------------------- # Test 11 — PARITY against PRIME-RL upstream's default_loss_fn. # Skip-marked when prime-rl is not installable. # --------------------------------------------------------------------- @pytest.mark.skipif( not _HAS_PRIME_RL, reason="prime-rl not installed; skipping upstream parity test", ) def test_parity_with_prime_rl_default_loss_fn(): """Run identical inputs through ours and PRIME-RL's; loss must match.""" seq = 32 g = torch.Generator().manual_seed(42) trainer_lp = -(0.1 + 2.0 * torch.rand(seq, generator=g)).to(torch.float32) inference_lp = (trainer_lp + 0.05 * torch.randn(seq, generator=g)).to(torch.float32) advantages = torch.randn(seq, generator=g, dtype=torch.float32) loss_mask = torch.ones(seq, dtype=torch.bool) # Use PRIME-RL's defaults (dppo_mask_high=0.2, etc.) directly. cfg = PrimeRLDefaultLossConfig() # type: ignore[name-defined] upstream_inputs = PrimeRLLossInputs( # type: ignore[name-defined] trainer_logprobs=trainer_lp, inference_logprobs=inference_lp, teacher_logprobs=None, advantages=advantages, loss_mask=loss_mask, ) upstream_out = prime_rl_default_loss_fn(upstream_inputs, cfg) # type: ignore[name-defined] ours = _loss_value(loss_fn( FakeLossInputs( trainer_logprobs=trainer_lp.clone(), inference_logprobs=inference_lp.clone(), advantages=advantages.clone(), loss_mask=loss_mask.clone(), ), alpha_sdpo=0.0, beta_dpo=0.0, dppo_mask_high=cfg.dppo_mask_high, dppo_mask_low=cfg.dppo_mask_low, adv_tau=cfg.adv_tau, kl_tau=cfg.kl_tau, )) assert torch.isclose(ours, upstream_out.loss, atol=1e-5, rtol=1e-5), ( f"Parity mismatch with PRIME-RL upstream: ours={ours.item()}, " f"upstream={upstream_out.loss.item()}" )