Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/nni.git" did not exist on "cd3a912a07c2187fcd9e81205da35b97746b76f3"
Commit 1a4edd19 authored by comfyanonymous's avatar comfyanonymous
Browse files

Fix overflow issue with inplace softmax.

parent 509c7dfc
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
comfy/ldm/modules/sub_quadratic_attention.py
View file @ 1a4edd19
...@@ -158,6 +158,7 @@ def _get_attention_scores_no_kv_chunking( ...@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment