Keep overrides of the deprecated gradient checkpointing methods working

@@ -3129,6 +3129,19 @@ def activation_checkpointing_enable(             activation_checkpointing_kwargs (dict, *optional*):                 Additional keyword arguments passed along to the `torch.utils.checkpoint.checkpoint` function.         """+        # Subclasses and wrappers of the deprecated `gradient_checkpointing_enable` stay the entry point until v6+        if self._is_overridden("gradient_checkpointing_enable"):+            self.gradient_checkpointing_enable(+                gradient_checkpointing_kwargs=activation_checkpointing_kwargs,+                every_n_layers=every_n_layers,+                offload=offload,+            )+        else:+            self._activation_checkpointing_enable(activation_checkpointing_kwargs, every_n_layers, offload)++    def _activation_checkpointing_enable(+        self, activation_checkpointing_kwargs=None, every_n_layers: int = 1, offload: bool = False+    ):         if not self.supports_gradient_checkpointing:             raise ValueError(f"{self.__class__.__name__} does not support activation checkpointing.") @@ -3184,7 +3197,17 @@ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kw             "`activation_checkpointing_enable` instead.",             FutureWarning,         )-        self.activation_checkpointing_enable(activation_checkpointing_kwargs=gradient_checkpointing_kwargs, **kwargs)+        # Reached from the override itself, which `activation_checkpointing_enable` would call again+        if self._is_overridden("gradient_checkpointing_enable"):+            self._activation_checkpointing_enable(gradient_checkpointing_kwargs, **kwargs)+        else:+            self.activation_checkpointing_enable(+                activation_checkpointing_kwargs=gradient_checkpointing_kwargs, **kwargs+            )++    def _is_overridden(self, name: str) -> bool:+        # A subclass override is another function and a wrapper set on the instance is not a bound method+        return getattr(getattr(self, name), "__func__", None) is not getattr(PreTrainedModel, name)      def _set_gradient_checkpointing(         self,@@ -3227,6 +3250,13 @@ def activation_checkpointing_disable(self):         """         Deactivates activation checkpointing for the current model.         """+        # Subclasses and wrappers of the deprecated `gradient_checkpointing_disable` stay the entry point until v6+        if self._is_overridden("gradient_checkpointing_disable"):+            self.gradient_checkpointing_disable()+        else:+            self._activation_checkpointing_disable()++    def _activation_checkpointing_disable(self):         if self.supports_gradient_checkpointing:             # For old GC format (transformers < 4.35.0) for models that live on the Hub             # we will fall back to the overwritten `_set_gradient_checkpointing` method@@ -3252,7 +3282,11 @@ def gradient_checkpointing_disable(self):             "`activation_checkpointing_disable` instead.",             FutureWarning,         )-        self.activation_checkpointing_disable()+        # Reached from the override itself, which `activation_checkpointing_disable` would call again+        if self._is_overridden("gradient_checkpointing_disable"):+            self._activation_checkpointing_disable()+        else:+            self.activation_checkpointing_disable()      @property     def is_activation_checkpointing(self) -> bool:

@@ -4265,6 +4265,47 @@ def test_aliases_warn_and_forward(self):             model.gradient_checkpointing_disable()         self.assertFalse(model.is_activation_checkpointing) +    def test_subclass_overrides_of_the_aliases_are_called(self):+        from transformers import LlamaModel++        class LegacyLlamaModel(LlamaModel):+            def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):+                self.calls.append("enable")+                super().gradient_checkpointing_enable(gradient_checkpointing_kwargs, **kwargs)++            def gradient_checkpointing_disable(self):+                self.calls.append("disable")+                super().gradient_checkpointing_disable()++        config = LlamaConfig(num_hidden_layers=2, hidden_size=32, intermediate_size=64, num_attention_heads=4)+        model = LegacyLlamaModel(config)+        model.calls = []+        with self.assertWarns(FutureWarning):+            model.activation_checkpointing_enable()+        self.assertTrue(model.is_activation_checkpointing)+        with self.assertWarns(FutureWarning):+            model.activation_checkpointing_disable()+        self.assertFalse(model.is_activation_checkpointing)+        self.assertEqual(model.calls, ["enable", "disable"])++    def test_wrappers_of_the_aliases_are_called(self):+        from transformers import LlamaModel++        config = LlamaConfig(num_hidden_layers=2, hidden_size=32, intermediate_size=64, num_attention_heads=4)+        model = LlamaModel(config)+        original_enable = model.gradient_checkpointing_enable+        received_kwargs = []++        def gradient_checkpointing_enable(gradient_checkpointing_kwargs=None, **kwargs):+            received_kwargs.append(gradient_checkpointing_kwargs)+            original_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs, **kwargs)++        model.gradient_checkpointing_enable = gradient_checkpointing_enable+        with self.assertWarns(FutureWarning):+            model.activation_checkpointing_enable({"use_reentrant": False}, every_n_layers=2)+        self.assertEqual(received_kwargs, [{"use_reentrant": False}])+        self.assertEqual([layer.gradient_checkpointing for layer in model.layers], [True, False])+  @require_torch_accelerator class GradientCheckpointingOffloadTest(unittest.TestCase):