diff --git a/single_shot_inference.py b/single_shot_inference.py index 0be94b1d..8233b7fe 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -424,6 +424,13 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id): 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}", flush=True) + # Check weight integrity + if hasattr(se_runner, '_l1_mat_b') and se_runner._l1_mat_b is not None: + wb = se_runner._l1_mat_b.view(torch.uint8) + print(f" L{li} SE l1 weight: shape={list(se_runner._l1_mat_b.shape)} dtype={se_runner._l1_mat_b.dtype} uint8_range=[{wb.min().item()},{wb.max().item()}]", flush=True) + if hasattr(se_runner, '_l1_scale_b') and se_runner._l1_scale_b is not None: + sb = se_runner._l1_scale_b + print(f" L{li} SE l1 scale: shape={list(sb.shape)} dtype={sb.dtype} range=[{sb.min().item():.6f},{sb.max().item():.6f}] has_nan={torch.isnan(sb).any().item()}", 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