Я загружаю модель Hugging Face с плавающей запятой32, привожу ее к float16 и сохраняю. Как я могу загрузить его как float16?

Я загружаю модель float32 Huggingface-Transformers, привожу ее к float16 и сохраняю. Как я могу загрузить его как float16?

Пример:

# pip install transformers
from transformers import AutoModelForTokenClassification, AutoTokenizer

# Load model
model_path = 'huawei-noah/TinyBERT_General_4L_312D'
model = AutoModelForTokenClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Convert the model to FP16
model.half()

# Check model dtype
def print_model_layer_dtype(model):
    print('\nModel dtypes:')
    for name, param in model.named_parameters():
        print(f"Parameter: {name}, Data type: {param.dtype}")

print_model_layer_dtype(model)
save_directory = 'temp_model_SE'
model.save_pretrained(save_directory)

model2 = AutoModelForTokenClassification.from_pretrained(save_directory, local_files_only=True)
print('\n\n##################')
print(model2)
print_model_layer_dtype(model2)

В этом примере model2 загружается как модель float32 (как показано print_model_layer_dtype(model2)), хотя model2 было сохранено как float16 (как показано в config.json). Как правильно загрузить его как float16?

Протестировано с transformers==4.36.2 и Python 3.11.7 в Windows 10.