From 262f844e2e4d6240777920f64e39c7684cff3d10 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 06:10:33 +0000 Subject: [PATCH] =?UTF-8?q?PART=20A:=20add=20detailed=20blowup=20diagnosti?= =?UTF-8?q?cs=20=E2=80=94=20capture=20mHC=20intermediate=20values=20when?= =?UTF-8?q?=20|X|=20>=201e6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit/test_part_a_decode_diagnostics.py | 49 ++++++++++++++++---- 1 file changed, 39 insertions(+), 10 deletions(-) diff --git a/tests/unit/test_part_a_decode_diagnostics.py b/tests/unit/test_part_a_decode_diagnostics.py index 4b0e6109..4c91f461 100644 --- a/tests/unit/test_part_a_decode_diagnostics.py +++ b/tests/unit/test_part_a_decode_diagnostics.py @@ -240,6 +240,7 @@ def main(): if X.device != torch.device(dev): X = X.to(dev) torch.cuda.set_device(gpu) + X_prev = X.clone() # Save for blowup diagnostics X_in_mag = X.abs().max().item() X = forward_layer(X, layer_w[li], li, cfg, *rope_caches[gpu], attn_mhcs.get(li), ffn_mhcs.get(li), attn_norms.get(li), ffn_norms.get(li), @@ -254,16 +255,44 @@ def main(): if pi < 3 or pi == len(input_ids) - 1: print(f" {pi:>3} {li:>3} {X_in_mag:>12.2f} {X_out_mag:>12.2f} {ratio:>5} {kc.n_comp:>6} {kc.swa_len:>4}", flush=True) - # Early abort if |X| blows up - if X_out_mag > 1e10: - print(f" *** BLOWUP at token {pi} layer {li}: |X|={X_out_mag:.2e} — ABORTING ***", flush=True) - print(f" This means the production pipeline is numerically unstable.", flush=True) - print(f" Check: mHC residual growth, NVFP4 quantization, MoE scaling.", flush=True) - # Print KV cache state at this point - for l2 in range(li + 1): - kc2 = kv_caches[l2] - r2 = cr[l2] if l2 < len(cr) else 128 - print(f" L{l2} (ratio={r2}): n_comp={kc2.n_comp} swa_len={kc2.swa_len}", flush=True) + # Early abort if |X| blows up — run detailed diagnostics on THIS layer + if X_out_mag > 1e6: + print(f" *** BLOWUP at token {pi} layer {li}: |X|={X_out_mag:.2e} ***", flush=True) + print(f" Re-running layer {li} with detailed diagnostics...", flush=True) + # Re-run the SAME input through forward_layer but capture intermediates + X_diag = X_prev.clone() # X before this layer + attn_mhc_d = attn_mhcs.get(li) + ffn_mhc_d = ffn_mhcs.get(li) + A_l_a, B_l_a, C_l_a = attn_mhc_d._dynamic_params(X_diag) + ctx_a_d = mHCContext(B_l=B_l_a, C_l=C_l_a) + x_quant_attn = mhc_rmsnorm_quantize_nvfp4( + X_diag, A_l_a, attn_norms.get(li).to(dev, torch.float32)) + x_normed = dequantize_nvfp4(x_quant_attn.x_fp4, x_quant_attn.x_sf, x_quant_attn.gsa) + print(f" |x_normed|={x_normed.abs().max().item():.2f} gsa={x_quant_attn.gsa}", flush=True) + F_attn_d, q_a_d = forward_attention( + x_normed, layer_w[li], li, cfg, *rope_caches[gpu], + kv_caches[li], pos, compressors.get(li), indexers.get(li), prod_lins.get(li), + x_quant=x_quant_attn) + print(f" |F_attn|={F_attn_d.abs().max().item():.2f}", flush=True) + X_mid_d = attn_mhc_d.post_block(X_diag, F_attn_d, ctx_a_d) + print(f" |X_mid|={X_mid_d.abs().max().item():.2f} B_l_row=[{B_l_a.sum(-1).min().item():.4f},{B_l_a.sum(-1).max().item():.4f}] C_l=[{C_l_a.min().item():.4f},{C_l_a.max().item():.4f}]", flush=True) + A_l_f, B_l_f, C_l_f = ffn_mhc_d._dynamic_params(X_mid_d) + ctx_f_d = mHCContext(B_l=B_l_f, C_l=C_l_f) + x_quant_ffn = mhc_rmsnorm_quantize_nvfp4( + X_mid_d, A_l_f, ffn_norms.get(li).to(dev, torch.float32)) + x_ffn = dequantize_nvfp4(x_quant_ffn.x_fp4, x_quant_ffn.x_sf, x_quant_ffn.gsa) + F_ffn_d = moe_forward(x_ffn, li, moe_runners.get(li), se_runners.get(li), + routers.get(li), tid32.to(dev)) + print(f" |F_ffn|={F_ffn_d.abs().max().item():.2f}", flush=True) + X_next_d = ffn_mhc_d.post_block(X_mid_d, F_ffn_d, ctx_f_d) + print(f" |X_next|={X_next_d.abs().max().item():.2e}", flush=True) + # Check per-component magnitudes + BX = torch.bmm(ctx_a_d.B_l.transpose(-1, -2), X_diag.float()) + CF = ctx_a_d.C_l.unsqueeze(-1) * F_attn_d.unsqueeze(1) + print(f" |B@X|={BX.abs().max().item():.2f} |C*F|={CF.abs().max().item():.2f}", flush=True) + BX_f = torch.bmm(ctx_f_d.B_l.transpose(-1, -2), X_mid_d.float()) + CF_f = ctx_f_d.C_l.unsqueeze(-1) * F_ffn_d.unsqueeze(1) + print(f" FFN: |B@X|={BX_f.abs().max().item():.2f} |C*F|={CF_f.abs().max().item():.2f}", flush=True) return 1 if pi % 5 == 0: