From bf99ad49ecbd0d6f978532db4c5b59d506fab2e9 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 22:56:56 +0000 Subject: [PATCH] Print both MoE and residual cosine --- tests/test_multilayer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_multilayer.py b/tests/test_multilayer.py index 9a952877..1f83d27d 100644 --- a/tests/test_multilayer.py +++ b/tests/test_multilayer.py @@ -142,10 +142,11 @@ def main(): ref_out[t] += w * (act @ down_bf16.T) ref_hidden = ref_out + ref_hidden_saved # Residual + cos_moe = F.cosine_similarity(ref_out.flatten().unsqueeze(0), run_out.flatten().unsqueeze(0)).item() cos = F.cosine_similarity(ref_hidden.flatten().unsqueeze(0), run_hidden.flatten().unsqueeze(0)).item() has_nan = torch.isnan(run_hidden).any().item() has_inf = torch.isinf(run_hidden).any().item() - print(f"Layer {layer}: cosine={cos:.6f} ref_amax={ref_hidden.amax().item():.4f} run_amax={run_hidden.amax().item():.4f} NaN={has_nan} Inf={has_inf}") + print(f"Layer {layer}: MoE_cosine={cos_moe:.6f} Residual_cosine={cos:.6f} ref_amax={ref_hidden.amax().item():.4f} run_amax={run_hidden.amax().item():.4f} NaN={has_nan} Inf={has_inf}") if has_nan: print(f" ❌ NaN detected after layer {layer}! Stopping.")