add gradient checkpointing for transformer token embedders (#4544)

* add gradient checkpointing for transformer token embedders

* Adds test for gradient checkpointing

Co-authored-by: Dirk Groeneveld <[email protected]>

Author: ArneBinder

CHANGELOG.md

@@ -35,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added the option to specify `requires_grad: false` within an optimizer's parameter groups.
 - Added the `file-friendly-logging` flag back to the `train` command. Also added this flag to the `predict`, `evaluate`, and `find-learning-rate` commands.
 - Added an `EpochCallback` to track current epoch as a model class member. 
+- Added the option to enable or disable gradient checkpointing for transformer token embedders via boolean parameter `gradient_checkpointing`.
 
 ### Removed
 

allennlp/modules/token_embedders/pretrained_transformer_embedder.py

@@ -40,6 +40,8 @@ class PretrainedTransformerEmbedder(TokenEmbedder):
         When `True` (the default), only the final layer of the pretrained transformer is taken
         for the embeddings. But if set to `False`, a scalar mix of all of the layers
         is used.
+    gradient_checkpointing: `bool`, optional (default = `None`)
+        Enable or disable gradient checkpointing.
     """
 
     def __init__(
@@ -51,14 +53,19 @@ def __init__(
         train_parameters: bool = True,
         last_layer_only: bool = True,
         override_weights_file: Optional[str] = None,
-        override_weights_strip_prefix: Optional[str] = None
+        override_weights_strip_prefix: Optional[str] = None,
+        gradient_checkpointing: Optional[bool] = None,
     ) -> None:
         super().__init__()
         from allennlp.common import cached_transformers
 
         self.transformer_model = cached_transformers.get(
             model_name, True, override_weights_file, override_weights_strip_prefix
         )
+
+        if gradient_checkpointing is not None:
+            self.transformer_model.config.update({"gradient_checkpointing": gradient_checkpointing})
+
         self.config = self.transformer_model.config
         if sub_module:
             assert hasattr(self.transformer_model, sub_module)

allennlp/modules/token_embedders/pretrained_transformer_mismatched_embedder.py

@@ -31,6 +31,8 @@ class PretrainedTransformerMismatchedEmbedder(TokenEmbedder):
         When `True` (the default), only the final layer of the pretrained transformer is taken
         for the embeddings. But if set to `False`, a scalar mix of all of the layers
         is used.
+    gradient_checkpointing: `bool`, optional (default = `None`)
+        Enable or disable gradient checkpointing.
     """
 
     def __init__(
@@ -39,6 +41,7 @@ def __init__(
         max_length: int = None,
         train_parameters: bool = True,
         last_layer_only: bool = True,
+        gradient_checkpointing: Optional[bool] = None,
     ) -> None:
         super().__init__()
         # The matched version v.s. mismatched
@@ -47,6 +50,7 @@ def __init__(
             max_length=max_length,
             train_parameters=train_parameters,
             last_layer_only=last_layer_only,
+            gradient_checkpointing=gradient_checkpointing,
         )
 
     @overrides

tests/modules/token_embedders/pretrained_transformer_embedder_test.py

@@ -26,10 +26,22 @@ def test_forward_runs_when_initialized_from_params(self):
         assert tuple(output.size()) == (1, 4, 768)
 
     @pytest.mark.parametrize(
-        "train_parameters, last_layer_only",
-        [(True, True), (False, True), (True, False), (False, False)],
+        "train_parameters, last_layer_only, gradient_checkpointing",
+        [
+            (True, True, False),
+            (False, True, False),
+            (True, False, False),
+            (False, False, False),
+            (
+                True,
+                False,
+                True,
+            ),  # checkpointing only makes sense when we're actually training the layers
+        ],
     )
-    def test_end_to_end(self, train_parameters: bool, last_layer_only: bool):
+    def test_end_to_end(
+        self, train_parameters: bool, last_layer_only: bool, gradient_checkpointing: bool
+    ):
         tokenizer = PretrainedTransformerTokenizer(model_name="bert-base-uncased")
         token_indexer = PretrainedTransformerIndexer(model_name="bert-base-uncased")
 
@@ -53,6 +65,7 @@ def test_end_to_end(self, train_parameters: bool, last_layer_only: bool):
                         "model_name": "bert-base-uncased",
                         "train_parameters": train_parameters,
                         "last_layer_only": last_layer_only,
+                        "gradient_checkpointing": gradient_checkpointing,
                     }
                 }
             }