Repository: transformers
Issue: Add gradient_checkpointing_segment_size
Labels: Feature request

### Feature request

Currently when we enable gradient checkpointing, e.g. in `LlamaModel`, we call `torch.utils.checkpoint.checkpoint` on every `LlamaDecoderLayer`. As per [Training Deep Nets with Sublinear Memory Cost](https://arxiv.org/pdf/1604.06174.pdf), https://github.com/cybertronai/gradient-checkpointing and hinted by [torch.utils.checkpoint.checkpoint_sequential](https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint_sequential), we can strike a balance between memory and compute by supporting `gradient_checkpointing_segment_size`.

### Motivation

Gradient checkpointing makes training slow and people enable it mainly to avoid VRAM OOM. It would be great if we can leverage `gradient_checkpointing_segment_size` and make gradient checkpointing configurable.

### Your contribution

Here is a diff to start off with:

```diff
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index bbf229b58..685400ffd 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -1745,6 +1745,16 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
         """
         return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
 
+    def set_gradient_checkpointing_segment_size(self, value: int):
+        """
+        Set gradient checkpointing segment size for the current model.
+        """
+        if not self.supports_gradient_checkpointing:
+            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
+        if not hasattr(self.model, "_set_gradient_checkpointing_segment_size"):
+            raise ValueError(f"{self.__class__.__name__} does not support configuring gradient checkpointing segment size.")
+        self.apply(partial(self._set_gradient_checkpointing_segment_size, value=value))
+
     def save_pretrained(
         self,
         save_directory: Union[str, os.PathLike],
diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py
index 5e7a879c0..682e49b21 100644
--- a/src/transformers/models/llama/modeling_llama.py
+++ b/src/transformers/models/llama/modeling_llama.py
@@ -491,6 +491,10 @@ class LlamaPreTrainedModel(PreTrainedModel):
         if isinstance(module, LlamaModel):
             module.gradient_checkpointing = value
 
+    def _set_gradient_checkpointing_segment_size(self, module, value=1):
+        if isinstance(module, LlamaModel):
+            module.gradient_checkpointing_segment_size = value
+
 
 LLAMA_INPUTS_DOCSTRING = r"""
     Args:
@@ -578,6 +582,7 @@ class LlamaModel(LlamaPreTrainedModel):
         self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
         self.gradient_checkpointing = False
+        self.gradient_checkpointing_segment_size = 1
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -689,7 +694,7 @@ class LlamaModel(LlamaPreTrainedModel):
 
             past_key_value = past_key_values[idx] if past_key_values is not None else None
 
-            if self.gradient_checkpointing and self.training:
+            if self.gradient_checkpointing and self.training and (idx % self.gradient_checkpointing_segment_size == 0):
 
                 def create_custom_forward(module):
                     def custom_forward(*inputs):
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 2b8cb013f..f7d5e87ea 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1687,6 +1687,9 @@ class Trainer:
         if args.gradient_checkpointing:
             self.model.gradient_checkpointing_enable()
 
+            if args.gradient_checkpointing_segment_size > 1:
+               self.model.set_gradient_checkpointing_segment_size(args.gradient_checkpointing_segment_size)
+
         model = self._wrap_model(self.model_wrapped)
 
         if (is_sagemaker_mp_enabled() or self.is_fsdp_enabled) and resume_from_checkpoint is not None:
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
index 2f72622d9..ecffc5660 100644
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -586,6 +586,8 @@ class TrainingArguments:
             Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not finished.
         gradient_checkpointing (`bool`, *optional*, defaults to `False`):
             If True, use gradient checkpointing to save memory at the expense of slower backward pass.
+        gradient_checkpointing_segment_size (`int`, *optional*, defaults to `1`):
+            If gradient_checkpointing is True, use gradient checkpointing for every segment size
         include_inputs_for_metrics (`bool`, *optional*, defaults to `False`):
             Whether or not the inputs will be passed to the `compute_metrics` function. This is intended for metrics
             that need inputs, predictions and references for scoring calculation in Metric class.
@@ -1144,6 +1146,12 @@ class TrainingArguments:
             "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass."
         },
     )
+    gradient_checkpointing_segment_size: int = field(
+        default=1,
+        metadata={
+            "help": "If gradient_checkpointing is True, use gradient checkpointing for every segment size."
+        },
+    )
     include_inputs_for_metrics: bool = field(
         default=False, metadata={"help": "Whether or not the inputs will be passed to the `compute_metrics` function."}
     )
@@ -2137,6 +2145,7 @@ class TrainingArguments:
         gradient_accumulation_steps: int = 1,
         seed: int = 42,
         gradient_checkpointing: bool = False,
+        gradient_checkpointing_segment_size: int = 1,
     ):
         """
         A method that regroups all basic arguments linked to the training.
@@ -2179,6 +2188,8 @@ class TrainingArguments:
                 parameters.
             gradient_checkpointing (`bool`, *optional*, defaults to `False`):
                 If True, use gradient checkpointing to save memory at the expense of slower backward pass.
+            gradient_checkpointing_segment_size (`int`, *optional*, defaults to `1`):
+                If gradient_checkpointing is True, use gradient checkpointing for every segment size
 
         Example:
 
```