riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-02 19:21:31 +06:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

Fix new FA2 if is_causal is passed explicitly (#35390)

Browse Source 
* fix

* Update modeling_decision_transformer.py

* Update flash_attention.py
This commit is contained in:
Cyril Vallez 2024-12-22 20:00:07 +01:00 committed by GitHub
parent 8f38f58f3d
commit 05260a1fc1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/integrations/flash_attention.py
View File

@ -44,6 +44,9 @@ def flash_attention_forward(
else:
target_dtype = next(layer for layer in module.modules() if isinstance(layer, torch.nn.Linear)).weight.dtype
# FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice
kwargs.pop("is_causal", None)
attn_output = _flash_attention_forward(
query,
key,

6
src/transformers/models/decision_transformer/modeling_decision_transformer.py
View File

@ -285,9 +285,9 @@ class DecisionTransformerGPT2Attention(nn.Module):
shape_q = (*query_states.shape[:-1], -1, self.head_dim)
shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
query_states = query_states.reshape(shape_q).transpose(1, 2)
key_states = key_states.reshape(shape_kv).transpose(1, 2)
value_states = value_states.reshape(shape_kv).transpose(1, 2)
query_states = query_states.view(shape_q).transpose(1, 2)
key_states = key_states.view(shape_kv).transpose(1, 2)
value_states = value_states.view(shape_kv).transpose(1, 2)
if layer_past is not None:
past_key, past_value = layer_past

6
src/transformers/models/gpt2/modeling_gpt2.py
View File

@ -295,9 +295,9 @@ class GPT2Attention(nn.Module):
shape_q = (*query_states.shape[:-1], -1, self.head_dim)
shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
query_states = query_states.reshape(shape_q).transpose(1, 2)
key_states = key_states.reshape(shape_kv).transpose(1, 2)
value_states = value_states.reshape(shape_kv).transpose(1, 2)
query_states = query_states.view(shape_q).transpose(1, 2)
key_states = key_states.view(shape_kv).transpose(1, 2)
value_states = value_states.view(shape_kv).transpose(1, 2)
if layer_past is not None:
past_key, past_value = layer_past