| """CISPO objective (MiniMax-M1 arXiv:2506.13585 Eq. 4-5) + FP32 LM head. |
| |
| Every oracle here is an independent fp64 implementation written IN the test |
| (explicit loops / hand-derived softmax gradient), never a call back into the |
| function under test. The properties that define CISPO -- a clipped token still |
| gets a gradient; the weight is a stop-gradient constant; the default GRPO loss |
| is untouched -- each have a test that fails on the corresponding mutant |
| (scripts in .tmp/g2/p4a/mutants.py run them). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import pytest |
| import torch |
|
|
| from align.rl.cispo import ( |
| P4CispoConformanceError, |
| assert_p4_cispo_conformant, |
| cispo_policy_loss, |
| fp32_head_forward, |
| fp32_head_logprobs, |
| ) |
| from align.rl.grpo_loss import ( |
| A3ConformanceError, |
| GRPOLossConfig, |
| GRPOLossType, |
| assert_a3_conformant, |
| compute_grpo_loss, |
| ) |
| from align.rl.grpo_module import GRPOTrainModule |
| from align.rl.types import RolloutBatch |
|
|
| D = torch.float64 |
|
|
|
|
| |
|
|
|
|
| def ref_cispo_grad_logits(logits, labels, adv, old_lp, mask, eps_high, eps_low, denom): |
| """Hand-derived fp64 gradient of -(1/denom) sum_t mask*A*w_t*log pi(y_t) |
| wrt the logits, with w_t = clip(exp(logp_t - old_t)) a CONSTANT (Eq. 4-5). |
| d logp_y / d logits = onehot(y) - softmax(logits).""" |
| logits = logits.to(D) |
| B, T, V = logits.shape |
| grad = torch.zeros(B, T, V, dtype=D) |
| loss = 0.0 |
| for b in range(B): |
| for t in range(T): |
| z = logits[b, t] |
| m = z.max() |
| p = torch.exp(z - m) |
| p = p / p.sum() |
| logp = torch.log(p[labels[b, t]]) |
| r = math.exp(float(logp) - float(old_lp[b, t])) |
| w = min(r, 1 + eps_high) |
| if eps_low is not None: |
| w = max(w, 1 - eps_low) |
| if not mask[b, t]: |
| continue |
| loss += -float(adv[b]) * w * float(logp) / denom |
| onehot = torch.zeros(V, dtype=D) |
| onehot[labels[b, t]] = 1.0 |
| grad[b, t] = -float(adv[b]) * w * (onehot - p) / denom |
| return loss, grad |
|
|
|
|
| def _setup(B=3, T=5, V=11, seed=0): |
| g = torch.Generator().manual_seed(seed) |
| logits = torch.randn(B, T, V, generator=g, dtype=D) |
| labels = torch.randint(0, V, (B, T), generator=g) |
| adv = torch.tensor([1.0, -0.7, 0.4][:B], dtype=D) |
| mask = torch.ones(B, T, dtype=torch.bool) |
| mask[:, 0] = False |
| return logits, labels, adv, mask |
|
|
|
|
| def _old_from_ratios(logits, labels, ratios): |
| """old = logp - log(ratio) so that exp(logp-old) == ratios exactly.""" |
| logp = logits.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) |
| return (logp - torch.log(ratios)).detach() |
|
|
|
|
| def _cispo_cfg(**kw): |
| base = dict(loss_fn=GRPOLossType.cispo, clip_higher=0.272, beta=0.0, load_ref_policy=False, |
| loss_denominator=64.0, num_samples_per_prompt_rollout=8) |
| base.update(kw) |
| return GRPOLossConfig(**base) |
|
|
|
|
| def _run_loss(logits, labels, adv, old, mask, cfg, loss_fn=None): |
| logits = logits.clone().requires_grad_(True) |
| lp = logits.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) |
| ratio = torch.exp(lp - old) |
| a = adv[:, None].expand_as(lp) |
| pg, clipfrac, _ = compute_grpo_loss(lp, ratio, a, None, cfg, torch.ones_like(lp)) |
| loss = (pg * mask).sum() / cfg.loss_denominator |
| loss.backward() |
| return loss.detach(), logits.grad, clipfrac |
|
|
|
|
| |
|
|
|
|
| @pytest.mark.parametrize("eps_low", [None, 0.5]) |
| def test_cispo_gradient_matches_fp64_reference(eps_low): |
| logits, labels, adv, mask = _setup() |
| B, T, _ = logits.shape |
| |
| ratios = torch.tensor([[0.2, 0.9, 1.0, 1.1, 3.0], [0.1, 0.6, 1.2, 1.3, 2.0], [0.95, 1.0, 5.0, 0.3, 1.25]], |
| dtype=D) |
| old = _old_from_ratios(logits, labels, ratios) |
| cfg = _cispo_cfg(cispo_clip_lower=eps_low) |
| loss, grad, _ = _run_loss(logits, labels, adv, old, mask, cfg) |
| ref_loss, ref_grad = ref_cispo_grad_logits(logits, labels, adv, old, mask, 0.272, eps_low, 64.0) |
| assert abs(float(loss) - ref_loss) < 1e-12 |
| torch.testing.assert_close(grad, ref_grad, rtol=1e-10, atol=1e-12) |
|
|
|
|
| def test_cispo_matches_explicit_eq4_loop_value(): |
| """Eq. 4 with sum over (i,t) written as a plain python double loop.""" |
| logits, labels, adv, mask = _setup(seed=3) |
| ratios = torch.rand(3, 5, dtype=D) * 3 + 0.05 |
| old = _old_from_ratios(logits, labels, ratios) |
| lp = logits.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) |
| pg, _ = cispo_policy_loss(lp, torch.exp(lp - old), adv[:, None].expand_as(lp), 0.272) |
| total = 0.0 |
| for i in range(3): |
| for t in range(5): |
| if mask[i, t]: |
| w = min(float(torch.exp(lp[i, t] - old[i, t])), 1.272) |
| total += -float(adv[i]) * w * float(lp[i, t]) |
| assert abs(float((pg * mask).sum()) - total) < 1e-12 |
|
|
|
|
| |
|
|
|
|
| def test_clipped_tokens_keep_gradient_ppo_clip_zeroes_them(): |
| """The defining difference. One sequence, A>0, ratio >> 1+eps_high on every |
| token: DAPO/PPO-clip gradient is EXACTLY zero, CISPO gradient is |
| A*(1+eps_high)*(onehot-softmax) per token.""" |
| logits, labels, adv, mask = _setup(B=1, T=6) |
| adv = torch.tensor([1.0], dtype=D) |
| mask = torch.ones(1, 6, dtype=torch.bool) |
| ratios = torch.full((1, 6), 4.0, dtype=D) |
| old = _old_from_ratios(logits, labels, ratios) |
|
|
| _, g_dapo, cf_dapo = _run_loss(logits, labels, adv, old, mask, GRPOLossConfig(loss_fn=GRPOLossType.dapo, |
| loss_denominator=6.0)) |
| assert torch.count_nonzero(g_dapo) == 0 and bool(cf_dapo.all()) |
|
|
| cfg = _cispo_cfg(loss_denominator=6.0) |
| _, g_cispo, cf = _run_loss(logits, labels, adv, old, mask, cfg) |
| assert bool(cf.all()) |
| p = logits.softmax(-1) |
| oh = torch.nn.functional.one_hot(labels, logits.shape[-1]).to(D) |
| expected = -1.0 * 1.272 * (oh - p) / 6.0 |
| torch.testing.assert_close(g_cispo, expected, rtol=1e-10, atol=1e-13) |
|
|
|
|
| def test_low_side_negative_advantage_also_keeps_gradient(): |
| """A<0 with ratio << 1-eps_low: PPO-clip zeroes it (min-side clip). CISPO |
| with no lower bound keeps weight = r (small, nonzero); with a floor it keeps |
| weight = 1-eps_low. Neither is zero.""" |
| logits, labels, adv, mask = _setup(B=1, T=4) |
| adv = torch.tensor([-1.0], dtype=D) |
| mask = torch.ones(1, 4, dtype=torch.bool) |
| ratios = torch.full((1, 4), 0.05, dtype=D) |
| old = _old_from_ratios(logits, labels, ratios) |
| _, g_dapo, _ = _run_loss(logits, labels, adv, old, mask, |
| GRPOLossConfig(loss_fn=GRPOLossType.dapo, loss_denominator=4.0)) |
| assert torch.count_nonzero(g_dapo) == 0 |
| for eps_low, w in ((None, 0.05), (0.5, 0.5)): |
| _, g, _ = _run_loss(logits, labels, adv, old, mask, _cispo_cfg(cispo_clip_lower=eps_low, |
| loss_denominator=4.0)) |
| oh = torch.nn.functional.one_hot(labels, logits.shape[-1]).to(D) |
| expected = 1.0 * w * (oh - logits.softmax(-1)) / 4.0 |
| torch.testing.assert_close(g, expected, rtol=1e-10, atol=1e-13) |
|
|
|
|
| def test_weight_is_stop_gradient(): |
| """If `ratio` is not detached, d loss/d ratio-source is nonzero (an extra |
| A*log pi*grad r term). Assert the gradient flows ONLY through log pi.""" |
| lp = torch.tensor([[-1.0, -2.0]], dtype=D, requires_grad=True) |
| ratio_src = torch.tensor([[1.1, 0.9]], dtype=D, requires_grad=True) |
| pg, _ = cispo_policy_loss(lp, ratio_src * 1.0, torch.ones(1, 2, dtype=D), 0.272) |
| pg.sum().backward() |
| assert ratio_src.grad is None or torch.count_nonzero(ratio_src.grad) == 0 |
| torch.testing.assert_close(lp.grad, -torch.tensor([[1.1, 0.9]], dtype=D)) |
|
|
|
|
| def test_eps_cap_changes_weight_exactly_at_boundary(): |
| lp = torch.zeros(1, 3, dtype=D) |
| ratio = torch.tensor([[1.271, 1.272, 1.2721]], dtype=D) |
| a = torch.ones(1, 3, dtype=D) |
| lp1 = lp - 1.0 |
| pg, cf = cispo_policy_loss(lp1, ratio, a, 0.272) |
| torch.testing.assert_close(pg, torch.tensor([[1.271, 1.272, 1.272]], dtype=D)) |
| assert cf.tolist() == [[0.0, 0.0, 1.0]] |
|
|
|
|
| def test_clipfrac_counts_both_sides_when_floor_set(): |
| lp = torch.full((1, 4), -1.0, dtype=D) |
| ratio = torch.tensor([[0.1, 0.6, 1.0, 2.0]], dtype=D) |
| _, cf = cispo_policy_loss(lp, ratio, torch.ones(1, 4, dtype=D), 0.272, eps_low=0.5) |
| assert cf.tolist() == [[1.0, 0.0, 0.0, 1.0]] |
| _, cf = cispo_policy_loss(lp, ratio, torch.ones(1, 4, dtype=D), 0.272) |
| assert cf.tolist() == [[0.0, 0.0, 0.0, 1.0]] |
|
|
|
|
| def test_invalid_eps_refused(): |
| x = torch.zeros(1, 1) |
| for kw in ({"eps_high": 0.0}, {"eps_high": float("nan")}, {"eps_high": 0.2, "eps_low": 0.0}, |
| {"eps_high": 0.2, "eps_low": float("inf")}): |
| with pytest.raises(ValueError): |
| cispo_policy_loss(x, x + 1, x, **kw) |
|
|
|
|
| |
|
|
|
|
| def _legacy_compute_grpo_loss(new_logprobs, ratio, advantages, config): |
| """VERBATIM pre-change bodies of the dapo and cispo branches (HEAD 780ce3d).""" |
| if config.loss_fn == GRPOLossType.dapo: |
| pg_losses = -advantages * ratio |
| pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - config.clip_lower, 1.0 + config.clip_higher) |
| clipfrac = (pg_losses2 > pg_losses).float() |
| pg_loss = torch.max(pg_losses, pg_losses2) |
| else: |
| clipfrac = (ratio > 1.0 + config.clip_higher).float() |
| pg_loss = -advantages * torch.clamp(ratio.detach(), max=1.0 + config.clip_higher) * new_logprobs |
| return pg_loss, clipfrac |
|
|
|
|
| @pytest.mark.parametrize("fn", [GRPOLossType.dapo, GRPOLossType.cispo]) |
| @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| def test_default_loss_byte_identical_to_pre_change(fn, dtype): |
| g = torch.Generator().manual_seed(7) |
| lp = (-torch.rand(4, 9, generator=g) * 5).to(dtype) |
| ratio = (torch.rand(4, 9, generator=g) * 3).to(dtype) |
| adv = torch.randn(4, 1, generator=g).to(dtype).expand(4, 9) |
| cfg = GRPOLossConfig(loss_fn=fn) |
| assert cfg.cispo_clip_lower is None and cfg.fp32_head is False |
| pg, cf, _ = compute_grpo_loss(lp, ratio, adv, None, cfg, torch.ones_like(lp)) |
| rpg, rcf = _legacy_compute_grpo_loss(lp, ratio, adv, cfg) |
| assert torch.equal(pg, rpg) and torch.equal(cf, rcf) |
| assert pg.dtype == rpg.dtype |
|
|
|
|
| def test_a3_gate_still_refuses_bare_cispo_and_p4_gate_accepts(): |
| base = dict(beta=0.0, load_ref_policy=False, loss_denominator=1024.0, |
| num_samples_per_prompt_rollout=16, record_entropy=True) |
| with pytest.raises(A3ConformanceError): |
| assert_a3_conformant(GRPOLossConfig(loss_fn=GRPOLossType.cispo, **base)) |
| assert_p4_cispo_conformant(GRPOLossConfig(loss_fn=GRPOLossType.cispo, clip_higher=2.0, **base)) |
| assert_a3_conformant(GRPOLossConfig(**base)) |
|
|
|
|
| @pytest.mark.parametrize( |
| "bad", |
| [ |
| {"loss_denominator": "token"}, |
| {"beta": 0.05}, |
| {"load_ref_policy": True}, |
| {"num_samples_per_prompt_rollout": 4}, |
| {"use_rho_correction": False}, |
| {"filter_zero_std_samples": False}, |
| {"advantage_normalization_type": "standard"}, |
| {"clip_higher": 0.0}, |
| {"clip_higher": float("inf")}, |
| {"cispo_clip_lower": 0.0}, |
| {"cispo_clip_lower": -1.0}, |
| {"fp32_head_chunk_tokens": 0}, |
| {"loss_fn": GRPOLossType.dapo}, |
| ], |
| ) |
| def test_p4_gate_refuses_each_violation(bad): |
| cfg = dict(loss_fn=GRPOLossType.cispo, beta=0.0, load_ref_policy=False, loss_denominator=1024.0, |
| num_samples_per_prompt_rollout=16, clip_higher=2.0, record_entropy=True) |
| cfg.update(bad) |
| with pytest.raises(P4CispoConformanceError): |
| assert_p4_cispo_conformant(GRPOLossConfig(**cfg)) |
|
|
|
|
| |
|
|
|
|
| def _hw(seed=0, B=2, T=7, H=16, V=40, scale=12.0, dtype=torch.bfloat16): |
| g = torch.Generator().manual_seed(seed) |
| h = (torch.randn(B, T, H, generator=g) * scale).to(dtype) |
| W = (torch.randn(V, H, generator=g) * 0.5).to(dtype) |
| labels = torch.randint(0, V, (B, T), generator=g) |
| return h, W, labels |
|
|
|
|
| def ref_head_logprobs(h, W, labels, temperature=1.0): |
| """fp64 oracle on the bf16-VALUED operands (exact upcast, no rounding).""" |
| z = (h.to(D) @ W.to(D).t()) / temperature |
| out = torch.empty(labels.shape, dtype=D) |
| for b in range(labels.shape[0]): |
| for t in range(labels.shape[1]): |
| zz = z[b, t] |
| lse = zz.max() + torch.log(torch.exp(zz - zz.max()).sum()) |
| out[b, t] = zz[labels[b, t]] - lse |
| return out |
|
|
|
|
| @pytest.mark.parametrize("temperature", [1.0, 0.7]) |
| def test_fp32_head_matches_fp64_oracle(temperature): |
| h, W, labels = _hw() |
| lp = fp32_head_logprobs(h, W, labels, temperature, chunk_tokens=5) |
| assert lp.dtype == torch.float32 |
| err = (lp.to(D) - ref_head_logprobs(h, W, labels, temperature)).abs().max() |
| assert float(err) < 5e-5, float(err) |
|
|
|
|
| def test_fp32_head_chunk_invariance_and_gradient_oracle(): |
| h, W, labels = _hw(seed=2) |
| outs = [fp32_head_logprobs(h, W, labels, chunk_tokens=c) for c in (1, 3, 100)] |
| for o in outs[1:]: |
| torch.testing.assert_close(o, outs[0], rtol=0, atol=2e-5) |
| |
| hh = h.clone().requires_grad_(True) |
| Wp = W.clone().requires_grad_(True) |
| w = torch.randn(labels.shape, generator=torch.Generator().manual_seed(9)) |
| (fp32_head_logprobs(hh, Wp, labels, chunk_tokens=4) * w).sum().backward() |
| h64 = h.to(D).requires_grad_(True) |
| W64 = W.to(D).requires_grad_(True) |
| z = h64 @ W64.t() |
| ref = (z.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) * w.to(D)).sum() |
| ref.backward() |
| assert hh.grad.dtype == torch.bfloat16 and Wp.grad.dtype == torch.bfloat16 |
| |
| torch.testing.assert_close(hh.grad.to(D), h64.grad, rtol=2 ** -7, atol=2 ** -7 * float(h64.grad.abs().max())) |
| torch.testing.assert_close(Wp.grad.to(D), W64.grad, rtol=2 ** -7, atol=2 ** -7 * float(W64.grad.abs().max())) |
|
|
|
|
| def test_fp32_head_entropy_matches_oracle(): |
| h, W, labels = _hw(seed=4, scale=2.0) |
| lp, ent = fp32_head_logprobs(h, W, labels, with_entropy=True, chunk_tokens=3) |
| p = (h.to(D) @ W.to(D).t()).softmax(-1) |
| ref = -(p * p.log()).sum(-1) |
| torch.testing.assert_close(ent.to(D), ref, rtol=0, atol=1e-5) |
| assert not ent.requires_grad |
|
|
|
|
| def test_fp32_head_disables_tf32_and_restores(): |
| |
| |
| prev = torch.get_float32_matmul_precision() |
| torch.set_float32_matmul_precision("high") |
| try: |
| h, W, labels = _hw() |
| seen = [] |
| import align.rl.cispo as c |
|
|
| import contextlib |
|
|
| orig = c._strict_fp32_matmul |
|
|
| @contextlib.contextmanager |
| def spy(): |
| with orig(): |
| seen.append(torch.backends.cuda.matmul.allow_tf32) |
| yield |
|
|
| c._strict_fp32_matmul = spy |
| try: |
| fp32_head_logprobs(h, W, labels) |
| finally: |
| c._strict_fp32_matmul = orig |
| assert seen and not any(seen) |
| assert torch.backends.cuda.matmul.allow_tf32 is True |
| finally: |
| torch.set_float32_matmul_precision(prev) |
|
|
|
|
| |
|
|
|
|
| class _Bf16LM(torch.nn.Module): |
| """bf16 embedding -> bf16 body -> bf16 lm_head, with the hidden_consumer seam |
| (the Qwen3/Olmoe contract). Scaled so head logits reach |z|~30-60: the |
| "high-magnitude activation" regime MiniMax identified.""" |
|
|
| def __init__(self, V=48, H=32, scale=6.0): |
| super().__init__() |
| torch.manual_seed(0) |
| self.emb = torch.nn.Embedding(V, H).to(torch.bfloat16) |
| self.body = torch.nn.Linear(H, H, bias=False).to(torch.bfloat16) |
| self.lm_head = torch.nn.Linear(H, V, bias=False).to(torch.bfloat16) |
| with torch.no_grad(): |
| self.emb.weight.mul_(scale) |
| self.lm_head.weight.mul_(1.5) |
|
|
| def forward(self, input_ids, attention_mask=None, hidden_consumer=None, return_hidden_states=False): |
| h = self.body(self.emb(input_ids)) |
| if hidden_consumer is not None: |
| return hidden_consumer(h, self.lm_head.weight) |
| return h if return_hidden_states else self.lm_head(h) |
|
|
|
|
| def _split_k_logits(h, W, dtype, parts=4): |
| """An INFERENCE engine's head: split-K accumulation, each partial rounded to |
| `dtype`, partials summed in `dtype` -- a different but legitimate summation |
| order than the trainer's single GEMM.""" |
| H = h.shape[-1] |
| acc = None |
| for k in range(parts): |
| sl = slice(k * H // parts, (k + 1) * H // parts) |
| part = (h[..., sl].to(dtype) @ W[:, sl].to(dtype).t()).to(dtype) |
| acc = part if acc is None else (acc + part).to(dtype) |
| return acc |
|
|
|
|
| def _sampler_lp(h, W, labels, dtype): |
| z = _split_k_logits(h, W, dtype).float() |
| return z.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) |
|
|
|
|
| def test_fp32_head_reduces_sampler_trainer_logprob_mismatch(capsys): |
| m = _Bf16LM() |
| g = torch.Generator().manual_seed(5) |
| ids = torch.randint(0, 48, (16, 24), generator=g) |
| labels = ids[:, 1:] |
| with torch.no_grad(): |
| h = m(ids, return_hidden_states=True)[:, :-1] |
| W = m.lm_head.weight |
| truth = ref_head_logprobs(h, W, labels).float() |
| |
| trainer_bf16 = m(ids)[:, :-1].log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1).float() |
| sampler_bf16 = _sampler_lp(h, W, labels, torch.bfloat16) |
| |
| trainer_fp32 = fp32_head_forward(m, ids, torch.ones_like(ids), chunk_tokens=64) |
| sampler_fp32 = _sampler_lp(h, W, labels, torch.float32) |
| mm_bf16 = (trainer_bf16 - sampler_bf16).abs() |
| mm_fp32 = (trainer_fp32 - sampler_fp32).abs() |
| err_bf16 = (trainer_bf16 - truth).abs() |
| err_fp32 = (trainer_fp32 - truth).abs() |
| print( |
| f"\nMISMATCH trainer-vs-sampler |dlogp|: bf16 head mean={mm_bf16.mean():.3e} max={mm_bf16.max():.3e} | " |
| f"fp32 head mean={mm_fp32.mean():.3e} max={mm_fp32.max():.3e}\n" |
| f"ERROR vs fp64 truth: bf16 mean={err_bf16.mean():.3e} max={err_bf16.max():.3e} | " |
| f"fp32 mean={err_fp32.mean():.3e} max={err_fp32.max():.3e}" |
| ) |
| |
| r_bf16, r_fp32 = torch.exp(trainer_bf16 - sampler_bf16), torch.exp(trainer_fp32 - sampler_fp32) |
| print(f"max |r-1|: bf16 {float((r_bf16 - 1).abs().max()):.3e} fp32 {float((r_fp32 - 1).abs().max()):.3e}") |
| assert mm_bf16.mean() > 1e-3 |
| assert mm_fp32.mean() * 100 < mm_bf16.mean() |
| assert mm_fp32.max() * 100 < mm_bf16.max() |
| assert err_fp32.max() < 1e-4 < err_bf16.max() |
|
|
|
|
| |
|
|
|
|
| class _Sgd: |
| def __init__(self, params, lr=0.01): |
| self.params = [p for p in params if p.requires_grad] |
| self.inner = torch.optim.SGD(self.params, lr=lr) |
| self.last_grad_norm = 0.0 |
| self.skipped_steps = 0 |
|
|
| def step(self): |
| self.last_grad_norm = float(torch.linalg.vector_norm(torch.stack( |
| [p.grad.float().norm() for p in self.params if p.grad is not None]))) |
| self.inner.step() |
| return True |
|
|
| def zero_grad(self): |
| self.inner.zero_grad(set_to_none=True) |
|
|
| def state_dict(self): |
| return self.inner.state_dict() |
|
|
| def load_state_dict(self, sd): |
| self.inner.load_state_dict(sd) |
|
|
|
|
| def _p4_cfg(**kw): |
| base = dict(loss_fn=GRPOLossType.cispo, beta=0.0, load_ref_policy=False, loss_denominator=256.0, |
| num_samples_per_prompt_rollout=8, clip_higher=0.272, record_entropy=True) |
| base.update(kw) |
| return GRPOLossConfig(**base) |
|
|
|
|
| def _batch(tm, model, B=4, T=12, prompt=4, old_shift=0.0): |
| g = torch.Generator().manual_seed(1) |
| ids = torch.randint(1, 48, (B, T), generator=g) |
| attn = torch.ones(B, T, dtype=torch.int64) |
| resp = torch.zeros(B, T, dtype=torch.int64) |
| resp[:, prompt:] = 1 |
| old = tm.compute_logprobs_nograd(ids, attn) + old_shift |
| return RolloutBatch(input_ids=ids, attention_mask=attn, response_mask=resp, |
| advantages=torch.tensor([1.0, -1.0, 0.5, -0.5])[:B], old_logprobs=old, |
| vllm_logprobs=old.clone(), group_ids=torch.arange(B) // 2, |
| rewards={"reward/total": torch.ones(B)}) |
|
|
|
|
| def test_module_cispo_fp32_head_end_to_end(): |
| model = _Bf16LM() |
| tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, _p4_cfg(fp32_head=True, fp32_head_chunk_tokens=7)) |
| batch = _batch(tm, model) |
| |
| out = tm._grpo_loss_adapter(model, batch) |
| assert out.metrics["grpo/ratio"] == pytest.approx(1.0, abs=1e-6) |
| assert math.isfinite(out.metrics["entropy"]) and out.metrics["entropy"] >= 0 |
| out.loss.backward() |
| assert all(p.grad is not None and torch.isfinite(p.grad.float()).all() for p in model.parameters()) |
| |
| ids, attn = batch.input_ids, batch.attention_mask |
| with torch.no_grad(): |
| h = model(ids, return_hidden_states=True)[:, :-1] |
| ref = ref_head_logprobs(h, model.lm_head.weight, ids[:, 1:]) |
| torch.testing.assert_close(batch.old_logprobs.to(D), ref, rtol=0, atol=1e-4) |
|
|
|
|
| def test_module_fp32_head_off_is_bit_identical_to_logits_path(): |
| """Default (fp32_head=False) takes the untouched branch: same loss bits as a |
| manual replay of the pre-change adapter math.""" |
| model = _Bf16LM() |
| cfg = _p4_cfg() |
| assert cfg.fp32_head is False |
| tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, cfg) |
| batch = _batch(tm, model) |
| out = tm._grpo_loss_adapter(model, batch) |
| logits = model(batch.input_ids)[:, :-1] |
| from align.rl.ops import log_softmax_and_gather |
|
|
| nl = log_softmax_and_gather(logits, batch.input_ids[:, 1:]) |
| ratio = torch.exp(nl - batch.old_logprobs) |
| adv = batch.advantages[:, None].expand_as(nl) |
| pg, _ = _legacy_compute_grpo_loss(nl, ratio, adv, cfg) |
| mask = batch.shifted_response_mask.bool() |
| expect = (pg * mask).sum() / 256.0 |
| assert float((out.loss.detach() - expect.to(out.loss.dtype)).abs()) == 0.0 |
|
|
|
|
| def test_module_clipped_tokens_still_update_where_dapo_does_not(): |
| """End to end through the module: push old_logprobs far below new so every |
| ratio >> 1+eps (positive advantage rows): the DAPO module's grad on those rows |
| is zero, the CISPO module's is not.""" |
| shift = -3.0 |
| grads = {} |
| for name, cfg in (("cispo", _p4_cfg()), |
| ("dapo", GRPOLossConfig(beta=0.0, load_ref_policy=False, loss_denominator=256.0, |
| num_samples_per_prompt_rollout=8, record_entropy=True))): |
| model = _Bf16LM() |
| tm = GRPOTrainModule(model, _Sgd(model.parameters()), 1, cfg) |
| batch = _batch(tm, model, old_shift=shift) |
| batch.advantages = torch.tensor([1.0, 1.0, 1.0, 1.0]) |
| batch.vllm_logprobs = batch.old_logprobs.clone() |
| tm._grpo_loss_adapter(model, batch).loss.backward() |
| grads[name] = sum(float(p.grad.float().abs().sum()) for p in model.parameters()) |
| assert grads["dapo"] == 0.0 |
| assert grads["cispo"] > 0.0 |
|
|
|
|
| def test_module_p4_gate_wired_and_a3_default_kept(): |
| model = _Bf16LM() |
| with pytest.raises(P4CispoConformanceError): |
| GRPOTrainModule(model, _Sgd(model.parameters()), 1, _p4_cfg(beta=0.05)) |
| with pytest.raises(A3ConformanceError): |
| GRPOTrainModule(model, _Sgd(model.parameters()), 1, |
| GRPOLossConfig(beta=0.0, load_ref_policy=False, loss_denominator="token")) |
|
|
|
|
| def test_fp32_head_needs_the_seam(): |
| class NoSeam(torch.nn.Module): |
| def forward(self, input_ids, attention_mask=None): |
| return torch.zeros(*input_ids.shape, 4) |
|
|
| with pytest.raises(TypeError, match="hidden_consumer"): |
| fp32_head_forward(NoSeam(), torch.zeros(1, 3, dtype=torch.long), torch.ones(1, 3)) |
|
|
