composer_replication/recipes/prime_rl/_parity_harness.py

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