Gradient Checkpointing with Transformers BERT model
Author: sajastu
Pull request: no
Labels: name: wontfix

Hi @sshleifer. I'm doing a summarization task (finetuning bert for extractive summarization). I'd like to use gradient checkpointing to save GPU memory. Here, I'm posting a description of the problem that I'm facing. I'd be grateful if you'd help me get around this.

- `transformers` version: 3.0.2
- Platform: Linux
- Python version: 3.6
- PyTorch version (GPU?): 1.1.0
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Single-GPU
 
Model I am using (Bert, XLNet ...): Bert

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## The problem:

I'm trying to apply gradient checkpointing to the huggingface's Transformers BERT model. I'm skeptical if I'm doing it right, though! Here is my code snippet wrapped around the BERT class:


```
class Bert(nn.Module):
    def __init__(self, large, temp_dir, finetune=False):
        super(Bert, self).__init__()
            self.model = BertModel.from_pretrained('allenai/scibert_scivocab_uncased', cache_dir=temp_dir)

        self.finetune = finetune # either the bert should be finetuned or not... default(True)

    def custom_bert_forward(self, module):
        def custom_forward(*inputs):
            output = module(inputs[0], attention_mask=inputs[1], token_type_ids=inputs[2])
            return output

        return custom_forward
    
    def forward(self, x, segs, mask):
        if (self.finetune):
            
            ## (1) without checkpointing
            top_vec, _ = self.model(x.long(), attention_mask=mask.long(), token_type_ids=segs.long())
            
            ## (2) with checkpointing
            # top_vec = checkpoint.checkpoint(
            #     self.custom_bert_forward(self.model),
            #     x, mask, segs,
            # )

        else:
            self.eval()
            with torch.no_grad():
                top_vec, _ = self.model(x, attention_mask=mask, token_type_ids=segs)
        return top_vec
```
As I'm checkpointing the BERT's forward function, the memory usage drops significantly (~1/5), but  I'm getting relatively inferior performance compared to non-checkpointing, in terms of the metrics (for my task, which is summarization) that I'm calculating on the validation set. Another observation: while checkpointing, the model's training speed also increases considerably which is totally odd to what I have learned from gradient checkpointing. Is there any problem with the implementation? or have I done any part wrong?


-----------------------

P.S. *Update.* I found out that this has been recently implemented by Hugginface guys! There's an argument in `BertConfig` constructor which takes in a boolean `gradient_checkpointing` (default set to False). Then, seemingly, the `bertConfig` object must be sent into `BertModel` class. I'm doing this procedure, but yet, I'm having the aforementioned problem. Following is exactly what I'm doing –updated code:

```
class Bert(nn.Module):
    def __init__(self, large, temp_dir, finetune=False):
        super(Bert, self).__init__()
        if (large):
            self.model = BertModel.from_pretrained('bert-large-uncased', cache_dir=temp_dir)
        else:
            # added line below for BertConfig
            config = BertConfig.from_pretrained('allenai/scibert_scivocab_uncased')
            config.gradient_checkpointing = True
            
            self.model = BertModel.from_pretrained('allenai/scibert_scivocab_uncased', cache_dir=temp_dir, config=config)

        self.finetune = finetune

    def forward(self, x, segs, mask):
        if (self.finetune):
            top_vec, _ = self.model(x.long(), attention_mask=mask.long(), token_type_ids=segs.long())

        else:
            self.eval()
            with torch.no_grad():
                top_vec, _ = self.model(x, attention_mask=mask, token_type_ids=segs)
        return top_vec
```

 



Comments:
> Another observation: while checkpointing, the model's training speed also increases considerably which is totally odd to what I have learned from gradient checkpointing. Is there any problem with the implementation? or have I done any part wrong?

from torch docs
> Checkpointing works by trading compute for memory. Rather than storing all intermediate activations of the entire computation graph for computing backward, the checkpointed part does not save intermediate activations, and instead recomputes them in backward pass. It can be applied on any part of a model.

As the mode; needs to recompute the activations during backward pass it adds up training time considerably.

@patil-suraj Thx for your answer. I totally agree with you and, of course, the doc :). The problem is that while monitoring the training process, I see that the model's training **speed** increases considerably, resulting in faster training (decreasing training time) which is unconventional to the concept of Gradient Checkpointing. That's what made me skeptical about the code snippet (above) that I wrote to use gradient checkpointing. In addition, when comparing checkpointed with non-checkpointed model, I see quite different performance. I'm wondering if there is/are some other considerations that I need to take to make it work...

What I expect with the checkpointed model: I expect that it adds up training time (i.e., lower training speed) due to the recomputation of activation functions during the backward pass, and achieve the same performance/scores as to the non-checkpointed model. 

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.