# Freezing layers when using gradient checkpointing

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffreezing-layers-when-using-gradient-checkpointing%2F15886%3E%5C
**Category:** 🤗Transformers\
**Created:** [March 20, 2022, 4:57am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffreezing-layers-when-using-gradient-checkpointing%2F15886 "2022-03-20T04:57:55Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![entslscheia](https://avatars.discourse-cdn.com/v4/letter/e/47e85d/32.png) [@entslscheia](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fentslscheia)%5C
**Post date:** [March 20, 2022, 4:57am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffreezing-layers-when-using-gradient-checkpointing%2F15886%2F1 "2022-03-20T04:57:55Z")

</div>

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])` in `BertEncoder`. 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?