File size: 1,671 Bytes
bcc8d35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | """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 # noqa: F401
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)) # a healthy micro first (arms the tripwire)
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 # graph-connected, non-finite
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)"
)
|