Upload 137 files

hugging/td_fuse/merge.py

@@ -903,12 +903,20 @@ def run_single_merge(
     # --- Step 8.5: Extract post-merge activations for ARM/OTMF ---
     print(f"\n[merge] Step 8.5/10: Post-merge activations + ARM/OTMF prep..."); sys.stdout.flush()
     step_t = time.time()
-    post_merge_activations = extract_activations(target_model, calibration_data[:100])
+    arm_sample_size = 100  # Use a small subset for speed
+    post_merge_activations = extract_activations(target_model, calibration_data[:arm_sample_size])
+
+    # Slice pre_merge_target_activations to match post_merge sample count
+    # (pre_merge used all 1500 samples, post_merge uses 100 — ARM needs same shape)
+    pre_merge_activations_subset = {}
+    for key in pre_merge_target_activations:
+        act = pre_merge_target_activations[key]
+        pre_merge_activations_subset[key] = act[:arm_sample_size]
 
     # Record this merge's delta + compute ARM/OTMF for next merge
     protection.after_merge(
         target_model, pre_merge_state,
-        pre_merge_activations=pre_merge_target_activations,
+        pre_merge_activations=pre_merge_activations_subset,
         post_merge_activations=post_merge_activations,
     )