SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[Mistral] Add Flash Attention-2 support for mistral (#26464)

Browse Source 
* add FA-2 support for mistral

* fixup

* add sliding windows

* fixing few nits

* v1 slicing cache - logits do not match

* add comment

* fix bugs

* more mem efficient

* add warning once

* add warning once

* oops

* fixup

* more comments

* copy

* add safety checker

* fixup

* Update src/transformers/models/mistral/modeling_mistral.py

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>

* copied from

* up

* raise when padding side is right

* fixup

* add doc + few minor changes

* fixup

---------

Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This commit is contained in:
Younes Belkada
2023-10-03 13:44:46 +02:00
committed by GitHub
parent 1a2e966cfe
commit ae9a344cce
5 changed files with 435 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_modeling_common.py
View File

@@ -2926,7 +2926,7 @@ class ModelTesterMixin:
model.save_pretrained(tmpdirname)
dummy_input = torch.LongTensor([[0, 2, 3, 4], [0, 2, 3, 4]]).to(torch_device)
dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1], [1, 1, 1, 0]]).to(torch_device)
dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1], [0, 1, 1, 1]]).to(torch_device)
model = model_class.from_pretrained(
tmpdirname, torch_dtype=torch.float16, use_flash_attention_2=True, low_cpu_mem_usage=True