[Longformer, Bert, Roberta, ...] Fix multi gpu training
Author: patrickvonplaten
Pull request: yes

<!-- This line specifies which issue to close after the pull request is merged. -->
Fixes #6256.

Issue #6256 shows that distributed training is not possible when the model has layers that are not used at all.
Bert, Roberta and Longformer all have a "pooling_layer" added to their BaseModel (`BertModel`, `RobertaModel`, `LongformerModel`). This pooling layer is however not needed by all the "higher" model extensions. The `LongformerForSequenceClassification`, `...ForTokenClassification` model extensions e.g. simply disregards the pooled output of the BaseModel. 
Looking back, I think the pooling layer should not have been part of the BaseModel classes, but is now not really possible to revert the design choice anymore given the number of `BertModel`s already trained with this layer.

Ideally, we should both in TF and PT automatically never instantiate a layer that is never needed. 
This is now implemented in PyTorch. In TF loading pretrained weights depends on the `load_weights` function here: https://github.com/huggingface/transformers/blob/38f17037957d325b5540a8031f065e6f23c9e265/src/transformers/modeling_tf_utils.py#L614
The load_weights function **cannot** load a layer which has a different number of weights in the `resolved_archive_file` as `model`. 
This means that if we would remove the `pooler` layer for `TFBertForMaskedLM` because it is not needed, the following command:
```
python -c 'from transformers import TFBertForMaskedLM; TFBertForMaskedLM.from_pretrained("bert-base-cased")'
```
would throw the following error:

```
ValueError: Layer #0 (named "bert") expects 197 weight(s), but the saved weights have 199 element(s).
```

 At the moment I don't see how we could keep backwards compatibility in TF without completely replacing the `from_pretrained` functionality and even then I don't think it will be easy to load hdf5 files with more layers than expected in TF.


So I propose the following solution:
Only implement the feature for PT and leave the layer in TF for now because a) It solves the issue b) is 100% backwards compatible.
It would be nice to remove the pooling layers for TF in a future PR, but I think we would have to rewrite `from_pretrained` for TF models in this case (cc @jplu )

@sgugger @LysandreJik - what do you think? 



Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/7272?src=pr&el=h1) Report
> Merging [#7272](https://codecov.io/gh/huggingface/transformers/pull/7272?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/2c8ecdf8a87019c438262d8c692e1bdffe05149f?el=desc) will **increase** coverage by `1.10%`.
> The diff coverage is `92.92%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/7272/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/7272?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #7272      +/-   ##
==========================================
+ Coverage   77.58%   78.69%   +1.10%     
==========================================
  Files         181      181              
  Lines       35725    35781      +56     
==========================================
+ Hits        27719    28157     +438     
+ Misses       8006     7624     -382     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/7272?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/configuration\_longformer.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX2xvbmdmb3JtZXIucHk=) | `81.81% <0.00%> (-18.19%)` | :arrow_down: |
| [src/transformers/modeling\_longformer.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19sb25nZm9ybWVyLnB5) | `18.69% <45.45%> (-55.46%)` | :arrow_down: |
| [src/transformers/configuration\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3V0aWxzLnB5) | `97.35% <100.00%> (+0.68%)` | :arrow_up: |
| [src/transformers/modeling\_albert.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19hbGJlcnQucHk=) | `84.33% <100.00%> (+0.17%)` | :arrow_up: |
| [src/transformers/modeling\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iZXJ0LnB5) | `88.48% <100.00%> (+0.11%)` | :arrow_up: |
| [src/transformers/modeling\_mobilebert.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19tb2JpbGViZXJ0LnB5) | `89.43% <100.00%> (+0.04%)` | :arrow_up: |
| [src/transformers/modeling\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yb2JlcnRhLnB5) | `81.51% <100.00%> (-15.05%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_albert.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9hbGJlcnQucHk=) | `97.35% <100.00%> (+0.01%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9iZXJ0LnB5) | `98.91% <100.00%> (+<0.01%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_longformer.py](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9sb25nZm9ybWVyLnB5) | `17.46% <100.00%> (-81.13%)` | :arrow_down: |
| ... and [30 more](https://codecov.io/gh/huggingface/transformers/pull/7272/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/7272?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/7272?src=pr&el=footer). Last update [2c8ecdf...0d1f787](https://codecov.io/gh/huggingface/transformers/pull/7272?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


@patrickvonplaten:

I am just curious about this statement:

    This pooling layer is however not needed by all of the "higher" model extensions. The ...ForSequenceClassification model extensions e.g. simply disregards the pooled output of the BaseModel.

If I remembered correctly, at least for `Bert`, the extension `TFBertForSequenceClassification` does use the pooler output. Since Roberta is similar to Bert, I think it also uses pooler output for sequence classification. Not sure about LongFormer though.

> @patrickvonplaten:
> 
> I am just curious about this statement:
> 
> ```
> This pooling layer is however not needed by all of the "higher" model extensions. The ...ForSequenceClassification model extensions e.g. simply disregards the pooled output of the BaseModel.
> ```
> 
> If I remembered correctly, at least for `Bert`, the extension `TFBertForSequenceClassification` does use the pooler output. Since Roberta is similar to Bert, I think it also uses pooler output for sequence classification. Not sure about LongFormer though.

True, thanks for the remark -> I updated the description above :-) 

I'm not 100% happy with adding a new config parameter the user will have to learn about instead of a solution where it just works out of the box. I would be happier with a solution where there is a parameter at init for the base model to choose whether or not to add that layer, and set it to False for the models that don't need it.

I understand this might break the `from_pretrained` method with existing models, but we can probably add a test that removes any warning that might appear because of that.

What do you think?

> I'm not 100% happy with adding a new config parameter the user will have to learn about instead of a solution where it just works out of the box. I would be happier with a solution where there is a parameter at init for the base model to choose whether or not to add that layer, and set it to False for the models that don't need it.
> 
> I understand this might break the `from_pretrained` method with existing models, but we can probably add a test that removes any warning that might appear because of that.
> 
> What do you think?

Yes I thought about this as well, but was a bit afraid since some people seem to rely on `load_state_dict` in strict mode (issue here: https://github.com/huggingface/transformers/issues/6882). In strict mode there would actually be an error instead of just a warning. 

But I agree that it's cleaner to not expose another config param! 
@LysandreJik and @sgugger - what do you think? I think I would also prefer @sgugger proposition and a slight breaking change here (one that we already had before anyways for all models)

This is just a slight breaking change as we would have a workaround with `strict=False`, and it would only be for user who stubbornly don't want to use our methods, since `from_pretrained` would have no breaking change ;-) Also for those users, it would only be a one-time thing since I imagine they download the model by themselves and store it somewhere. Just one load with `from_pretrained` then one save to the place they want to save the model would solve the potential bug, since the model saved would not have the weights anymore.





Agree! Is this fine for you as well @LysandreJik ?

I think that what @sgugger proposes is cleaner for the end-user. Maybe we should specify somewhere that even if we aim for 0 breaking changes between minor (and minimize them for major) version changes, we can't expect to have no breaking changes for users that do not leverage our methods but rely on methods like `torch.load` and `torch.save` between versions.

I'm talking here only for the TF part.

For TF I'm not really in favor of ignoring a layer mostly because you can easily mess up a lot of things. Let me explain.

The error raised by @patrickvonplaten `ValueError: Layer #0 (named "bert") expects 197 weight(s), but the saved weights have 199 element(s).` beyond saying that you try to load saved weights into the wrong layer, it simply means that the order where the weights are saved and the order where the layers are instantiated in the `__init__` must be exactly the same, otherwise it fails. Including the total number of weights.

When training, if a layer is not used its gradients will be None and will be properly handled by the `.fit()` method anyway. This is also exactly for this reason of missing layer that when I started to develop the TF trainer I handled it with:
```
gradients = [
            g if g is not None else tf.zeros_like(v) for g, v in zip(gradients, self.model.trainable_variables)
        ]
```

So at the end, for TF having this layer is really not an issue. The main issue will be when someone would like to load a PT model in TF with `TFBertModel.from_pretrained("bert-base-cased", from_pt=True)` and the opposite, we have to be sure it works properly.

The quickest solution I can propose for TF is to change the way we load the models, from:
```
model.load_weights(resolved_archive_file, by_name=True)
````
to:
```
model.load_weights(resolved_archive_file, by_name=True, skip_mismatch=True)
````

The `skip_mismatch` parameters is a boolean that skip loading the layers where there is a mismatch in the number of weights, or a mismatch in the shape of the weight. But it has an inconvenient, if it cannot load the weights of a layer it might not be able to load all the weights after the one that has failed and might results to an "empty" model.

BTW @patrickvonplaten I cannot reproduce your error, can you tell me more on how you did it please?

> I'm talking here only for the TF part.
> 
> For TF I'm not really in favor of ignoring a layer mostly because you can easily mess up a lot of things. Let me explain.
> 
> The error raised by @patrickvonplaten `ValueError: Layer #0 (named "bert") expects 197 weight(s), but the saved weights have 199 element(s).` beyond saying that you try to load saved weights into the wrong layer, it simply means that the order where the weights are saved and the order where the layers are instantiated in the `__init__` must be exactly the same, otherwise it fails. Including the total number of weights.
> 
> When training, if a layer is not used its gradients will be None and will be properly handled by the `.fit()` method anyway. This is also exactly for this reason of missing layer that when I started to develop the TF trainer I handled it with:
> 
> ```
> gradients = [
>             g if g is not None else tf.zeros_like(v) for g, v in zip(gradients, self.model.trainable_variables)
>         ]
> ```
> 
> So at the end, for TF having this layer is really not an issue. The main issue will be when someone would like to load a PT model in TF with `TFBertModel.from_pretrained("bert-base-cased", from_pt=True)` and the opposite, we have to be sure it works properly.
> 
> The quickest solution I can propose for TF is to change the way we load the models, from:
> 
> ```
> model.load_weights(resolved_archive_file, by_name=True)
> ```
> 
> to:
> 
> ```
> model.load_weights(resolved_archive_file, by_name=True, skip_mismatch=True)
> ```
> 
> The `skip_mismatch` parameters is a boolean that skip loading the layers where there is a mismatch in the number of weights, or a mismatch in the shape of the weight. But it has an inconvenient, if it cannot load the weights of a layer it might not be able to load all the weights after the one that has failed and might results to an "empty" model.
> 
> BTW @patrickvonplaten I cannot reproduce your error, can you tell me more on how you did it please?

Just to be clear to problem is the following: `TFBertModel` and `BertModel` have a pooling layer that is not needed by classes such as `BertForMaskedLM` and `TFBertForMaskedLM`. What happens in practice is that during the forward pass of `BertForMaskedLM` the `BertModel` pools the hidden states, but `BertForMaskedLM` does not need the pooled output and simply ignores it. In PyTorch this leads to problems with multi-gpu training as shown in the issue linked above. Now, the best and obvious solution here is to remove this layer for all Bert models that don't need it. Removing a layer in a model architecture is dangerous however as all previously serialized weights for this architecture include the removed layer. This can lead to problems with backwards compatibility when loading serialized weight files of previous transformer versions.

Now, in PyTorch this is no problem as such single layers "[bert.pooler]" are simply ignored, but in TF there does not seem to be such an option. Using `skip_mismatch=True` would simply lead to not loading any weights of the whole `TFBertMainLayer` which is obviously not possible. So I do not see a possibility to use `model.load_weights` in TF **and** removing the unnecessary pooling layer for TFBert, ... 

I was wondering if there is a better method or way to load model weights than `model.load_weights`, but could not find one and it seems like to big of a change for such a small issue.

Therefore, I think the best option is to only remove it for PT and leave it for TF which is done in this PR currently because:
- 100% backwards compatibility
- Solves the issue, as TF does not seem to have a problem with multi-gpu inference in its current version.
- Verified than loading into PT from TF and vice-versa works

>Now, in PyTorch this is no problem as such single layers "[bert.pooler]" are simply ignored, but in TF there does not seem to be such an option. Using skip_mismatch=True would simply lead to not loading any weights of the whole TFBertMainLayer which is obviously not possible. So I do not see a possibility to use model.load_weights in TF and removing the unnecessary pooling layer for TFBert, ...
> I was wondering if there is a better method or way to load model weights than model.load_weights, but could not find one and it seems like to big of a change for such a small issue.

Indeed, there is not such way unfortunately, or at least to do it properly in a as nice as in the PT manner.

>Therefore, I think the best option is to only remove it for PT and leave it for TF which is done in this PR currently because:
> 100% backwards compatibility
> Solves the issue, as TF does not seem to have a problem with multi-gpu inference in its current version.
> Verified than loading into PT from TF and vice-versa works

100% agree