huggingface / huggingface/diffusers
XFormer fails when passing attention mask while using bfloat and key's sequence length not being a multiple of 8
- Langage dominant
- Python
- Étoiles
- 34.5k
- Forks
- 7.3k
- Merge moyen
- 3 j 3 h
- PR mergées (30 j)
- 91
Description
### 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
Guide de contribution
Ouvrir le guide de contribution
Piste de recherche
Commencez dans diffusers/models/attention_processor.py, au niveau de XFormersAttnProcessor.__call__, puis exécutez la reproduction CUDA fournie avec bfloat16 et une clé de 700 tokens. Suivez la manière dont le masque d’attention est transmis à xformers et vérifiez que le même cas réussit lorsque la longueur de séquence de la clé n’est pas un multiple de 8, avec une couverture de régression pour l’échec signalé.
Rédigé par le modèle d'indexation à partir du texte de l'issue.
Évaluation
- Stack technique
- python, pytorch
- Domaine
- machine-learning
- Type d'issue
- Bug
- Difficulté
- 3/5
- Temps estimé
- 1-2 jours
- Activité
- À l'abandon
- Clarté
- Plutôt claire
- Accessibilité débutants
- 45/100