Repository: huggingface/transformers
Pull request: Add __iter__ to DynamicCache

Currently, DDP is broken when there is a `DynamicCache` because it has no `__iter___` method and so it cannot be concatenated after the distributed forward. This PR adds back and `__iter__` and adapts the way ddp data is consumed to properly initialize sliding windows.

Changed files:
filename: src/transformers/cache_utils.py
patch: @@ -176,6 +176,11 @@ def __init__(self, sliding_window: int):
         super().__init__()
         self.sliding_window = sliding_window
         self.cumulative_length = 0
+        self._sliding_window_tensor = torch.tensor(self.sliding_window, dtype=torch.long)
+
+    def lazy_initialization(self, key_states: torch.Tensor) -> None:
+        super().lazy_initialization(key_states)
+        self._sliding_window_tensor = self._sliding_window_tensor.to(self.device)
 
     def update(
         self,
@@ -932,7 +937,7 @@ class DynamicCache(Cache):
 
     def __init__(
         self,
-        ddp_cache_data: Optional[Iterable[tuple[torch.Tensor, torch.Tensor]]] = None,
+        ddp_cache_data: Optional[Iterable[tuple[Optional[torch.Tensor], torch.Tensor, torch.Tensor]]] = None,
         config: Optional[PreTrainedConfig] = None,
         offloading: bool = False,
         offload_only_non_sliding: bool = False,
@@ -965,10 +970,15 @@ def __init__(
         # In this case, use the passed data to already fill in the Cache
         if ddp_cache_data is not None:
             # Init all the layers with the data
-            for layer_idx, (key_states, value_states) in enumerate(ddp_cache_data):
-                # If the config was not passed above, initialize a DynamicLayer for each entry of the ddp_data
+            for layer_idx, (sliding_window_tensor, key_states, value_states) in enumerate(ddp_cache_data):
+                # If the config was not passed above, initialize a new cache layer for each entry of the ddp_data
                 if config is None:
-                    layers.append(DynamicLayer())
+                    if sliding_window_tensor is not None:
+                        # Since the same layer is dispatched across replicas, sliding_window is the same for all
+                        sliding_window = sliding_window_tensor[0].item()
+                        layers.append(DynamicSlidingWindowLayer(sliding_window=sliding_window))
+                    else:
+                        layers.append(DynamicLayer())
                 # Update the layer with the data
                 _, _ = layers[layer_idx].update(key_states, value_states)
 
@@ -982,6 +992,10 @@ def __init__(
         else:
             super().__init__(layers=layers, offloading=offloading, offload_only_non_sliding=offload_only_non_sliding)
 
+    def __iter__(self):
+        for layer in self.layers:
+            yield getattr(layer, "_sliding_window_tensor", None), layer.keys, layer.values
+
 
 class StaticCache(Cache):
     """

filename: tests/test_modeling_common.py
patch: @@ -1598,6 +1598,7 @@ def _create_and_check_torch_fx_tracing(self, config, inputs_dict, output_loss=Fa
                         cache_shape = (batch_size, num_heads, cache_length, head_dim)
                         non_empty_pkv = tuple(
                             (
+                                None,
                                 torch.rand(cache_shape, dtype=torch.float, device=torch_device),
                                 torch.rand(cache_shape, dtype=torch.float, device=torch_device),
                             )

filename: tests/utils/test_modeling_utils.py
patch: @@ -1806,7 +1806,10 @@ def test_cache_when_needed_at_train_time(self):
 
         # simulate injecting virtual tokens like in prefix tuning
         num_virtual_tokens = 3
-        past_key_values = [torch.randn(2, 1, 2, num_virtual_tokens, 8)] * 2
+        past_key_values = [
+            (None, torch.randn(1, 2, num_virtual_tokens, 8), torch.randn(1, 2, num_virtual_tokens, 8)),
+            (None, torch.randn(1, 2, num_virtual_tokens, 8), torch.randn(1, 2, num_virtual_tokens, 8)),
+        ]
         past_key_values = DynamicCache(past_key_values)
         model_inputs["attention_mask"] = torch.cat(
             (