riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-30 17:52:35 +06:00
Code Issues Actions 30 Packages Projects Releases Wiki Activity

Patch GPTNeoX to use adequate FA2 if position_ids is provided (#35318)

Browse Source
This commit is contained in:
Taha Yassine 2024-12-23 13:45:55 +01:00 committed by GitHub
parent 5e7aedebeb
commit 2bb60982ac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/models/gpt_neox/modeling_gpt_neox.py
View File

@ -148,6 +148,7 @@ def flash_attention_forward(
norm_factor,
attention_dropout,
training,
position_ids=None,
target_dtype=None,
**_kwargs,
):
@ -173,6 +174,7 @@ def flash_attention_forward(
attention_mask,
query_length,
dropout=attention_dropout,
position_ids=position_ids,
softmax_scale=norm_factor,
is_causal=True,
use_top_left_mask=flash_attn_uses_top_left_mask,
@ -353,6 +355,7 @@ class GPTNeoXAttention(nn.Module):
key,
value,
attention_mask=attention_mask,
position_ids=position_ids,
head_mask=head_mask,
norm_factor=self.norm_factor,
attention_dropout=self.config.attention_dropout,