Code simplification

mergekit/scripts/ABM/activations_based_merge.py

@@ -16,8 +16,8 @@
 
 
 @click.command("mergekit-activation-based-merge")
-@click.argument("model_path", type=str)
-@click.argument("secondary_model_path", type=str)
+@click.argument("model_path", type=str, help="Path to the anchor model")
+@click.argument("secondary_model_path", type=str, help="Path to the secondary model")
 @click.argument("merge_unmerge_directory", type=str)
 @click.option("--out-path", "-o", required=True, type=str, help="Output model path")
 @click.option(
@@ -121,8 +121,8 @@ def main(
 
         if merge_matrix is not None:
             if weight_info.is_embed:
-                w = (merge_matrix[0] @ w.T).T
-                w2 = (merge_matrix[1] @ w2.T).T
+                w = w @ merge_matrix[0].T
+                w2 = w2 @ merge_matrix[1].T
             else:
                 w = merge_matrix[0] @ w
                 w2 = merge_matrix[1] @ w2
@@ -151,10 +151,7 @@ def main(
             )
 
         # average weights and save them
-        if merge_matrix:
-            w = w + w2
-        else:
-            w = (w + w2) / 2
+        w = (w + w2) / 2
         writer.save_tensor(weight_info.name, w)
     writer.finalize()
 

mergekit/scripts/ABM/extract_permutation_matrices.py

@@ -39,14 +39,10 @@ def match_tensors_permute(
     new_mat = torch.eye(Om, device=device)[torch.tensor(col_ind).long().to(device)]
     mats.append(new_mat.T)
 
-    unmerge_mats = mats
+    unmerge = torch.cat(mats, dim=0)
+    merge = unmerge.clone().T
 
-    unmerge = torch.cat(unmerge_mats, dim=0)
-
-    merge = torch.cat(mats, dim=0)
-    merge = merge / (merge.sum(dim=0, keepdim=True) + 1e-5)
-
-    return merge.T, unmerge
+    return merge, unmerge
 
 
 def match_tensors_permute_MHA(
@@ -111,13 +107,10 @@ def match_tensors_permute_MHA(
     ]
     mats.append(new_mat.T)
 
-    unmerge_mats = mats
-
-    unmerge = torch.cat(unmerge_mats, dim=0)
-    merge = torch.cat(mats, dim=0)
-    merge = merge / (merge.sum(dim=0, keepdim=True) + 1e-5)
+    unmerge = torch.cat(mats, dim=0)
+    merge = unmerge.clone().T
 
-    return merge.T, unmerge
+    return merge, unmerge
 
 
 @click.command("mergekit-abm-extract-permutations")