Skip to content

[lora]feat: use exclude modules to loraconfig. #11806

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Jun 30, 2025
Merged

Conversation

sayakpaul
Copy link
Member

What does this PR do?

If we try to do:

from diffusers import AutoModel, DiffusionPipeline 
import torch

model_id = "Wan-AI/Wan2.1-VACE-14B-diffusers"
vae = AutoModel.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe = DiffusionPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights(
   "vrgamedevgirl84/Wan14BT2VFusioniX", 
   weight_name="FusionX_LoRa/Wan2.1_T2V_14B_FusionX_LoRA.safetensors",
)

It prints:

Unfold
Loading adapter weights from state_dict led to missing keys in the model: vace_blocks.0.proj_out.lora_A.default_0.weight, vace_blocks.0.proj_out.lora_B.default_0.weight, vace_blocks.1.proj_out.lora_A.default_0.weight, vace_blocks.1.proj_out.lora_B.default_0.weight, vace_blocks.2.proj_out.lora_A.default_0.weight, vace_blocks.2.proj_out.lora_B.default_0.weight, vace_blocks.3.proj_out.lora_A.default_0.weight, vace_blocks.3.proj_out.lora_B.default_0.weight, vace_blocks.4.proj_out.lora_A.default_0.weight, vace_blocks.4.proj_out.lora_B.default_0.weight, vace_blocks.5.proj_out.lora_A.default_0.weight, vace_blocks.5.proj_out.lora_B.default_0.weight, vace_blocks.6.proj_out.lora_A.default_0.weight, vace_blocks.6.proj_out.lora_B.default_0.weight, vace_blocks.7.proj_out.lora_A.default_0.weight, vace_blocks.7.proj_out.lora_B.default_0.weight.

It happens because target_modules in LoraConfig treats the values in target_modules as suffixes. In this case, proj_out is a part of target_modules which is why vace_blocks_* get targeted here (and hence the missing warning).

LoraConfig allows us to specify exclude_modules, too. This PR introduces support to add it when initializing the LoraConfig.

@apolinario does it work for you?

@sayakpaul sayakpaul requested a review from BenjaminBossan June 25, 2025 12:20
@HuggingFaceDocBuilderDev

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.

Copy link
Member

@BenjaminBossan BenjaminBossan left a comment

Choose a reason for hiding this comment

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

Generally looks good. As already discussed, for the final PR, we need a version guard on PEFT, as exclude_modules was only added in v0.14.0. For the finished PR, there should also be unit tests.

@sayakpaul sayakpaul requested a review from BenjaminBossan June 27, 2025 12:29
@sayakpaul sayakpaul changed the title [wip][lora]feat: use exclude modules to loraconfig. [lora]feat: use exclude modules to loraconfig. Jun 27, 2025
@sayakpaul sayakpaul marked this pull request as ready for review June 27, 2025 12:29
@sayakpaul sayakpaul requested a review from a-r-r-o-w June 27, 2025 12:29
Copy link
Member

@a-r-r-o-w a-r-r-o-w left a comment

Choose a reason for hiding this comment

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

Nice, looks good to me.

pipe_cp, _ = self.check_if_adapters_added_correctly(
pipe_cp, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config
)
denoiser_exclude_modules = self._get_exclude_modules(pipe_cp)
Copy link
Member

Choose a reason for hiding this comment

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

I'm afraid I don't fully understand how this test works but that shouldn't be a blocker. If this tests the exact same condition we're facing in the FusionX lora, then it should be good :)

Copy link
Member Author

Choose a reason for hiding this comment

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

No, not the cleanest test this one. I added a description in the test, see if that helps?

Currently, we don't have:

If this tests the exact same condition we're facing in the FusionX lora

I will think of a way to include a test for that, too.

string_to_replace = f"{adapter_name}." if adapter_name else ""

for name in model_state_dict.keys():
if string_to_replace:
Copy link
Member

Choose a reason for hiding this comment

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

maybe the if-statement is not needed here because string_to_replace will be an empty string, but no problem keeping as micro optimization

Comment on lines 163 to 164
denoiser_lora_config.target_modules = ["to_q", "to_k", "to_v", "out"]
denoiser_lora_config.exclude_modules = ["proj_out"]
Copy link
Member

Choose a reason for hiding this comment

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

I wonder: Would this test not pass even without the fix of this PR? When I tested it on main, it passed, even when I commented out the denoiser_lora_config.exclude_modules = ["proj_out"] line.

We would need an example with something like:

target_modules = ["proj_out"]
exclude_modules = ["foo.proj_out"]

Copy link
Member Author

Choose a reason for hiding this comment

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

True. But we will have to then use the Wan VACE model as the current model being tested doesn't meet the module naming criterion we're looking for. I will delete this test case and add it in a future PR then if that sounds good?

Copy link
Member

Choose a reason for hiding this comment

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

SGTM

@sayakpaul sayakpaul requested a review from BenjaminBossan June 30, 2025 13:41
Copy link
Member

@BenjaminBossan BenjaminBossan left a comment

Choose a reason for hiding this comment

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

Thanks, LGTM.

@sayakpaul sayakpaul merged commit bc34fa8 into main Jun 30, 2025
32 checks passed
@sayakpaul sayakpaul deleted the exclude-modules-lora branch June 30, 2025 14:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants