deepspeedai / deepspeedai/DeepSpeed
[BUG] Sparse Attention with block sizes < 16
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
If we run (sparse) BigBird with block sizes < 16, we get the following error:
File "deepspeed/ops/sparse_attention/matmul.py", line 472, in make_dxx_lut
xincs[:, 0] -= (div - 1) * step
IndexError: index 0 is out of bounds for dimension 1 with size 0
It seems that step is pre set to 16 (here), and then div will be set to 0 when block < 16 (here), which means that xincs = xincs.view(-1, 1).repeat(1, div) will have a shape of (-1, 0), so xincs[:, 0] raises an IndexError.
Is it possible to use smaller blocks or is this behavior a limitation of deepspeed?
To Reproduce
from deepspeed.ops.sparse_attention import SparseSelfAttention, BigBirdSparsityConfig
# create a dummy config
config = BigBirdSparsityConfig(
num_heads=8
block=8
different_layout_per_head=False,
num_random_blocks=3,
num_sliding_window_blocks=3,
num_global_blocks=1
)
# define the attention module
bigbird_attn = SparseSelfAttention(sparsity_config=bigbird_config, max_seq_length=1024)
# call the attention module with fake data:
x = torch.randn(8, 1024, 64).cuda() # (num_heads, seq_len, hidden_size)
bigbird_attn(x, x, x) # yields IndexError
ds_report output
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
runtime if needed. Op compatibility means that your system
meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
[WARNING] async_io requires the dev libaio .so object and headers but these were not found.
[WARNING] async_io: please install the libaio-dev package with apt
[WARNING] If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/mtreviso/test-fast-transformers/env/lib/python3.8/site-packages/torch']
torch version .................... 1.10.0+cu102
torch cuda version ............... 10.2
nvcc version ..................... 11.1
deepspeed install path ........... ['/home/mtreviso/test-fast-transformers/env/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.5.8, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.10, cuda 10.2
System info
- OS: Ubuntu 20.04
- GPU: RTX 2080
- Python version: 3.8.5
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in deepspeed/ops/sparse_attention/matmul.py at the div and xincs handling around lines 469-472, then inspect the step initialization near line 795. Reproduce the BigBird sparse-attention case with block size 8 and verify that supported smaller blocks no longer raise the reported IndexError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100