298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-25 15:55:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Run softmax in place if it OOMs.

Browse Source
This commit is contained in:
comfyanonymous 2023-01-30 19:55:01 -05:00
parent 0d8ad93852
commit 1daccf3678
1 changed files with 11 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
comfy/ldm/modules/sub_quadratic_attention.py
View File

@ -146,8 +146,17 @@ def _get_attention_scores_no_kv_chunking(
alpha=scale,
beta=0,
)
attn_probs = attn_scores.softmax(dim=-1)
del attn_scores
try:
attn_probs = attn_scores.softmax(dim=-1)
del attn_scores
except torch.cuda.OutOfMemoryError:
print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
torch.exp(attn_scores, out=attn_scores)
summed = torch.sum(attn_scores, dim=-1, keepdim=True)
attn_scores /= summed
attn_probs = attn_scores
hidden_states_slice = torch.bmm(attn_probs, value)
return hidden_states_slice