bug of SlicedAttnProcessor
shinetzh opened this issue · comments
Describe the bug
![image](https://private-user-images.githubusercontent.com/18564791/347691362-e98ed9ef-b874-4502-9cd4-f775e279ad9a.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjEyODMwNzksIm5iZiI6MTcyMTI4Mjc3OSwicGF0aCI6Ii8xODU2NDc5MS8zNDc2OTEzNjItZTk4ZWQ5ZWYtYjg3NC00NTAyLTljZDQtZjc3NWUyNzlhZDlhLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA3MTglMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNzE4VDA2MDYxOVomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWFhOTU2MzEwNmJjNTVjMDVmMjA3Nzg0MzZlNjQ5MTg0MDNiMmVlOTM5YTU1YmJiMzU2MDEyNjg5ZGE5ZWMyNjgmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.4YwtLr74cRby2wZcwER469RPzVKtWxlV3ItMz8AZIBE)
"batch_size_attention // self.slice_size" will forget to compute the last batch of attention。
should be this:
"(batch_size_attention - 1) // self.slice_size + 1"
I have submit a pull request for this issue:
#8836
Reproduction
import torch
from diffusers.models.attention_processor import Attention, SlicedAttnProcessor
attn = Attention(320, None, heads=5, dim_head=64)
slice_proc = SlicedAttnProcessor(slice_size=2)
hidden_states = torch.randn(1, 4096, 320)
result1 = slice_proc(attn=attn, hidden_states=hidden_states)
hidden_states = torch.cat([hidden_states]*2, dim=0)
result2 = slice_proc(attn=attn, hidden_states=hidden_states)
print(torch.sum(torch.abs(result1 - result2[:1]))) ### Ideally, the numerical difference should be 0
Logs
No response
System Info
- __ Diffusers version: 0.29.0
- Platform: Linux-4.14.105-1-tlinux3-0013-x86_64-with-glibc2.28
- Running on a notebook?: No
- Running on Google Colab?: No
- Python version: 3.9.19
- PyTorch version (GPU?): 2.1.2+cu118 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.23.3
- Transformers version: 4.41.2
- Accelerate version: 0.31.0
- PEFT version: not installed
- Bitsandbytes version: not installed
- Safetensors version: 0.4.3
- xFormers version: 0.0.23.post1+cu118
- Accelerator: Tesla V100-SXM2-32GB, 32510 MiB VRAM
- Using GPU in script?:
- Using distributed or parallel set-up in script?:
Who can help?
import torch
from diffusers.models.attention_processor import Attention, SlicedAttnProcessor, AttnProcessor2_0
attn = Attention(320, None, heads=5, dim_head=64)
slice_proc1 = SlicedAttnProcessor(slice_size=1)
slice_proc2 = SlicedAttnProcessor(slice_size=2)
slice_proc3 = SlicedAttnProcessor(slice_size=3)
attn_proc = AttnProcessor2_0()
hidden_states = torch.randn(2, 4096, 320)
result1 = slice_proc1(attn=attn, hidden_states=hidden_states)
result2 = slice_proc2(attn=attn, hidden_states=hidden_states)
result3 = slice_proc3(attn=attn, hidden_states=hidden_states)
result = attn_proc(attn=attn, hidden_states=hidden_states)
print(torch.sum(torch.abs(result1 - result))) ### Ideally, the numerical difference should be 0
print(torch.sum(torch.abs(result2 - result)))
print(torch.sum(torch.abs(result3 - result)))