Gradient checkpointing for GPT-2
Author: cemilcengiz
Pull request: no

# 🚀 Feature request

Gradient checkpointing for GPT-2 would be very useful especially for the larger models in the GPT-2 family. I tried to train the largest release, i.e. "gpt2-xl", on AWS's p3dn.24xlarge instance (has 8 V100 GPUs with 32 GB VRAM on each) with no success. Even making the batch size 1 and using mixed precision training (through torch.cuda.amp) failed with OOM. By the way, I was able to train the second largest model ("gpt2-large") with that configuration i.e. batch size of 1 and enabling mixed precision).


Comments:
Did you run into the issue of imbalanced GPU usage? I have been trying to fine-tune the gpt2-xl model myself on two Titan RTX GPUs (24 GB RAM each) but the imbalanced GPU usage seems to be the main culprit to me. Not sure if gpt2-xl would fit even if I wasn't facing this issue. Why don't they provide some details about the memory requirements for fine-tuning the models. 

> Did you run into the issue of imbalanced GPU usage? I have been trying to fine-tune the gpt2-xl model myself on two Titan RTX GPUs (24 GB RAM each) but the imbalanced GPU usage seems to be the main culprit to me. Not sure if gpt2-xl would fit even if this I wasn't facing this issue. Why don't they provide some details about the memory requirements for fine-tuning the models.

No, my GPU usage was balanced among multiple GPUs.

For anyone interested, I was able to train the "gpt2-xl" model by implementing the gradient checkpointing myself by looking at the suggestions in the previous issues.
Specifically, I added


```
            gradient_checkpointing = kwargs.pop("gradient_checkpointing", None)
            if gradient_checkpointing is not None:
                config_dict["gradient_checkpointing"] = gradient_checkpointing
```

to https://github.com/huggingface/transformers/blob/eb074af75e2fc64c9ec2f5f80637885c455ee15b/src/transformers/configuration_utils.py#L358
. Additionally, I replaced the Block output with 
`return tuple(outputs)`
in https://github.com/huggingface/transformers/blob/eb074af75e2fc64c9ec2f5f80637885c455ee15b/src/transformers/modeling_gpt2.py#L322
. Finally, I replaced the block() call with the 
 ```
            if getattr(self.config, "gradient_checkpointing", False):

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs,
                                      layer_past,
                                      attention_mask,
                                      head_mask[i],
                                      encoder_hidden_states,
                                      encoder_attention_mask,
                                      use_cache,
                                      output_attentions)

                    return custom_forward

                outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states)

            else:
                outputs = block(
                    hidden_states,
                    layer_past=layer_past,
                    attention_mask=attention_mask,
                    head_mask=head_mask[i],
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                )

```
in https://github.com/huggingface/transformers/blob/eb074af75e2fc64c9ec2f5f80637885c455ee15b/src/transformers/modeling_gpt2.py#L611
. To activate the gradient checkpointing, I construct the model by supplying `gradient_checkpointing=True` to the constructor, e.g. 
`model = GPT2LMHeadModel.from_pretrained(model_checkpoint_directory, gradient_checkpointing=True)`. 
I made similar changes to `openai-gpt` just to test if checkpointing works since I can train it with and without gradient checkpointing due to its small size. I verified that both the normal and checkpointed training yielded the same losses in the same iterations. Moreover, I checked the memory utilization of the GPUs and verified that the checkpointed training needs considerably lower memory. Thus, I concluded that the gradient checkpointing is working. 

By the way, I could not train the xlarge model even with the gradient checkpointing if I use AdamW as the optimizer, even when the batch size is 1! Therefore, I switched to the RMSprop since its memory requirement is much smaller. I went ahead and performed a crude memory benchmark for training "gpt2-large" and "gpt2-xl" with sequence length = 1024. You can see the results below.
![gpt2-memory](https://user-images.githubusercontent.com/32267027/93687651-ed9ac300-fac7-11ea-8f31-175157a52966.png)


This is great. Will be implementing this. So with gradient checkpointing, gpt2-xl only requires about 32.47 GB of RAM with RMSProp? 

> This is great. Will be implementing this. So with gradient checkpointing, gpt2-xl only requires about 32.47 GB of RAM with RMSProp?

That is correct for batch size of 3 according to my manual measurements.

When I implement this, I get the error : 
```
File "hyde/src/transformers/modeling_utils.py", line 912, in from_pretrained
TypeError: __init__() got an unexpected keyword argument 'gradient_checkpointing'

```
Did you do anything else besides the steps you mentioned? 
In the configuration_utils.py file, I have added the lines you mentioned after this piece of code:

```
if resolved_config_file is None:
                raise EnvironmentError
            config_dict = cls._dict_from_json_file(resolved_config_file)
			gradient_checkpointing = kwargs.pop("gradient_checkpointing", None)
            if gradient_checkpointing is not None:
                config_dict["gradient_checkpointing"] = gradient_checkpointing      

```

Thank you so much.


Hey, could you check this?

Hi @cemilcengiz, could you tell me which version of transformers is used here?

> Hi @cemilcengiz, could you tell me which version of transformers is used here?

3.1.0

> When I implement this, I get the error :
> 
> ```
> File "hyde/src/transformers/modeling_utils.py", line 912, in from_pretrained
> TypeError: __init__() got an unexpected keyword argument 'gradient_checkpointing'
> ```
> 
> Did you do anything else besides the steps you mentioned?
> In the configuration_utils.py file, I have added the lines you mentioned after this piece of code:
> 
> ```
> if resolved_config_file is None:
>                 raise EnvironmentError
>             config_dict = cls._dict_from_json_file(resolved_config_file)
> 			gradient_checkpointing = kwargs.pop("gradient_checkpointing", None)
>             if gradient_checkpointing is not None:
>                 config_dict["gradient_checkpointing"] = gradient_checkpointing      
> ```
> 
> Thank you so much.

No, that was all I did.

Can someone please make a PR and add this feature?
I think with this added, finetuning would become a lot better with this library

Well, apparently this feature was added to the library 2 weeks after I opened the issue. Therefore, I am closing it.