mrq/vall-e
5
9
Fork 0
Code Issues 8 Pull Requests Packages Projects Releases Wiki Activity

do not include SDPA attention if there's no available SDPA backends

Browse Source
This commit is contained in:
mrq 2024-08-06 20:42:39 -05:00
parent 613024ec0d
commit 949339a3fa
2 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
vall_e/models/arch/llama.py
View File

@ -9,7 +9,7 @@ from transformers.cache_utils import Cache
from transformers import LlamaModel, LlamaConfig, LlamaForCausalLM
from transformers.models.llama.modeling_llama import LlamaAttention, apply_rotary_pos_emb, repeat_kv
AVAILABLE_ATTENTIONS = ["sdpa"]
AVAILABLE_ATTENTIONS = []
if torch.backends.cuda.flash_sdp_enabled():
AVAILABLE_ATTENTIONS.append("flash")
@ -20,6 +20,12 @@ if torch.backends.cuda.mem_efficient_sdp_enabled():
if torch.backends.cuda.math_sdp_enabled():
AVAILABLE_ATTENTIONS.append("math")
if torch.backends.cuda.cudnn_sdp_enabled():
AVAILABLE_ATTENTIONS.append("cudnn")
if AVAILABLE_ATTENTIONS:
AVAILABLE_ATTENTIONS.append("sdpa")
try:
from xformers.ops import LowerTriangularMask
from xformers.ops.fmha import memory_efficient_attention

4
vall_e/models/base.py
View File

@ -522,8 +522,10 @@ class Base(nn.Module):
attention_backend = "mem_efficient"
elif "math" in AVAILABLE_ATTENTIONS:
attention_backend = "math"
else:
elif "sdpa" in AVAILABLE_ATTENTIONS:
attention_backend = "sdpa"
else:
attention_backend = "eager"
if attention_backend == "xformers":
attention_backend = "mem_efficient"