| """Adversarial-review additions for the P4 CISPO objective + FP32 head. |
| |
| Each test pins a failure mode the builder's suite left open (see the docstring |
| of each). Oracles are independent fp64 computations or exact-bit comparisons. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import pytest |
| import torch |
|
|
| import align.rl.cispo as cispo_mod |
| from align.rl.cispo import ( |
| P4CispoConformanceError, |
| assert_p4_cispo_conformant, |
| cispo_policy_loss, |
| fp32_head_logprobs, |
| ) |
| from align.rl.grpo_loss import A3ConformanceError, GRPOLossConfig, GRPOLossType, assert_a3_conformant |
| from align.rl.grpo_module import GRPOTrainModule |
| from tests.test_cispo import _Bf16LM, _Sgd, _batch, _hw, _p4_cfg |
|
|
| D = torch.float64 |
|
|
|
|
| |
|
|
|
|
| @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) |
| def test_ratio_exactly_at_clip_boundary(dtype): |
| """r == 1+eps_high is NOT clipped (weight == r, clipfrac 0); the next |
| representable value above IS (weight pinned at the cap, clipfrac 1), and the |
| weight is continuous across the boundary. Same on the low side with a floor. |
| Failure mode: `>=` vs `>` flips the clip metric at the single most common |
| value (ratio==1+eps after eps-sized drift) without changing the loss.""" |
| eps_hi, eps_lo = 0.25, 0.5 |
| hi = torch.tensor(1.0 + eps_hi, dtype=dtype) |
| lo = torch.tensor(1.0 - eps_lo, dtype=dtype) |
| ratios = torch.stack([torch.nextafter(hi, torch.tensor(0.0, dtype=dtype)), hi, |
| torch.nextafter(hi, torch.tensor(9.0, dtype=dtype)), |
| torch.nextafter(lo, torch.tensor(9.0, dtype=dtype)), lo, |
| torch.nextafter(lo, torch.tensor(0.0, dtype=dtype))]) |
| lp = torch.full_like(ratios, -2.0, requires_grad=True) |
| loss, clipfrac = cispo_policy_loss(lp, ratios, torch.ones_like(ratios), eps_hi, eps_lo) |
| weight = (-loss / lp.detach()).detach() |
| assert clipfrac.tolist() == [0, 0, 1, 0, 0, 1] |
| assert float(weight[2]) == float(hi) and float(weight[5]) == float(lo) |
| assert float(weight[1]) == float(hi) and float(weight[4]) == float(lo) |
| assert float((weight[0] - hi).abs()) <= 2 * torch.finfo(dtype).eps |
|
|
|
|
| def test_all_zero_advantages_give_exactly_zero_gradient_and_loss(): |
| """A=0 (every group dead / centered to zero) must be an exact no-op even with |
| wild ratios. Failure mode: weight*logp leaking through when A is zero.""" |
| lp = torch.randn(3, 6, requires_grad=True) |
| ratio = torch.tensor([[0.0, 1e-9, 0.5, 1.0, 5.0, 1e9]] * 3) |
| loss, _ = cispo_policy_loss(lp, ratio, torch.zeros(3, 6), 0.272, 0.5) |
| loss.sum().backward() |
| assert float(loss.abs().sum()) == 0.0 and float(lp.grad.abs().sum()) == 0.0 |
|
|
|
|
| def test_inf_ratio_is_clipped_not_propagated_but_nan_is_not_laundered(): |
| """exp(new-old) overflows to +inf on a runaway token. CISPO must turn that |
| into the capped weight (finite loss and gradient); a NaN ratio (a real bug |
| upstream) must surface as NaN, not be silently clamped into a finite weight. |
| Failure modes: inf*0 -> NaN poisoning every grad; nan_to_num hiding a broken |
| logprob path behind a healthy-looking loss.""" |
| lp = torch.tensor([-1.0, -1.0], requires_grad=True) |
| loss, cf = cispo_policy_loss(lp, torch.tensor([float("inf"), 1.0]), torch.tensor([1.0, 1.0]), 0.272) |
| loss.sum().backward() |
| assert torch.isfinite(loss).all() and torch.isfinite(lp.grad).all() and float(cf[0]) == 1.0 |
| assert float(lp.grad[0]) == pytest.approx(-1.272) |
| for eps_low in (None, 0.5): |
| lp2 = torch.tensor([-1.0], requires_grad=True) |
| loss2, _ = cispo_policy_loss(lp2, torch.tensor([float("nan")]), torch.tensor([1.0]), 0.272, eps_low) |
| assert torch.isnan(loss2).all(), eps_low |
|
|
|
|
| @pytest.mark.parametrize("eps_low", [1.0, 5.0, 1e6]) |
| def test_eps_low_ge_one_is_the_papers_no_lower_bound(eps_low): |
| """Paper Sec 3.1: no lower bound is realised by setting eps_low to a LARGE |
| value, so 1-eps_low <= 0 and the floor can never bind a positive ratio. |
| `None` and any eps_low >= 1 must therefore give identical loss, gradient and |
| clipfrac, including for ratio -> 0. Failure mode: treating a large eps_low |
| as a real floor (e.g. abs()/clamp to >=eps) would re-introduce a bound.""" |
| ratio = torch.tensor([1e-12, 1e-3, 0.3, 1.0, 3.0, 50.0]) |
| adv = torch.tensor([1.0, -1.0, 1.0, -1.0, 1.0, -1.0]) |
| out = [] |
| for e in (None, eps_low): |
| lp = torch.linspace(-3, -1, 6).requires_grad_(True) |
| loss, cf = cispo_policy_loss(lp, ratio, adv, 0.272, e) |
| loss.sum().backward() |
| out.append((loss.detach(), lp.grad.clone(), cf)) |
| for a, b in zip(out[0], out[1]): |
| assert torch.equal(a, b) |
|
|
|
|
| |
|
|
|
|
| def test_head_is_checkpointed_not_stored(): |
| """Memory contract: backward RECOMPUTES each chunk's [chunk, V] logits |
| (checkpoint) instead of saving them. Counted via the chunk function: n calls |
| forward + n recompute in backward; with no grad enabled exactly n and no |
| checkpoint machinery. Failure mode: dropping the checkpoint silently turns |
| the head into a [B*T, V] fp32 resident tensor (10 GB at V=152k, 16k tok).""" |
| h, W, labels = _hw(B=2, T=8) |
| calls = [] |
| orig = cispo_mod._head_chunk_logprobs |
| cispo_mod._head_chunk_logprobs = lambda *a, **k: (calls.append(1), orig(*a, **k))[1] |
| try: |
| hh = h.clone().requires_grad_(True) |
| lp = fp32_head_logprobs(hh, W, labels, chunk_tokens=4) |
| assert len(calls) == 4 |
| lp.sum().backward() |
| assert len(calls) == 8 |
| calls.clear() |
| with torch.no_grad(): |
| fp32_head_logprobs(h, W, labels, chunk_tokens=4) |
| assert len(calls) == 4 |
| finally: |
| cispo_mod._head_chunk_logprobs = orig |
|
|
|
|
| def test_strict_fp32_matmul_survives_every_precision_start_state(): |
| """On torch>=2.9 flipping the legacy `matmul.allow_tf32` flag while the |
| global precision is 'high' (what most launchers set) leaves a mixed-API |
| state in which torch.get_float32_matmul_precision() RAISES. The context must |
| (a) never raise from any start state, (b) report 'highest' inside, and |
| (c) restore the exact start precision.""" |
| start = torch.get_float32_matmul_precision() |
| try: |
| for s in ("highest", "high", "medium"): |
| torch.set_float32_matmul_precision(s) |
| with cispo_mod._strict_fp32_matmul(): |
| assert torch.get_float32_matmul_precision() == "highest" |
| assert torch.backends.cuda.matmul.allow_tf32 is False |
| assert torch.get_float32_matmul_precision() == s |
| |
| with pytest.raises(RuntimeError, match="boom"): |
| with cispo_mod._strict_fp32_matmul(): |
| raise RuntimeError("boom") |
| assert torch.get_float32_matmul_precision() == s |
| finally: |
| torch.set_float32_matmul_precision(start) |
|
|
|
|
| |
|
|
|
|
| def test_chunk_tokens_is_plumbed_to_both_module_paths(monkeypatch): |
| """fp32_head_chunk_tokens is the memory knob; it must reach the head from |
| BOTH the training forward and compute_logprobs_nograd. Failure mode: a |
| hardcoded chunk silently ignoring the config on one path (OOM on 7B).""" |
| seen = [] |
| orig = cispo_mod.fp32_head_logprobs |
|
|
| def spy(hidden, weight, labels, temperature=1.0, chunk_tokens=2048, with_entropy=False): |
| seen.append(chunk_tokens) |
| return orig(hidden, weight, labels, temperature, chunk_tokens, with_entropy) |
|
|
| monkeypatch.setattr(cispo_mod, "fp32_head_logprobs", spy) |
| model = _Bf16LM() |
| tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, _p4_cfg(fp32_head=True, fp32_head_chunk_tokens=13)) |
| batch = _batch(tm, model) |
| tm._grpo_loss_adapter(model, batch) |
| assert seen and set(seen) == {13}, seen |
|
|
|
|
| def test_record_entropy_off_skips_entropy_pass(monkeypatch): |
| """record_entropy=False must not pay the extra no-grad head pass.""" |
| seen = [] |
| orig = cispo_mod.fp32_head_logprobs |
|
|
| def spy(hidden, weight, labels, temperature=1.0, chunk_tokens=2048, with_entropy=False): |
| seen.append(with_entropy) |
| return orig(hidden, weight, labels, temperature, chunk_tokens, with_entropy) |
|
|
| monkeypatch.setattr(cispo_mod, "fp32_head_logprobs", spy) |
| model = _Bf16LM() |
| tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, |
| _p4_cfg(fp32_head=True, record_entropy=False)) |
| tm._grpo_loss_adapter(model, _batch(tm, model)) |
| assert seen and not any(seen) |
|
|
|
|
| def test_inert_clip_lower_knobs_are_refused(): |
| """`clip_lower` is never read by the cispo branch and `cispo_clip_lower` is |
| never read by dapo; setting either to a non-default value used to be |
| accepted and silently ignored (a config that LOOKS floored but is not).""" |
| base = dict(beta=0.0, load_ref_policy=False, loss_denominator=1024.0, num_samples_per_prompt_rollout=16) |
| with pytest.raises(P4CispoConformanceError, match="clip_lower"): |
| assert_p4_cispo_conformant(GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=2.0, clip_lower=0.5, **base)) |
| with pytest.raises(A3ConformanceError, match="cispo_clip_lower"): |
| assert_a3_conformant(GRPOLossConfig(cispo_clip_lower=0.5, **base)) |
| |
| assert_p4_cispo_conformant( |
| GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=2.0, cispo_clip_lower=0.5, **base) |
| ) |
|
|
|
|
| def test_gate_message_names_every_independent_violation(): |
| """The refusal collects, not short-circuits: a config with three problems |
| must list all three (operators fix them in one pass).""" |
| cfg = GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=-1.0, cispo_clip_lower=-2.0, beta=0.5, |
| load_ref_policy=False, loss_denominator=64.0, num_samples_per_prompt_rollout=16) |
| with pytest.raises(P4CispoConformanceError) as ei: |
| assert_p4_cispo_conformant(cfg) |
| msg = str(ei.value) |
| assert "clip_higher" in msg and "cispo_clip_lower" in msg and "beta" in msg |
|
|
