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?