如何正确加载已转换并保存为float16的Hugging Face模型？

问题原因 #

调用model.half()仅会转换当前模型参数的数据类型，但默认不会将torch_dtype配置写入模型的config.json文件。from_pretrained方法在未指定数据类型时，会默认使用float32加载权重——哪怕权重文件本身是float16格式。

两种解决方法 #

方法1：加载时显式指定float16数据类型

在调用from_pretrained时，通过torch_dtype参数强制指定使用float16加载模型：

import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer

# 加载时指定数据类型
model2 = AutoModelForTokenClassification.from_pretrained(
    save_directory, 
    local_files_only=True,
    torch_dtype=torch.float16
)
# 若有GPU可将模型移至CUDA以发挥float16性能
# model2 = model2.to("cuda")
print_model_layer_dtype(model2)

方法2：保存模型时写入float16配置

在保存前修改模型配置的torch_dtype字段，让后续加载自动沿用float16：

import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer

# 转换为float16后修改配置
model.half()
model.config.torch_dtype = torch.float16

# 保存模型
save_directory = 'temp_model_SE'
model.save_pretrained(save_directory)

# 加载时无需额外指定，自动使用float16
model2 = AutoModelForTokenClassification.from_pretrained(
    save_directory, 
    local_files_only=True
)
print_model_layer_dtype(model2)

补充说明 #

• 若使用CPU环境，float16的兼容性可能受限，可考虑替换为torch.bfloat16；
• 有NVIDIA GPU时，配合device_map="auto"参数能自动分配模型到GPU，进一步优化float16的运行效率。

内容的提问来源于stack exchange，提问作者Franck Dernoncourt