From 08db916009a28f366b697e621f4566a1b6b46f8a Mon Sep 17 00:00:00 2001 From: Edwin Fennell Date: Wed, 6 Nov 2024 20:23:09 +0000 Subject: [PATCH] CLI now handles user input strings for dtype correctly (#1235) Co-authored-by: root --- unsloth/models/loader.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index db7259b1d9..8dcdebab12 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -44,6 +44,25 @@ if SUPPORTS_GEMMA2: from .gemma2 import FastGemma2Model pass +def get_dtype_from_input( + dtype +): + '''Converts user-defined dtype input string to a usable dtype''' + TORCH_FLOAT16_SYNONYMS = {"torch.float16"} + TORCH_BFLOAT16_SYNONYMS = {"torch.bfloat16"} + TORCH_FLOAT32_SYNONYMS = {"torch.float32"} + if dtype in TORCH_FLOAT16_SYNONYMS: + return torch.float16 + if dtype in TORCH_BFLOAT16_SYNONYMS: + return torch.bfloat16 + if dtype in TORCH_FLOAT32_SYNONYMS: + return torch.float32 + if dtype != "None": + print(f"--------------------------------------------------\n"\ + f"User-specified dtype not recognised. Defaulting to dtype = None\n"\ + f"--------------------------------------------------") + return None + def __get_model_name( model_name, @@ -332,7 +351,7 @@ class FastLanguageModel(FastLlamaModel): model, tokenizer = dispatch_model.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = dtype, + dtype = get_dtype_from_input(dtype), load_in_4bit = load_in_4bit, token = token, device_map = device_map,