Fix Seasame

This commit is contained in:
Daniel Han 2025-05-15 01:45:15 -07:00
commit 3a9c10c33c
2 changed files with 12 additions and 8 deletions

View file

@ -543,7 +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)"
os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"] = "torch.floa16;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

View file

@ -303,13 +303,6 @@ 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":
@ -317,6 +310,17 @@ class FastBaseModel:
bnb_compute_dtype = torch.float16
do_forced_float32 = True
pass
# Check for custom data-types
custom_datatype = None
if os.environ.get("UNSLOTH_FORCE_CUSTOM_DTYPE", "") != "":
custom_datatype = os.environ["UNSLOTH_FORCE_CUSTOM_DTYPE"]
assert custom_datatype.count(";") == 1
bnb_compute_dtype, custom_datatype = custom_datatype.split(";", 1)
dtype = torch.float32
bnb_compute_dtype = eval(bnb_compute_dtype)
pass
# Stop SDPA for some archs like Pixtral / Mistral3
if not ("attn_implementation" in kwargs):
kwargs["attn_implementation"] = "sdpa"