MoE LoRA: auto-target per-expert Linear experts (gpt-oss 4bit) instead of leaving them frozen - #6936
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a check to verify if MoE (Mixture of Experts) parameters actually exist on a model before attempting to enable LoRA on them. This prevents silent failures on certain quantized checkpoints (such as gpt-oss bnb-4bit) where experts are stored as per-expert modules rather than fused parameters. A warning is now logged when these parameters are missing, advising the user to use a 16-bit checkpoint if they wish to train the experts. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@codex review |
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 845dc43030
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| f"as per-expert modules (for example gpt-oss bnb-4bit). LoRA is applied to the " | ||
| f"attention and other targeted layers only; the experts are left frozen. Use a " | ||
| f"16bit checkpoint to LoRA-train the experts." |
There was a problem hiding this comment.
When an MoE keeps experts as per-expert Linear modules named experts.<i>.gate_proj/up_proj/down_proj rather than as fused gate_up_proj parameters, _moe_parameter_exists is false and this warning fires, but PEFT still receives target_modules=final_modules with the default gate_proj, up_proj, and down_proj leaves from FastLlamaModel.get_peft_model, so module-based LoRA can still attach to those expert modules. In that context the warning tells users the experts are left frozen even though only fused target_parameters were skipped; please gate it on the absence of matching expert modules or reword it to avoid claiming all expert LoRA is disabled.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Agreed, the flat 'experts are left frozen' over-claimed. Reworded: the warning now says fused-parameter LoRA on the experts is skipped, expert Linears your target_modules match are still adapted, and only experts in a non-standard layout (like gpt-oss bnb-4bit) get no LoRA. Behavior is unchanged (still returns None for the fused target_parameters); message-accuracy fix only.
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…bit) MoE checkpoints whose experts are stored as per-expert nn.Linear ModuleLists could not receive expert LoRA. gpt-oss bnb-4bit is the canonical case: its experts live at mlp.experts.gate_up_projs.<i> and mlp.experts.down_projs.<i> as per-expert Linear4bit modules, not a fused nn.Parameter. The target_parameters path only handles the fused nn.Parameter layout, and the plain gate_proj/up_proj/down_proj leaf names do not match the per-expert indices, so get_peft_model attached LoRA to attention only and left every expert frozen (0 of 1536 on gpt-oss-20b) even though the grouped bnb-4bit training forward exists. Add get_moe_target_modules, the module-LoRA counterpart of get_moe_target_parameters: it detects per-expert Linear ModuleLists under an experts container and returns their suffix target_modules names (gate_up_projs.<i> / down_projs.<i>). get_peft_model in both llama.py and vision.py extends target_modules with these, handling the explicit leaf-list form and the regex form (auto / all-linear / scoped). It is gated on the same MLP-in-scope condition as the parameter path, so an attention-only request still skips the experts. Also gate get_moe_target_parameters on the fused parameter actually existing, so a per-expert-Linear layout no longer produces a dead target_parameters path or a misleading "Enabling LoRA on MoE parameters" line; those experts are handled through target_modules instead. Validated on gpt-oss-20b-unsloth-bnb-4bit (transformers 5.5.0): experts attach (1536 modules, trainable 0.036 percent to 1.65 percent) across the default, None and all-linear paths; training memorizes and the LoRA adapter reproduces exactly after a cold reload in a fresh process. No regression: fused-parameter MoEs (Qwen3-30B-A3B-4bit), non-MoE models, and attention-only requests are unaffected (get_moe_target_modules returns an empty list). Merging these per-expert adapters into a merged_16bit checkpoint is handled by a companion unsloth-zoo change (saving_utils folds each per-expert delta into the fused gate_up_proj / down_proj tensor). With both, the LoRA adapter and the merged_16bit checkpoint reload the trained behavior identically.
495af96 to
2129fc7
Compare
danielhanchen
changed the title
|
Expanded this PR from the honest-gating-only version to the complete fix. It keeps Folding these per-expert adapters into a |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2129fc7ed4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for expert_index in range(len(module)): | ||
| targets.add(f"{leaf}.{expert_index}") |
There was a problem hiding this comment.
When a per-expert Linear MoE is called with a narrowed target list such as target_modules=["down_proj"] (or a regex that names only down_proj), this loop still adds every Linear ModuleList under experts, including gate_up_projs.*. That makes explicit down-only or gate/up-only LoRA requests train extra expert projections, unlike get_moe_target_parameters, which keeps the fused-parameter path scoped to the requested leaves.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 5cf6e3c: get_moe_target_modules now scopes to the requested projection leaves (gate/up map to gate_up_projs, down maps to down_projs), matching get_moe_target_parameters, so a down-only or gate/up-only request no longer trains the other projection.
| _moe_module_targets = get_moe_target_modules(model, target_modules) | ||
| if _moe_module_targets: | ||
| _added = [t for t in _moe_module_targets if t not in final_modules] | ||
| final_modules.extend(_added) |
There was a problem hiding this comment.
On per-expert Linear MoEs this appends hidden gate_up_projs.N/down_projs.N targets to the config saved by PEFT, but the existing-adapter pass-through check earlier in this function compares peft_config["target_modules"] against the unaugmented user target_modules before this block runs. A second get_peft_model call with the same arguments will therefore see the saved extra expert targets as a mismatch and raise instead of taking the intended idempotent skip path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 5cf6e3c: the existing-adapter check now recomputes the auto-added expert targets (and detection is base_layer-aware), so a repeat get_peft_model with the same arguments stays idempotent. gpt-oss routes through vision.py, whose get_peft_model has no pass-through, so this was only reachable for a per-expert-Linear MoE on FastLlamaModel; fixed defensively.
| for expert_index in range(len(module)): | ||
| targets.add(f"{leaf}.{expert_index}") |
There was a problem hiding this comment.
This starts producing adapter keys like gate_up_projs.N/down_projs.N, but merged saves are still delegated to unsloth_zoo.saving_utils.merge_and_overwrite_lora while pyproject.toml continues to allow unsloth_zoo>=2026.6.7. In environments that keep an older allowed zoo version without the companion logic that folds those per-expert deltas back into fused gate_up_proj/down_proj tensors, save_pretrained_merged("merged_16bit") can silently drop or fail to merge the newly trained expert LoRA weights; bump the minimum zoo version or add a runtime guard before emitting this layout.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 5cf6e3c: get_peft_model now warns when the installed unsloth_zoo cannot fold these per-expert experts into a merged_16bit checkpoint, so the expert LoRA is not silently dropped on save. The fold is unslothai/unsloth-zoo#885; the zoo floor will be bumped once it releases.
|
|
||
| targets = set() | ||
| for name, module in model.named_modules(): | ||
| if not isinstance(module, torch.nn.ModuleList) or len(module) == 0: |
There was a problem hiding this comment.
When patch_regional_compilation() has run, _utils.py replaces torch.nn.ModuleList with the UnslothModuleList factory, so this isinstance(..., torch.nn.ModuleList) check raises TypeError instead of just skipping or detecting the expert lists. In that configuration, per-expert MoE get_peft_model fails before LoRA injection; compare against the original container class (for example torch.nn.modules.container.ModuleList) rather than the patched factory.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
patch_regional_compilation is defined but never called anywhere in unsloth or unsloth_zoo, so torch.nn.ModuleList is never rebound to the UnslothModuleList factory. After import, torch.nn.ModuleList.name == 'ModuleList' and the isinstance check detects the expert ModuleLists without error, so no change is needed here.
…potent, warn on old zoo Address review of the per-expert Linear MoE targeting: - Scope get_moe_target_modules to the requested projection leaves (gate/up map to the gate_up ModuleList, down maps to the down ModuleList), so a narrowed request such as target_modules=["down_proj"] no longer also trains gate_up_projs, matching get_moe_target_parameters. - Detect experts through a PEFT-wrapped base_layer as well, and recompute the auto-added expert targets in the llama.py existing-adapter check, so a repeat get_peft_model call with the same arguments stays idempotent instead of raising on the saved expert targets. - Warn when the installed unsloth_zoo cannot fold these per-expert experts into a merged_16bit checkpoint (older releases keep the fused gate_up_proj / down_proj tensors and drop the per-expert deltas), so the expert LoRA is not silently lost on save_pretrained_merged; the fold lands in unsloth-zoo #885. The LoRA adapter itself is unaffected.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Can't wait for the next one! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Summary
get_moe_target_parametershands PEFT the fused expertnn.Parameters (mlp.experts.gate_up_proj/down_proj) astarget_parameters. Some quantized checkpoints have no fused expert Parameter at all: gpt-oss bnb-4bit stores its experts as a per-expertLinear4bitModuleList (mlp.experts.gate_up_projs.<i>/down_projs.<i>). Those experts matched neither the fusedtarget_parameterspath nor the plaingate_proj/up_proj/down_projleaf names, so LoRA attached to attention only and left every expert frozen, while still printing "Enabling LoRA on MoE parameters".Fix
Keep the honest gating and additionally attach the per-expert experts so they actually train:
_moe_parameter_exists(model, name): gate each fused expert name on a matching parameter actually existing, so we never hand PEFT a deadtarget_parameterspath or print a misleading "Enabling LoRA on MoE parameters" line.get_moe_target_modules(model, target_modules): the module-LoRA counterpart ofget_moe_target_parameters. It detects per-expertnn.LinearModuleLists under anexpertscontainer and returns their suffixtarget_modules(gate_up_projs.<i>/down_projs.<i>), so PEFT attaches LoRA to every expert Linear via ordinary suffix matching.get_peft_model(llama.py and vision.py) extendstarget_moduleswith these, handling both the explicit leaf-list form and the regex form (auto / all-linear / scoped). Gated on the same MLP-in-scope condition as the parameter path, so an attention-only request still skips the experts.get_moe_target_modulesreturns an empty list for non-MoE models, fused-parameter MoEs (Qwen3 / Mixtral / LFM2 / BF16 gpt-oss, handled byget_moe_target_parameters), and absent per-expert layouts, so nothing else changes.Validation
gpt-oss-20b-unsloth-bnb-4bit (transformers 5.5.0): experts attach (3072 modules, trainable 1.65 percent). A memorize overfit trains and the LoRA adapter reloads exactly in a fresh process (teacher-forced CE 4e-6). No regression: fused-parameter MoEs (Qwen3-30B-A3B-4bit), non-MoE models, and attention-only requests are unaffected (
get_moe_target_modulesreturns[]).Folding these per-expert adapters into a
merged_16bitcheckpoint is handled by the companion PR unslothai/unsloth-zoo#885. With both, the LoRA adapter and the merged_16bit checkpoint reload the trained behaviour identically.This replaces the honest-gating-only version this PR previously carried: the per-expert Linear4bit layout is now attached and trained rather than left frozen.