Gradient Checkpointing with Transformers BERT model

username: sajastu
cooked: <p>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:</p>
<pre><code class="lang-auto">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
</code></pre>
<p>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>

username: ptrblck
cooked: <p>Based on the code snippet it seems you are calling the complete model via <code>checkpoint</code>, which should not save any memory, or is this module just part of a larger model?<br>
However, based on the issue description is seems that parts of the model might not be executed at all.</p>

username: sajastu
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thanks for your response!</p>
<p>Yes, this model is just part of a larger network, i.e., <code>top_vec</code> which is the output of this model is being used by another model. I see <code>top_vec</code> as a vector that has the encoded version of vector <code>x</code> (i…e, src) by the BERT. I a sense, the weights associated with this class should be updated (i.e., learned) during training. Upon my investigations, I noticed that this part of the model consumes much of a memory, so that I thought it’d be better to checkpoint it. Isn’t that true argument, wherever we have learning (i.e., the update of model parameters), we can use gradient checkpointing??</p>
<hr>
<p>I also noticed that there’s a recently implemented option in Huggingface’s BERT which allows us to apply gradient checkpointing easily. That’s an argument that is specified in <code>BertConfig</code> and then the object is passed to   <code>BertModel.from_pretrained</code>. I also tried that, but have the same above issues that I mentioned: 1) the performance does not yield to that of setting without gradient-checkpointing. 2) the training is much faster. While I expect it to be slower, in expense of having memory usage reduced, as Gradient Checkpointing makes the model recompute activations wherever necessary instead of retrieving them, which totally results in more computations, and more training time.</p>

username: ptrblck
cooked: <p>I would have the same assumptions and also think that something might not work as expected.<br>
Did you just recently encountered this issue and was it working before?<br>
Also, which PyTorch version are you using?</p>

username: sajastu
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I’m having this issue since the first time I implemented it, and was not working before as my assumptions. The code that you see runs on torch 1.1.0 (BERTSUM main implementation), but I also tested it on BART Huggingface which uses Pytorch &gt; 1.4.0.</p>
<p>Does pytorch version affect the checkpointing? I thought this before, and searched if gradient checkpointing has been added since a certain version of pytorch, but couldn’t find anything useful.</p>
<p>Another observation: what if the returned value has Gradient None? I’ve had this issue: <a href="https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117" class="inline-onebox">Checkpoint with no grad requiring inputs PROBLEM</a> as well, which is solved when the <code>required_gradient</code> is set to True for the input arguments. I should mention that in HuggingFace’s <code>BertModel</code>, there is <code>src</code> argument that does not need gradients at all, but to fade the warning, I have to set its gradients to True. I this also an issue?</p>

username: ptrblck
cooked: <p>The second issue might be disabling the gradient calculation.<br>
After you’ve set <code>requires_grad=True</code> in the input, are you still seeing a speedup and bad model performance?</p>