Wave 21c: verify PRIME-RL adapter parity against upstream source (byte-for-byte)

The recipe's only automated parity check (test_parity_with_prime_rl_default_loss_fn)
is skip-marked whenever prime-rl isn't importable in the framework venv — i.e.
essentially always, since we deliberately keep prime-rl's heavy deps (vLLM,
pydantic config trees, flash-attn) out of our test env. The fallback was an
in-file reimplementation (_reference_default_loss), so a shared bug could pass
silently. This closes that gap against the ACTUAL upstream source.

## What was done
Built an out-of-band parity harness that:
- clones PrimeIntellect-ai/prime-rl (shallow)
- builds an isolated venv with ONLY torch+beartype+jaxtyping+numpy
- loads upstream src/prime_rl/trainer/rl/loss.py directly by path, stubbing the
two modules it imports (prime_rl.configs.trainer, prime_rl.utils.utils) so we
skip the vLLM/pydantic dependency tree entirely
- runs identical inputs through our loss_fn and upstream default_loss_fn

## Result
24/24 cases match (12 seeds × 2 regimes: tiny-perturbation + wide-divergence,
the latter exercising both DPPO masking branches), with partial loss masks.
**Max absolute difference 0.00e+00** — bit-identical, not merely within tolerance.
Upstream rev: f510ef6 (2026-05-28).

## Added
- composer_replication/recipes/prime_rl/verify_parity.sh: one-command reproducible
check (clones + isolated venv + sweep). Exit 0 = parity confirmed.
- _parity_harness.py: the sweep harness it runs.
- PARITY_VERIFIED.md: result + provenance + reproduce instructions.
- composer_loss.py docstring: notes parity is verified and that upstream
refactored the importance-ratio into compute_importance_ratio_and_mismatch_kl
(math unchanged); re-run verify_parity.sh after any upstream bump.

## Tests
Recipe unit tests: 15 passed, 1 skipped (the in-venv parity test still skips by
design — the out-of-band script is its reproducible counterpart).

composer_replication/recipes/prime_rl/PARITY_VERIFIED.md

@@ -0,0 +1,50 @@
+# PRIME-RL upstream parity — VERIFIED
+
+**Status:** PASS ✅ — our adapter's Channel-1 loss matches PrimeIntellect-ai/prime-rl's
+upstream `default_loss_fn` **byte-for-byte** (max absolute difference `0.00e+00`).
+
+## What was verified
+
+`composer_replication/recipes/prime_rl/composer_loss.py::loss_fn` (Channel 1:
+DPPO + KL on the importance-sampling ratio, with the advantage-sign-conditioned
+DPPO mask) produces numerically identical loss to upstream
+`prime_rl.trainer.rl.loss.default_loss_fn` across:
+
+- **12 random seeds × 2 regimes** (24 cases total)
+  - `tiny_perturb`: inference ≈ trainer + small noise → no DPPO masking (the
+    common on-policy regime)
+  - `wide_diff`: large trainer/inference divergence → exercises both the
+    `dppo_invalid_mask_high` (positive-advantage) and `dppo_invalid_mask_low`
+    (negative-advantage) branches hard
+- partial loss masks (~10% of tokens masked out)
+- PRIME-RL's own default config (`dppo_mask_low=0.2`, `dppo_mask_high=0.2`,
+  `adv_tau=1.0`, `kl_tau=1e-3`)
+
+Result: **24/24 exact matches**, max abs diff `0.00e+00` (not merely within
+`atol=1e-5` — bit-identical for these inputs).
+
+## Provenance
+
+- Upstream: `PrimeIntellect-ai/prime-rl` @ `f510ef6` (2026-05-28)
+- Verified by loading upstream `src/prime_rl/trainer/rl/loss.py` directly by path
+  in an isolated venv (torch+beartype+jaxtyping+numpy only — no vLLM, no pydantic
+  config tree), with `prime_rl.configs.trainer` / `prime_rl.utils.utils` stubbed.
+- Reproduce: `bash composer_replication/recipes/prime_rl/verify_parity.sh`
+
+## Why this matters
+
+Previously the only automated check was `test_parity_with_prime_rl_default_loss_fn`,
+which is skip-marked whenever prime-rl isn't importable in the framework venv —
+i.e. essentially always, because we deliberately keep prime-rl's heavy deps out of
+our test env. The fallback `_reference_default_loss` in the unit tests is an *in-file
+reimplementation*, so a shared bug between it and `loss_fn` would pass silently.
+This out-of-band check closes that gap against the **actual upstream source**.
+
+## Note on upstream drift
+
+Upstream refactored the importance-ratio computation into a helper
+(`compute_importance_ratio_and_mismatch_kl`) since the line-references in
+`composer_loss.py`'s docstring were written. The **math is unchanged** — the helper
+just extracts `log_importance_ratio / importance_ratio / mismatch_kl`. Our adapter
+remains exact against current `f510ef6`. Re-run `verify_parity.sh` after any
+upstream bump to catch a real divergence early.

composer_replication/recipes/prime_rl/_parity_harness.py

@@ -0,0 +1,132 @@
+"""Isolated PRIME-RL parity harness — runs OUR adapter vs UPSTREAM default_loss_fn
+byte-for-byte, without installing the full prime-rl package (which drags vLLM,
+pydantic config trees, etc.).
+
+Strategy: stub the two modules upstream loss.py imports (`prime_rl.configs.trainer`
+for DefaultLossConfig + CustomLossConfig + LossConfig, and `prime_rl.utils.utils`
+for import_object), then load loss.py by file path. Compare on random inputs.
+
+Run with the throwaway venv that has torch+beartype+jaxtyping+numpy:
+    /tmp/prime-parity-venv/bin/python this_file.py /path/to/prime-rl /path/to/framework
+"""
+import importlib.util
+import sys
+import types
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+PRIME_RL = Path(sys.argv[1])
+FRAMEWORK = Path(sys.argv[2])
+
+# --- Stub the config + utils modules loss.py needs at import time -----------
+cfg_mod = types.ModuleType("prime_rl.configs.trainer")
+
+
+@dataclass
+class DefaultLossConfig:
+    # Exact upstream defaults (trainer.py lines 412-425).
+    dppo_mask_low: float = 0.2
+    dppo_mask_high: float = 0.2
+    adv_tau: float = 1.0
+    kl_tau: float = 1e-3
+
+
+class CustomLossConfig:  # only referenced in type hints / isinstance paths
+    pass
+
+
+class LossConfig:
+    pass
+
+
+cfg_mod.DefaultLossConfig = DefaultLossConfig
+cfg_mod.CustomLossConfig = CustomLossConfig
+cfg_mod.LossConfig = LossConfig
+
+utils_mod = types.ModuleType("prime_rl.utils.utils")
+utils_mod.import_object = lambda path: None  # unused by default_loss_fn
+
+# Register stub package tree so `from prime_rl.configs.trainer import ...` resolves.
+for name in ("prime_rl", "prime_rl.configs", "prime_rl.utils"):
+    sys.modules.setdefault(name, types.ModuleType(name))
+sys.modules["prime_rl.configs.trainer"] = cfg_mod
+sys.modules["prime_rl.utils.utils"] = utils_mod
+
+# --- Load upstream loss.py by path ------------------------------------------
+loss_path = PRIME_RL / "src" / "prime_rl" / "trainer" / "rl" / "loss.py"
+spec = importlib.util.spec_from_file_location("prime_rl.trainer.rl.loss", loss_path)
+upstream = importlib.util.module_from_spec(spec)
+sys.modules["prime_rl.trainer.rl.loss"] = upstream
+spec.loader.exec_module(upstream)
+print(f"loaded upstream loss.py from {loss_path}")
+
+# --- Load our adapter -------------------------------------------------------
+sys.path.insert(0, str(FRAMEWORK))
+from composer_replication.recipes.prime_rl.composer_loss import loss_fn as ours  # noqa: E402
+
+
+@dataclass
+class FakeLossInputs:
+    trainer_logprobs: torch.Tensor
+    inference_logprobs: torch.Tensor
+    teacher_logprobs: object
+    advantages: torch.Tensor
+    loss_mask: torch.Tensor
+
+
+# --- Parity sweep across seeds + regimes ------------------------------------
+cfg = DefaultLossConfig()
+n_pass = 0
+n_total = 0
+max_abs_diff = 0.0
+for seed in range(12):
+    for regime in ("tiny_perturb", "wide_diff"):
+        g = torch.Generator().manual_seed(seed)
+        seq = 32
+        trainer_lp = -(0.1 + 2.0 * torch.rand(seq, generator=g)).to(torch.float32)
+        if regime == "tiny_perturb":
+            inference_lp = (trainer_lp + 0.05 * torch.randn(seq, generator=g)).to(torch.float32)
+        else:
+            # Large divergence -> exercises the DPPO masking branches hard.
+            inference_lp = -(0.1 + 2.0 * torch.rand(seq, generator=g)).to(torch.float32)
+        advantages = torch.randn(seq, generator=g, dtype=torch.float32)
+        loss_mask = (torch.rand(seq, generator=g) > 0.1)  # ~10% masked out
+
+        up_inputs = upstream.LossInputs(
+            trainer_logprobs=trainer_lp,
+            inference_logprobs=inference_lp,
+            teacher_logprobs=None,
+            advantages=advantages,
+            loss_mask=loss_mask,
+        )
+        up_out = upstream.default_loss_fn(up_inputs, cfg)
+
+        our_out = ours(
+            FakeLossInputs(
+                trainer_logprobs=trainer_lp.clone(),
+                inference_logprobs=inference_lp.clone(),
+                teacher_logprobs=None,
+                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,
+        )
+        our_loss = our_out.loss if hasattr(our_out, "loss") else our_out
+        diff = abs(float(our_loss) - float(up_out.loss))
+        max_abs_diff = max(max_abs_diff, diff)
+        ok = torch.isclose(our_loss, up_out.loss, atol=1e-5, rtol=1e-5).item()
+        n_total += 1
+        n_pass += int(ok)
+        if not ok:
+            print(f"  MISMATCH seed={seed} {regime}: ours={float(our_loss):.6f} up={float(up_out.loss):.6f} diff={diff:.2e}")
+
+print(f"\nPARITY: {n_pass}/{n_total} cases match upstream (max abs diff {max_abs_diff:.2e})")
+print("RESULT:", "PASS ✅" if n_pass == n_total else "FAIL ❌")
+sys.exit(0 if n_pass == n_total else 1)

composer_replication/recipes/prime_rl/composer_loss.py

@@ -85,6 +85,14 @@ 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
 

composer_replication/recipes/prime_rl/verify_parity.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# Verify our PRIME-RL composer-loss adapter matches UPSTREAM default_loss_fn
+# byte-for-byte, WITHOUT installing the full prime-rl package (which pulls vLLM,
+# pydantic config trees, flash-attn, etc.). We clone prime-rl, build a throwaway
+# venv with only torch+beartype+jaxtyping+numpy, load upstream loss.py by path
+# with stubbed config/utils modules, and run identical inputs through both.
+#
+# Usage:
+#   bash composer_replication/recipes/prime_rl/verify_parity.sh
+#
+# Exit 0 = byte-for-byte parity confirmed; non-zero = mismatch or setup failure.
+#
+# This is the reproducible counterpart to the skip-marked
+# test_parity_with_prime_rl_default_loss_fn unit test: that test only runs when
+# prime-rl is importable in the framework venv (it usually isn't, by design —
+# we don't want prime-rl's heavy deps in our test env). This script provides the
+# real upstream check out-of-band.
+set -euo pipefail
+
+PRIME_RL_REPO="${PRIME_RL_REPO:-https://github.com/PrimeIntellect-ai/prime-rl.git}"
+WORK="${WORK:-/tmp/prime-rl-parity-check}"
+FRAMEWORK="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
+CLONE="$WORK/prime-rl"
+VENV="$WORK/venv"
+HARNESS="$WORK/harness.py"
+
+mkdir -p "$WORK"
+
+echo "==> Cloning prime-rl (shallow) into $CLONE"
+if [ ! -d "$CLONE/.git" ]; then
+    git clone --depth 1 "$PRIME_RL_REPO" "$CLONE"
+fi
+PRIME_REV="$(cd "$CLONE" && git rev-parse --short HEAD)"
+echo "    upstream rev: $PRIME_REV"
+
+echo "==> Building isolated venv (torch+beartype+jaxtyping+numpy only)"
+if [ ! -x "$VENV/bin/python" ]; then
+    python3 -m venv "$VENV"
+    "$VENV/bin/pip" install --quiet --upgrade pip
+    # CPU torch is plenty for a loss-numerics parity check.
+    "$VENV/bin/pip" install --quiet torch --index-url https://download.pytorch.org/whl/cpu
+    "$VENV/bin/pip" install --quiet beartype jaxtyping numpy
+fi
+
+echo "==> Writing parity harness"
+cp "$FRAMEWORK/composer_replication/recipes/prime_rl/_parity_harness.py" "$HARNESS"
+
+echo "==> Running parity sweep"
+"$VENV/bin/python" "$HARNESS" "$CLONE" "$FRAMEWORK"