Debugging NaN Activations in vLLM on V100 GPUs
What we're building: A minimal reproduction script to isolate NaN activations when running quantized models (like MiniMax-M2.7-AWQ) through vLLM on V100 GPUs with flash-attention-v100.
Prerequisites:
- NVIDIA V100 GPU (16GB+ VRAM)
- Python 3.10+
- vLLM installed (
pip install vllm) - flash-attention-v100 installed (
pip install flash-attention-v100) - A quantized model (we'll use MiniMax-M2.7-AWQ)
The Problem
V100 GPUs don't have native support for FlashAttention-2, so projects like flash-attention-v100 provide custom kernels. But quantized models + custom attention kernels = NaN activations with certain inputs. Let's build a script to reproduce and debug this.
Step 1: Minimal Reproduction Script
Create reproduce_nan.py:
import torch
from vllm import LLM, SamplingParams
import traceback
MODEL_PATH = "/home/youruser/models/MiniMax-M2.7-AWQ"
PROMPTS = [
"Explain quantum computing in simple terms.",
"Write a Python function to reverse a linked list.",
"What is the capital of France?",
"Translate 'hello world' to Spanish.",
"Summarize the plot of Dune.",
]
def test_single_prompt(llm, prompt: str, max_tokens: int = 50):
"""Test a single prompt and catch NaN issues."""
print(f"\n{'='*60}")
print(f"Testing prompt: {prompt[:50]}...")
try:
sampling_params = SamplingParams(
temperature=0.0,
max_tokens=max_tokens,
top_p=0.95,
)
outputs = llm.generate([prompt], sampling_params)
output_text = outputs[0].outputs[0].text
# Check for NaN in output (indicates problem)
if "nan" in output_text.lower():
print("❌ NaN detected in output!")
return False
# Also check tokens
token_ids = outputs[0].outputs[0].token_ids
if any(t == float('nan') for t in token_ids):
print("❌ NaN in token IDs!")
return False
print(f"✅ OK: {output_text[:80]}...")
return True
except Exception as e:
print(f"❌ Exception: {e}")
traceback.print_exc()
return False
def main():
print("Loading model...")
llm = LLM(
model=MODEL_PATH,
tensor_parallel_size=1,
gpu_memory_utilization=0.85,
max_model_len=2048,
enforce_eager=True, # Avoid CUDA graph issues
dtype=torch.float16,
)
print("Model loaded successfully.")
results = []
for prompt in PROMPTS:
success = test_single_prompt(llm, prompt)
results.append(success)
print(f"\n{'='*60}")
print(f"Results: {sum(results)}/{len(results)} passed")
if not all(results):
print("❌ NaN issues detected. See above for failing prompts.")
if __name__ == "__main__":
main()
Step 2: Add Activation Monitoring
Now let's add hooks to catch NaNs before they become garbage output:
import torch.nn as nn
def add_nan_monitor(model):
"""Add hooks to detect NaN in activations."""
nan_locations = []
def hook_fn(module, input, output):
# Check output for NaN
if isinstance(output, torch.Tensor):
if torch.isnan(output).any():
nan_locations.append(f"{module.__class__.__name__} at {id(module)}")
# Check input
if isinstance(input, tuple):
for i, inp in enumerate(input):
if isinstance(inp, torch.Tensor) and torch.isnan(inp).any():
nan_locations.append(f"{module.__class__.__name__} input[{i}]")
return output
# Register hooks on attention layers
for name, module in model.named_modules():
if "attention" in name.lower() or "layernorm" in name.lower():
module.register_forward_hook(hook_fn)
return nan_locations
# In main():
# After loading model, access the underlying model
# vLLM doesn't expose this directly - you'd need to patch vllm or use engine internals
# Instead, let's check intermediate states via vLLM's internal API
Step 3: Use vLLM's Internal State
Since vLLM abstracts the model, let's use its debug features:
import os
from vllm import LLM
def debug_nan_vllm():
# Enable vLLM's debug logging
os.environ["VLLM_LOG_LEVEL"] = "DEBUG"
os.environ["VLLM_ATTENTION_BACKEND"] = "FLASH_ATTN" # Force flash-attn
# Or try XFORMERS as alternative
# os.environ["VLLM_ATTENTION_BACKEND"] = "XFORMERS"
llm = LLM(
model=MODEL_PATH,
enforce_eager=True,
max_num_seqs=1, # Reduce batch size to isolate
max_num_batched_tokens=512, # Smaller batches
)
# Test with minimal input
result = llm.generate(
["Hello"],
SamplingParams(max_tokens=10, temperature=0.0)
)
print(result[0].outputs[0].text)
Step 4: Binary Search for the Problem
Create binary_search.py to find the exact failing input length:
from vllm import LLM, SamplingParams
def find_failing_length(llm, base_prompt="The quick brown fox jumps over the lazy dog. "):
"""Binary search for input length that triggers NaN."""
def test_length(length):
prompt = base_prompt * (length // len(base_prompt) + 1)
prompt = prompt[:length]
try:
out = llm.generate([prompt], SamplingParams(max_tokens=5, temperature=0.0))
text = out[0].outputs[0].text
return not ("nan" in text.lower() or "inf" in text.lower())
except Exception as e:
if "nan" in str(e).lower() or "inf" in str(e).lower():
return False
return True
# Binary search
lo, hi = 1, 2048
while lo < hi:
mid = (lo + hi) // 2
if test_length(mid):
lo = mid + 1
else:
hi = mid
return lo
# Usage
llm = LLM(model=MODEL_PATH, enforce_eager=True)
failing_len = find_failing_length(llm)
print(f"First NaN at input length: {failing_len}")
Adding Observability with TracePilot
Now let's add proper tracing to see where in the pipeline things break:
pip install tracepilot-sdk
python
from tracepilot_sdk import TracePilot
from vllm import LLM, SamplingParams
tp = TracePilot('tp_live_YOUR_KEY') # Get free key at tracepilotai.com
def traced_generate(llm, prompt):
"""Generate with full TracePilot visibility."""
async def run():
await tp.startTrace('vllm-nan-debug')
# Trace the input
messages = [{"role": "user", "content": prompt}]
# Wrap the generation call
result = await tp.wrapOpenAI(
lambda: llm.generate(
[prompt],
SamplingParams(max_tokens=50, temperature=0.0)
),
messages,
---
**Debugging AI agents shouldn't feel like reading The Matrix.**
Join other engineers who are building reliable autonomous workflows in our community: [TracePilot Discord](https://discord.gg/KzXRAXFM8)
Top comments (0)