Repository: huggingface/transformers
Pull request: Fix Layer device placement in Caches

# What does this PR do?

Fixes https://github.com/huggingface/transformers/issues/39730!
If the Layer were initialized on the correct device in the beginning, this should not be needed. However, as explained in the comment and as was the case before, sometimes `generate` does not initialize them correctly (this should be fixable in the future and would be the preffered way, I will investigate this week)

Changed files:
filename: src/transformers/cache_utils.py
patch: @@ -263,6 +263,14 @@ def update(
         key_states = key_states.to(self.keys.dtype)
         value_states = value_states.to(self.values.dtype)
 
+        # This may be needed if the Layer was not created with the right device in the beginning, i.e. if it did not respect
+        # the device_map. However, even if it is the case, this will only run once, because then the new states received
+        # will always have the same device
+        if self.device != key_states.device:
+            self.device = key_states.device
+            self.keys = self.keys.to(self.device)
+            self.values = self.values.to(self.device)
+
         if cache_position is None:
             # Prefill phase where seq_len potentially equals max_cache_len. Directly copy.
             self.keys.copy_(key_states)
@@ -341,6 +349,14 @@ def update(
         if cache_position is None:
             raise ValueError("`cache_position` must be provided for SlidingWindowLayer.")
 
+        # This may be needed if the Layer was not created with the right device in the beginning, i.e. if it did not respect
+        # the device_map. However, even if it is the case, this will only run once, because then the new states received
+        # will always have the same device
+        if self.device != key_states.device:
+            self.device = key_states.device
+            self.keys = self.keys.to(self.device)
+            self.values = self.values.to(self.device)
+
         key_states = key_states.to(self.keys.dtype)
         value_states = value_states.to(self.values.dtype)
 
@@ -354,7 +370,7 @@ def update(
             return key_states, value_states
 
         # Sliding window logic for generation phase or prefill < window
-        slicing = torch.arange(self.max_cache_len, device=value_states.device)
+        slicing = torch.arange(self.max_cache_len, device=self.device)
         current_seq_len = cache_position[-1] + 1  # Use last position to determine current length
         to_shift = current_seq_len > self.max_cache_len
         indices = (slicing + to_shift.sum()) % self.max_cache_len
@@ -411,6 +427,14 @@ def update(
         if cache_position is None:
             raise ValueError("`cache_position` must be provided for ChunkedSlidingLayer.")
 
+        # This may be needed if the Layer was not created with the right device in the beginning, i.e. if it did not respect
+        # the device_map. However, even if it is the case, this will only run once, because then the new states received
+        # will always have the same device
+        if self.device != key_states.device:
+            self.device = key_states.device
+            self.keys = self.keys.to(self.device)
+            self.values = self.values.to(self.device)
+
         cumulative_length = self.cumulative_length
         self.cumulative_length += key_states.shape[-2]
         is_full = cumulative_length >= self.max_cache_len