Update debug_ee.py

debug_ee.py

@@ -1,5 +1,6 @@
 """
-EE Sanity Check + Layer Diagnostics
+EE Sanity Check Script
+Run this locally (not on HF Spaces) to verify the transform is correct.
 Usage:
   python debug_ee.py --original Qwen/Qwen3-0.6B --ee your/model-dp-ee --seed 424242
 """
@@ -16,97 +17,93 @@ def get_sigma(hidden_size, seed):
 
 def run_check(original_name, ee_name, seed, prompt="Hello, how are you?"):
     print(f"\n{'='*60}")
+    print(f"Original : {original_name}")
+    print(f"EE model : {ee_name}")
+    print(f"Seed     : {seed}")
+    print(f"Prompt   : {prompt}")
+    print('='*60)
+
     tokenizer = AutoTokenizer.from_pretrained(original_name, trust_remote_code=True)
     inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs.input_ids
+
+    print("\n[1] Loading original model...")
+    orig = AutoModelForCausalLM.from_pretrained(
+        original_name, torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True
+    )
+    orig.eval()
 
-    print("[1] Loading models...")
-    orig = AutoModelForCausalLM.from_pretrained(original_name, torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True)
-    ee   = AutoModelForCausalLM.from_pretrained(ee_name,       torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True)
-    orig.eval(); ee.eval()
+    print("[2] Loading EE model...")
+    ee = AutoModelForCausalLM.from_pretrained(
+        ee_name, torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True
+    )
+    ee.eval()
 
     hidden_size = orig.config.hidden_size
     sigma, sigma_inv = get_sigma(hidden_size, seed)
-    print(f"hidden_size={hidden_size}, seed={seed}")
 
-    # --- CHECK 1: Embed layers ---
-    embed_match = torch.allclose(orig.model.embed_tokens.weight.data, ee.model.embed_tokens.weight.data, atol=1e-3)
+    # --- Check 1: Does the EE embed layer match original? ---
+    orig_embed = orig.model.embed_tokens.weight.data
+    ee_embed   = ee.model.embed_tokens.weight.data
+    embed_match = torch.allclose(orig_embed, ee_embed, atol=1e-3)
     print(f"\n[CHECK 1] Embed layers identical: {embed_match}")
+    if not embed_match:
+        diff = (orig_embed - ee_embed).abs().max().item()
+        print(f"  ⚠️  Max diff: {diff:.6f} — EE embed was permuted, this BREAKS client-side encryption")
+        print(f"  → Re-run transform with the embed layer skipped (see transform_fix.py)")
+
+    # --- Check 2: Run plain forward on original ---
+    print("\n[CHECK 2] Running plain forward on original...")
+    with torch.no_grad():
+        plain_embeds = orig.model.embed_tokens(input_ids)
+        orig_out = orig(inputs_embeds=plain_embeds, output_hidden_states=False)
+    orig_logits = orig_out.logits  # (1, seq, vocab)
 
-    # --- LAYER DIFF: print every layer that differs and HOW ---
-    print("\n[LAYER DIFF] Comparing every named parameter...")
-    ROPE_OUTPUT_LAYERS = {"q_proj", "k_proj"}
-    issues = []
-    for (name_o, param_o), (name_e, param_e) in zip(orig.named_parameters(), ee.named_parameters()):
-        assert name_o == name_e
-        if torch.allclose(param_o.data, param_e.data, atol=1e-3):
-            continue  # unchanged — skip
-
-        basename = name_o.split(".")[-1]  # "weight", "bias"
-        layer    = name_o.split(".")[-2]  # "q_proj", "embed_tokens", etc.
-        shape    = tuple(param_o.shape)
-
-        # Check what the transform DID to this param
-        changed_cols = changed_rows = False
-        if param_o.dim() == 2:
-            if not torch.allclose(param_o.data, param_e.data[:, np.argsort(sigma_inv)], atol=1e-3):
-                pass
-            # Did it permute cols?
-            reconstructed_cols = param_e.data[:, np.argsort(sigma_inv)]
-            changed_cols = torch.allclose(param_o.data, reconstructed_cols, atol=1e-3)
-            # Did it permute rows?
-            reconstructed_rows = param_e.data[np.argsort(sigma_inv), :]
-            changed_rows = torch.allclose(param_o.data, reconstructed_rows, atol=1e-3)
-            # Did it permute both?
-            reconstructed_both = param_e.data[np.argsort(sigma_inv), :][:, np.argsort(sigma_inv)]
-            changed_both = torch.allclose(param_o.data, reconstructed_both, atol=1e-3)
-
-            what = []
-            if changed_both: what = ["BOTH rows+cols"]
-            elif changed_cols: what = ["cols only"]
-            elif changed_rows: what = ["rows only"]
-            else: what = ["UNKNOWN permutation"]
-
-            flag = ""
-            if layer in ROPE_OUTPUT_LAYERS and ("BOTH" in what[0] or "rows" in what[0]):
-                flag = "  ⚠️  BAD: RoPE layer has rows permuted!"
-                issues.append(f"{name_o}: rows permuted on RoPE layer")
-            elif layer not in ROPE_OUTPUT_LAYERS and shape[0] == hidden_size and shape[1] == hidden_size and "BOTH" not in what[0]:
-                flag = f"  ⚠️  BAD: square hidden layer should have BOTH permuted"
-                issues.append(f"{name_o}: square layer missing full permutation")
-
-            print(f"  {layer:20s} {str(shape):20s} → {what[0]}{flag}")
-
-        elif param_o.dim() == 1:
-            print(f"  {layer:20s} {str(shape):20s} → 1D (norm/bias)")
-
-    # --- CHECK 4: Logits ---
-    print("\n[CHECK 4] Equivariance test...")
+    # --- Check 3: Run encrypted forward on EE model ---
+    print("[CHECK 3] Running encrypted forward on EE model...")
     with torch.no_grad():
-        plain_embeds = orig.model.embed_tokens(inputs.input_ids)
         encrypted_embeds = plain_embeds[..., sigma]
-        orig_logits = orig(inputs_embeds=plain_embeds).logits
-        ee_logits   = ee(inputs_embeds=encrypted_embeds).logits
+        ee_out = ee(inputs_embeds=encrypted_embeds, output_hidden_states=False)
+    ee_logits = ee_out.logits
 
+    # --- Check 4: Do logits match? ---
+    logit_match = torch.allclose(orig_logits, ee_logits, atol=1e-1)
     max_diff = (orig_logits - ee_logits).abs().max().item()
-    match = max_diff < 0.5
-    print(f"  Max logit diff: {max_diff:.4f}  →  {'✅ PASS' if match else '❌ FAIL'}")
-
-    # --- CHECK 5: Decode ---
-    print("\n[CHECK 5] Greedy decode (10 tokens)...")
+    print(f"\n[CHECK 4] Logits match (atol=0.1): {logit_match}")
+    print(f"  Max logit diff: {max_diff:.4f}")
+    if not logit_match:
+        print("  ⚠️  Logits differ — equivariance is BROKEN")
+        # Find where it breaks — check RoPE suspicion
+        print("\n  Diagnosing: checking if RoPE is the culprit...")
+        print("  RoPE applies rotation in head_dim space (64), not hidden space (1024)")
+        print("  If q_proj/k_proj output is permuted (because output==hidden_size),")
+        print("  the head_dim slices fed to RoPE will be scrambled → broken attention")
+
+    # --- Check 5: Greedy decode comparison ---
+    print("\n[CHECK 5] Greedy decode comparison (10 tokens)...")
     with torch.no_grad():
-        orig_ids = orig.generate(inputs.input_ids, max_new_tokens=10, do_sample=False)
+        orig_ids = orig.generate(input_ids, max_new_tokens=10, do_sample=False)
         ee_ids   = ee.generate(inputs_embeds=encrypted_embeds,
                                attention_mask=inputs.attention_mask,
                                max_new_tokens=10, do_sample=False,
                                pad_token_id=tokenizer.eos_token_id)
-    print(f"  Original : {repr(tokenizer.decode(orig_ids[0], skip_special_tokens=True))}")
-    print(f"  EE model : {repr(tokenizer.decode(ee_ids[0],   skip_special_tokens=True))}")
-
-    if issues:
-        print(f"\n⚠️  {len(issues)} issue(s) found:")
-        for i in issues: print(f"  - {i}")
-    else:
-        print("\n✅ No layer issues detected")
+
+    orig_text = tokenizer.decode(orig_ids[0], skip_special_tokens=True)
+    ee_text   = tokenizer.decode(ee_ids[0],   skip_special_tokens=True)
+    print(f"  Original output : {repr(orig_text)}")
+    print(f"  EE model output : {repr(ee_text)}")
+    print(f"  Match: {orig_text == ee_text}")
+
+    if orig_text != ee_text:
+        print("\n  ⚠️  OUTPUTS DIFFER. Most likely causes in order:")
+        print("  1. Embed layer was permuted in EE model (Check 1 above)")
+        print("  2. RoPE disruption — q_proj/k_proj output rows were permuted")
+        print("     FIX: do NOT permute output rows of q_proj and k_proj")
+        print("     because their outputs are split into heads for RoPE rotation")
+        print("  3. Model on Hub is stale — re-run transform and re-push")
+
+    print(f"\n{'='*60}\n")
+    return embed_match and logit_match
 
 if __name__ == "__main__":