biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

diag: check SE weight and scale integrity for first 3 layers

Browse Source
This commit is contained in:
biondizzle
2026-06-01 04:08:21 +00:00
parent 8ad617e2ff
commit bcd7a0cf0d
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
single_shot_inference.py
View File

@@ -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