Codeseys's picture
Wave 21c: verify PRIME-RL adapter parity against upstream source (byte-for-byte)
c98928e
|
Raw
History Blame Contribute Delete
2.53 kB
# 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.