Repository: huggingface/transformers
Pull request: feat: add support for gradient checkpointing for TimmWrapperModel and TimmWrapperForImageClassification

# What does this PR do?

This enables gradient checkpointing for timm wrapped models.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

@qubvel 


Changed files:
filename: src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
patch: @@ -70,6 +70,10 @@ def __init__(self, *args, **kwargs):
         requires_backends(self, ["vision", "timm"])
         super().__init__(*args, **kwargs)
 
+    def post_init(self):
+        self.supports_gradient_checkpointing = self._timm_model_supports_gradient_checkpointing()
+        super().post_init()
+
     @staticmethod
     def _fix_state_dict_key_on_load(key) -> tuple[str, bool]:
         """
@@ -107,6 +111,24 @@ def _init_weights(self, module):
             if module.bias is not None:
                 module.bias.data.zero_()
 
+    def _timm_model_supports_gradient_checkpointing(self):
+        """
+        Check if the timm model supports gradient checkpointing by checking if the `set_grad_checkpointing` method is available.
+        Some timm models will have the method but will raise an AssertionError when called so in this case we return False.
+        """
+        if not hasattr(self.timm_model, "set_grad_checkpointing"):
+            return False
+
+        try:
+            self.timm_model.set_grad_checkpointing(enable=True)
+            self.timm_model.set_grad_checkpointing(enable=False)
+            return True
+        except Exception:
+            return False
+
+    def _set_gradient_checkpointing(self, enable: bool = True, *args, **kwargs):
+        self.timm_model.set_grad_checkpointing(enable)
+
 
 class TimmWrapperModel(TimmWrapperPreTrainedModel):
     """

filename: tests/models/timm_wrapper/test_modeling_timm_wrapper.py
patch: @@ -170,6 +170,16 @@ def test_mismatched_shapes_have_properly_initialized_weights(self):
     def test_model_is_small(self):
         pass
 
+    def test_gradient_checkpointing(self):
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+        model = TimmWrapperModel._from_config(config)
+        self.assertTrue(model.supports_gradient_checkpointing)
+
+    def test_gradient_checkpointing_on_non_supported_model(self):
+        config = TimmWrapperConfig.from_pretrained("timm/hrnet_w18.ms_aug_in1k")
+        model = TimmWrapperModel._from_config(config)
+        self.assertFalse(model.supports_gradient_checkpointing)
+
     def test_forward_signature(self):
         config, _ = self.model_tester.prepare_config_and_inputs_for_common()