Skip to content

Fix adapter KV cache reuse during training - #2327

Open
Excelius-Wang wants to merge 1 commit into
Lightning-AI:mainfrom
Excelius-Wang:fix-adapter-training-cache
Open

Excelius-Wang wants to merge 1 commit into
Lightning-AI:mainfrom
Excelius-Wang:fix-adapter-training-cache

Conversation

@Excelius-Wang

Copy link
Copy Markdown

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.

This branch has not been deployed

No deployments
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.

Is it correct to keep using adapter_kv_cache during training in litgpt/adapter.py?

1 participant