Merge pull request #703 from TengJiao33/codex/fix-grace-wise-ultraedit-multigpu

Fix multi-GPU placement for GRACE, WISE, and ULTRAEDIT

Author: HaomingX

easyeditor/models/grace/GRACE.py

@@ -53,6 +53,7 @@ def __init__(self, config, model, device):
         edit_module = parent_module(self.model, brackets_to_periods(self.layer))
         layer_name = self.layer.rsplit(".", 1)[-1]
         original_layer = getattr(edit_module, layer_name)
+        self.device = original_layer.weight.device
         if type(original_layer) is not GRACEAdapter:
             setattr(edit_module, layer_name, GRACEAdapter(config, original_layer, transpose=transpose).to(self.device))
             self.original_layer = copy.deepcopy(original_layer)

easyeditor/models/wise/WISE.py

@@ -39,6 +39,7 @@ def euc(query, key, config, act_mask=None, infer=False):
         return topk.values.mean()
 
     if act_mask is not None:
+        act_mask = act_mask.to(device=l2_norm.device, dtype=l2_norm.dtype)
         return torch.sum(l2_norm * act_mask, dim=1) / torch.sum(act_mask, dim=1)
     else:
         return torch.mean(l2_norm, dim=-1)

easyeditor/trainer/algs/malmen/util.py

@@ -125,14 +125,16 @@ def forward_hook(
             inputs: Tuple[torch.FloatTensor],
             outputs: Tuple[torch.FloatTensor]
         ):
-            self.keys = inputs[0][cache_indices].detach()
+            layer_cache_indices = tuple(index.to(inputs[0].device) for index in cache_indices)
+            self.keys = inputs[0][layer_cache_indices].detach()
 
         def backward_hook(
             module: nn.Module,
             inputs_grad: Tuple[torch.FloatTensor],
             outputs_grad: Tuple[torch.FloatTensor]
         ):
-            self.values_grad = outputs_grad[0][cache_indices].detach()
+            layer_cache_indices = tuple(index.to(outputs_grad[0].device) for index in cache_indices)
+            self.values_grad = outputs_grad[0][layer_cache_indices].detach()
 
         self.handles = [
             module.register_forward_hook(forward_hook),