Repository: huggingface/transformers (Python)
Task type: Feature

Issue:
Support gradient checkpointing for ESM models
Would you please add `gradient_checkpointing_enable()` feature for ESM models?
These models currently are the best available pre-trained protein language models for researchers.
Many thanks.


Discussion:
cc @Rocketknight1 
Any updates?
It's on the to-do list, but I'm afraid there are competing priorities at the moment!
Let's open it up for anyone in the community who might want to tackle it :) 
Hi @amyeroberts @Rocketknight1 I would like to work on this 
@sanjeevk-os Great! Once you have the code ready, open a PR and ping both @Rocketknight1 and me. Looking forward to reviewing! 
Hi @sanjeevk-os, I actually took a look at the ESM code - it actually looks like some of the supports for gradient checkpointing are already there, in which case you just need to make a one-line change to set `supports_gradient_checkpointing = True`
Hi @Rocketknight1 Thank you for taking a look. I also noticed that the ESM model has the _create_custom_forward_ passed to torch checkpoint function. I will do some more checks and will raise a PR soon.
Hi @sanjeevk-os - we're getting even more requests for this, so we'd like to try to add it soon! If you're having trouble, just let us know. We can take over the PR internally to try to get it through, and we appreciate your effort regardless.

Solution patch:
diff --git a/src/transformers/models/esm/modeling_esm.py b/src/transformers/models/esm/modeling_esm.py
--- a/src/transformers/models/esm/modeling_esm.py
+++ b/src/transformers/models/esm/modeling_esm.py
@@ -690,6 +690,7 @@ class EsmPreTrainedModel(PreTrainedModel):
 
     config_class = EsmConfig
     base_model_prefix = "esm"
+    supports_gradient_checkpointing = True
     _no_split_modules = ["EsmLayer", "EsmFoldTriangularSelfAttentionBlock", "EsmEmbeddings"]
 
     # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
@@ -709,6 +710,10 @@ def _init_weights(self, module):
             module.bias.data.zero_()
             module.weight.data.fill_(1.0)
 
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, EsmEncoder):
+            module.gradient_checkpointing = value
+
 
 ESM_START_DOCSTRING = r"""
 
@@ -785,8 +790,6 @@ class EsmModel(EsmPreTrainedModel):
     `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
     """
 
-    supports_gradient_checkpointing = False
-
     def __init__(self, config, add_pooling_layer=True):
         super().__init__(config)
         self.config = config
@@ -803,10 +806,6 @@ def __init__(self, config, add_pooling_layer=True):
         # Initialize weights and apply final processing
         self.post_init()
 
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, EsmEncoder):
-            module.gradient_checkpointing = value
-
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
 


Test patch:
diff --git a/tests/models/esm/test_modeling_esm.py b/tests/models/esm/test_modeling_esm.py
--- a/tests/models/esm/test_modeling_esm.py
+++ b/tests/models/esm/test_modeling_esm.py
@@ -151,6 +151,24 @@ def create_and_check_for_token_classification(
         result = model(input_ids, attention_mask=input_mask, labels=token_labels)
         self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))
 
+    def create_and_check_forward_and_backwards(
+        self,
+        config,
+        input_ids,
+        input_mask,
+        sequence_labels,
+        token_labels,
+        choice_labels,
+        gradient_checkpointing=False,
+    ):
+        model = EsmForMaskedLM(config)
+        if gradient_checkpointing:
+            model.gradient_checkpointing_enable()
+        model.to(torch_device)
+        result = model(input_ids, attention_mask=input_mask, labels=token_labels)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
+        result.loss.backward()
+
     def prepare_config_and_inputs_for_common(self):
         config_and_inputs = self.prepare_config_and_inputs()
         (
@@ -219,6 +237,10 @@ def test_for_token_classification(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*config_and_inputs)
 
+    def test_esm_gradient_checkpointing(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_forward_and_backwards(*config_and_inputs, gradient_checkpointing=True)
+
     @slow
     def test_model_from_pretrained(self):
         for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]: