From b8df4a8cc5cd5c914e285d0b6c349a647f34af7e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 18 May 2026 02:20:14 +0000 Subject: [PATCH] Fix NaN check: use os.environ gate instead of is_current_stream_capturing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit torch.cuda.is_current_stream_capturing() returns bool, which breaks Dynamo FX tracing (non-Tensor output). Switch to env var gate: CLAWMINE_NAN_CHECK=1 enables NaN/Inf detection. Dynamo evaluates os.environ at trace time — if the env var is not set, the entire NaN check block is compiled away. Set it before first inference to get NaN detection during prefill only. --- vllm/patches/deepseek_v4.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 886149ff..60a040d2 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os import typing from collections.abc import Callable, Iterable from itertools import islice @@ -1315,18 +1316,21 @@ class DeepseekV4Model(nn.Module): positions, input_ids, ) - # NaN detection (prefill only — no CPU-GPU syncs during cudagraph) - if not torch.cuda.is_current_stream_capturing(): - if layer_idx % 10 == 0: - print(f"[CLAWMINE] Layer {layer_idx}: amax={hidden_states.amax().item():.4f} mean={hidden_states.mean().item():.6f}") - if torch.isnan(hidden_states).any(): - nan_pct = torch.isnan(hidden_states).float().mean().item() * 100 - print(f"[CLAWMINE] NaN after layer {layer_idx}! {nan_pct:.2f}% NaN, amax={hidden_states.amax().item():.4f}") - break - if torch.isinf(hidden_states).any(): - inf_pct = torch.isinf(hidden_states).float().mean().item() * 100 - print(f"[CLAWMINE] Inf after layer {layer_idx}! {inf_pct:.2f}% Inf, amax={hidden_states.amax().item():.4f}") - break + # NaN detection — only during prefill. Disabled via env var during cudagraph. + # os.environ is evaluated at trace time by Dynamo, so the entire + # NaN check block is skipped during compilation. + if os.environ.get('CLAWMINE_NAN_CHECK', '0') == '1': + with torch.no_grad(): + if torch.isnan(hidden_states).any(): + nan_pct = torch.isnan(hidden_states).float().mean().item() * 100 + print(f"[CLAWMINE] NaN after layer {layer_idx}! {nan_pct:.2f}% NaN, amax={hidden_states.amax().item():.4f}") + break + if torch.isinf(hidden_states).any(): + inf_pct = torch.isinf(hidden_states).float().mean().item() * 100 + print(f"[CLAWMINE] Inf after layer {layer_idx}! {inf_pct:.2f}% Inf, amax={hidden_states.amax().item():.4f}") + break + if layer_idx % 10 == 0: + print(f"[CLAWMINE] Layer {layer_idx}: amax={hidden_states.amax().item():.4f} mean={hidden_states.mean().item():.6f}") # Stash pre-hc_head residual for the MTP draft (captured copy_). num_tokens = hidden_states.shape[0]