Skip to content

Fix batched generation which is dependent on left padding - BREAKING! 😮‍💨#5

Open
vasqu wants to merge 1 commit into
mainfrom
batching-fix-for-using-left-padding
Open

Fix batched generation which is dependent on left padding - BREAKING! 😮‍💨#5
vasqu wants to merge 1 commit into
mainfrom
batching-fix-for-using-left-padding

Conversation

@vasqu

@vasqu vasqu commented Sep 26, 2024

Copy link
Copy Markdown
Owner

I will probably not merge this PR but keep it as is. It fixes a huge issue within mamba-related models: Generation with padding. This is due to the fact that the eos token is used as pad token which results in unwanted repeated influence of it.

This fix requires the tokenizer to be left padding. But, for right padding, there is an interesting use case to enable correct last ssm states even with padding. See the tests...

…t's breaking the previous behaviour with the tokenizer

also supports a special case for right padding --> return the correct last ssm state even within a batch
Comment on lines +812 to +817
if (
model_kwargs.get("use_cache", True)
and "cache_position" in model_kwargs
and model_kwargs["cache_position"] is not None
):
model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could likely ignore cache_position as it's used for the cache to be torch compile compatible but since the mamba2 kernel can't really compile either way...

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant