From c3a2656c486b5f5612ab77b04ceb4c05c7bc70db Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 31 May 2026 03:12:52 +0000 Subject: [PATCH] diag: add FFN and pre_block diagnostics --- single_shot_inference.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/single_shot_inference.py b/single_shot_inference.py index bb2c6498..aa828864 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -425,6 +425,10 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, # -- mHC pre_block (attention) -- x_in, attn_ctx = attn_mhc.pre_block(X_l) # x_in: (T, H) + if li < 3: + A_l = None + B_l, C_l = attn_ctx + print(f" L{li} pre_attn: |X_l|={X_l.abs().max().item():.2f} |x_in|={x_in.abs().max().item():.2f}", flush=True) # -- RMSNorm (pre-norm before attention) -- x_normed = attn_norm.forward(x_in) # (T, H) BF16 @@ -525,6 +529,9 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, # -- mHC post_block (FFN) -- X_next = ffn_mhc.post_block(X_mid, F_ffn, ffn_ctx) # (T, n_hc, H) + if li < 3: + B_l, C_l = ffn_ctx + print(f" L{li} ffn: |X_mid|={X_mid.abs().max().item():.2f} |F_ffn|={F_ffn.abs().max().item():.2f} |B|={B_l.abs().max().item():.4f} |C|={C_l.abs().max().item():.4f} |X_next|={X_next.abs().max().item():.2f}", flush=True) return X_next