refactor transformer decoder and revamp the left padding attention mask #178
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
There was a problem hiding this comment.
Nicely done! This simplification had been overdue for quite some time.
(Did not test changes on my end yet, but I trust the existing tests should cover most of the changes here.)
Minor comments to address.
Also, for reference, what is the rationale behind embedding the left_pad flag within the batch?
When using left_padding we need to mask the pad tokens at EACH time step when decoding. Also, previously it was ugly hardcoded in the generator. |
vince62s
added a commit
to vince62s/eole
that referenced
this pull request
Jan 14, 2025
vince62s
added a commit
that referenced
this pull request
Jan 14, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The initial goal was to gather transformer_lm and transformer decoder type in a single class
After carefull review I realized that left padding was not handled 100% properly.
After this PR:
For training:
we rely on pytorch
scaled_dot_product_attentioninmemory_efficientmode.we may need to investigate the
flash_varlen_func()(which seems slower for now). see Dao-AILab/flash-attention#1125==>
self_attn_backend: has no impact at training at the moment.For inference:
For EncoderDecoder models we rely also on pytorch
scaled_dot_product_attentioninmemory_efficientmode since it is much faster thanflash_attn_with_kvcache()due to short sequences / large batches.For LLMs (Decoder only models) we rely on
flash_attn_with_kvcache()which much faster.When we use SDPA we always buid ourselves the attention_mask and never rely on is_causal=True
The main reason is that we need to handle padding properly and not attend padded tokens.
With
flash_attn_with_kvcache()we added a "recent" functionalitycache_leftpadwe requires flash-attn > 2.6.1Results are much more consistent between the two options now.
==>
self_attn_backend: if "flash" thenflash_attn_with_kvcache()will be used, if "pytorch" then sdpa in memory_efficient mode will be used.