Skip to content

Add KeyRerotationPress#31

Merged
maxjeblick merged 25 commits intomainfrom
max/rerotate_keys_2
Dec 12, 2024
Merged

Add KeyRerotationPress#31
maxjeblick merged 25 commits intomainfrom
max/rerotate_keys_2

Conversation

@maxjeblick
Copy link
Copy Markdown
Collaborator

This PR adds KeyRerotationPress which allows to apply key rerotation to any press that inherits from ScorerPress .

I added a relatively lengthy test to ensure that RoPE * prune * RoPE^-1 * RoPE * keys is numerically equivalent to
RoPE * prune * keys (which wasn't obvious to me).

Fixes #16

@maxjeblick maxjeblick requested a review from SimJeg December 11, 2024 11:39
@maxjeblick
Copy link
Copy Markdown
Collaborator Author

@giulio98 feel free to comment on this PR if you have any observations.

@maxjeblick
Copy link
Copy Markdown
Collaborator Author

I should mention that
context_length=(cache.get_seq_length() if isinstance(press, KeyRerotationPress) else context_length),
is currently not correct for quantized cache. This is fixed in #30.

@maxjeblick maxjeblick mentioned this pull request Dec 11, 2024
Comment thread kvpress/presses/key_rerotation_press.py
Comment thread kvpress/presses/key_rerotation_press.py
Comment thread kvpress/presses/key_rerotation_press.py
Comment thread kvpress/presses/key_rerotation_press.py
@SimJeg SimJeg assigned SimJeg and unassigned SimJeg Dec 11, 2024
Comment thread tests/presses/test_key_rerotation_press_rope.py
Comment thread README.md Outdated
Comment thread kvpress/presses/key_rerotation_press.py
Comment thread tests/presses/test_key_rerotation_press_rope.py
Comment thread tests/presses/test_key_rerotation_press_rope.py
@maxjeblick maxjeblick merged commit 2b350b0 into main Dec 12, 2024
@maxjeblick maxjeblick deleted the max/rerotate_keys_2 branch December 12, 2024 09:18
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.

Adjust Positional Embedding For Compression

2 participants