[Improvement] Formalize KV sharing#47290
Open
remi-or wants to merge 13 commits into
Open
Conversation
|
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. |
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: diffusion_gemma, gemma3n, gemma4, gemma4_unified |
Contributor
CI recapDashboard: View test results in Grafana |
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
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.
This PR aims to formalize / standardize KV cache sharing in
transformersin case more models add KV cache sharing patterns.In recent years, there has been more efforts to reduce the amount of KV cache generated by a model, in the hope of reducing the memory footprint of long-context generation. One technique is to share KV cache between layers, like in YOCO (https://arxiv.org/abs/2405.05254) and CLA (https://arxiv.org/html/2405.12981v1).
Recent models, like those from the Gemma family, were implemented in
transformerswith KV cache sharing. Although the actual implementation is simple and correct, it would not allow for complex KV cache sharing patterns like CLA.To make this possible, we implement the following:
KVCacheSharingStoragewhich function as an alternative to theUserDictthat was previously used to share KV cache across layers. Like the UserDict, it is not an instance of a dict object, so it does not get picked up by torch's FSDP2 helpers, which would lead to a crash. But it also implements atomethod, soacceleratecan dispatch it across device using the same hooks as the rest of the inputs. Furthermore, thetomethod acts in place (documented) so once this container crosses a device boundary, it does not need to be moved again if next layer is on the same device (not the case atm). Finally, this means we can remove the manual.tocalls from the modeling code, making it lighter.DummyLayer, which is used for layers that use KV cache sharing and read their cache from another layer. Before, since only the last layers were shared (YOCO scheme), what was done was to simply truncate the number of cache layers, so that KV sharing layers would not have a corresponding cache layer that would store cache. This led to a discrepancy between the number of layers and the number of cache layers, and it was not comaptible with KV sharing schemes where some layers in the middle of the model share their KV cache (CLA). The DummyLayer solves both these problems.@propertyon the configs with KV cache sharing (and only those), namedkv_sharing_roles, which the Cache object can read when constructing a new cache. There are 3 roles:consumer(the layer reads KV cache from another layer),producer, (the cache produced by this layer is going to be read by another layer), andindependant(the layer does not participate in cache sharing, ie. a "normal" layer). When aCacheobject is created, if the config has this property, then the Cache object createDummyLayerobjects for consumer layers. The distinction betweenproducerandindependentis needed becauseproducerlayer store their KV cache in theKVCacheSharingStoragefor other layers to reuse (important for sliding window layers). This property centralizes everything related to KV cache sharing: to know the exact behavior of all layers, it is the only thing needed.Quite a long description, sorry! The idea here is to standardize KV sharing-related attributes and mechanism so that we have a simple yet adaptable interface so that when more KV-sharing models and schemes get added to
trasnformers, we can support them out of the box without having to change this API. Hence feedback is very much appreciated!