[WIP] Update `LoraConfig` for KaSA implementation #2698

nsbg · 2025-08-02T05:41:15Z

I was delayed in updating the code because I was focusing on company work, but now I'm planning to resume the project in earnest. If I have any questions about implementing the code, may I continue to ask you?

I apologize for opening a new pull request, as the previous one was closed 🥲 Thank you for your understanding.

BenjaminBossan

Thank you for resuming your work on KaSA.

Implementation-wise, we need to take a different approach. Right now, KaSA is just added to the normal LoRA code, but we only want to activate it if the user opts in. Therefore, it should be implemented in a separate class, something like KasaVariant, in peft/tuners/lora/variants.py. Please check how DoRA is implemented and use a similar approach, as I have detailed in my previous comment. If anything is unclear, feel free to ask.

github-actions · 2025-09-01T15:04:02Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

BenjaminBossan · 2025-09-01T15:12:53Z

gentle ping @nsbg

… variants.py

nsbg · 2025-09-02T12:14:03Z

Thank you for your alert!

I spent some time looking over the KaSA paper and code to get ready for more serious work, but it does seem pretty difficult 🥲 My goal is to upload code that's ready for review before the end of September, so I'm going to try even harder.

Right now, I'm stuck at the 'Extend LoRA variant resolution' stage you mentioned. Honestly, this seems like the most important part, but it's hard for me to figure out where to start—specifically, which file and class I should work on first. Could you help me with this?

BenjaminBossan · 2025-09-02T14:04:09Z

That's great to see, thanks for picking this back up.

Right now, I'm stuck at the 'Extend LoRA variant resolution' stage you mentioned. Honestly, this seems like the most important part, but it's hard for me to figure out where to start—specifically, which file and class I should work on first. Could you help me with this?

You're already on the right track, you added KasaLinearVariant, which is the most important step. There are definitely some changes required there, as there is some code that is only relevant for DoRA and can be removed for KaSA. But we can leave that as is for now.

Next about resolving the variants. As a first step, let's revert the changes you made to lora/layer.py and start fresh. We don't need a self.use_kasa attribute, we only have self.use_dora for backwards compatibility, as we didn't have LoRA variants when we first implemented DoRA.

Then let's look at these lines in lora.Linear:

peft/src/peft/tuners/lora/layer.py

Lines 636 to 642 in a3197b1

    
           def resolve_lora_variant(self, *, use_dora: bool, **kwargs) -> Optional[LoraVariant]: 
        
               if not use_dora: 
        
                   return None 
        
               from .variants import DoraLinearVariant 
        
               return DoraLinearVariant()

Here we need to extend the functionality to add KaSA. The updated method could be something like:

    def resolve_lora_variant(self, *, use_dora: bool, use_kasa: bool, **kwargs) -> Optional[LoraVariant]:
        if use_dora and use_kasa:
            raise ValueError("Cannot use DoRA and KaSA at the same time, please choose only one.")

        variant = None
        if use_dora:
            from .variants import DoraLinearVariant

            variant = DoraLinearVariant()
        elif use_kasa:
            ...

        return variant

Does that make sense? Similarly, we'd have to update the resolve_lora_variant methods of other LoRA layers, depending on whether they work with KaSA or not (I'm not sure if KaSA works with Conv2d etc.).

I would suggest that you work on this as a next step, then we'll see what else needs to be done.

nsbg · 2025-09-04T14:57:54Z

wow I really appreciate your sincere feedback. I'll read your advice carefully and then move forward 🤗

nsbg · 2025-09-08T18:39:23Z

@BenjaminBossan I modified the code in the files below based on what you explained. Please give me feedback if there are parts that still need fixing, and then we can discuss the next steps.

1. variants.py

Completed updates to methods in the KasaLinearVariants class

2. layer.py

In the LoraLayer class, added self.use_kasa[adapter_name] = use_kasa inside the update_layer method
In the Linear class, added KaSA handling logic inside the get_delta_weight method

BenjaminBossan

Thanks for integrating my feedback. I gave this another review and noted the next few changes that are necessary. Please check my comments.

Apart from this, the branch is now encountering merge conflicts. Could you please bring your fork up-to-date with the remote and then merge with, or rebase on, the latest main branch from PEFT? If you have questions on how to resolve the merge conflicts, don't hesitate to ask.

Furthermore, please always run make style on your changes before pushing to make our linter happy.

More of a note for myself: Since KaSA updates the base weights of the model, we will have to take extra care to ensure that it works correctly when saving and loading the adapter.

BenjaminBossan · 2025-09-09T13:44:26Z

src/peft/tuners/lora/layer.py


        """
-        return None
+        if use_dora and use_kasa:


Let's undo the changes in this method body and return None. Instead, since this KaSA layer is implemented for Linear only, add the logic to lora.Linear.resolve_lora_variant instead.

Also, we should update the resolve_lora_variant methods of the other layer types like lora.Embedding.resolve_lora_variant to accept the use_kasa argument but raise an error if it's True. Otherwise, users may add it to non-supported layers and not notice that it doesn't actually do anything there.

BenjaminBossan · 2025-09-09T13:45:08Z

src/peft/tuners/lora/layer.py

+        ############ kasa ############# 
+        self.lora_diag[adapter_name] = nn.Parameter(torch.randn(r), requires_grad=True)
+
+        weight = self.get_base_layer().weight
+        dtype = weight.dtype
+        svd_rank = self.in_features - r
+        weight = weight.to(torch.float32)
+        U, S, Vh = torch.linalg.svd(weight.data, full_matrices=False)
+        U_principle, S_principle, Vh_principle = U[:, :svd_rank], S[:svd_rank], Vh[:svd_rank, :]
+        self.get_base_layer().weight.data = (U_principle @ torch.diag(S_principle) @ Vh_principle).to(dtype)

+        #########################


All of this can be removed, since it's part of KasaLinearVariant.init, right?

BenjaminBossan · 2025-09-09T13:47:30Z

src/peft/tuners/lora/variants.py

+        # initialize lora_diag
+        module.lora_diag[adapter_name] = nn.Parameter(torch.randn(module.r[adapter_name]), requires_grad=True)
+
+        # SVD


Let's add a reference here, so that we know the origin:
# see https://github.com/juyongjiang/KaSA/blob/f85e88c22d0fa4cb8ab2923d7c2bf1bbec152da3/peft/src/peft/tuners/lora/layer.py#L132

# initialize lora_diag module.lora_diag[adapter_name] = nn.Parameter(torch.randn(module.r[adapter_name]), requires_grad=True) # see https://github.com/juyongjiang/KaSA/blob/f85e88c22d0fa4cb8ab2923d7c2bf1bbec152da3/peft/src/peft/tuners/lora/layer.py#L132 # SVD

I put it in here, how is it?

BenjaminBossan · 2025-09-09T13:51:30Z

src/peft/tuners/lora/variants.py

+    @staticmethod
+    def merge_safe(module: Linear, active_adapter: str, orig_weight: torch.Tensor) -> torch.Tensor:
+        delta_weight = module.get_delta_weight(active_adapter)
+        return orig_weight + delta_weight
+
+    @staticmethod
+    def merge_unsafe(module: Linear, active_adapter: str, orig_weight: torch.Tensor) -> None:
+        delta_weight = module.get_delta_weight(active_adapter)
+        orig_weight.data += delta_weight
+
+    @staticmethod
+    def unmerge(module: Linear, active_adapter: str, orig_weight: torch.Tensor) -> torch.Tensor:
+        delta_weight = module.get_delta_weight(active_adapter)
+        return orig_weight - delta_weight


KaSA should have an influence on the merged weights, should it not?

Although this PR is closed, it seems I've incorporated everything else except for this comment (of course, you'd have to look at the code). Could you explain this question in more detail?

BenjaminBossan · 2025-09-09T13:53:07Z

src/peft/tuners/lora/variants.py

+            x = dropout(x)
+
+        # KaSA calculation
+        lora_output = lora_B(torch.einsum('ijk,kl->ijl', lora_A(x), diag)) * scaling


Again, let's add a reference:
# see https://github.com/juyongjiang/KaSA/blob/f85e88c22d0fa4cb8ab2923d7c2bf1bbec152da3/peft/src/peft/tuners/lora/layer.py#L602C21-L602C110

# KaSA calculation # see https://github.com/juyongjiang/KaSA/blob/f85e88c22d0fa4cb8ab2923d7c2bf1bbec152da3/peft/src/peft/tuners/lora/layer.py#L602C21-L602C110 lora_output = lora_B(torch.einsum('ijk,kl->ijl', lora_A(x), diag)) * scaling return result + lora_output

I inserted this near where the actual calculation logic begins, rather than just in an empty space. I think this is a bit better.

nsbg · 2025-09-16T17:02:08Z

@BenjaminBossan oh I didn't mean to close the branch, but it seems to have closed while I was merging with the main branch. I guess I'll have to open a new PR, right? 😰

+) when I tried to sync with the main branch, I ended up discarding all my commits, so did that cause it to close?

BenjaminBossan · 2025-09-17T08:55:18Z

oh I didn't mean to close the branch, but it seems to have closed while I was merging with the main branch. I guess I'll have to open a new PR, right? 😰

+) when I tried to sync with the main branch, I ended up discarding all my commits, so did that cause it to close?

I don't know what happened, but I could re-open the PR and there are some changes visible. Can you double check that everything looks as expected? If for some reason it's not what it's expected, you can create a new PR and push your local branch.

nsbg · 2025-09-17T09:08:15Z

I usually handle merges in the terminal, and I suspect the pull request was closed because I accidentally wiped the commit history while using the 'Sync fork' feature on GitHub. I'll be more careful in the future. Thanks for reopening it.

I'll review the changes and open a new PR if needed. Sorry to keep bothering you with this.

BenjaminBossan · 2025-09-18T09:35:02Z

I'll review the changes and open a new PR if needed. Sorry to keep bothering you with this.

No worries. If the diff on this PR looks good, let me know and I'll do a review. Only open a new PR if for some reason, the code here does not correspond to what it should be.

nsbg · 2025-09-20T07:25:12Z

@BenjaminBossan I checked layer.py/variants.py and KasaLinearVariants class in variants.py was removed. I added it again and I updated file based on your minor feedback, so I think we can discuss in this PR continually.

BTW I ran make style command and got this error.

make style
ruff check --fix src tests examples docs scripts docker
process_begin: CreateProcess(NULL, ruff check --fix src tests examples docs scripts docker, ...) failed.

I ran pip install -e .[test] command in https://huggingface.co/docs/peft/install#source, but I got same error. Do I just run that command directly without needing to set up a virtual environment?

nsbg · 2025-09-20T08:08:42Z

maybe make style related error was fixed. After applying this command, quite a few files have changed. Is it okay to just push them? Also, what exactly does make style do?

BenjaminBossan · 2025-09-23T12:28:36Z

maybe make style related error was fixed. After applying this command, quite a few files have changed. Is it okay to just push them? Also, what exactly does make style do?

No, let's not push any changes to unrelated files. If make style changes unrelated files, it's often to one of these reasons:

Wrong ruff version: check that v0.12.12 is installed in your virtual environment.
Not picking up the config: There are some settings for ruff in the pyproject.toml file, ensure that it's there when you run make style.

nsbg · 2025-09-27T06:52:18Z

@BenjaminBossan

Also, we should update the resolve_lora_variant methods of the other layer types like lora.Embedding.resolve_lora_variant to accept the use_kasa argument but raise an error if it's True. Otherwise, users may add it to non-supported layers and not notice that it doesn't actually do anything there.

I referred to your explanation and added the use_kasa parameter to the resolve_lora_variant method in the classes below.

Linear (line 702)
Embedding (line 934)
_ConvNd (line 1263)
Conv2d (line 1507)
Conv1d (line 1524)
Conv3d (line 1541)

The logic for raising errors in each layer hasn’t been applied yet, but I committed first to check whether adding the parameter in this way matches what you meant. Excluding the Linear class, it seems like an error should be raised when use_kasa is true in the other classes. However, I might be mistaken, so please feel free to give me feedback anytime. Also, I noticed there’s no part that calls KasaLinearVariant—should this be called inside the linear class? I’m a bit confused about this part.

BenjaminBossan

Thanks for the updates, we're making good progress here.

The logic for raising errors in each layer hasn’t been applied yet, but I committed first to check whether adding the parameter in this way matches what you meant. Excluding the Linear class, it seems like an error should be raised when use_kasa is true in the other classes.

Nice, this looks correct, please raise the error in the unsupported layers as indicated.

Also, I noticed there’s no part that calls KasaLinearVariant—should this be called inside the linear class? I’m a bit confused about this part.

Don't worry, it is being called. E.g. in lora.Linear.forward we have this code:

peft/src/peft/tuners/lora/layer.py

Lines 806 to 816 in 190f987

    
           if active_adapter not in self.lora_variant:  # vanilla LoRA 
        
               result = result + lora_B(lora_A(dropout(x))) * scaling 
        
           else: 
        
               result = self.lora_variant[active_adapter].forward( 
        
                   self, 
        
                   active_adapter=active_adapter, 
        
                   x=x, 
        
                   result=result, 
        
                   **variant_kwargs, 
        
                   **kwargs, 
        
               )

So if the KaSA variant is found, KasaLinearVariant.forward will be used here. Same for the other methods.

src/peft/tuners/lora/layer.py

BenjaminBossan · 2025-09-29T12:58:29Z

src/peft/tuners/lora/variants.py

+        U, S, Vh = torch.linalg.svd(weight.data, full_matrices=False)
+        U_principle, S_principle, Vh_principle = U[:, :svd_rank], S[:svd_rank], Vh[:svd_rank, :]
+        module.get_base_layer().weight.data = (U_principle @ torch.diag(S_principle) @ Vh_principle).to(dtype)
+


Let's add a new @staticmethod called _get_delta_weight here. This method should implement the KaSA delta weight logic from above:

diag = torch.diag(module.lora_diag[adapter]) output_tensor = transpose(weight_B @ diag @ weight_A, module.fan_in_fan_out) * module.scaling[adapter]

Then, the merge_safe, merge_unsafe, and unmerge methods below can call KasaLinearVariant._get_delta_weight(...).

src/peft/tuners/lora/variants.py

nsbg added 2 commits May 18, 2025 01:37

Add KaSA implementation to layer.py

387ad2a

Add use_kasa argument to LoraConfig

ae00a34

nsbg marked this pull request as draft August 2, 2025 05:45

BenjaminBossan requested changes Aug 5, 2025

View reviewed changes

nsbg added 3 commits September 2, 2025 20:54

Add use_kasa parameter to Linear class

6588f1a

Add KasaLinearVariant class (just copy of DoraLinearVariant class) in…

a824ac9

… variants.py

Add kasa description

05e4e07

nsbg added 5 commits September 5, 2025 02:37

Remove unnecessary self.kasa

d1e7e43

[WIP] update KasaLinearVariant class with SVD implementation

9e53b9f

Modify merge/unmerge method in KasaLinearVariant class

aa37111

update KasaLinearVariant class with SVD implementation

9cfe65c

fix type in init method

f9d7cc7

BenjaminBossan requested changes Sep 9, 2025

View reviewed changes

nsbg added 2 commits September 16, 2025 22:19

delete unnecessary part in layer.py

84813a3

add original reference in layer.py

39abcad

nsbg closed this Sep 16, 2025

nsbg force-pushed the peft-kasa branch from 39abcad to 20a9829 Compare September 16, 2025 16:53

merge main to peft-kasa

06f76d8

BenjaminBossan reopened this Sep 17, 2025

re-add KaSA implementation to variants.py

0043ae3

add use_kasa param to resolve_lora_variant in other layers

ea59432

BenjaminBossan requested changes Sep 29, 2025

View reviewed changes

nsbg added 2 commits September 29, 2025 22:46

delete unnecessary part in layer.py

574c1b8

delete unnecessary part in variants.py

73fa58f

	if active_adapter not in self.lora_variant: # vanilla LoRA
	result = result + lora_B(lora_A(dropout(x))) * scaling
	else:
	result = self.lora_variant[active_adapter].forward(
	self,
	active_adapter=active_adapter,
	x=x,
	result=result,
	**variant_kwargs,
	**kwargs,
	)

[WIP] Update LoraConfig for KaSA implementation #2698

Are you sure you want to change the base?

[WIP] Update LoraConfig for KaSA implementation #2698

Uh oh!

Conversation

nsbg commented Aug 2, 2025

Uh oh!

BenjaminBossan left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Sep 1, 2025

Uh oh!

BenjaminBossan commented Sep 1, 2025

Uh oh!

nsbg commented Sep 2, 2025

Uh oh!

BenjaminBossan commented Sep 2, 2025

Uh oh!

nsbg commented Sep 4, 2025

Uh oh!

nsbg commented Sep 8, 2025

1. variants.py

2. layer.py

Uh oh!

BenjaminBossan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nsbg commented Sep 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

BenjaminBossan commented Sep 17, 2025

Uh oh!

nsbg commented Sep 17, 2025

Uh oh!

BenjaminBossan commented Sep 18, 2025

Uh oh!

nsbg commented Sep 20, 2025

Uh oh!

nsbg commented Sep 20, 2025

Uh oh!

BenjaminBossan commented Sep 23, 2025

Uh oh!

nsbg commented Sep 27, 2025

Uh oh!

BenjaminBossan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[WIP] Update `LoraConfig` for KaSA implementation #2698

[WIP] Update `LoraConfig` for KaSA implementation #2698

nsbg commented Sep 16, 2025 •

edited

Loading