ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

CLI now handles user input strings for dtype correctly (#1235)

Browse source 
Co-authored-by: root <root@ieeres.chu.cam.ac.uk>
This commit is contained in:
Edwin Fennell 2024-11-06 20:23:09 +00:00 • committed by GitHub
commit 08db916009
1 changed files with 20 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

21
unsloth/models/loader.py
View file

@ -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,