Repository: transformers
Issue: Supporting Selective Activation Checkpointing and CPU Offloading Option.
Labels: Feature request

### Feature request

- references
  - [deepspeed docs](https://deepspeed.readthedocs.io/en/latest/activation-checkpointing.html)

### Motivation

Hi, first of all, thank you for your great contribution for open-source community, HF team !
I really enjoy training with HF transformers but found that activation checkpointing (a.k.a gradient checkpointing) is applied for every layers and checkpoints activation in GPU memory where it can easily leads to GPU OOM.
However, there are well-known techniques to checkpoint activations selectively and even in CPU.
Implementations of these features can be found in Megatron and DeepSpeed, and torch team also release [the blog](https://pytorch.org/blog/maximizing-training/#selective-activation-checkpointing) about this technique today.  

So i tested selective activation checkpointing with CPU offloading by modifying HF model (in my case, gpt-2 xl).
i just changed two things, `model._gradient_checkpointing_func` and forward pass like below.

```python
from transformers import AutoConfig, AutoModelForCausalLM
from deepspeed.runtime.activation_checkpointing import checkpointing

# init model and activate grad ckpt
model = AutoModelForCausalLM(**kwargs)
model.gradient_checkpointing_enable()

# change grad ckpt function
checkpointing. configure(
    mpu_=None,
    deepspeed_config=None,
    partition_activations=False,
    contiguous_checkpointing=False,
    num_checkpoints=num_checkpoints, # only use these for now
    checkpoint_in_cpu=True, # only use these for now
    synchronize=False,
    profile=True,
)
model._gradient_checkpointing_func = checkpointing.checkpoint # use deepspeed activation checkpointing function
```

```python
def forward(**inputs):
    ...
    def custom(start, end):
        def custom_forward(*inputs):
            layers_ = self.h[start:end]
            x_ = inputs[0] # hidden_states
            head_mask_ = inputs[3] # hidden_states
            for i, layer in enumerate(layers_):
                x_ = layer(
                    x_, 
                    None,
                    inputs[2],
                    head_mask_[i],
                    inputs[4],
                    inputs[5],
                    False,
                    None
                )[0] # tuple to tensor
                # Tra()
            return x_
        return custom_forward

    if self.gradient_checkpointing and self.training:
        # print('use_deepspeed_activation_checkpointing is True !')
        # Tra()
        l = 0
        total_num_layers = len(self.h)
        chunk_length = num_layers ## how many checkpoint do you want?
        # print(f'chunk_length : {chunk_length}')

        while l < total_num_layers:
            if self.model_parallel:
                raise NotImplementedError("Idc model parallel :)")
            if output_hidden_states:
                raise NotImplementedError("Idc layerwise output :)")

            hidden_states = self._gradient_checkpointing_func(
                custom(l, l+chunk_length),
                hidden_states,
                None,
                attention_mask,
                head_mask[l:l+chunk_length],
                encoder_hidden_states,
                encoder_attention_mask,
                False,
                None,
            )

            if output_attentions:
                raise NotImplementedError("We use efficient SDPA and it does not support to output attention score map")
            if self.model_parallel:
                raise NotImplementedError("Idc model parallel :)")

            l += chunk_length
```

and i could see significant improvement in GPU memory (trading off throughput).

<img width="1238" alt="cpu_gpu_elapsed_from_gpt4" src="https://github.com/huggingface/transformers/assets/48202736/9f9d102c-4c35-4216-b8b7-440fdc127ff8">

The quantitative results (CPU/GPU memory and elapsed time for 3 times forwarding) are recorded by torch profiler and figure was generated by GPT-4. 
My setting was like this.

- 2x A100 80GB
- gpt2-xl (near 1.5B)
- fixed batch size with same random seed (`[128, 512] # B, T`)
- zero-3 without CPU offloading (using accelerate)
- flash attn 2 is applied (monkey patched with xformers)

Results show that it could save about 10GB for each GPU device.
The reason why reserved GPU memory for checkpointing every 2, 3, 6 layers are same is that peak memory is dominated by computing loss i guess.
here is my detailed result for checkpointing every 3 layers with CPU offloading,

<img width="1047" alt="every3_cpu" src="https://github.com/huggingface/transformers/assets/48202736/4a9ee4e7-7b7e-4e51-96e7-78868943dda4">

<img width="1047" alt="every3_gpu" src="https://github.com/huggingface/transformers/assets/48202736/df40e4b6-f1fb-4cff-aebb-f85aa0e39c64">

and HF default version (using torch.utils.checkpoint for every layers)

<img width="1057" alt="default_cpu" src="https://github.com/huggingface/transformers/assets/48202736/8455b32e-2410-4faa-b093-b918c1ebe4d0">

<img width="1063" alt="default_gpu" src="https://github.com/huggingface/transformers/assets/48202736/61ab67b6-2d54-4388-8fe1-5f7ebdc06967">

i didn't add profiling result for larger model but there were more improvement and i can find sweet spot for GPU memory VS throughput trade off.  


### Your contribution

Didnt PR yet because in HF transformers, every model has its own forward method and applying it for all would be a huge update.
i just add profiling results and naive code snippet for implementation.
just want to discuss with you guys !