Skip to content

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

Open
ErenAta16 wants to merge 5 commits 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 5 commits 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
Comment thread tests/lora/utils.py Outdated
Comment on lines +1721 to +1722
if "text_encoder" not in self.pipeline_class._lora_loadable_modules:
return

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should use skip rather than returning so that the test properly gets skipped.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, fixed in cacdeee. The old return was reporting a pass for pipelines that don't support text encoder LoRAs at all, instead of skipping. Switched to self.supports_text_encoder_loras + pytest.skip(...), matching the convention already used elsewhere in this file.

@sergereview sergereview Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤗 Serge says:

The core fix in lora_base.py is correct and well-reasoned: tracking unmerge_candidates and only dropping an adapter from the pipeline-wide _merged_adapters once _is_adapter_merged_in_any_component confirms it's gone from every loadable component correctly resolves the partial-unfuse_lora bookkeeping desync described in #14214. The None-guard from the old code (if adapter and ...) is safely subsumed by the if adapter not in self._merged_adapters: continue filter, so no regression there.

Tests

  • The new test uses a bare if "text_encoder" not in ... : return, which makes the test report as passed (green) on pipelines without a text encoder rather than skipped. The established convention in this file is the supports_text_encoder_loras flag plus pytest.skip(...) (see lines 385, 521). Using that keeps the skip visible and consistent, and correctly reflects pipelines that carry a text_encoder module but don't support text-encoder LoRAs.

Style (minor)

  • _is_adapter_merged_in_any_component has a single caller. Per the repo's coding-style guidance (inline single-use private helpers), this could be inlined into unfuse_lora, though the helper is self-contained and the current form is readable — not blocking.

No correctness or security concerns; the description's claims match the diff.

serge v0.1.0 · model: claude-opus-4-8 · 12 LLM turns · 11 tool calls · 59.1s · 180773 in / 3307 out tokens

Comment thread tests/lora/utils.py Outdated
component and then unfusing only *some* of them should not report the adapter as fully unfused,
it's still merged into the untouched component(s).
"""
if "text_encoder" not in self.pipeline_class._lora_loadable_modules:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Prefer the existing skip convention over a bare return. As written, this test reports as passed for pipelines that don't have a text_encoder loadable module, which hides the fact that it never exercised the scenario. The rest of this file gates on the supports_text_encoder_loras flag with pytest.skip(...) (e.g. lines 385, 521), which is also more accurate — a pipeline can have a text_encoder module but not support text-encoder LoRAs.

Suggested change
if "text_encoder" not in self.pipeline_class._lora_loadable_modules:
if not self.supports_text_encoder_loras:
pytest.skip("Skipping test as text encoder LoRAs are not currently supported.")

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I looked at this again against AGENTS.md's single-caller inlining guidance. I went back and forth, but decided to leave _is_adapter_merged_in_any_component as a separate method: inlining it means a two-level nested loop with a boolean flag sitting inside the unfuse_lora unmerge-candidate loop, which reads worse than a named, docstring'd predicate, even with one caller. Happy to inline it if you'd still prefer that on a second look, but wanted to explain the call rather than silently leave it.

@sayakpaul
sayakpaul requested a review from BenjaminBossan July 18, 2026 02:59
@ErenAta16

Copy link
Copy Markdown
Author

Process note, read after opening this: I found .ai/AGENTS.md's coordination guidance (added in #14185, a couple days before this PR) after the fact, not before, so this didn't go through the "wait for explicit maintainer acknowledgment on the issue before opening a PR" step it asks for. Issue #14214 has the full trace and reproduction but no prior acknowledgment. Apologies for the process miss, sharing the self-review now to at least close that part of the gap.

Self-review against .ai/review-rules.md:

  • Ephemeral context: no PR-review-specific comments, no debug printouts, no scratch/dev-path scripts in the diff. The two added comments (on the unmerge_candidates tracking and the _is_adapter_merged_in_any_component docstring) explain the underlying design tension, self-contained, not "per discussion" references.
  • Documentation impact: num_fused_loras/fused_loras's existing docstrings ("Returns the number of LoRAs that have been fused" / "Returns names of the LoRAs that have been fused") already describe the correct, intended contract, this PR fixes the implementation to actually match what was already documented, not a new or changed public surface. No doc updates needed.
  • Dead code analysis: not applicable, this isn't a new-model PR.

One thing I did not fix and want to flag rather than silently leave: unfuse_lora's existing docstring doesn't mention that _merged_adapters/num_fused_loras is pipeline-wide rather than per-component, so a caller reading only the docstring wouldn't know that fusing into multiple components and unfusing just one still shows the adapter as fused. Happy to add a short note there if that'd be useful, left it out of this PR to keep the diff scoped to the actual bug rather than bundling in a docstring rewrite, but can add it as a follow-up commit here if you'd rather have it in the same PR.

Matches the established skip pattern used elsewhere in this file
instead of a silent early return, per review feedback.

@BenjaminBossan BenjaminBossan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generally looks good and makes sense. Just some small comments.

unfuse_lora's existing docstring doesn't mention that _merged_adapters/num_fused_loras is pipeline-wide rather than per-component

Yes, that's worth documenting.

per-component PEFT merge state after a partial `unfuse_lora(components=...)` call."""
for component in self._lora_loadable_modules:
model = getattr(self, component, None)
if model is None or not issubclass(model.__class__, (ModelMixin, PreTrainedModel)):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why issubclass(model.__class__, ...) and not isinstance(model, ...)?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's copied from the existing pattern right above it in the same unfuse_lora loop (and used throughout this file, 16 occurrences), not a deliberate choice on my part - happy to switch to isinstance() here if you'd rather not perpetuate it, but left it matching the surrounding code for now.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. IMO, isinstance makes more sense but maybe I'm missing something. I'll leave that decision up to the Diffusers maintainers.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds like the consensus is to leave it as-is for now and revisit library-wide later, so no change here.

Comment thread tests/lora/utils.py

# Now unfuse the remaining component; bookkeeping should correctly drop to 0.
pipe.unfuse_lora(components=[denoiser_component_name])
self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For completeness, I would also check "adapter-1" is not in pipe.fused_loras.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in f6c54bc - now also asserts adapter-1 is out of fused_loras once it's unfused from every component.

… test

- unfuse_lora docstring now notes num_fused_loras/fused_loras are
  tracked pipeline-wide, not per component.
- Regression test now also asserts the adapter is absent from
  fused_loras after it's been unfused from every component.
@ErenAta16

Copy link
Copy Markdown
Author

Pushed f6c54bc addressing the review: added the pipeline-wide vs per-component note to unfuse_lora's docstring, and replied inline on the two other points.

@BenjaminBossan BenjaminBossan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the changes, just one nit.

per-component PEFT merge state after a partial `unfuse_lora(components=...)` call."""
for component in self._lora_loadable_modules:
model = getattr(self, component, None)
if model is None or not issubclass(model.__class__, (ModelMixin, PreTrainedModel)):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. IMO, isinstance makes more sense but maybe I'm missing something. I'll leave that decision up to the Diffusers maintainers.

Comment thread src/diffusers/loaders/lora_base.py Outdated
Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
LoRA parameters then it won't have any effect.

Note that `num_fused_loras`/`fused_loras` (backed by `self._merged_adapters`) are tracked pipeline-wide,

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's not mention implementation details here, as those don't matter to the user. It's sufficient to mention the part about fused_loras, as this is what's exposed to the uesr.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's fine to keep it as is and it can be replaced library-wide to fix on a common pattern.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair, trimmed in b52aa28 - just the fused_loras behavior now, dropped the _merged_adapters mention.

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your work!

@sayakpaul

Copy link
Copy Markdown
Member

@bot /style

@github-actions

github-actions Bot commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Style bot fixed some files and pushed the changes.

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

Drop the _merged_adapters implementation detail per review feedback;
fused_loras is what's actually exposed to callers.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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

4 participants