[feat]Adding a simple blockwise attention

[xwhan]

What does this PR do?

This attention is used in fairseq to build a long-doc summarization model.

More efficient implementations should be possible with Triton or blocksparse. However, it needs to be compatible with CPU inference for the sake of deployment and demo.

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.
@dianaml0

[blefaudeux]

hey @xwhan thanks for the PR ! FYI sparse attention (just pass a block mask) would work also (it's CPU compatible, although not very optimized as of now, cc @fmassa), curious to see a benchmark in between the two. There are pattern preset here if that helps https://github.com/facebookresearch/xformers/blob/main/xformers/components/attention/attention_patterns.py

[blefaudeux]

on another front I've just added some more explanations on pre-commit here, in case that helps ? sorry about all these steps, we're trying to keep the repo coherent and easy to read

[xwhan]

hey @xwhan thanks for the PR ! FYI sparse attention (just pass a block mask) would work also (it's CPU compatible, although not very optimized as of now, cc @fmassa), curious to see a benchmark in between the two. There are pattern preset here if that helps https://github.com/facebookresearch/xformers/blob/main/xformers/components/attention/attention_patterns.py

Thanks, I did not know those are already CPU compatible. I'll try it out with block masks and see if it's faster and stable in my experiments.

[xwhan]

@blefaudeux I spent some time trying out passing a blockwise mask to scaled_dot_product_attention.

        repeat = seq_len // self.block_size
        block_mask = torch.ones(self.block_size, self.block_size)
        mask = torch.kron(torch.eye(repeat), block_mask).bool()

The fp16 support is the main pain point. It uses 1/4 more memory and is 2x slower than the PR with fp16. When compared them in fp32 settings, the sputnik one saves a little memory but is still much slower. IMO, this fp16 bottleneck of sputnik might block most NLP use cases, especially for large-scale pretraining.

[blefaudeux]

@blefaudeux I spent some time trying out passing a blockwise mask to scaled_dot_product_attention.

        repeat = seq_len // self.block_size
        block_mask = torch.ones(self.block_size, self.block_size)
        mask = torch.kron(torch.eye(repeat), block_mask).bool()

The fp16 support is the main pain point. It uses 1/4 more memory and is 2x slower than the PR with fp16. When compared them in fp32 settings, the sputnik one saves a little memory but is still much slower. IMO, this fp16 bottleneck of sputnik might block most NLP use cases, especially for large-scale pretraining.

fair enough, thanks for checking it out ! fp16 support is a known issue for sparse (see #15), we don't have that many cycles or manpower to work on that unfortunately.

Blocksparse handles it though, but in turn it's not CPU compatible (give it a try, it should be pretty fast ! Some example here, this can also be a helper), so IMO your PR has value in this in-between space. Ideally we should have stronger primitives with sparse and blocksparse though, and @fmassa is working towards a single interface for all these

[fmassa]

As @blefaudeux mentioned, I'm working on refactoring the sparse abstractions so that they are unified.
I can add a naive CPU implementation that will make it possible to run things on both CPU and well as GPU (through Triton). A WIP branch is in https://github.com/facebookresearch/xformers/tree/blocksparse_refactoring_v2

I will look into it today

[fmassa]

Hi @xwhan

With #202 merged, we now have native support for blocksparse tensors with both CPU and CUDA support.

[fmassa]

Hi @xwhan

Let me know if you would like some help getting the new BlockSparse Tensor integrated in your pipeline. I believe it might be preferable over implementing multiple (similar) variants using the same base code, as all changes needed would just be on defining the sparse matrix

[xwhan]

Thanks @fmassa for the update. That would be a great feature. I am closing the PR for now