| """Audit 2026-10-14 (G4-rl-batch) regression for `align/ppo/trainer.py` (U05-2).""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from test_ppo import _build_module, _make_rollout, _micro, _perturb, _ppo_cfg, tiny |
|
|
|
|
| def test_u05_2_a_skipped_micro_with_non_finite_logprobs_still_reaches_every_parameter(tiny, monkeypatch): |
| """FSDP2 reduce-scatter fires from the parameters' grad hooks. A rank whose `new_lp` is |
| non-finite used to skip on a DETACHED leaf: backward touched no parameter, the rank never |
| joined the collective its finite-`new_lp` peers sat in, and the job hung until the watchdog. |
| The skip must back-propagate through the model whatever `new_lp` holds.""" |
| import align.ppo.trainer as T |
|
|
| cfg = _ppo_cfg() |
| mod = _build_module(tiny, cfg) |
| exp = mod.prepare_experience(_make_rollout(tiny, cfg)) |
| _perturb(mod.model, "policy", seed=5, scale=0.1) |
| mod.set_train() |
| mod.train_batch(_micro(mod, exp, iteration=1)) |
| for p in mod.model.parameters(): |
| p.grad = None |
|
|
| real = T.policy_forward |
|
|
| def nan_forward(*a, **k): |
| new_lp, values = real(*a, **k) |
| return new_lp * float("nan"), values |
|
|
| monkeypatch.setattr(T, "policy_forward", nan_forward) |
| out = mod.train_batch(_micro(mod, exp, iteration=1, rows=(2, 3))) |
| assert out.metrics["ppo/skipped_micro"] == 1.0 |
| grads = [p.grad for p in mod.model.parameters() if p.requires_grad] |
| assert grads and all(g is not None for g in grads), ( |
| "a non-finite-new_lp skip must still back-propagate through every adapter (FSDP2 lockstep)" |
| ) |
|
|
