298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-20 03:13:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix overflow issue with inplace softmax.

Browse Source
This commit is contained in:
comfyanonymous 2023-02-10 11:47:41 -05:00
parent 509c7dfc6d
commit 1a4edd19cd
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -158,6 +158,7 @@ def _get_attention_scores_no_kv_chunking(
del attn_scores del attn_scores
except OOM_EXCEPTION: except OOM_EXCEPTION:
print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead") print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
attn_scores -= attn_scores.max(dim=-1, keepdim=True).values
torch.exp(attn_scores, out=attn_scores) torch.exp(attn_scores, out=attn_scores)
summed = torch.sum(attn_scores, dim=-1, keepdim=True) summed = torch.sum(attn_scores, dim=-1, keepdim=True)
attn_scores /= summed attn_scores /= summed