Remove NaN check — incompatible with Dynamo fullgraph compilation

Dynamo fullgraph mode rejects BOTH data-dependent branching AND
torch.compiler.disable as graph breaks. The NaN check cannot coexist
with vLLM's AOT compilation. Use layertest/cudagraph_test for debugging.
This commit is contained in:
2026-05-18 12:17:25 +00:00
parent 65763a200c
commit 2d1e9f42b1

View File

@@ -1210,24 +1210,6 @@ class DeepseekV4DecoderLayer(nn.Module):
return x
@torch.compiler.disable
def _clawmine_nan_check(hidden_states: torch.Tensor, layer_idx: int):
"""NaN/Inf detection — only active when CLAWMINE_NAN_CHECK=1.
Decorated with @torch.compiler.disable so Dynamo never traces
the data-dependent branching (if tensor.any())."""
if os.environ.get('CLAWMINE_NAN_CHECK', '0') != '1':
return
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}")
elif 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}")
elif layer_idx % 10 == 0:
print(f"[CLAWMINE] Layer {layer_idx}: amax={hidden_states.amax().item():.4f} mean={hidden_states.mean().item():.6f}")
@support_torch_compile
class DeepseekV4Model(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -1334,9 +1316,7 @@ class DeepseekV4Model(nn.Module):
positions,
input_ids,
)
# NaN detection — guarded by env var, wrapped in torch.compiler.disable
# so Dynamo never traces the data-dependent branching.
_clawmine_nan_check(hidden_states, layer_idx)
# Stash pre-hc_head residual for the MTP draft (captured copy_).
num_tokens = hidden_states.shape[0]