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

check SDPA for Mistral 3, Pixtral

Browse source
This commit is contained in:
Daniel Han 2025-03-26 04:11:27 -07:00
commit cb7157c4bf
3 changed files with 14 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/models/_utils.py
View file

@ -1176,9 +1176,10 @@ def unsloth_compile_transformers(
"so turning off some optimizations!"
)
return
if disable: return
model_types = list(dict().fromkeys(model_types).keys())
if disable: return model_types, False
supports_sdpa = [True]
for model_type in model_types:
_unsloth_compile_transformers(
model_type,
@ -1206,12 +1207,13 @@ def unsloth_compile_transformers(
import_from_cache = import_from_cache,
disable = disable,
return_logits = return_logits,
supports_sdpa = supports_sdpa,
)
pass
# Redo patches which override compiler
for temporary_patch in TEMPORARY_PATCHES:
temporary_patch()
return model_types
return model_types, supports_sdpa[0]
pass
# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie

4
unsloth/models/loader.py
View file

@ -642,7 +642,6 @@ class FastModel(FastBaseModel):
trust_remote_code = trust_remote_code,
)
model_types = ["siglip"] + model_types
print("model_types", model_types)
# Set forced float32 env flag
os.environ["UNSLOTH_FORCE_FLOAT32"] = "0"
@ -664,7 +663,7 @@ class FastModel(FastBaseModel):
with redirector:
patch_loss_functions(torch_compile = False)
model_types = unsloth_compile_transformers(
model_types, supports_sdpa = unsloth_compile_transformers(
dtype = dtype,
model_name = model_name,
model_types = model_types,
@ -727,6 +726,7 @@ class FastModel(FastBaseModel):
tokenizer_name = tokenizer_name,
auto_model = auto_model,
use_gradient_checkpointing = use_gradient_checkpointing,
supports_sdpa = supports_sdpa,
*args, **kwargs,
)

25
unsloth/models/vision.py
View file

@ -66,11 +66,6 @@ __all__ = [
"FastBaseModel",
]
global FORCE_EAGER_ATTENTION
FORCE_EAGER_ATTENTION = [
"pixtral", # Pixtral SDPA not implemented
]
global NUM_LOGITS_TO_KEEP
NUM_LOGITS_TO_KEEP = dict()
global PROMPT_LOOPKUP
@ -240,6 +235,7 @@ class FastBaseModel:
tokenizer_name = None,
auto_model = AutoModelForVision2Seq,
use_gradient_checkpointing = "unsloth",
supports_sdpa = True,
**kwargs,
):
if model_types is None:
@ -307,16 +303,11 @@ class FastBaseModel:
bnb_compute_dtype = torch.float16
do_forced_float32 = True
pass
global FORCE_EAGER_ATTENTION
attn_implementation = "sdpa"
for disable_name in FORCE_EAGER_ATTENTION:
if (disable_name.lower() == model_type_arch.lower() or \
disable_name.lower() in model_name.lower()):
print(f"Unsloth: {model_type_arch} does not support SDPA - switching to eager!")
attn_implementation = "eager"
break
# Stop SDPA for some archs like Pixtral / Mistral3
kwargs["attn_implementation"] = "sdpa"
if not supports_sdpa:
print(f"Unsloth: {model_type_arch} does not support SDPA - switching to eager!")
del kwargs["attn_implementation"]
pass
bnb_config = None
@ -355,8 +346,6 @@ class FastBaseModel:
os.environ["UNSLOTH_ENABLE_FULL_FINETUNING"] = "0"
pass
kwargs.pop("attn_implementation", None); # No need since we auto call it
# Cannot be None, since HF now checks for the config
if load_in_4bit: kwargs["quantization_config"] = bnb_config
@ -370,7 +359,7 @@ class FastBaseModel:
# quantization_config = bnb_config,
token = token,
trust_remote_code = trust_remote_code,
attn_implementation = attn_implementation,
# attn_implementation = attn_implementation,
**kwargs,
)
# Return old flag