Fix MoE target_parameters module_count alignment (#3405, #3701)

[GoldenGrapeGentleman]

Summary

Resolves the TODO: handle MoE target_parameters to align counts at line 410 of saving_utils.py.

Problem

When using PEFT target_parameters for MoE expert LoRA (e.g., gpt-oss), the create_lora_statistics function reports a misleading diagnostic warning:

[Unsloth merge debug] LoRA count mismatch: modules=120, lora_A=144, lora_B=144, scaling=144

This happens because MoE ParamWrapper entries have lora_A/lora_B/scaling counted, but lack a .base_layer module, so module_count is never incremented for them.

Fix

After the LoRA collection loop, count MoE expert entries that have lora_A/lora_B but no .module, aligning module_count with the other counts. Only applies to .mlp.experts entries to avoid affecting non-MoE models.

Changes

• saving_utils.py: +12 lines / -3 lines (removes TODO comment)

Verified on 8× AMD Instinct MI355X (gfx950), ROCm 7.1

• gpt-oss-20b BF16 + MoE expert LoRA: 46.2M trainable params (vs 2.0M attention-only)
• save_pretrained_merged: ✅ success, no count mismatch warning
• Before fix: modules=120, lora_A=144 warning printed
• After fix: counts aligned, clean merge

Related: unslothai/unsloth#3405, unslothai/unsloth#3701

cc @danielhanchen @Datta0

[gemini-code-assist]

Summary of Changes

Hello @GoldenGrapeGentleman, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses and resolves a misleading "LoRA count mismatch" diagnostic warning that occurred when using PEFT target_parameters for Mixture-of-Experts (MoE) expert LoRA. The fix ensures that the module_count accurately reflects the number of LoRA modules, particularly for ParamWrapper entries within MoE experts, leading to clean merges and correct statistics without false warnings.

Highlights

• MoE LoRA Count Mismatch Resolution: Resolved an issue in the create_lora_statistics function where Mixture-of-Experts (MoE) expert LoRA configurations would incorrectly trigger a "LoRA count mismatch" warning.
• Accurate Module Counting for MoE: Implemented logic to correctly increment module_count for MoE ParamWrapper entries that possess lora_A and lora_B but lack a .base_layer module, specifically targeting .mlp.experts.
• TODO Comment Addressed: Addressed and removed the TODO: handle MoE target_parameters to align counts comment at line 410 in saving_utils.py, indicating the completion of the planned fix.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• unsloth_zoo/saving_utils.py
  • Added a loop to iterate through lora_weights and increment module_count for MoE ParamWrapper entries that have lora_A and lora_B but no associated .module, specifically for keys containing ".mlp.experts".
  • Removed a TODO comment related to handling MoE target parameters for count alignment.

Activity

• The fix was verified on 8× AMD Instinct MI355X (gfx950) with ROCm 7.1.
• Successfully eliminated the "LoRA count mismatch" warning for gpt-oss-20b BF16 + MoE expert LoRA, achieving aligned counts and clean merges.
• This PR is related to issues [Bug] GPT OSS: No adapters for the experts? unsloth#3405 and [Bug] RuntimeError: Unsloth: Saving LoRA finetune failed since # of LoRAs = 9360 does not match # of saved modules = 144. Please file a bug report! unsloth#3701.
• Reviewers @danielhanchen and @Datta0 were CC'd.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request provides a clean and effective fix for a LoRA module count mismatch that occurred with MoE models using target_parameters. The added logic correctly identifies uncounted MoE expert modules by checking for LoRA weights on modules without a registered base layer, and appropriately increments the module count. The change is well-targeted to .mlp.experts to avoid side effects. The removal of the corresponding TODO comment is also appropriate. The implementation is straightforward and directly resolves the issue described.

[Datta0]

Hey @GoldenGrapeGentleman thanks for the contribution. I do remember that there's only one base_layer in the experts module (referring to gate_up_proj iirc?). This does effect transformers v5.
But if I also remember, in transformers v4, the experts are nn.ModuleList(nn.Linear) which should play well with LoRA.

Can we handle this specifically for nn.Parameter instead of doing a blanket addition? Also can you test your fix for the following models : unsloth/Qwen3-30B-A3B-Instruct-2507,'unsloth/gpt-oss-20b-BF16',"unsloth/GLM-4.7-Flash",imdatta0/tiny_qwen3_moe_2.8B_0.7B,unsloth/Qwen3-VL-30B-A3B-Instruct? IF you can do that this would be of great help and would ease the review process...

Thanks a lot :)

[GoldenGrapeGentleman]

Hi @Datta0 Thanks for your quick response, I will do all you mentioned, maybe after the Chinese Spring Festival. BTW, happy new year to you and unsloth team if anyone is from China(❁´◡`❁)

[Datta0]

Happy new year @GoldenGrapeGentleman. Wishing you a great time :)

[Datta0]

Hey @GoldenGrapeGentleman
I tried your patch with Qwen3.5 family and I see that there's still some issues.
You don't need to test on the full model. I made a small dummy ~20B scale model for testing purposes. You can give it a try.
Also did you try verifying that merged model is same as the intended model aka training learnings are intact after merging...

[FAIL] save_pretrained_merged: Unsloth: Saving LoRA finetune failed since # of LoRAs = 200 does not match # of saved modules = 0. Please file a bug report!
PS: LoRA is only applied to MoE params not on attn at all