Upload 141 files

hugging/td_fuse/transport.py

@@ -394,6 +394,9 @@ def _compute_plans_fallback(
             print("[transport] This may take 2-5 minutes...")
         sys.stdout.flush()
 
+        # Track which block indices already have permutations (avoid computing twice)
+        block_perms = {}  # block_index -> perm array
+
         for i, (sl, tl) in enumerate(zip(source_layers, target_layers)):
             S = source_act[sl].numpy()
             T = target_act[tl].numpy()
@@ -403,26 +406,45 @@ def _compute_plans_fallback(
                     # Neurons already correspond (e.g. DeepSeek) — identity Q
                     Q_matrices[(sl, tl)] = np.eye(S.shape[1]) / S.shape[1]
                 else:
-                    # Neurons are SCRAMBLED (e.g. MiMo) — find the permutation
-                    # 1. Compute correlation matrix between source and target neurons
-                    S_norm = (S - S.mean(0)) / (S.std(0) + 1e-8)
-                    T_norm = (T - T.mean(0)) / (T.std(0) + 1e-8)
-                    corr = S_norm.T @ T_norm / S.shape[0]  # [hidden_dim, hidden_dim]
-
-                    # 2. Run Sinkhorn on cost matrix to get soft transport plan
-                    cost = 1.0 - corr
-                    Q_soft = _sinkhorn(cost, reg=0.05, max_iter=cfg.sinkhorn_max_iter)
-
-                    # 3. Extract hard permutation: for each source neuron, which target neuron?
-                    perm = np.argmax(Q_soft, axis=1)  # source_neuron -> target_neuron
-
-                    # 4. Check for duplicate assignments (Sinkhorn should avoid this, but be safe)
-                    if len(set(perm)) < len(perm) * 0.9:
-                        # Too many collisions — fall back to Hungarian-style greedy
-                        perm = _greedy_permutation(corr)
-
-                    permutations[(sl, tl)] = perm
-                    Q_matrices[(sl, tl)] = Q_soft
+                    # Extract block index (e.g. "model.layers.5.mlp" -> 5)
+                    block_idx = None
+                    for part_j, part in enumerate(tl.split(".")):
+                        if part == "layers":
+                            try:
+                                block_idx = int(tl.split(".")[part_j + 1])
+                            except (ValueError, IndexError):
+                                pass
+                            break
+
+                    # Reuse permutation if we already computed it for this block
+                    if block_idx is not None and block_idx in block_perms:
+                        perm = block_perms[block_idx]
+                        permutations[(sl, tl)] = perm
+                        Q_matrices[(sl, tl)] = np.eye(S.shape[1]) / S.shape[1]  # placeholder
+                    else:
+                        # Neurons are SCRAMBLED (e.g. MiMo) — find the permutation
+                        # 1. Compute correlation matrix between source and target neurons
+                        S_norm = (S - S.mean(0)) / (S.std(0) + 1e-8)
+                        T_norm = (T - T.mean(0)) / (T.std(0) + 1e-8)
+                        corr = S_norm.T @ T_norm / S.shape[0]  # [hidden_dim, hidden_dim]
+
+                        # 2. Run Sinkhorn on cost matrix to get soft transport plan
+                        # Use reg=0.1 and 30 iters (faster — we only need argmax, not precision)
+                        cost = 1.0 - corr
+                        Q_soft = _sinkhorn(cost, reg=0.1, max_iter=30)
+
+                        # 3. Extract hard permutation: for each source neuron, which target neuron?
+                        perm = np.argmax(Q_soft, axis=1)  # source_neuron -> target_neuron
+
+                        # 4. Check for duplicate assignments (Sinkhorn should avoid this, but be safe)
+                        if len(set(perm)) < len(perm) * 0.9:
+                            # Too many collisions — fall back to Hungarian-style greedy
+                            perm = _greedy_permutation(corr)
+
+                        permutations[(sl, tl)] = perm
+                        Q_matrices[(sl, tl)] = Q_soft
+                        if block_idx is not None:
+                            block_perms[block_idx] = perm
             else:
                 # Different dims -- do lightweight Sinkhorn on this pair only
                 print(f"  Layer {i}: dim mismatch ({S.shape[1]} vs {T.shape[1]}), using Sinkhorn...")
@@ -438,8 +460,8 @@ def _compute_plans_fallback(
                 print(f"  Matched layer {i + 1}/{n_source}: {sl} -> {tl}")
                 sys.stdout.flush()
 
-            # Timeout: 15 min (permutation takes longer than identity)
-            tracker.check_timeout(timeout_seconds=900)
+            # Timeout: 90 min (Sinkhorn on 4096x4096 is slow on CPU)
+            tracker.check_timeout(timeout_seconds=5400)
 
         if permutations:
             print(f"[transport] Computed {len(permutations)} neuron permutations")