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}
}
24 Nov 01, 2022
2.1k Dec 28, 2022
5k Dec 31, 2022
24 Nov 29, 2022
63 Dec 16, 2022
46 Dec 12, 2022
102 Dec 26, 2022
23 Nov 29, 2022
607 Dec 31, 2022
104 Dec 07, 2022
6.7k Dec 31, 2022
3 Jul 16, 2022
10 Nov 30, 2022
108 Dec 25, 2022
146 Dec 22, 2022
41 Dec 10, 2022
5 Jun 08, 2022
3.5k Jan 01, 2023
10 Aug 08, 2022
26 Nov 29, 2022