mirror of
https://github.com/vladmandic/automatic
synced 2026-08-26 06:30:44 +02:00
Vladimir Mandic
40 lines
1.6 KiB
Python
40 lines
1.6 KiB
Python
from modules.logger import log
|
|
|
|
|
|
ATTENTIONS = ['eager', 'sdpa', 'flash_attention_2', 'flash_attention_3', 'flex_attention', 'paged|eager', 'paged|sdpa', 'paged|flash_attention_2', 'paged|flash_attention_3']
|
|
|
|
|
|
def get_first_attention_block(model): # for models still using transformers==4 internal architecture
|
|
m = model.model
|
|
if hasattr(m, "layers"):
|
|
return m.layers[0].self_attn if hasattr(m.layers[0], "self_attn") else m.layers[0].attention
|
|
if hasattr(m, "transformer") and hasattr(m.transformer, "layers"):
|
|
return m.transformer.layers[0].attention
|
|
if hasattr(m, "block"):
|
|
return m.block[0].attention
|
|
return None
|
|
|
|
|
|
def set_attention(model):
|
|
if not hasattr(model, 'set_attn_implementation'):
|
|
return
|
|
supported = []
|
|
unsupported = []
|
|
if hasattr(model, "_attn_implementation"):
|
|
default = model._attn_implementation # pylint: disable=protected-access
|
|
elif hasattr(model, "_get_attn_implementation"):
|
|
default = model._get_attn_implementation() # pylint: disable=protected-access
|
|
else:
|
|
default = get_first_attention_block(model)
|
|
log.debug(f"LLM attention: cls={model.__class__.__name__} default={default} fixed")
|
|
return
|
|
|
|
for name in ATTENTIONS: # type: ignore
|
|
try:
|
|
model.set_attn_implementation(name)
|
|
supported.append(name)
|
|
except Exception:
|
|
unsupported.append(name)
|
|
log.debug(f"LLM attention: cls={model.__class__.__name__} default={default} supported={supported} unsupported={unsupported}")
|
|
model.set_attn_implementation(default) # restore default
|