Repository: huggingface/transformers (Python) — Bug Fix

Issue:
Difference in LlamaAttention & LlamaFlashAttention2 attn_output
### System Info

- `transformers` version: 4.34.1
- Platform: Linux-5.15.0-86-generic-x86_64-with-glibc2.31
- Python version: 3.11.5
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker and @younesbelkada

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

We notice `LlamaFlashAttention2._flash_attention_forward` returns a different `attn_output` than `LlamaAttention` computes.

`flash_attn_non_determinism.py`:
```python
import argparse

import torch
import torch.backends.cudnn
import transformers
from transformers.models import llama


def main() -> None:
    torch.backends.cudnn.deterministic = True

    parser = argparse.ArgumentParser()
    parser.add_argument("--use-flash-attention-2", action="store_true")
    args = parser.parse_args()
    use_flash_attention_2 = args.use_flash_attention_2

    tokenizer = transformers.AutoTokenizer.from_pretrained(
        "/models/huggingface/meta-llama/llama-2-7b-chat-hf", local_files_only=True, use_safetensors=True, device_map=torch.device("cuda")
    )
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    text = "Hello world!"
    tokenized_text = tokenizer(text)
    tokenized_text = {key: torch.tensor(value).unsqueeze(dim=0).to(torch.device("cuda")) for key, value in tokenized_text.items()}
    tokenized_text["labels"] = tokenized_text["input_ids"].clone()

    torch.manual_seed(0)
    model = llama.LlamaForCausalLM.from_pretrained(
        "/models/huggingface/meta-llama/llama-2-7b-chat-hf",
        local_files_only=True,
        use_safetensors=True,
        device_map=torch.device("cuda"),
        use_flash_attention_2=use_flash_attention_2,
        torch_dtype=torch.bfloat16,
    )
    assert isinstance(model, llama.LlamaForCausalLM)
    model.eval()
    for param in model.parameters():
        param.requires_grad = False

    model.model.layers[0].train()
    for param in model.model.layers[0].parameters():
        param.requires_grad = True

    optim = torch.optim.AdamW(model.parameters())

    torch.manual_seed(0)

    for i in range(10):
        output = model(**tokenized_text)
        loss = output["loss"]
        if i in (0, 9):
            print(loss)
        loss.backward()
        optim.step()
        optim.zero_grad()


if __name__ == "__main__":
    main()
```

```console
$ python flash_attn_non_determinism.py --use-flash-attention-2
tensor(5.6612, device='cuda:0', grad_fn=<NllLossBackward0>)
tensor(0.3542, device='cuda:0', grad_fn=<NllLossBackward0>)
$ python flash_attn_non_determinism.py
tensor(5.6589, device='cuda:0', grad_fn=<NllLossBackward0>)
tensor(0.2275, device='cuda:0', grad_fn=<NllLossBackward0>)
```

### Expected behavior

I am not expecting the magnitude of the difference between the 2 implementations. A difference of `0.1267` compared to `0.3542` seems very large.


Discussion:
Hey, I think this is related to flash attention version, could you have a look at #26697? 
We are currently using `flash-attn==2.3.2`. There was a minor version release of flash attention literally yesterday.

The problem persists with `flash-attn==2.3.3`.

Are you able to reproduce on your end with the supplied script?
cc @younesbelkada if you can have a look 😉 
hi @KyleMylonakisProtopia !
I think that difference is expected, I am not sure if flash-attn guarantees full reproducibility for gradient computation, note also that some slight differences in logits are expected between FA-2 and non FA-2 models. 
The code demonstrates non-trivial differences in the loss prior to even the first backwards call. Flash attention and flash attention 2 are supposed to be exact algorithms for computing attention. 

From the Flash attention 2 paper "To speed up attention on hardware accelerators such as GPU, [5] proposes an algorithm to reduce the memory
reads/writes while maintaining the same output (without approximation)." That seems pretty unambiguous to me. 

The slight differences from whatever parallelization differences are happening should not be manifesting at the third significant digit on the first loss call. This points to some other kind of issue.
> Flash attention and flash attention 2 are supposed to be exact algorithms for computing attention.

yes, but in the script above you are comparing vanilla attention vs FA-2 no?
That sentence is referring to Flash attention (and implicitly flash attention 2) to "vanilla" attention. That is what our script is showing.
ah correct yes you are right, sorry for the confusion, I'll have a deeper look !
I also encountered the same problem at inference. Environment: `transformers==4.34.0`, `flash-attn==2.3.3`, `torch==2.0.1+cu117`.

```python
seed = 42
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
prompt = """

Tell me the story about a dog.
"""
d_model = "/path/to/CodeLlama-13b-Instruct-hf"
tokenizer = CodeLlamaTokenizer.from_pretrained(d_model)
model = LlamaForCausalLM.from_pretrained(d_model, device_map="auto", torch_dtype=torch.bfloat16)
tokenized = tokenizer(prompt, return_tensors="pt", truncation=False).to("cuda")
generated_ids = model.generate(**tokenized, max_new_tokens=1024, do_sample=True, streamer=TextStreamer(tokenizer, skip_prompt=True))
```

use-flash-attention-2=False:

Once upon a time, there was a dog named Max. Max was a lovable golden retriever who loved nothing more than to go for walks with his owner, Sarah. One day, while they were out on **a walk**,

use-flash-attention-2=True:

Once upon a time, there was a dog named Max. Max was a lovable golden retriever who loved nothing more than to go for walks with his owner, Sarah. One day, while they were out on **their usual stroll**,
Here is my minimal reproducible script:
```python
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaAttention, LlamaModel, _make_causal_mask

device = torch.device("cuda")
dtype = torch.float16

config_ori = LlamaConfig(
    hidden_size=1024,
    intermediate_size=128,
    num_hidden_layers=1,
    num_attention_heads=8,
    max_position_embeddings=16,
    _flash_attn_2_enabled=False
)

config_new = LlamaConfig(
    hidden_size=1024,
    intermediate_size=128,
    num_hidden_layers=1,
    num_attention_heads=8,
    max_position_embeddings=16,
    _flash_attn_2_enabled=True
)

model_ori = LlamaModel(config_ori)
model_new = LlamaModel(config_new)

model_new.load_state_dict(model_ori.state_dict())

model_ori.to(dtype).to(device)
model_new.to(dtype).to(device)

attn_ori = model_ori.layers[0].self_attn
attn_new = model_new.layers[0].self_attn

bsz, hs, seqlen = 2, config_ori.hidden_size, 4
inputs_embeds = torch.randn((bsz, seqlen, hs), dtype=dtype, device=device)

padding_mask = torch.full((bsz, seqlen), 1, dtype=torch.long, device=device)
# or pad a part
# padding_mask[0, 2:] = 0

out_ori = model_ori(attention_mask=padding_mask, inputs_embeds=inputs_embeds, use_cache=False)['last_hidden_state']
out_new = model_new(attention_mask=padding_mask, inputs_embeds=inputs_embeds, use_cache=False)['last_hidden_state']

out_ori.sum(), out_new.sum(), (out_ori - out_new).mean().item(), (out_ori - out_new).abs().max().item(), (out_ori - out_new).abs().mean().item()
```
I noticed that the numerical difference mainly comes from the padding_mask. If the padding_mask is None, it means we only use the causal mask, and the difference is small. However, if we set the padding_mask, we cannot ignore the difference.
![image](https://github.com/huggingface/transformers/assets/13296106/c67111ae-3f1d-4128-91f4-cfa4bac70ba6)
![image](https://github.com/huggingface/transformers/assets/13296106/bede7232-8028-493b-ac9f-6819fe427a25)
If we run pytest from the offical flash-attn repo, the diff.abs().max().item() is always small:
![image](https://github.com/huggingface/transformers/assets/13296106/38f018b9-6af8-4c7f-b211-251dae3ba23b)

The diff comes from the attention module. A more fine-grained code:
```python
bsz, hs, seqlen = 2, config_ori.hidden_size, 4
hidden = torch.rand((bsz, seqlen, hs), dtype=dtype, device=device)

padding_mask = torch.full((bsz, seqlen), 1, dtype=torch.long, device=device)
# padding_mask[0, 2:] = 0

past_key_values_length = 0
key_value_length = seqlen + past_key_values_length

position_ids = torch.arange(past_key_values_length, key_value_length, dtype=torch.long, device=device)
position_ids = position_ids.unsqueeze(0)

if padding_mask is not None:
    attention_mask_ori = model_ori.attn_mask_converter.to_4d(
        padding_mask, seqlen, key_value_length, dtype=hidden.dtype
    )
else:
    attention_mask_ori = model_ori.attn_mask_converter.to_causal_4d(
        bsz, seqlen, key_value_length, dtype=hidden.dtype, device=hidden.device
    )

out_ori, _, _ = attn_ori.forward(
    hidden, attention_mask=attention_mask_ori, position_ids=position_ids, 
)

out_new, _, _ = attn_new.forward(
    hidden, attention_mask=padding_mask, position_ids=position_ids
)

out_ori.sum(), out_new.sum(), (out_ori - out_new).mean().item(), (out_ori - out_new).abs().max().item(), (out_ori - out_new).abs().mean().item()
```

UPDATE: It seems the diff lies in the padded part in the final attn weights? So maybe this should not affect the final training loss and the inference results?

my env:
- `transformers` version: 4.35.0.dev0 (from commit aa4198a at 2023.10.27 main branch)
- Platform: Linux-4.14.0_1-0-0-43-x86_64-with-glibc2.27
- Python version: 3.9.17
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.2
- Accelerate version: 0.22.0
- Accelerate config:    not found
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

hope this helps!
Thanks for the deep dive @wizyoung! This thread already shows differences in the loss and the inference results, so something is afoot. 
cc @younesbelkada If I remember correctly when we debugged the flash attention tests, we found out that the attention mask was not properly taken into account and the attention weights for pad tokens was non zero in vanilla and zero for flash attention. This came from the way we create our attention mask, which adds two inf values, creating overflows. We should be able to easily fix! cc @patrickvonplaten as we talked about this 

Resolving patch:
diff --git a/src/transformers/modeling_attn_mask_utils.py b/src/transformers/modeling_attn_mask_utils.py
--- a/src/transformers/modeling_attn_mask_utils.py
+++ b/src/transformers/modeling_attn_mask_utils.py
@@ -11,11 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import torch
 
 
+@dataclass
 class AttentionMaskConverter:
     """
     A utility attention mask class that allows one to:
@@ -24,6 +26,21 @@ class AttentionMaskConverter:
         - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length,
           key_value_length) that can be multiplied with attention scores
 
+    Examples:
+
+    ```python
+    >>> import torch
+    >>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+
+    >>> converter = AttentionMaskConverter(True)
+    >>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, 5)
+    tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
+            [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
+            [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
+            [-3.4028e+38, -3.4028e+38, -3.4028e+38,  0.0000e+00, -3.4028e+38],
+            [-3.4028e+38, -3.4028e+38, -3.4028e+38,  0.0000e+00,  0.0000e+00]]]])
+    ```
+
     Parameters:
         is_causal (`bool`):
             Whether the attention mask should be a uni-directional (causal) or bi-directional mask.
@@ -32,6 +49,9 @@ class AttentionMaskConverter:
             Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer.
     """
 
+    is_causal: bool
+    sliding_window: int
+
     def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
         self.is_causal = is_causal
         self.sliding_window = sliding_window
@@ -112,7 +132,11 @@ def to_4d(
         expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
             attention_mask_2d.device
         )
-        expanded_4d_mask = expanded_attn_mask if causal_4d_mask is None else expanded_attn_mask + causal_4d_mask
+        if causal_4d_mask is not None:
+            expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
+
+        # expanded_attn_mask + causal_4d_mask can cause some overflow
+        expanded_4d_mask = expanded_attn_mask
 
         return expanded_4d_mask
 


Test patch:
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -1266,6 +1266,9 @@ def check_to_4d(self, mask_converter, q_len, kv_len, additional_mask=None, bsz=3
 
         assert mask_4d.shape == (bsz, 1, q_len, kv_len)
 
+        # make sure there are no overflows
+        assert mask_4d.min() != float("-inf")
+
         context = mask_converter.sliding_window
         if mask_converter.is_causal and context is None:
             # k * (k+1) / 2 tokens are masked in triangualar masks
@@ -1341,6 +1344,9 @@ def test_2d_to_4d_causal(self):
         self.check_to_4d(mask_converter, q_len=3, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
         self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 2), (1, 3), (2, 0)])
 
+        # check that the mask does not overflow on causal masked tokens
+        self.check_to_4d(mask_converter, q_len=7, kv_len=7, additional_mask=[(0, 0), (1, 0), (1, 1)])
+
     def test_2d_to_4d(self):
         mask_converter = AttentionMaskConverter(is_causal=False)