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

Seasame force float16 / float32

Browse source
This commit is contained in:
Daniel Han 2025-05-15 01:27:14 -07:00
commit 8f42a99e49
2 changed files with 15 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
unsloth/models/loader.py
View file

@ -543,6 +543,7 @@ class FastModel(FastBaseModel):
elif "csm-1b" in model_name.lower():
os.environ["UNSLOTH_COMPILE_DISABLE"] = "1"
os.environ["UNSLOTH_DISABLE_FAST_GENERATION"] = "1"
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "if name.endswith(('_proj', 'fc1', 'fc2', 'codebook', 'head')): module.to(torch.float16)"
elif "olmo-2" in model_name.lower() and transformers_version < Version("4.50.0.dev0"):
raise RuntimeError("Unsloth: OLMo-2 only works on transformers >= 4.50.0." + NIGHTLY)
pass

14
unsloth/models/vision.py
View file

@ -303,6 +303,13 @@ class FastBaseModel:
pass
assert(dtype in (torch.float16, torch.bfloat16, torch.float32))
# Check for custom data-types
custom_datatype = None
if os.environ.get("UNSLOTH_FORCE_CUSTOM_DTYPE", "") != "":
custom_datatype = os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"]
dtype = torch.float32
pass
bnb_compute_dtype = dtype
do_forced_float32 = False
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
@ -374,6 +381,13 @@ class FastBaseModel:
# Return old flag
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer
# Edit data-types
if custom_datatype is not None:
with torch.inference_mode():
for name, module in model.named_modules():
exec(custom_datatype)
pass
# Counteract saved tokenizers
tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
is_vlm = (auto_model is AutoModelForVision2Seq)