diff --git a/single_shot_inference.py b/single_shot_inference.py index b4949944..bb2c6498 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -505,6 +505,10 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, # -- mHC post_block (attention) -- X_mid = attn_mhc.post_block(X_l, F_attn, attn_ctx) # (T, n_hc, H) + # Diagnostic: check mHC is stabilizing the residual + if li < 3: + B_l, C_l = attn_ctx + print(f" L{li} attn: |X_l|={X_l.abs().max().item():.2f} |F_attn|={F_attn.abs().max().item():.2f} |B|={B_l.abs().max().item():.4f} |C|={C_l.abs().max().item():.4f} |X_mid|={X_mid.abs().max().item():.2f}", flush=True) # ================================================================== # FFN SUB-BLOCK