deepspeedai / deepspeedai/DeepSpeed

[BUG] Sparse Attention with block sizes < 16

Open
#1,639 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.