ENH Add KaSA (Knowledge-aware Singular-value Adaptation) as a LoRA variant - #3446
Conversation
Implement KaSA (Knowledge-aware Singular-value Adaptation, arXiv:2412.06071) using the LoRA-variant framework, following the SVD-based variants (CorDA/DoRA). KaSA changes vanilla LoRA in two ways: - A one-time, destructive SVD truncation of the frozen base weight that drops its r smallest singular components, leaving the rank-(k-r) approximation as the new frozen base (k = min(in_features, out_features)). - A learnable diagonal of singular values (lora_diag) inserted between the LoRA A and B factors, so the update is ΔW = scaling * B @ diag(lora_diag) @ A. - New KasaConfig sub-config (beta, gamma) and LoraConfig.kasa_config field; selection is driven by kasa_config being non-None via resolve_lora_variant, with explicit guards rejecting KaSA on embedding/conv/MHA/ParamWrapper and fan_in_fan_out layers. - KasaLinearVariant implements init (SVD truncation + lora_diag), forward, merge_safe/merge_unsafe/unmerge. lora_diag is registered in adapter_layer_names so it is saved/loaded. - get_kasa_regularization_loss helper exposes the paper's two auxiliary terms (L2 singular-value penalty + L3 orthogonal regularization), since the variant forward has no channel to inject an extra loss into the training loop. - Tests in tests/test_kasa.py (SVD-truncation faithfulness, lora_diag shape, zero-init update, merge/unmerge round-trip, delta-weight formula, save/load, regularization closed-form checks) plus wiring in tests/test_lora_variants.py. Faithfulness notes: - The base-weight truncation is destructive; disabling/unloading does not restore the original weight and merge/unmerge round-trips to the truncated base. This is inherent to the method and documented. - The paper's L2/L3 regularizers are required for the SVD interpretation to hold but cannot be auto-injected; users must add get_kasa_regularization_loss to their loss.
Re-express the KaSA integration in the new declarative LoRA-variant framework introduced by huggingface#3219: - Register KaSA in Linear.lora_variants as ("kasa_config",) -> KasaLinearVariant instead of overriding resolve_lora_variant. - Drop the explicit embedding/conv guards; unsupported layers are now rejected generically by the base resolve_lora_variant (the variant is absent from those layers' lora_variants mappings). - Keep explicit guards for MultiheadAttention and ParamWrapper, mirroring VeLoRA. - Tag kasa_config with is_lora_variant metadata, as the new framework requires. - Re-apply the KaSA test additions in tests/test_lora_variants.py on top of the refactored test module.
- docs/source/package_reference/lora_variant_kasa.md following the existing 'Variant: ...' page pattern, plus the _toctree.yml entry. - examples/kasa_finetuning with a runnable SFT script; it subclasses SFTTrainer to add get_kasa_regularization_loss to the task loss, since the paper's auxiliary regularizers cannot be injected by PEFT. - method_comparison MetaMathQA experiment config (rank32 + KasaConfig defaults), mirroring the existing rank32 LoRA experiment.
…orward Three bugs in the low_cpu_mem_usage=True (meta-device) path, where the destructive SVD truncation of the base weight is deferred past init: 1. Merging before the first forward (e.g. merge_and_unload right after from_pretrained) merged the delta into the un-truncated base weight, and because merged layers skip the variant forward, the truncation then never ran at all - silently wrong outputs. 2. The first forward crashed on mixed dtypes (bf16 base with the default fp32-autocast adapter): the deferred branch recomputed the base forward with the input already cast to the adapter dtype. 3. The recompute overwrote the accumulated result, discarding the contributions of other active adapters processed earlier in the same forward. Fix: extract _apply_deferred_truncation (shared by forward, merge_safe and merge_unsafe; merge_safe refreshes its stale pre-truncation clone), and have _truncate_base_weight return the dropped low-rank factors so the first forward subtracts exactly the dropped component in fp32 instead of recomputing the base forward - preserving other adapters' contributions and avoiding the dtype mismatch. Each fix has a regression test that fails without it.
|
@BenjaminBossan when this merges, could you add Co-authored-by: Bogeum Kim 53206051+iambogeumkim@users.noreply.github.com to the squash commit, so @iambogeumkim gets co-author credit for their groundwork? |
There was a problem hiding this comment.
🤗 Serge says:
KaSA is implemented cleanly and follows the existing variant framework conventions well. The deferred-truncation path for low_cpu_mem_usage=True is carefully handled across forward, merge-safe, and merge-unsafe, with dedicated regression tests. The auxiliary regularization helper is correctly computed in fp32 and matches the paper's formulas. Below are two issues worth addressing before merge.
Correctness
get_kasa_regularization_lossreturnstorch.zeros(())(a CPU tensor) when the model contains no KaSA layers. If the user's task loss lives on GPU,task_loss + get_kasa_regularization_loss(model)will raise a device-mismatchRuntimeError. Since the function explicitly handles the "no KaSA layers" case, it should return a tensor on the right device.convert_to_lorawill silently droplora_diagfor KaSA layers becauseLinear.get_delta_weightcomputes the vanillaB @ Adelta. The PR description acknowledges this and notes DoRA has the same behavior onmain, so it is not a regression — but users who runconvert_to_loraon a KaSA model will get a silently incorrect approximation with no warning. Consider overridingsupports_lora_conversionto returnFalsefor KaSA-adapted layers, or at least emitting a warning, in a follow-up.
Tests
- Test coverage is strong: config wiring, SVD truncation rank, merge/unmerge round-trip, save/load with re-truncation, deferred-truncation with bf16 base, merge-before-forward, multi-adapter interaction, and regularization gradients. The four regression tests for the
low_cpu_mem_usagedeferred path are particularly valuable.
serge v0.1.0 · model: zai-org/GLM-5.2 · 40 LLM turns · 46 tool calls · 397.9s · 2030530 in / 22370 out tokens
|
@robbiebusinessacc Thanks for the PR. I didn't have time for a review yet, but I let our review bot take a look and the feedback seems to be helpful, so I posted its review here. From my side, I'd add that you could add a KaSA experiment to the image-gen benchmark as well, but that's optional. |
…en config - get_kasa_regularization_loss now returns its zero on the model's device when there are no KaSA layers, so adding it to a GPU task loss does not raise a device-mismatch error. - Adding a KaSA adapter now makes supports_lora_conversion return False: the update depends on lora_diag and on the destructive base truncation, neither of which is representable in a vanilla LoRA adapter. Implemented as an opt-out flag on LoraVariant so other variants keep their current behavior. - Add a KaSA experiment config to the image-gen benchmark, mirroring the DoRA rank64 experiment.
|
Thanks. Both of sergereview's points are addressed in 3844f16: the regularization helper now returns its zero on the model's device, and rather than leaving the |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for reviving the effort to add KaSA to PEFT. The PR is in a good state. I focused in this review on the pure integration, so I haven't checked the example, docs, and KaSA algorithm in detail yet.
Besides my comments, I think we also need to adjust LoraModel._check_new_adapter_config: We can't really have multiple adapters when using KaSA (unless they all use KaSA), right? Please add checks there and also the appropriate tests.
…st reorganization - LoraModel._check_new_adapter_config now rejects mixing KaSA and non-KaSA adapters on one model (multiple KaSA adapters remain allowed), since the destructive truncation changes the base weights under the other adapters' feet. - Replace the standalone get_kasa_regularization_loss helper with LoraModel._get_kasa_loss, mirroring _get_monteclora_loss: coefficients are read from each adapter's KasaConfig (removing the layer-side _lora_kasa_config stash) and 0.0 is returned when KaSA is not used. Public exports, docs, and the example are updated accordingly. - Remove the supports_lora_conversion attribute from the LoraVariant base class; Linear.supports_lora_conversion now defaults to True and only honors an opt-out set on a variant class (KaSA). - Move the user-facing behavior and config-validation tests from tests/test_kasa.py into TestKasaInitialization (tests/test_initialization.py), the regularization tests into TestKasaRegularization (tests/test_lora_variants.py), and the conversion test into tests/test_lora_conversion.py; delete tests/test_kasa.py. Add KaSA to the test_custom_models.py matrix with a skip for test_disable_adapters (disabled output intentionally differs from the base model). The multi-adapter deferred-truncation test now uses two KaSA adapters, matching the new mixing constraint. - Docstring cleanups: single backticks, user-facing notes moved from KasaLinearVariant to KasaConfig, and a comment clarifying when the deferred truncation triggers in forward. Reword the docs' 'spectral domain' phrasing to the paper's framing.
|
Thanks for the thorough review. All points are addressed in e7e6e76:
Ran |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the updates, the PR is in a good state. My comments mostly concern details of the implementation and test, please check.
I haven't checked the example and experiments yet, this will come next once the integration is finished.
|
|
||
| # The KaSA update depends on lora_diag and on the destructive truncation of the base weight, neither of which is | ||
| # representable in a vanilla scaling * B @ A adapter on the unmodified base. | ||
| supports_lora_conversion = False |
There was a problem hiding this comment.
Let's make this a normal function, same as on the LoraLayer, for consistency. Moreover, on the LoraVariant parent class, this method should be added (with default True) to make it clear that it can be subclassed.
Adjust the call site accordingly.
There was a problem hiding this comment.
Done — it's a method on LoraVariant returning True by default, KasaLinearVariant overrides it, and the call site calls it as a method.
| if adapter_name not in lora_diag: | ||
| continue | ||
| kasa_config = self.peft_config[adapter_name].kasa_config | ||
| layer_loss = _kasa_layer_regularization_loss(module, adapter_name, kasa_config.beta, kasa_config.gamma) |
There was a problem hiding this comment.
Let's use kwargs after module for clarity.
| if num_kasa_layers == 0: | ||
| return 0.0 |
There was a problem hiding this comment.
If num_kasa_layers == 0, total will already be 0.0, right? So there is no need to track this.
There was a problem hiding this comment.
Right, removed.
| # Give the adapter a non-trivial value so merge/forward differences are observable (B and diag both non-zero). | ||
| with torch.no_grad(): | ||
| for name, param in model.named_parameters(): | ||
| if "lora_B" in name: |
There was a problem hiding this comment.
lora_B can already be made non-zero with init_lora_weights=False. As for lora_diag, it is already randomly initialized, is it not?
There was a problem hiding this comment.
Right on both counts — dropped the helper, the tests that need a non-trivial adapter now use init_lora_weights=False.
| # forward (see KasaLinearVariant.forward), which mirrors the deferral pattern used by Monteclora. Without | ||
| # this re-trigger the SVD truncation would be silently skipped on the low_cpu_mem_usage path and the model | ||
| # would compute the wrong thing (full base + an adapter trained against the truncated base). | ||
| module.lora_diag[adapter_name] = nn.Parameter(torch.randn(r, device=device, dtype=dtype)) |
There was a problem hiding this comment.
Why initialize to random in case of meta device?
There was a problem hiding this comment.
No good reason — a meta tensor has no values anyway. Changed it to torch.empty; the randn now only happens when the deferred truncation is applied and the values actually materialize (or they come from the loaded state dict), mirroring the regular init path.
| out2 = reloaded(x) # subsequent forwards: steady state | ||
| assert torch.allclose(out1, out2, atol=1e-6) | ||
|
|
||
| def test_kasa_mixing_with_non_kasa_adapter_raises(self): |
There was a problem hiding this comment.
Let's parametrize the test to also test the other order
There was a problem hiding this comment.
Done, parametrized over both orders.
- supports_lora_conversion is now a method on the LoraVariant base class (default True) that KasaLinearVariant overrides; the Linear call site calls it as a method. - _get_kasa_loss: keyword arguments for the layer-helper call and drop the redundant num_kasa_layers tracking (total already starts at 0.0). - Use torch.empty for the meta-device lora_diag placeholder; the values either come from a loaded state dict or are randn-initialized when the deferred truncation is applied. - Tests: drop the randomize_adapter helper in favor of init_lora_weights=False, extend the merge-before-forward test with merge_and_unload, and parametrize the adapter-mixing test over both orders.
|
Thanks, all seven points are addressed in 65998dc. Same test results as before: |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the new push, the PR is almost done. There is a failing test, could you please check?
I could run the finetuning example and also started the benchmarks. When it comes to the MetaMathQA benchmark, although it's accuracy wasn't particularly high (47%), the forgetting score was 0.013, which is the lowest (and thus best) score among all LoRA experiments, which is a nice confirmation. I had some issues with the image-gen benchmark, but it's probably just my machine, so no worries. Could you run it successfully?
- test_get_base_model_state_dict (new on main) compares the extracted base state dict to the original base weights, which KaSA's destructive truncation intentionally changes; skip it for KaSA like test_disable_adapters. - Move the image-gen KaSA experiment to experiments/lora/ flux2-klein-rank64-kasa, matching the MetaMathQA naming scheme and the new upstream flux2-klein-rank64 LoRA experiment.
|
The failing test is I also moved the image-gen experiment to |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the update, LGTM.
I revisited the image-gen issue and found the problem. At first it seemed like it was hanging for me and I thought it was an unrelated issue, e.g. HF Hub problems. But it was not actually hanging, the problem was that the SVD for KaSA was running on CPU!
In the image-gen benchmark, the transformer module is moved to the GPU late so that we can first precompute with, and then offload, other components without having a memory peak. This means that when get_peft_model is called, the transformer is still on CPU, causing KaSA to run SVD on the weights on CPU which took forever.
The solution is to move the transformer to GPU before calling get_peft_model and then offload it again, before moving it again to GPU later. It's a bit wasteful but better than hitting a memory peak. I will create a separate PR to make these changes, for this PR, there is nothing more to do.
Regarding the results: I got DINO similarity of 0.64 and drift of 0.32. "Drift" is similar to "forgetting" in the MetaMathQA benchmark, so I expected it to be lower (i.e. better) than normal LoRA, but normal LoRA with rank 64 gets drift of 0.29 while similarity is 0.743 (go here, select task "image-gen", filter by peft_type=="LORA"). Qualitatively, sample images also don't look all that great.
Anyway, the PR is still ready to be merged, if you want to further investigate the mediocre image-gen results and find a way to improve it, just create another PR.
|
Thanks for digging into the image-gen issue, that explanation makes sense — the SVD at init is the one expensive step KaSA has, so on CPU it really drags. Moving the transformer to the GPU before On the image-gen results: one thing I'd like to look at in a follow-up is that the benchmark trains without KaSA's auxiliary regularizers (the harness doesn't add |
That makes sense, thanks. |
Supersedes #3298, which was auto-closed by the stale bot and can no longer be reopened now that the branch has moved on.
Implements KaSA (Knowledge-aware Singular-value Adaptation, arXiv:2412.06071, ICLR 2025) as a LoRA variant, following the SVD-based variants. Continues the work from #2516 with @iambogeumkim's blessing and maintainer approval — their earlier PRs (#2543, #2698, #3346) did the original groundwork here.
What KaSA changes vs vanilla LoRA
rsmallest singular components (destructive by design; documented in the caveats).lora_diag) between the LoRA A and B factors:ΔW = scaling * B @ diag(ΔΣ) @ A.get_kasa_regularization_lossfor users to add to their task loss, since the variant API has no loss-injection channel.Scope: core variant + config, tests (
tests/test_kasa.py, additions totests/test_lora_variants.py), docs page,examples/kasa_finetuning/, and a MetaMathQA benchmark config. Rebuilt on the new variant-resolution framework from #3219 (declarativelora_variantsmapping). KaSA is not added to thetest_custom_models.pymatrix for the same reason PiSSA/OLoRA aren't: the destructive base-weight change breaks the matrix's output-restoration invariants; dedicated tests cover save/load, retruncation, merge-before-forward, and thelow_cpu_mem_usagedeferred path.Tests run:
pytest tests/test_kasa.py tests/test_lora_variants.py(43 passed, incl. 4 regression tests for bugs found in thelow_cpu_mem_usagedeferred-truncation path — merge-before-forward, bf16 base with fp32 adapter, multi-adapter interaction),pytest tests/test_config.py(1193 passed), andmake style/make quality, all green locally on CPU.Note for reviewers:
convert_to_loracurrently produces the vanillaB @ Adelta for variant layers, which for KaSA dropslora_diag(DoRA onmainhas the same behavior). Happy to address that here or in a follow-up if you have a preferred direction.Disclosure per the repo's contribution policy: this PR was developed with AI assistance. I've reviewed the changes and I'm accountable for them.