Upload merge_composition_audit.py

merge_composition_audit.py

@@ -0,0 +1,94 @@
+import os
+import torch
+from safetensors.torch import safe_open
+import yaml
+
+# --- CONFIGURATION ---
+YAML_PATH = "B:/24B/karcher_stock_24b/mergekit_config.yml"
+FINAL_MERGE_DIR = "B:/24B/karcher_stock_24b"
+LAYERS_TO_SCAN =[
+    "model.layers.10.mlp.down_proj.weight" # "model.language_model.layers.10.mlp.gate_proj.weight"
+]
+# ---------------------
+
+def load_tensor(model_dir, tensor_name):
+    """Finds and loads a tensor from a directory of safetensors."""
+    for file in os.listdir(model_dir):
+        if file.endswith(".safetensors"):
+            filepath = os.path.join(model_dir, file)
+            with safe_open(filepath, framework="pt", device="cpu") as f:
+                if tensor_name in f.keys():
+                    return f.get_tensor(tensor_name).float()
+    raise ValueError(f"Tensor {tensor_name} not found in {model_dir}")
+
+def main():
+    print("Loading YAML config...")
+    with open(YAML_PATH, 'r') as f:
+        config = yaml.safe_load(f)
+    
+    base_path = config['base_model']
+    donor_paths = [m['model'] for m in config['models']]
+    
+    print(f"\nScanning {len(LAYERS_TO_SCAN)} MLP layers for structural influence...\n")
+    
+    for layer in LAYERS_TO_SCAN:
+        print(f"--- Layer: {layer} ---")
+        try:
+            base_w = load_tensor(base_path, layer)
+            final_w = load_tensor(FINAL_MERGE_DIR, layer)
+            
+            # Use float64 for norm calculations to prevent precision loss in energy ratios
+            final_norm = torch.norm(final_w.double()).item()
+            final_tv = final_w - base_w
+            final_tv_norm = torch.norm(final_tv.double()).item()
+            
+            results = []
+            
+            # 1. Collect raw magnitudes of the components
+            base_norm = torch.norm(base_w.double()).item()
+            donor_tvs = []
+            for donor in donor_paths:
+                dw = load_tensor(donor, layer)
+                donor_tvs.append(dw - base_w)
+            
+            donor_tv_norms = [torch.norm(dtv.double()).item() for dtv in donor_tvs]
+            
+            # 2. Calculate Total Component Energy (Base + all Donor Deltas)
+            total_component_energy = base_norm + sum(donor_tv_norms)
+            
+            results = []
+            # 3. Assign Share to Base Model
+            base_share = (base_norm / total_component_energy) * 100
+            results.append(("(Base Model)", -1.0, 0.0, base_share))
+
+            # 4. Assign Share to Donors
+            for i, donor in enumerate(donor_paths):
+                donor_tv = donor_tvs[i]
+                
+                cos_sim = torch.nn.functional.cosine_similarity(
+                    final_tv.flatten(), donor_tv.flatten(), dim=0
+                ).item()
+                
+                rel_mag = (torch.norm(donor_tv.double()).item() / final_tv_norm)
+                
+                # Compositional Share: How much of the total energy sum belongs to this donor's delta
+                comp_share = (donor_tv_norms[i] / total_component_energy) * 100
+                
+                name = donor.split("/")[-1][:50]
+                results.append((name, cos_sim, rel_mag, comp_share))
+            
+            # Sort by highest similarity (Donors first, Base at the very bottom)
+            results.sort(key=lambda x: x[1], reverse=True)
+            
+            print(f"{'Model Name':<55} | {'Alignment {Cos}':<12} | {'Rel Mag (TV)':<12} | {'Merge Composition'}")
+            print("-" * 105)
+            for name, sim, mag, energy in results:
+                sim_str = f"{sim:12.4f}" if sim >= 0 else "    N/A     "
+                mag_str = f"{mag:11.2f}x" if mag > 0 else "    N/A    "
+                print(f"{name:<55} | {sim_str} | {mag_str} | {energy:>13.2f}%")
+            
+        except Exception as e:
+            print(f"Skipping layer due to error: {e}")
+
+if __name__ == "__main__":
+    main()