diff --git a/single_shot_inference.py b/single_shot_inference.py index 4ef98edf..82b7c49c 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -499,7 +499,7 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id): if VERBOSE >= 2 and (li < 3 or li >= 58): if topk_ids.max().item() >= 384 or topk_ids.min().item() < 0: print(f" L{li} BAD topk_ids: min={topk_ids.min().item()} max={topk_ids.max().item()}", flush=True) - if li >= 58: + if VERBOSE >= 2 and li >= 58: print(f" L{li} MoE DIAG: topk_ids={topk_ids[0].tolist()} topk_w=[{','.join(f'{w:.3f}' for w in topk_w[0].tolist())}]", flush=True) # Also print gate logits for debugging if hasattr(router, '_gate_lin') and router._gate_lin is not None: @@ -509,7 +509,7 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id): print(f" L{li} MoE input: |x|={x.abs().max().item():.4f} has_nan={torch.isnan(x).any().item()}", flush=True) routed_out = moe_runner.run(x, topk_w, topk_ids) shared_out = se_runner.run(x) - if li >= 58: + if VERBOSE >= 2 and li >= 58: print(f" L{li} MoE DIAG: |routed|={routed_out.abs().max().item():.1f} |shared|={shared_out.abs().max().item():.1f} |x|={x.abs().max().item():.1f}", flush=True) if VERBOSE >= 2 and li < 3: has_nan = torch.isnan(shared_out).any().item() @@ -541,11 +541,11 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, x_in_f, ctx_f = ffn_mhc.pre_block(X_mid); x_ffn = rmsnorm(x_in_f, ffn_norm_w) F_ffn = moe_forward(x_ffn, li, moe_runner, se_runner, router, token_id) X_next = ffn_mhc.post_block(X_mid, F_ffn, ctx_f) - if VERBOSE >= 1 and (li < 3 or li >= 58): + if VERBOSE >= 2 and (li < 3 or li >= 58): print(f" L{li}: |X|={X_l.abs().max().item():.1f}->{X_next.abs().max().item():.1f} " f"|Fa|={F_attn.abs().max().item():.1f} |Ff|={F_ffn.abs().max().item():.1f}", flush=True) - # Detailed diagnostics for last 3 layers or any layer with explosive growth - if li >= 58 or (li > 0 and X_next.abs().max().item() > 200): + # Detailed diagnostics — only with VERBOSE >= 2 to avoid .item() syncs on hot path + if VERBOSE >= 2 and (li >= 58 or (li > 0 and X_next.abs().max().item() > 200)): A_a, B_a, C_a = attn_mhc._dynamic_params(X_l) A_f, B_f, C_f = ffn_mhc._dynamic_params(X_mid) print(f" L{li} DIAG: A_attn=[{A_a.min().item():.4f},{A_a.max().item():.4f}] "