Skip to content

Fix _merged_adapters bookkeeping desync on partial unfuse_lora(components=...)#14215

Open
ErenAta16 wants to merge 1 commit into
huggingface:mainfrom
ErenAta16:fix/unfuse-lora-partial-components-bookkeeping
Open

Fix _merged_adapters bookkeeping desync on partial unfuse_lora(components=...)#14215
ErenAta16 wants to merge 1 commit into
huggingface:mainfrom
ErenAta16:fix/unfuse-lora-partial-components-bookkeeping

Conversation

@ErenAta16

Copy link
Copy Markdown

Fixes #14214.

What does this PR do?

LoraBaseMixin._merged_adapters (backing num_fused_loras/fused_loras) is a single set shared across the whole pipeline, but merge state itself is tracked per component by PEFT. fuse_lora()/unfuse_lora() both accept a components argument to operate on a subset of _lora_loadable_modules. When the same adapter is fused into multiple components (the default when fuse_lora() is called with no adapter_names) and later only some of those components are unfused, unfuse_lora() dropped the adapter name from _merged_adapters unconditionally, even though it's still physically merged into the base weights of the untouched component(s).

pipe.fuse_lora(components=["unet", "text_encoder"], adapter_names=["my_adapter"])
pipe.num_fused_loras  # 1

pipe.unfuse_lora(components=["text_encoder"])  # only unfuse text_encoder
pipe.num_fused_loras  # 0, but unet still has my_adapter merged into its weights!

Fix

Track which adapters got unmerged during the per-component loop, then only drop an adapter from the pipeline-wide _merged_adapters set once it's confirmed unmerged in every _lora_loadable_modules component (new _is_adapter_merged_in_any_component helper), not just the ones passed to this particular call.

Testing

Added test_fuse_unfuse_partial_components_keeps_merged_adapter_bookkeeping to the shared PeftLoraLoaderMixinTests mixin in tests/lora/utils.py, so it runs across the pipeline-specific test files that inherit from it (skips pipelines without a text_encoder loadable module, since the scenario needs at least two independent components). Ran it against StableDiffusionLoRATests:

tests/lora/test_lora_layers_sd.py::StableDiffusionLoRATests::test_fuse_unfuse_partial_components_keeps_merged_adapter_bookkeeping PASSED

Also ran the full existing fuse/unfuse suite for the same pipeline to check for regressions in the common (full-components) case:

test_lora_fuse_nan XPASS
test_simple_inference_with_text_denoiser_lora_unfused PASSED
test_simple_inference_with_text_lora_denoiser_fused PASSED
test_simple_inference_with_text_lora_denoiser_fused_multi PASSED
test_simple_inference_with_text_lora_fused PASSED

5 passed, 1 skipped (torch_compile variant, unrelated), 1 xpassed (pre-existing marker, unrelated), no regressions.

I verified the underlying bug and the fix directly with the actual LoraBaseMixin/PeftAdapterMixin classes and a real peft.LoraConfig before writing the pipeline-level test, in case that's useful:

pipe.fuse_lora(components=["unet", "text_encoder"], adapter_names=["my_adapter"])
print(pipe.num_fused_loras)  # 1

pipe.unfuse_lora(components=["text_encoder"])
print(pipe.num_fused_loras)  # before fix: 0 (wrong); after fix: 1 (correct, unet still merged)

…ents=...)

_merged_adapters is a single set shared across the whole pipeline, while
actual merge state is tracked per component by PEFT. When the same adapter
is fused into multiple components (the default when fuse_lora() is called
with no adapter_names) and later only some of those components are unfused
via unfuse_lora(components=[...]), the old code removed the adapter name
from _merged_adapters unconditionally, even when it was still physically
merged into the base weights of the untouched component(s).

num_fused_loras/fused_loras would then report the pipeline as having
nothing (or less) fused than it actually does.

Track unmerge candidates during the per-component unmerge loop, then only
drop an adapter from the pipeline-wide set once confirmed unmerged in every
_lora_loadable_modules component via the new
_is_adapter_merged_in_any_component helper.

Also adds a regression test (test_fuse_unfuse_partial_components_keeps_merged_adapter_bookkeeping
in tests/lora/utils.py) covering the partial-components case, since the
existing fuse/unfuse tests only ever pass every loadable component at once.

Fixes huggingface#14214.
@github-actions github-actions Bot added fixes-issue lora tests size/M PR with diff < 200 LOC and removed fixes-issue labels Jul 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

lora size/M PR with diff < 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

unfuse_lora(components=[...]) desyncs _merged_adapters/num_fused_loras from actual per-component PEFT merge state

1 participant