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
Merged

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

Show file tree
Hide file tree
Changes from 7 commits
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
1 change: 1 addition & 0 deletions docs/source/en/perf_infer_gpu_one.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ FlashAttention-2 is currently supported for the following architectures:
* [GPTBigCode](https://huggingface.co/docs/transformers/model_doc/gpt_bigcode#transformers.GPTBigCodeModel)
* [GPTNeo](https://huggingface.co/docs/transformers/model_doc/gpt_neo#transformers.GPTNeoModel)
* [GPTNeoX](https://huggingface.co/docs/transformers/model_doc/gpt_neox#transformers.GPTNeoXModel)
* [GPT-J](https://huggingface.co/docs/transformers/model_doc/gptj#transformers.GPTJModel)
* [Falcon](https://huggingface.co/docs/transformers/model_doc/falcon#transformers.FalconModel)
* [Llama](https://huggingface.co/docs/transformers/model_doc/llama#transformers.LlamaModel)
* [Llava](https://huggingface.co/docs/transformers/model_doc/llava)
Expand Down
1 change: 1 addition & 0 deletions src/transformers/models/codegen/modeling_codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,7 @@ def forward(self, hidden_states: Optional[torch.FloatTensor]) -> torch.FloatTens

# Copied from transformers.models.gptj.modeling_gptj.GPTJBlock with GPTJ->CodeGen
younesbelkada marked this conversation as resolved.
Show resolved Hide resolved
class CodeGenBlock(nn.Module):
# Ignore copy
def __init__(self, config):
super().__init__()
inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
Expand Down
Loading