fix(lora): only drop adapter from _merged_adapters when unfused from all components - #14385
fix(lora): only drop adapter from _merged_adapters when unfused from all components#14385AloysJehwin wants to merge 7 commits into
Conversation
…all components Signed-off-by: Aloys Jehwin <aloysjehwin@gmail.com>
sayakpaul
left a comment
There was a problem hiding this comment.
Thanks for this PR. Should we also have a test for this?
|
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. |
Signed-off-by: Aloys Jehwin <aloysjehwin@gmail.com>
ff8dbce to
8e1ca01
Compare
|
Added a test in |
|
It shouldn't be a regression test. The test should live under |
…lone unit) Signed-off-by: Aloys Jehwin <aloysjehwin@gmail.com>
9085ff6 to
9145c5f
Compare
|
The test uses a minimal mock pipeline (no real weights needed), so it lives in |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the PR, generally looks good. I just added some smaller comments. As for where to put the test, please follow the guidance from the mtaintainers.
| remaining_merged: set[str] = set() | ||
| for component_name in self._lora_loadable_modules: | ||
| component_model = getattr(self, component_name, None) | ||
| if component_model is not None and issubclass(component_model.__class__, (ModelMixin, PreTrainedModel)): |
There was a problem hiding this comment.
Why check issubclass instead of isinstance? I also think isinstance(component_model, nn.Module) would be sufficient here, but I'll leave that to the maintainers to decide.
| module.unmerge() | ||
|
|
||
| # Only remove an adapter from _merged_adapters once it is no longer | ||
| # physically merged in any remaining loadable component. Removing it |
There was a problem hiding this comment.
IMO the second sentence is not needed.
| # Clean up the hooks to prevent state leak | ||
| if hasattr(denoiser, "_diffusers_hook"): | ||
| denoiser._diffusers_hook.remove_hook(_GROUP_OFFLOADING, recurse=True) | ||
|
|
There was a problem hiding this comment.
Remove unrelated changes.
| Unfusing only a subset of components must keep _merged_adapters in sync | ||
| with the adapters still physically fused in the remaining components. | ||
| """ | ||
| import torch.nn as nn |
There was a problem hiding this comment.
Why local imports? Only PEFT needs to be local or guarded behind is_peft_available().
Signed-off-by: Aloys Jehwin <aloysjehwin@gmail.com>
c7d8593 to
5ec11b4
Compare
|
Addressed all four points — swapped to |
Signed-off-by: Aloys Jehwin <aloysjehwin@gmail.com>
|
Fixed — restored utils.py to main. |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the updates, LGTM. About where to put the test, let's wait for @sayakpaul's response.
Fixes #14214
When
unfuse_lora(components=[...])is called with a subset of components, the current code removes the adapter from_merged_adaptersas soon as it's unfused from the first component — even if it's still physically fused into the others. Sonum_fused_loras/fused_lorasreport nothing fused while the remaining components silently still have the LoRA baked in.The fix: call
module.unmerge()as before, then recompute_merged_adaptersby scanning all loadable components for adapters that are still physically merged (viaBaseTunerLayer.merged_adapters). That way the set only loses an adapter once it's gone from everything.Verified with the reproduction script from the issue —
num_fused_lorascorrectly stays at 1 after unfusing onlytext_encoder, and the unet is confirmed still merged at the PEFT level.