memory-efficient-attention-pytorch - GitLite

KO | EN

gitlite — search

#javascript #python #react #ai #android #linux #macos #nodejs #windows #awesome #cli #security #ios #openai

memory-efficient-attention-pytorch

(★ 394)

Implementation of a memory efficient multi-head attention as proposed in the paper, "Self-attention Does Not Need O(n²) Memory"

파일 탐색기

최종 버전 다운로드 (.zip)

- - python-package.yml
  - python-publish.yml
- enwik8.gz
- README.md
- __init__.py
- autoregressive_wrapper.py
- cosine_sim_flash_attention.py
- flash_attention.py
- memory_efficient_attention.py
- memory_efficient_cosine_sim_attention.py
- reversible.py
- transformer.py
- test.py
.gitignore
LICENSE
README.md
setup.cfg
setup.py
train.py

# 설치 가이드

1. 코드 내려받기

git clone https://github.com/lucidrains/memory-efficient-attention-pytorch

깃허브에서 프로젝트 코드 전체를 내 컴퓨터로 내려받습니다.

cd memory-efficient-attention-pytorch

방금 내려받은 프로젝트 폴더 안으로 이동합니다.

2. Python

쉬움 추천

사전 준비물

Git GitHub에서 프로젝트 코드를 내려받으려면 필요합니다.
Python 3 설치 시 'Add Python to PATH' 옵션을 꼭 체크하세요 (Windows).

$ pip install memory-efficient-attention-pytorch

requirements.txt 등에 명시된 파이썬 라이브러리를 설치합니다.

✅ 에러 메시지 없이 실행되고 터미널에 안내 문구가 출력되면 정상입니다.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

## Memory Efficient Attention Pytorch (obsolete)

Implementation of a memory efficient multi-head attention as proposed in the paper, <a href="https://arxiv.org/abs/2112.05682">Self-attention Does Not Need O(n²) Memory</a>. In addition, the module will take care of masking, causal masking, as well as cross attention.

This repository also contains a <a href="https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py">naive non-CUDA implementation</a> of the improvements made by <a href="https://tridao.me/">Tri Dao</a> with his <a href="https://github.com/HazyResearch/flash-attention">Flash Attention 2</a> paper, for educational purposes. It is a game changer for attention and building long-context transformers.

Update: from now on, you should just be using the <a href="https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html?highlight=scaled_dot_product#torch.nn.functional.scaled_dot_product_attention">`F.scaled_dot_product_attention`</a> function in Pytorch 2.0 for built-in Flash Attention v1 support - or use Flash Attention v2 at the <a href="https://github.com/Dao-AILab/flash-attention">official repository</a>

## Install

```bash
$ pip install memory-efficient-attention-pytorch
```

## Usage

For autoregressive language model

```python
import torch
from memory_efficient_attention_pytorch import Attention

attn = Attention(
    dim = 512,
    dim_head = 64,                # dimension per head
    heads = 8,                    # number of attention heads
    causal = True,                # autoregressive or not
    memory_efficient = True,      # whether to use memory efficient attention (can be turned off to test against normal attention)
    q_bucket_size = 1024,         # bucket size along queries dimension
    k_bucket_size = 2048          # bucket size along key / values dimension
).cuda()

x = torch.randn(1, 65536, 512).cuda()
out = attn(x) # (1, 65536, 512)
```

Cross attention

```python
import torch
from memory_efficient_attention_pytorch import Attention

cross_attn = Attention(
    dim = 512,
    dim_head = 64,
    heads = 8,
    memory_efficient = True,
    q_bucket_size = 1024,
    k_bucket_size = 2048
).cuda()

x = torch.randn(1, 65536, 512).cuda()
context = torch.randn(1, 65536, 512).cuda()
mask = torch.ones(1, 65536).bool().cuda()

out = cross_attn(x, context = context, mask = mask) # (1, 65536, 512)
```

## Citations

```bibtex
@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}
}
```

```bibtex
@misc{liu2021swin,
    title   = {Swin Transformer V2: Scaling Up Capacity and Resolution},
    author  = {Ze Liu and Han Hu and Yutong Lin and Zhuliang Yao and Zhenda Xie and Yixuan Wei and Jia Ning and Yue Cao and Zheng Zhang and Li Dong and Furu Wei and Baining Guo},
    year    = {2021},
    eprint  = {2111.09883},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV}
}
```

```bibtex
@article{Dao2022FlashAttentionFA,
    title   = {FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness},
    author  = {Tri Dao and Daniel Y. Fu and Stefano Ermon and Atri Rudra and Christopher R'e},
    journal = {ArXiv},
    year    = {2022},
    volume  = {abs/2205.14135}
}
```

```bibtex
@article{dao2023flashattention2,
  title     = {Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning,
  author    = {Dao, Tri},
  year      = {2023}
}
```