huggingface / huggingface/diffusers
XFormer fails when passing attention mask while using bfloat and key's sequence length not being a multiple of 8
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
### Describe the bug
XFormer will fail when passing attention mask with its last dimension not being a multiple of 8 (i.e. key's sequence length) under bfloat16. This seems to be because xformer needs to attn_bias's stride to be a multiple of 8. Padding the attention mask on its last dimension will solve this problem.
### Reproduction
```python
from diffusers.models.attention_processor import Attention, XFormersAttnProcessor
import torch
attn_processer = XFormersAttnProcessor()
attn = Attention(
query_dim=256,
heads=8,
dim_head=64,
processor=attn_processer,
).to(device="cuda", dtype=torch.bfloat16)
q = torch.zeros((2, 350, 256), device="cuda", dtype=torch.bfloat16)
kv = torch.zeros((2, 700, 256), device="cuda", dtype=torch.bfloat16)
attn_mask = torch.zeros((2, 1, 700), device="cuda", dtype=torch.bfloat16)
out = attn(q, kv, attn_mask)
```
### Logs
```shell
Traceback (most recent call last):
File "/mnt/hcufs/home/youjunqi/GENAD2/test/test_diffusers.py", line 17, in
out = attn(q, kv, attn_mask)
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/diffusers/models/attention_processor.py", line 490, in forward
return self.processor(
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/diffusers/models/attention_processor.py", line 2156, in __call__
hidden_states = xformers.ops.memory_efficient_attention(
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/xformers/ops/fmha/__init__.py", line 276, in memory_efficient_attention
return _memory_efficient_attention(
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/xformers/ops/fmha/__init__.py", line 403, in _memory_efficient_attention
return _fMHA.apply(
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/torch/autograd/function.py", line 574, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/xformers/ops/fmha/__init__.py", line 74, in forward
out, op_ctx = _memory_efficient_attention_forward_requires_grad(
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/xformers/ops/fmha/__init__.py", line 428, in _memory_efficient_attention_forward_requires_grad
op = _dispatch_fw(inp, True)
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/xformers/ops/fmha/dispatch.py", line 119, in _dispatch_fw
return _run_priority_list(
File "/mnt/hcufs/home/youjunqi/miniconda3/envs/nuplan/lib/python3.9/site-packages/xformers/ops/fmha/dispatch.py", line 55, in _run_priority_list
raise NotImplementedError(msg)
NotImplementedError: No operator found for `memory_efficient_attention_forward` with inputs:
query : shape=(16, 350, 1, 64) (torch.bfloat16)
key : shape=(16, 700, 1, 64) (torch.bfloat16)
value : shape=(16, 700, 1, 64) (torch.bfloat16)
attn_bias :
p : 0.0
`flshattF@2.5.6-pt` is not supported because:
attn_bias type is
`cutlassF-pt` is not supported because:
attn_bias.stride(-2) % 8 != 0 (attn_bias.stride() = (700, 0, 0, 1))
HINT: To use an `attn_bias` with a sequence length that is not a multiple of 8, you need to ensure memory is aligned by slicing a bigger tensor. Example: use `attn_bias = torch.zeros([1, 1, 5, 8])[:,:,:,:5]` instead of `torch.zeros([1, 1, 5, 5])`
`smallkF` is not supported because:
max(query.shape[-1] != value.shape[-1]) > 32
dtype=torch.bfloat16 (supported: {torch.float32})
has custom scale
unsupported embed per head: 64
```
### System Info
- 🤗 Diffusers version: 0.30.2
- Platform: Linux-5.15.0-41-generic-x86_64-with-glibc2.35
- Running on Google Colab?: No
- Python version: 3.9.19
- PyTorch version (GPU?): 2.4.1+cu121 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.24.6
- Transformers version: 4.44.2
- Accelerate version: 0.21.0
- PEFT version: 0.12.0
- Bitsandbytes version: not installed
- Safetensors version: 0.4.4
- xFormers version: 0.0.27.post2
- Accelerator: NVIDIA H800, 81559 MiB
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No
### Who can help?
@DN6 @yiyixuxu @sayakpaul
Beitragsleitfaden
Bewertung
Dieses Issue wurde noch nicht bewertet.