# Flash Attention

### About Flash Attention-2

GPU는 메모리 대역폭과 병렬 처리에 최적화되어 있기 때문에 CPU와 달리 머신 러닝을 위한 표준 하드웨어로 선택됩니다. 최신 모델의 더 큰 크기를 따라잡거나 기존 및 구형 하드웨어에서 이러한 대형 모델을 실행하기 위해 GPU 추론 속도를 높이는 데 사용할 수 있는 몇 가지 최적화 방법이 있습니다.

Flash Attention(현재 FlashAttention-2)는 메모리 효율이 높은 주의 메커니즘으로 BetterTransformer(파이토치 기본 빠른 경로 실행)와 Bitsandbytes를 사용하여 모델을 더 낮은 정밀도로 정량화합니다.

FlashAttention-2는 standard attention 메커니즘을 더 빠르고 효율적으로 구현한 것으로, 추론 속도를 크게 높일 수 있습니다:

* 시퀀스 길이에 대한 주의 계산을 추가로 병렬화합니다.
* GPU 스레드 간 작업을 분할하여 스레드 간 통신 및 공유 메모리 읽기/쓰기를 감소시킵니다.

논문에서 40GB GPU의 타일형 플래시어텐션 계산 패턴과 메모리 계층구조를 보여줍니다. 오른쪽 차트는 어텐션 메커니즘의 여러 구성 요소를 융합하고 재정렬하여 얻을 수 있는 상대적인 속도 향상을 보여줍니다.

<figure><img src="https://developer-blogs.nvidia.com/wp-content/uploads/2023/11/flash-attention-computation-pattern-memory-hierarchy-gpu.png" alt=""><figcaption><p><a href="https://arxiv.org/abs/2205.14135"><em>FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness</em></a></p></figcaption></figure>

\


## Flash Attention

설치 방법은 아래와 같습니다. 설치를 하면 Transformer 라이브러리와 함께 사용할 수 있습니다.

```python
%pip install flash-attn --no-build-isolation
```

```python
import torch
import bitsandbytes as bnb
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
```

### Memory Test

Flash Attention-2를 사용했을 때 실제  GPU에서 메모리가 어떻게 최적화되고 소모하는지 결과를 비교하겠습니다. 아래는 테스트 조건 입니다:

* Model: Mistral-7B-v0.3
* Batchsize: 2
* Sequence Lenght: 3072

테스트는 3가지 조건으로 실행할 예정입니다.

1. `AutoModelForCausalML`
2. `AutoModelForCausalML` + `BitsAndBytes`
3. `AutoModelForCausalML` + `BitsAndBytes` + `flash_attention_2`

### 1. Standard Inference

```python
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.3",
    device_map="auto",
    use_flash_attention_2=False,
)

unit_gib = 1024 ** 3

curr_mem = torch.cuda.memory_reserved() / unit_gib
print(f"Initial: {curr_mem:.4f} GiB")

batch_size = 2
seq_len = 3072
inn = torch.LongTensor([[0] * seq_len] * batch_size)

try:
    with torch.no_grad():
        out = model(inn)

    infer_mem = torch.cuda.memory_reserved() / unit_gib
    print(f"Inference: {infer_mem:.4f} GiB")
except:
    print(f"Inference: OOM")
```

```
Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]

Initial: 7.5684 GiB
Inference: 8.7402 GiB
```

### 2. Inference with Bitsandbytes

```python
bnb_config = BitsAndBytesConfig(  
    load_in_4bit= True,
    bnb_4bit_quant_type= "nf4",
    bnb_4bit_compute_dtype= torch.bfloat16,
    bnb_4bit_use_double_quant= False,
)

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.3",
    device_map="auto",
    use_flash_attention_2=False,
)

unit_gib = 1024 ** 3

curr_mem = torch.cuda.memory_reserved() / unit_gib
print(f"Initial: {curr_mem:.4f} GiB")

batch_size = 2
seq_len = 3072
inn = torch.LongTensor([[0] * seq_len] * batch_size)

try:
    with torch.no_grad():
        out = model(inn)

    infer_mem = torch.cuda.memory_reserved() / unit_gib
    print(f"Inference: {infer_mem:.4f} GiB")
except:
    print(f"Inference: OOM")
```

```
Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]

Initial: 7.5527 GiB
Inference: 8.7246 GiB
```

### 3. Inference with Bitsandbytes, FlashAttention2

```python
bnb_config = BitsAndBytesConfig(  
    load_in_4bit= True,
    bnb_4bit_quant_type= "nf4",
    bnb_4bit_compute_dtype= torch.bfloat16,
    bnb_4bit_use_double_quant= False,
)

model: ModelCls = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.3",
    device_map="auto",
    quantization_config=bnb_config,
    use_flash_attention_2=True,
)

unit_gib = 1024 ** 3

curr_mem = torch.cuda.memory_reserved() / unit_gib
print(f"Initial: {curr_mem:.4f} GiB")

batch_size = 2
seq_len = 3072
inn = torch.LongTensor([[0] * seq_len] * batch_size)

try:
    with torch.no_grad():
        out = model(inn)

    infer_mem = torch.cuda.memory_reserved() / unit_gib
    print(f"Inference: {infer_mem:.4f} GiB")
except:
    print(f"Inference: OOM")
```

```
Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]

Initial: 9.3535 GiB
Inference: OOM
```



### Conclusion

첫번째 일반적인 `AutoModelForCausalML`로  모델을 불러와 Inference 했을 때 8.742G의 GPU 메모리를 사용했습니다. 두번째 의 조합을 사용했을 때 8.724G의 메모리를 사용했습니다. 중요한 사실은 Initial Memory로 초기에 GPU 메모리 확보 입니다. 두 조건 모두 차이가 없습니다.

이제 마지막으로 `flash_attention_2=True` 로 했을 때 입니다. Initial Memory는 9.3535G로 일반적인 조건보다 더 많이 확보했고, Inference 후에는 0으로 GPU 공간을 잡아먹지 않았습니다.&#x20;

Train이나 Inference 시 Flash Attention을 사용하면 메모리 최적화를 이룰 수 있습니다.
