From 22d7cc9b7ad8e88a3bf5b448a306bd6d38c30ad0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 02:56:28 +0000 Subject: [PATCH] diag: cuda sync check after shared expert for first 3 layers --- single_shot_inference.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 2d68dbf3..d57b4f1e 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -421,7 +421,13 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id): print(f" L{li} MoE routed: |out|={routed_out.abs().max().item():.4f} has_nan={torch.isnan(routed_out).any().item()}", flush=True) shared_out = se_runner.run(x) if li < 3: - print(f" L{li} MoE shared: |out|={shared_out.abs().max().item():.4f} has_nan={torch.isnan(shared_out).any().item()}", flush=True) + try: + torch.cuda.synchronize(x.device) + cuda_ok = True + 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) 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