from_pretrained torch_dtype DO NOT affect model buffers

[Chandler-Bing]

System Info

pass

Who can help?

@ArthurZucker

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

I recently encountered a problem with inconsistent model output. After debugging, I find why.

If I use torch_dtype=torch.bfloat16 in the from_pretrained method, it only affects the model parameters, but has no effect on the model buffer. The buffer is still float32.
But in the configuration of trainer --bf16_full_eval=True, model.to(bfloat16) had been used, so that all parameters and buffers of the model become bfloat16.
And that's why I see difference logits of same pretrained models

So the question is, is from_pretrained torch_dtype designed to only focus on the parameters and not the buffer? Or is this a bug?

Expected behavior

I'm not sure what behavior should be, but it will be nice to see some docs about it

[ArthurZucker]

Hey!
To answer that question I would need a simple reproducer, as well as the output of transformers-cli env to see which version you are running this one.
AFAIK, buffers should we affected. I recently specifically wanted to keep certain buffers in a precision and it did not work

[Chandler-Bing]

Hey! To answer that question I would need a simple reproducer, as well as the output of transformers-cli env to see which version you are running this one. AFAIK, buffers should we affected. I recently specifically wanted to keep certain buffers in a precision and it did not work

@ArthurZucker
sure, you could get the code here, colab

from  transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
        'Qwen/Qwen1.5-0.5B',torch_dtype=torch.bfloat16)

for name,buffer in model.named_buffers():
    if buffer.dtype==torch.float32:
        print(name,buffer.dtype)

the model is randomly chosen to test.

seems just this buffer rotary_emb.inv_freq still torch.float32

[ArthurZucker]

Actually we want the inv freq to be in float 32 to keep as much precision as possible.

because in the code we specify the dtype, the default torch dtype is not used

import torch.nn as nn
class myModule(nn.Module):
  def __init__(self):
    super().__init__()
    self.register_buffer("freq", torch.ones(12,12, dtype=torch.float32))
torch.set_default_dtype(torch.bfloat16)
module = myModule()
module.freq.dtype

vs

import torch.nn as nn
class myModule(nn.Module):
  def __init__(self):
    super().__init__()
    self.register_buffer("freq", torch.ones(12,12))
torch.set_default_dtype(torch.bfloat16)
module = myModule()
module.freq.dtype

[ArthurZucker]

Closing this as it's affecting buffers that don't specify a dtype, but not others.

[zamroni111]

Actually we want the inv freq to be in float 32 to keep as much precision as possible. because in the code we specify the dtype, the default torch dtype is not used

import torch.nn as nn
class myModule(nn.Module):
  def __init__(self):
    super().__init__()
    self.register_buffer("freq", torch.ones(12,12, dtype=torch.float32))
torch.set_default_dtype(torch.bfloat16)
module = myModule()
module.freq.dtype

vs

import torch.nn as nn
class myModule(nn.Module):
  def __init__(self):
    super().__init__()
    self.register_buffer("freq", torch.ones(12,12))
torch.set_default_dtype(torch.bfloat16)
module = myModule()
module.freq.dtype

which py files needs to be edited in order NOT to have this forced float32?

i suggest gives option in from_pretrained function to allow buffer has same dtype as torch_dtype.
since 2023, new PCs have npu that can handle bfloat16
amd zen4 laptop and desktop also has avx512 bfloat16 capability

[ArthurZucker]

Depending on the modeling, you can just change _keep_in_fp32_modules (From here:

transformers/src/transformers/models/t5/modeling_t5.py

Lines 785 to 797 in 0c8021c

	class T5PreTrainedModel(PreTrainedModel):
	"""
	An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
	models.
	"""
	config_class = T5Config
	load_tf_weights = load_tf_weights_in_t5
	base_model_prefix = "transformer"
	is_parallelizable = True
	supports_gradient_checkpointing = True
	_no_split_modules = ["T5Block"]
	_keep_in_fp32_modules = ["wo"]

)