Farm: Error in conversion of BertForMaskedLM with BertLMHead to HuggingFace transformers

[himanshurawlani]

Describe the bug
I started with the lm_finetuning.py script to pre-train BERT on a domain-specific text corpus. After completing training, the saved artifacts included the following:

• language_model.bin
• language_model_config.json (BertForMaskedLM)
• prediction_head_0.bin
• prediction_head_0_config.json (BertLMHead)
• prediction_head_1.bin
• prediction_head_1_config.json (NextSentenceHead)
• processor_config.json
• special_tokens_map.json
• tokenizer_config.json
• vocab.txt

I know FARM currently doesn't support converting an AdaptiveModel with two prediction heads to a HuggingFace transformers model. However, after going through the AdaptiveModel.load() classmethod used in the model conversion script, I figured out that removing the prediction_head_1_config.json (NextSentenceHead) and prediction_head_1.bin (PyTorch model) files should result in successful conversion to HuggingFace model as the conversion script would consider it as an AdaptiveModel with single PredictionHead.

Error message
After removing the files and running the conversion script, I get an error saying:

Traceback (most recent call last):
  File "conversion_huggingface_models.py", line 88, in <module>
    convert_to_transformers("./farm_saved_models/bert-english-lm", 
  File "conversion_huggingface_models.py", line 46, in convert_to_transformers
    transformer_model = model.convert_to_transformers()
  File "/home/himanshu/.conda/envs/tf2/lib/python3.8/site-packages/farm/modeling/adaptive_model.py", line 509, in convert_to_transformers
    elif len(self.prediction_heads[0].layer_dims) != 2:
  File "/home/himanshu/.conda/envs/tf2/lib/python3.8/site-packages/torch/nn/modules/module.py", line 771, in __getattr__
    raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
torch.nn.modules.module.ModuleAttributeError: 'BertLMHead' object has no attribute 'layer_dims'

I made some changes to BertLMHead class by referring to HuggingFace's BertLMPredictionHead and added the following lines to __init__() method:

self.layer_dims = [hidden_size, vocab_size]
self.decoder.bias = self.bias

I also updated the __foward__() function by changing the line:

lm_logits = self.decoder(hidden_states) + self.bias

to

lm_logits = self.decoder(hidden_states)

Expected behavior
After making the above changes and re-training the model, the conversion script ran successfully and I was able to import the model into HuggingFace pipelines with the following code:

from transformers import pipeline

fill_mask = pipeline(
    "fill-mask",
    model="./bert-english-lm",
    tokenizer="./bert-english-lm"
)

Correct me if I'm wrong, I expected the conversion script to run successfully after removing the NextSentenceHead files. I'm not sure if this a bug or a feature request as I'm new to FARM. However, I was able to solve the issue with a few changes as described above.

System:

• OS: Ubuntu 18.04.4 LTS (Bionic Beaver)
• CPU: Intel(R) Xeon(R) Platinum 8168 CPU @ 2.70GHz
• FARM version: 0.4.7

[Timoeller]

Hey @himanshurawlani thanks for the detailed report about what you did.

Your method of removing the PH files, adjusting the BertLMHead and conversion seem reasonable but I did not test them myself.

Did you test the outcome and does the transformers model predict meaningful words when you fill the masks?

[tholor]

Related to #517

[himanshurawlani]

Hey @himanshurawlani thanks for the detailed report about what you did.

Your method of removing the PH files, adjusting the BertLMHead and conversion seem reasonable but I did not test them myself.

Did you test the outcome and does the transformers model predict meaningful words when you fill the masks?

Yes, I've tested it and it works perfectly.

[Timoeller]

Hey, I just realized something. Do we need the change in the predictionhead forward function (removal of bias addition) for the conversion?
I think we just need the additional parameters layer dims and decoder.bias for the conversion to work, because the actual forward code is happening in transformers pipeline.
Is that correct?

If so, we could integrate this functionality easy within FARM. How about creating a PR with

1. these additional params in FARMs BertLMHead (maybe make a comment that they are only used for conversion to transformers).
2. adding a function in https://github.com/deepset-ai/FARM/blob/master/examples/conversion_huggingface_models.py that describes what you did to make the conversion work (delete NSP head weights and config)

[himanshurawlani]

I think we need to change the prediction head forward function (by removing bias addition). In the code, we initialize the linear layer with bias=False and then assign the bias separately. If we do not remove the bias addition, in forward function, then it will be added twice during FARM's training.

# The output weights are the same as the input embeddings, but there is
# an output-only bias for each token.
self.decoder = nn.Linear(hidden_size,
                                 vocab_size,
                                 bias=False)
self.bias = nn.Parameter(torch.zeros(vocab_size))

# Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
self.decoder.bias = self.bias

This is exactly what is done in HuggingFace's BertLMPredictionHead.

We should be easily able to integrate this within FARM. I'll work on creating a PR, following are the changes I see:

1. Adding additional params in FARM's BertLMHead class as described in the issue.
2. We can modify the conversion script to delete NSP head weights and config but what if someone needs it for testing or for continuing training? I was thinking, we can modify AdaptiveModel.convert_to_transformers() to accept len(self.prediction_heads) == 2 if one of the head is self.prediction_heads[x].model_type == "language_modelling" and use that to initialize ph_state_dict variable.

Let me know if this is possible or will it have cascading effects on other modules.

[Timoeller]

Ah ok, I understand now what you mean: if we set the self.decoder.bias parameter in FARM the bias is also added when the decoder (a linear pytorch layer) is called. Therefore we need to remove the explicit bias addition.

@tholor do you think changing how bias is used inside the BertLMHead will have an impact on other parts of FARM? I am slightly worried about the AWS sagemaker algo...

Another solution would be to add this self.decoder.bias just before conversion to HF. Then we do not need to change the forward pass. Do you think that is an option @himanshurawlani

[himanshurawlani]

@Timoeller I tried your suggested alternative and that works too (I got meaningful predictions for [MASK]). We can assign the bias, in AdaptiveModel.convert_to_transformers(), as follows:

self.prediction_heads[0].decoder.bias = self.prediction_heads[0].bias

We do not need to change the forward pass in this case.

[Timoeller]

Nice, thanks for testing. Would you be interested in creating a PR with your method? We would be happy to include it in the library. I linked a similar issue, so it seems more people would be happy about this functionality as well.

[himanshurawlani]

Sure @Timoeller, I'll create a PR soon.

[himanshurawlani]

Thanks @Timoeller for approving the fix. I'll close this issue.