composer_replication/recipes/prime_rl/PARITY_VERIFIED.md

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.