Update merge.py

merge.py

@@ -23,7 +23,7 @@ def merge_tensors(tensor1: torch.Tensor, tensor2: torch.Tensor, p: float) -> tor
         torch.Tensor: The merged tensor.
     """
     delta = tensor2 - tensor1
-    m = torch.from_numpy(np.random.binomial(1, p, delta.shape)).to(tensor1.dtype)
+    m = torch.from_numpy(np.random.binomial(1, p, delta.shape)).to(tensor1.device)
     delta_tilde = m * delta
     delta_hat = delta_tilde / (1 - p)
     return delta_hat
@@ -42,16 +42,17 @@ def merge_safetensors(file_path1: str, file_path2: str, p: float, lambda_val: fl
         dict: A dictionary of merged tensors.
     """
     merged_tensors = {}
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     with safe_open(file_path1, framework="pt", device="cpu") as f1, safe_open(file_path2, framework="pt", device="cpu") as f2:
         keys1 = set(f1.keys())
         keys2 = set(f2.keys())
         common_keys = keys1.intersection(keys2)
 
         for key in common_keys:
-            tensor1 = f1.get_tensor(key)
-            tensor2 = f2.get_tensor(key)
+            tensor1 = f1.get_tensor(key).to(device)
+            tensor2 = f2.get_tensor(key).to(device)
             tensor1, tensor2 = resize_tensors(tensor1, tensor2)
-            merged_tensors[key] = tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)
+            merged_tensors[key] = (tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)).cpu()
             logging.info(f"Merging {key}")
 
     return merged_tensors
@@ -131,6 +132,7 @@ def merge_folder(tensor_map: dict, directory_path: str, p: float, lambda_val: fl
     """
     keys1 = set(tensor_map.keys())
     ext = None
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     for filename in os.listdir(directory_path):
         if filename.endswith(".safetensors"):
             ext = ".safetensors"
@@ -146,14 +148,14 @@ def merge_folder(tensor_map: dict, directory_path: str, p: float, lambda_val: fl
             common_keys = keys1.intersection(keys2)
             for key in common_keys:
                 if "block_sparse_moe.gate" in key:
-                    tensor1 = tensor_map[key]['tensor']
-                    tensor2 = f.get_tensor(key)
+                    tensor1 = tensor_map[key]['tensor'].to(device)
+                    tensor2 = f.get_tensor(key).to(device)
                     tensor_map[key]['tensor'] = (tensor1 + tensor2) / 2.0
                     continue
-                tensor1 = tensor_map[key]['tensor']
-                tensor2 = f.get_tensor(key)
+                tensor1 = tensor_map[key]['tensor'].to(device)
+                tensor2 = f.get_tensor(key).to(device)
                 tensor1, tensor2 = resize_tensors(tensor1, tensor2)
-                tensor_map[key]['tensor'] = tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)
+                tensor_map[key]['tensor'] = (tensor1 + lambda_val * merge_tensors(tensor1, tensor2, p)).cpu()
     return tensor_map
 
 def map_tensors_to_files(directory_path: str) -> dict:
@@ -238,4 +240,4 @@ def main():
         save_file(merged, args.output_model)
 
 if __name__ == '__main__':
-    main()
+    main()