Add NaN/Inf detection in DeepseekV4Model.forward layer loop
- Checks every layer during prefill (not during cudagraph capture) - is_current_stream_capturing() gate prevents CPU-GPU syncs during capture - Prints amax every 10 layers for magnitude tracking - Breaks on first NaN/Inf to avoid wasting compute
This commit is contained in:
@@ -1309,12 +1309,24 @@ class DeepseekV4Model(nn.Module):
|
||||
hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
|
||||
if self.use_mega_moe:
|
||||
input_ids = input_ids.to(torch.int64)
|
||||
for layer in islice(self.layers, self.start_layer, self.end_layer):
|
||||
for layer_idx, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)):
|
||||
hidden_states = layer(
|
||||
hidden_states,
|
||||
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
|
||||
|
||||
# Stash pre-hc_head residual for the MTP draft (captured copy_).
|
||||
num_tokens = hidden_states.shape[0]
|
||||
|
||||
Reference in New Issue
Block a user