From 3e4eb8ec7639541d2eedf91d52e23f1a097dd811 Mon Sep 17 00:00:00 2001 From: hhhappyshow <1253332472@qq.com> Date: Mon, 9 Nov 2026 12:12:25 +0800 Subject: [PATCH] fix(training): reset NaN-guard device latches under inference mode rsl_rl wraps the whole PPO rollout in torch.inference_mode(). With deferred finite validation (the new default) the guard's device latches are typically first allocated inside that context, so they are inference tensors; poll sites such as the cycle boundary and close() run outside it, and the in-place zero_() in reset_verdict() then raises "Inplace update to inference tensor outside InferenceMode is not allowed". Run the latch reset under torch.inference_mode() explicitly. The reset touches no autograd bookkeeping, so behavior is unchanged for plain tensors, and inference tensors reset cleanly at every poll site. Discovered by the downstream perceptive navigation port: the first full PPO train smoke against current main crashed in close(). Commands: uv run pytest tests/training/test_tensor_diagnostics.py -q (11 passed; new test fails without the fix) Co-authored-by: Cursor --- src/unilab/training/tensor_diagnostics.py | 16 ++++++++++++---- tests/training/test_tensor_diagnostics.py | 17 +++++++++++++++++ 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/src/unilab/training/tensor_diagnostics.py b/src/unilab/training/tensor_diagnostics.py index 080ebb324..9e19c7579 100644 --- a/src/unilab/training/tensor_diagnostics.py +++ b/src/unilab/training/tensor_diagnostics.py @@ -128,10 +128,18 @@ def capture_due(self, step: int) -> bool: return self._cfg.enabled and interval > 0 and step % interval == 0 def reset_verdict(self) -> None: - if self._device_bad_rows is not None: - self._device_bad_rows.zero_() - if self._device_verdict is not None: - self._device_verdict.zero_() + # The device latches may be inference tensors: rsl_rl wraps the whole + # rollout in torch.inference_mode(), so with deferred validation the + # buffers are typically first allocated inside that context. In-place + # zeroing of an inference tensor is only legal under inference mode, + # and poll sites (cycle boundaries, close) run outside it, so run the + # reset under inference mode explicitly; it is a no-op on the + # autograd bookkeeping either way. + with torch.inference_mode(): + if self._device_bad_rows is not None: + self._device_bad_rows.zero_() + if self._device_verdict is not None: + self._device_verdict.zero_() def _validate_finite_float(self, value: torch.Tensor, label: str) -> None: if not isinstance(value, torch.Tensor): diff --git a/tests/training/test_tensor_diagnostics.py b/tests/training/test_tensor_diagnostics.py index e19fc83b3..87a2995ab 100644 --- a/tests/training/test_tensor_diagnostics.py +++ b/tests/training/test_tensor_diagnostics.py @@ -200,3 +200,20 @@ def test_training_routes_guard_and_shared_finite_boundary() -> None: cfg.training.finite_validation_interval = 0 with pytest.raises(ValueError, match="positive integer"): finite_validation_env_overrides(cfg) + + +def test_reset_verdict_allows_inference_tensor_latches(tmp_path: Path) -> None: + """Latches allocated inside an rsl_rl rollout must survive an outside reset. + + rsl_rl wraps the whole rollout in ``torch.inference_mode()``; with + deferred validation the guard's device buffers are typically first + allocated there, and ``close()`` resets them outside the context. The + reset must not fail on inference tensors. + """ + guard = _guard(tmp_path) + obs = {"obs": torch.zeros((3, 2))} + reward = torch.zeros(3) + with torch.inference_mode(): + assert guard.check(obs, reward, defer=True) is None + assert guard.device_verdict is not None + guard.reset_verdict()