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

Update loader.py

Browse source
This commit is contained in:
Daniel Han 2024-11-06 15:05:10 -08:00
commit 52e3a2bf9a
1 changed files with 17 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

35
unsloth/models/loader.py
View file

@ -44,24 +44,23 @@ 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:
def _get_dtype(dtype):
__DTYPE_MAP = {
"float32": torch.float32,
torch.float32: torch.float32,
"float16": torch.float16,
torch.float16: torch.float16,
"bfloat16": torch.bfloat16,
torch.bfloat16: torch.bfloat16,
}
if dtype in __DTYPE_MAP:
return __DTYPE_MAP[dtype]
else:
print(f"Unsloth: {dtype} is not recognized, so we'll default to torch.float16")
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
pass
pass
def __get_model_name(
@ -351,7 +350,7 @@ class FastLanguageModel(FastLlamaModel):
model, tokenizer = dispatch_model.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = get_dtype_from_input(dtype),
dtype = _get_dtype(dtype),
load_in_4bit = load_in_4bit,
token = token,
device_map = device_map,