[CI] Detect and fail evals on when NaNs appear in logits (#50323)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Tyler Michael Smith committed
Aug 4, 2026 at 19:02 UTC
12292d94b25869be2af6b6d4f8eea6c2445e935f
6 files changed
+60
-3
tests/basic_correctness/test_basic_correctness.py
+31
@@ -14,6 +14,7 @@ import torch
14
from packaging.version import Version
15
from transformers import __version__ as TRANSFORMERS_VERSION
16
17
+import vllm.envs as envs
18
from vllm import LLM
19
from vllm.platforms import current_platform
20
from vllm.v1.engine.llm_engine import LLMEngine
@@ -287,6 +288,36 @@ def test_failed_model_execution(vllm_runner, monkeypatch) -> None:
288
v1_test_failed_model_execution(vllm_model)
289
290
291
+@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
292
+def test_raise_on_logit_nans(
293
+ vllm_runner, monkeypatch, use_v2_model_runner: bool
294
+) -> None:
295
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
296
+ monkeypatch.setenv("VLLM_RAISE_ON_LOGIT_NANS", "1")
297
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
298
+ envs.disable_envs_cache()
299
+
300
+ try:
301
+ with vllm_runner("facebook/opt-125m", enforce_eager=True) as vllm_model:
302
+ engine_core = vllm_model.llm.llm_engine.engine_core.engine_core
303
+ model_runner = engine_core.model_executor.driver_worker.worker.model_runner
304
+ assert model_runner.vllm_config.use_v2_model_runner is use_v2_model_runner
305
+ model_cls = type(model_runner.model)
306
+ original_compute_logits = model_cls.compute_logits
307
+
308
+ def compute_nan_logits(self, *args, **kwargs):
309
+ logits = original_compute_logits(self, *args, **kwargs)
310
+ logits[0, 0] = float("nan")
311
+ return logits
312
+
313
+ monkeypatch.setattr(model_cls, "compute_logits", compute_nan_logits)
314
+
315
+ with pytest.raises(RuntimeError, match="NaNs detected in logits"):
316
+ vllm_model.generate_greedy(["Hello, my name is"], 1, use_tqdm=False)
317
+ finally:
318
+ envs.disable_envs_cache()
319
+
320
+
321
def v1_test_failed_model_execution(vllm_model):
322
engine = vllm_model.llm.llm_engine
323
mocked_execute_model = Mock(side_effect=RuntimeError("Mocked Critical Error"))
tests/v1/worker/test_gpu_model_runner.py
+2
-1
@@ -524,7 +524,8 @@ def test_update_states_request_resumed(model_runner, dist_init):
524
assert _is_req_state_block_table_match(model_runner, req_id)
525
526
527
-def test_get_nans_in_logits(model_runner, dist_init):
527
+def test_get_nans_in_logits(model_runner, dist_init, monkeypatch):
528
+ monkeypatch.setattr(gpu_model_runner_module.envs, "VLLM_RAISE_ON_LOGIT_NANS", False)
529
req_ids = ("req_0", "req_1")
530
531
scheduler_output = _schedule_new_request(*req_ids)
vllm/envs.py
+7
@@ -228,6 +228,7 @@ if TYPE_CHECKING:
228
VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
229
VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
230
VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
231
+ VLLM_RAISE_ON_LOGIT_NANS: bool = False
232
VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
233
"FP", "INT8", "INT6", "INT4", "INT3", "NONE"
234
] = "NONE"
@@ -1720,6 +1721,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1721
# or bad hardware but it may add compute overhead.
1722
"VLLM_COMPUTE_NANS_IN_LOGITS": lambda: bool(
1723
int(os.getenv("VLLM_COMPUTE_NANS_IN_LOGITS", "0"))
1724
+ or int(os.getenv("VLLM_RAISE_ON_LOGIT_NANS", "0"))
1725
+ ),
1726
+ # Raise an exception when generated logits contain NaNs. Enabling this
1727
+ # also enables the NaN computation required to detect them.
1728
+ "VLLM_RAISE_ON_LOGIT_NANS": lambda: bool(
1729
+ int(os.getenv("VLLM_RAISE_ON_LOGIT_NANS", "0"))
1730
),
1731
# Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
1732
"VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
vllm/v1/worker/gpu/async_utils.py
+4
@@ -5,6 +5,7 @@ import contextlib
5
import numpy as np
6
import torch
7
8
+import vllm.envs as envs
9
from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
10
from vllm.v1.outputs import (
11
AsyncModelRunnerOutput,
@@ -14,6 +15,7 @@ from vllm.v1.outputs import (
15
RoutedExpertsTensors,
16
)
17
from vllm.v1.worker.gpu.sample.output import SamplerOutput
18
+from vllm.v1.worker.utils import raise_if_nan_logits
19
20
21
class AsyncOutput(AsyncModelRunnerOutput):
@@ -80,6 +82,8 @@ class AsyncOutput(AsyncModelRunnerOutput):
82
self.model_runner_output.num_nans_in_logits = dict(
83
zip(self.model_runner_output.req_ids, self.num_nans.tolist())
84
)
85
+ if envs.VLLM_RAISE_ON_LOGIT_NANS:
86
+ raise_if_nan_logits(self.model_runner_output.num_nans_in_logits)
87
88
if self.logprobs_tensors is not None:
89
self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
vllm/v1/worker/gpu_model_runner.py
+3
-1
@@ -231,7 +231,7 @@ from vllm.v1.worker.ubatch_utils import (
231
maybe_create_ubatch_slices,
232
split_attn_metadata,
233
)
234
-from vllm.v1.worker.utils import is_residual_scattered_for_sp
234
+from vllm.v1.worker.utils import is_residual_scattered_for_sp, raise_if_nan_logits
235
from vllm.v1.worker.workspace import lock_workspace
236
237
from .utils import (
@@ -5723,6 +5723,8 @@ class GPUModelRunner(
5723
if num_nans_for_index is not None and req_index < logits.shape[0]
5724
else 0
5725
)
5726
+ if envs.VLLM_RAISE_ON_LOGIT_NANS:
5727
+ raise_if_nan_logits(num_nans_in_logits)
5728
return num_nans_in_logits
5729
except IndexError:
5730
return {}
vllm/v1/worker/utils.py
+13
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import math
4
from collections import defaultdict
5
-from collections.abc import Iterable, Sequence
5
+from collections.abc import Iterable, Mapping, Sequence
6
from dataclasses import dataclass, field
7
from itertools import product as iprod
8
from typing import Any
@@ -41,6 +41,18 @@ from vllm.v1.worker.block_table import get_block_table_width
41
logger = init_logger(__name__)
42
43
44
+def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None:
45
+ if not any(num_nans_in_logits.values()):
46
+ return
47
+
48
+ corrupted_requests = {
49
+ req_id: num_nans
50
+ for req_id, num_nans in num_nans_in_logits.items()
51
+ if num_nans > 0
52
+ }
53
+ raise RuntimeError(f"NaNs detected in logits: {corrupted_requests}")
54
+
55
+
56
@triton.jit(do_not_specialize=["n_blocks"])
57
def _zero_kv_blocks_kernel(
58
seg_addrs_ptr,