Align MLlama code with Transformers 4.55 #2319
Conversation
|
The code quality check failed, please run |
pbielak
force-pushed
the
dev/pbielak/update-mllama-implementation
branch
2 times, most recently
from
9710363 to
a270769
Compare
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
pbielak
force-pushed
the
dev/pbielak/update-mllama-implementation
branch
2 times, most recently
from
f5c5287 to
939e520
Compare
|
The code quality check failed, please run |
pbielak
force-pushed
the
dev/pbielak/update-mllama-implementation
branch
from
939e520 to
720f6a7
Compare
After the Transformers 4.55 update, one of the attention classes failed to compute the attention scores due to mismatches between arguments in the `torch.matmul` op. This commits updates the whole `Mllama` code base to be fully aligned with the code in Transformers 4.55. In particular, it: - uses the `_attn_implementation` instead of custom classes, - applies the changes from PR [1] - handles `_attn_implementation` passed to the model - fix argument preparation in `gaudi_fused_sdpa_attention` [1] huggingface/transformers#40083
pbielak
force-pushed
the
dev/pbielak/update-mllama-implementation
branch
from
720f6a7 to
79a9ebc
Compare
| ] | ||
| args.prompt = processor.apply_chat_template(conversation, add_generation_prompt=True) | ||
|
|
||
| if model_type == "mllama" and args.use_flash_attention: |
There was a problem hiding this comment.
why not just allow user to select attn_implementation + add readme section about it?
There was a problem hiding this comment.
This is done, because this script is also used by other models (such as llava), which are not yet aligned with the attn_implementation interface.
| ) -> tuple[torch.Tensor, None]: | ||
| bsz, num_heads, tgt_len, head_dim = query.shape | ||
|
|
||
| softmax_mode = "fast" if os.getenv("FLASH_ATTENTION_FAST_SOFTMAX") == "1" else "None" |
There was a problem hiding this comment.
since attention implementation is now separated from model we don't have to use env vars, can you explore if it's possible to use kwargs instead of env vars?
There was a problem hiding this comment.
Makes sense - I will have a look at it
3 participants
What does this PR do?
After the Transformers 4.55 update, one of the attention classes failed to compute the attention scores due to mismatches between arguments in the
torch.matmulop. This commits updates the wholeMllamacode base to be fully aligned with the code in Transformers 4.55. In particular, it uses the_attn_implementationinstead of custom classes.