SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Update config.torch_dtype correctly (#36679)

Browse Source 
* fix

* style

* new test
This commit is contained in:
Marc Sun
2025-03-13 12:08:02 +01:00
committed by GitHub
parent c4161238bd
commit fbb18ce68b
2 changed files with 11 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/modeling_utils.py
View File

@@ -1223,10 +1223,12 @@ def _get_torch_dtype(
) )
elif hasattr(torch, torch_dtype): elif hasattr(torch, torch_dtype):
torch_dtype = getattr(torch, torch_dtype) torch_dtype = getattr(torch, torch_dtype)
config.torch_dtype = torch_dtype
for sub_config_key in config.sub_configs.keys(): for sub_config_key in config.sub_configs.keys():
sub_config = getattr(config, sub_config_key) sub_config = getattr(config, sub_config_key)
sub_config.torch_dtype = torch_dtype sub_config.torch_dtype = torch_dtype
elif isinstance(torch_dtype, torch.dtype): elif isinstance(torch_dtype, torch.dtype):
config.torch_dtype = torch_dtype
for sub_config_key in config.sub_configs.keys(): for sub_config_key in config.sub_configs.keys():
sub_config = getattr(config, sub_config_key) sub_config = getattr(config, sub_config_key)
sub_config.torch_dtype = torch_dtype sub_config.torch_dtype = torch_dtype

9
tests/utils/test_modeling_utils.py
View File

@@ -36,6 +36,7 @@ from transformers import (
AutoModel, AutoModel,
AutoModelForImageClassification, AutoModelForImageClassification,
AutoModelForSequenceClassification, AutoModelForSequenceClassification,
CLIPTextModelWithProjection,
DynamicCache, DynamicCache,
LlavaForConditionalGeneration, LlavaForConditionalGeneration,
MistralForCausalLM, MistralForCausalLM,
@@ -617,6 +618,14 @@ class ModelUtilsTest(TestCasePlus):
model = AutoModel.from_pretrained(TINY_BERT_FOR_TOKEN_CLASSIFICATION, torch_dtype="auto") model = AutoModel.from_pretrained(TINY_BERT_FOR_TOKEN_CLASSIFICATION, torch_dtype="auto")
self.assertEqual(model.dtype, torch.float32) self.assertEqual(model.dtype, torch.float32)
# test model that init the model with _from_config
model = CLIPTextModelWithProjection.from_pretrained(
"hf-internal-testing/diffusers-stable-diffusion-tiny-all",
subfolder="text_encoder",
torch_dtype=torch.bfloat16,
)
self.assertEqual(model.dtype, torch.bfloat16)
def test_model_from_pretrained_attn_implementation(self): def test_model_from_pretrained_attn_implementation(self):
# test that the model can be instantiated with attn_implementation of either # test that the model can be instantiated with attn_implementation of either
# 1. explicit from_pretrained's attn_implementation argument # 1. explicit from_pretrained's attn_implementation argument