fp16 models getting auto converted to fp32 in .from_pretrained()

stas00 edited: this Issue has nothing to do with Deepspeed, but pure transformers

---

Environment info

• transformers version: 4.6.1
• Platform: Linux-3.10.0-1127.13.1.el7.x86_64-x86_64-with-centos-7.7.1908-Core
• Python version: 3.6.8
• PyTorch version (GPU?): 1.6.0+cu92 (True)
• Tensorflow version (GPU?): not installed (NA)
• Using GPU in script?: Yes (not essential)
• Using distributed or parallel set-up in script?: Yes (not essential)

Who can help

@LysandreJik @sgugger

Information

Model I am using (Bert, XLNet …): BertForMaskedLM

The problem arises when using:

• my own modified scripts: (give details below)

The tasks I am working on is:

• my own task or dataset: (give details below) Masked LM

To reproduce

Steps to reproduce the behavior:

1. Finetune a 16-bit low precision BertForMaskedLM model on any dataset using DeepSpeed and Trainer
2. Load the model and check the dtype using:

from transformers import BertTokenizer, BertForMaskedLM
tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
model = BertForMaskedLM.from_pretrained(model_path)
print(model.dtype)

Expected behavior

Outputs torch.float32 instead of the expected torch.float16. I was able to recover the original weights using model.half()

I think it would be helpful to highlight this behaviour of forced autoconversion either as a warning or as a part of from_pretrained() method’s documentation or provide an additional argument to help retain fp16 weights. Willing to pick this issue up. Please let me know what would be the most appropriate fix.

[ShadenSmith]

Hi @asit2898 , thanks for reporting your issue. I can help look at things from DeepSpeed’s side.

Was the model fine-tuned with ZeRO enabled? From the DS config above it seems not, unless it is enabled somewhere on the HF side of things.

@stas00 , does the from_pretrained codepath go through DeepSpeed’s load_checkpoint(), or is the checkpoint logic all on HF’s side?

To start, I did a quick experiment with DeepSpeed (without ZeRO) and examined model parameter dtypes before and after deepspeed.initialize(). So far I haven’t reproduced the issue:

• When FP16 is not enabled, the model’s dtype is unchanged (eg., fp32 stays fp32 and fp16 stays fp16).
• When fp16 is enabled, the model weights are fp16 after deepspeed.initialize() no matter the initial dtype of fp32 or fp16.

[sgugger]

Yes, from_config uses just 1d. For your question, I’m not aware of such a situation existing.