diag: handle NaN in shared expert output print

This commit is contained in:
2026-06-01 03:16:25 +00:00
parent 7000762309
commit ee8318ad58

View File

@@ -427,11 +427,9 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id):
except RuntimeError as e:
cuda_ok = False
print(f" L{li} SE CUDA ERROR: {e}", flush=True)
print(f" L{li} MoE shared: |out|={shared_out.abs().max().item():.4f} has_nan={torch.isnan(shared_out).any().item()} cuda_ok={cuda_ok}", flush=True)
# Check SE weight integrity
if hasattr(se_runner, '_l1_mat_b') and se_runner._l1_mat_b is not None:
w = se_runner._l1_mat_b
print(f" L{li} SE _l1_mat_b: |max|={w.abs().max().item():.4f} has_nan={torch.isnan(w.view(torch.float16)).any().item()} shape={w.shape}", flush=True)
has_nan = torch.isnan(shared_out).any().item()
out_max = shared_out.abs().max().item() if not has_nan else float('nan')
print(f" L{li} MoE shared: |out|={out_max:.4f} has_nan={has_nan} cuda_ok={cuda_ok}", flush=True)
print(f" L{li} SE gsa: l1={se_runner._l1_activation_global_scale:.6f} l2={se_runner._l2_activation_global_scale:.6f} gsb: l1={se_runner._l1_gsb[0].item():.6f} l2={se_runner._l2_gsb[0].item():.6f}", flush=True)
return routed_out + shared_out