Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Flash Attention 2] Add flash attention 2 for GPT-J #28295

Merged
merged 11 commits into from
Mar 13, 2024
+350 −23
Merged

[Flash Attention 2] Add flash attention 2 for GPT-J #28295

Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
6419b04
initial implementation of flash attention for gptj
bytebarde Jan 1, 2024
3a9e31f
modify flash attention and overwrite test_flash_attn_2_generate_paddi…
bytebarde Jan 2, 2024
cb65655
update flash attention support list
bytebarde Jan 2, 2024
e47ef13
remove the copy line in the `CodeGenBlock`
bytebarde Jan 4, 2024
0c31cb3
address copy mechanism
bytebarde Jan 28, 2024
def626e
Update src/transformers/models/gptj/modeling_gptj.py
younesbelkada Jan 30, 2024
6e9c707
Merge branch 'huggingface:main' into flash_attn_gptj
bytebarde Jan 31, 2024
af0752e
Add GPTJ attention classes
bytebarde Feb 9, 2024
cd73e33
Merge branch 'huggingface:main' into flash_attn_gptj
bytebarde Mar 8, 2024
cb265c7
add expected outputs in the gptj test
bytebarde Mar 8, 2024
2b489b0
Ensure repo consistency with 'make fix-copies'
bytebarde Mar 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Update src/transformers/models/gptj/modeling_gptj.py 
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
@younesbelkada @ArthurZucker
younesbelkada and ArthurZucker authored Jan 30, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
GPG key ID: 4AEE18F83AFDEB23
Expired
Verified
Learn about vigilant mode
commit def626ef1699b58e78507161fe32433290454e59
2 changes: 1 addition & 1 deletion src/transformers/models/gptj/modeling_gptj.py
View file
Original file line number Diff line number Diff line change
@@ -375,7 +375,7 @@ def forward(
else:
present = None

# The Falsh attention requires the input to have the shape
# The Flash attention requires the input to have the shape
# batch_size x seq_length x head_dim x hidden_dim
# therefore we need to keep the original shape for query and key, and reshape value
# to have the correct shape.