Fix adapter KV cache reuse during training - #2327
Open
Excelius-Wang wants to merge 1 commit into
Open
Excelius-Wang wants to merge 1 commit into
Excelius-Wang wants to merge 1 commit into
Conversation
Excelius-Wang
requested review from
andyland,
k223kim,
lianakoleva and
t-vi
as code owners
September 9, 2026 13:57
This branch has not been deployed
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.
Adapter attention caches its projected K/V even when gradients are enabled. A second training backward can therefore reuse a freed graph. If validation runs first under
no_grad, training instead reuses detached K/V and the adapter embeddings receive no gradients. This affects both Adapter and Adapter V2.Only cache adapter K/V when the model KV cache is enabled and gradients are disabled. Recompute and discard old adapter K/V otherwise. Extend the Adapter model's existing
clear_kv_cache()to clear adapter K/V as well; Adapter V2 inherits the correction.Regression coverage includes repeated optimization, validation/inference-to-training transitions, gradients in eval mode, cached/uncached output equivalence, cache reuse and cache clearing for both adapter variants. CPU tests use small original models with nonzero gates for output comparisons. No pretrained quality or performance improvement is claimed.
Fixes #1287. This is a focused correction to the existing cache behavior; the previously linked broad KV-cache refactor #2061 was closed without merging.
Validation: final Adapter/V2 suite 253 passed, 8 skipped on CPU (PyTorch 2.13, Transformers 4.56.2). All 14 new regression cases fail on baseline and pass with the fix. A separate probe calls the official finetune validate functions before training: baseline loses the embedding gradient, while both fixed variants update their embeddings. Changed-file Ruff checks, formatting and git diff --check pass. GPU-specific checks were skipped.