Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(cache): StaticCache uses index_copy_ to avoid useless copy #31857

Merged
merged 7 commits into from
Jul 23, 2024
Merged

feat(cache): StaticCache uses index_copy_ to avoid useless copy #31857

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d329ad2
feat(cache): StaticCache uses index_copy_ to avoid useless copy
tengomucho Jul 9, 2024
53e99d1
feat(cache): SlidingWindowCache uses index_copy_ to avoid useless copy
tengomucho Jul 10, 2024
8e622ec
fix(cache): fallback of index_copy_ when not implemented
tengomucho Jul 11, 2024
aba28b5
fix(cache): in index_copy_ ensure tensors are on same device
tengomucho Jul 12, 2024
02608dd
[run slow] llama
tengomucho Jul 12, 2024
4c818f3
fix(cache): add move of cache_position to same device in SlidingWindo…
tengomucho Jul 12, 2024
a76cfb7
Revert "[run slow] llama"
tengomucho Jul 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 20 additions & 4 deletions src/transformers/cache_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -862,8 +862,18 @@ def update(
k_out.copy_(key_states)
v_out.copy_(value_states)
else:
k_out[:, :, cache_position] = key_states
v_out[:, :, cache_position] = value_states
# Note: here we use `tensor.index_copy_(dim, index, tensor)` that is equivalent to
# `tensor[:, :, index] = tensor`, but the first one is compile-friendly and it does explicitly an in-place
# operation, that avoids copies and uses less memory.
try:
# If using several devices (e.g.: multiple GPUs), we need to ensure everything is on the same one
cache_position.to(device=k_out.device)
k_out.index_copy_(2, cache_position, key_states)
v_out.index_copy_(2, cache_position, value_states)
except NotImplementedError:
# The operator 'aten::index_copy.out' is not currently implemented for the MPS device.
k_out[:, :, cache_position] = key_states
v_out[:, :, cache_position] = value_states
Comment on lines +868 to +876
Copy link
Collaborator

Choose a reason for hiding this comment

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

hope this is not slowing anything down!

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I didn't notice anything when trying it...

Copy link
Collaborator

Choose a reason for hiding this comment

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

then good to go once check quality is green!


return k_out, v_out

Expand Down Expand Up @@ -958,8 +968,14 @@ def update(
k_out = k_out[:, :, indices]
v_out = v_out[:, :, indices]

k_out[:, :, cache_position] = key_states
v_out[:, :, cache_position] = value_states
try:
Copy link
Member

Choose a reason for hiding this comment

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

in the final commit, ensure that the two sets of changes are the same (in the current commit, this one is missing some logic)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

you are right, I forgot it!

cache_position.to(device=k_out.device)
k_out.index_copy_(2, cache_position, key_states)
v_out.index_copy_(2, cache_position, value_states)
except NotImplementedError:
# The operator 'aten::index_copy.out' is not currently implemented for the MPS device.
k_out[:, :, cache_position] = key_states
v_out[:, :, cache_position] = value_states

# `_.zero()` followed by `+=` is equivalent `=`, but compile-friendly (without graph breaks due to assignment)
self.key_cache[layer_idx].zero_()
Expand Down
Loading