text: 
  
  TITLE
  Freezing layers does not work with torch.utils.checkpoint
  
  COMMENTS
  1
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  Hi,
  
  I am using BERT with `pytorch_transformers V1.1.0`. In this version, gradients checkpointing was not implemented, so I do it by myself simply by replacing `layer_outputs = layer_module(hidden_states, attention_mask, head_mask[i])` with `layer_outputs = checkpoint(layer_module, hidden_states, attention_mask, head_mask[i])`. 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:
  ```python
          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 = False
  ```
  The training procedures behave the same no matter what `layers_to_freeze` I set (i.e., the trace of loss is exactly the same for different `layers_to_freeze`), while after I disable gradient checkpointing, it works as expected. I think this suggest that using `torch.utils.checkpoint` might interfere with `param.requires_grad = False` for freezing layers.
  
  I notice that in your official implementation for gradient checkpointing, you [define ` create_custom_forward` first](https://github.com/huggingface/transformers/blob/f466936476ed10410b07999fa7dd65b1dd0a961a/src/transformers/models/bert/modeling_bert.py#L570-L574) instead of directly call `layer_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 calling `layer_module` lead to any serious issue?