Attention for PyTorch with Linear Memory Footprint
Unofficially implements https://arxiv.org/abs/2112.05682 to get Linear Memory Cost on Attention (+ some sidekick speedup on the GPU when compared to reference implementation in JAX)
Usage:
git clone https://github.com/CHARM-Tx/linear_mem_attention_pytorch
cd linear_mem_attention_pytorch
python setup.py install
Usage:
High Level
from linear_mem_attention_torch.fast_attn import Attention
batch, length, features = 2, 2**8, 64
x, ctx = torch.randn(2, batch, length, features)
mask = torch.randn(batch, length) < 1.
attn = Attention(dim=features, heads = 8, dim_head = 64, bias=False)
# self-attn
v_self = attn(x, x, mask, query_chunk_size=1024, key_chunk_size=4096)
# cross-attn
v_cross = attn(x, ctx, mask, query_chunk_size=1024, key_chunk_size=4096)
Low level
from linear_mem_attention_torch import attention
batch, length, heads, features = 2, 2**8, 8, 64
mask = torch.randn(batch, length) < 1.
q, k, v = torch.randn(3, batch, length, heads, features)
v_ = attention(q, k, v, mask, query_chunk_size=1024, key_chunk_size=4096)
Benchmarks
See examples/example_benchamrk.ipynb for more information.
Citations:
@misc{rabe2021selfattention,
title={Self-attention Does Not Need $O(n^2)$ Memory},
author={Markus N. Rabe and Charles Staats},
year={2021},
eprint={2112.05682},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
85 Dec 30, 2022
4.1k Dec 28, 2022
6 Jan 26, 2022
7 Sep 29, 2022
4 Jul 11, 2022
0 May 26, 2022
139 Jan 01, 2023
91 Jan 06, 2023
1.4k Jan 01, 2023
336 Dec 26, 2022
2.5k Dec 30, 2022
4 Mar 08, 2022
5 Dec 14, 2022
75 Dec 22, 2022
328 Jan 09, 2023
59 Dec 16, 2022
62 Dec 30, 2022
6 Nov 29, 2022
6.2k Jan 07, 2023
35 Aug 21, 2022