Freezing layers when using gradient checkpointing
Hi,I am using BERT withpytorch_transformers V1.1.0. In this version, gradients checkpointing was not implemented, so I do it by myself simply by replacinglayer_outputs = layer_module(hidden_states, attention_mask, head_mask[i])withlayer_outputs = checkpoint(layer_module, hidden_states, attention_mask, head_mask[i])inBertEncoder. This works fine when I try to fine-tune the entire BERT model. However, when I tried to freeze several layers of BERT with the code below:if layers_to_freeze is not None:
            modules = [self.transformer_model.embeddings,
                       *self.transformer_model.encoder.layer[:layers_to_freeze]]
            for module in modules:
                for param in module.parameters():
                    param.requires_grad = FalseThe training procedures behave the same no matter whatlayers_to_freezeI set (i.e., the trace of loss is exactly the same for differentlayers_to_freeze), while after I disable gradient checkpointing, it works as expected. I think this suggest that usingtorch.utils.checkpointmight interfere withparam.requires_grad = Falsefor freezing layers.I notice that in your official implementation for gradient checkpointing, youdefinecreate_custom_forwardfirstinstead of directly calllayer_module. I haven’t tested whether doing this would avoid the issue I mentioned so far. Before I test it, I am curious about out of what concern you are implementing it in this way? Will directly callinglayer_modulelead to any serious issue?