From 66a66f8244e0f4580b313a6f48970184c7ff504e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 23:48:32 +0000 Subject: [PATCH] Add per-layer NaN tracking for mHC debug --- single_shot_inference.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/single_shot_inference.py b/single_shot_inference.py index b7479c85..afc749ef 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -437,6 +437,17 @@ def main(): ffn_mhc = ffn_mhc_blocks.get(li) rc, rs = rope_caches[gpu] X = forward_layer(X, layer_weights[li], li, cfg, rc, rs, attn_mhc, ffn_mhc) + + # NaN check every 5 layers or on first NaN + if li % 10 == 0 or (X.float().abs().max().item() > 1e4 or torch.isnan(X.float()).any().item()): + has_nan = torch.isnan(X.float()).any().item() + xmax = X.float().abs().max().item() + print(f" L{li}: nan={has_nan}, max_abs={xmax:.2f}") + if has_nan: + # Debug: check mHC outputs + x_out = X[:, 0, :] + print(f" L{li} stream0: nan={torch.isnan(x_out.float()).any().item()}, max={x_out.float().abs().max().item():.2f}") + break # Back to gpu0 X = X.to('cuda:0')