diff --git a/single_shot_inference.py b/single_shot_inference.py index 622741b8..8f11ec17 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -366,6 +366,8 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, a_grp = a_flat.reshape(T, o_groups, gid); oa_3d = oa_bf.reshape(o_groups, o_rank, gid) g_out = torch.bmm(a_grp.permute(1, 0, 2), oa_3d.transpose(1, 2)) g_flat = g_out.permute(1, 0, 2).reshape(T, o_groups * o_rank) + if li < 3: + print(f" L{li} wo_a: |g_flat|={g_flat.abs().max().item():.6f} shape={g_flat.shape}", flush=True) F_attn = prod_lin['o_b'](g_flat) else: # o_a_proj as full-rank BF16 linear (no low-rank) @@ -375,6 +377,8 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, else: log.warning(f"L{li}: No o_a_proj weight, zero attention output") F_attn = torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev) + if li < 3: + print(f" L{li} F_attn: |F_attn|={F_attn.abs().max().item():.6f}", flush=True) return F_attn, q_a # =====================================================================