Scope MoE expert LoRA detection to actual MLP projection targets#6849
Conversation
_moe_target_set_from_string treated any regex containing the substring mlp or ffn as targeting the expert MLP projections. Unsloth's auto-generated attention-only regex lists mlp, ffn and feed_forward as allowed intermediate path segments while its final group matches only q_proj/k_proj/v_proj/o_proj, so attention-only finetuning on MoE models silently enabled expert LoRA as well: the experts were trained and every MoE layer paid the extra expert LoRA grouped matmuls. Detect expert intent from the projection names themselves (gate_proj/up_proj/down_proj/gate_up_proj) instead of the mlp substring.
for more information, see https://pre-commit.ci
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
|
@codex review |
|
/gemini review |
![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 refactors the expert-LoRA detection logic in unsloth/models/_utils.py to check for specific projection names (gate_proj, up_proj, down_proj, gate_up_proj) instead of generic mlp or ffn substrings, preventing accidental activation during attention-only finetuning. The reviewer noted that the new substring check is fragile when target_modules is a regular expression, as literal substring matches will fail. They suggested compiling the regex and testing it against dummy MoE expert module names to ensure robust matching.
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.
| targets_mlp_proj = any( | ||
| name in target_modules for name in ("gate_proj", "up_proj", "down_proj", "gate_up_proj") | ||
| ) | ||
| if is_regex and targets_mlp_proj: |
There was a problem hiding this comment.
The current substring check any(name in target_modules for name in ...) is fragile when target_modules is specified as a regular expression. For example, if a user specifies target_modules as .*(gate|up|down)_proj, none of the literal strings "gate_proj", "up_proj", "down_proj", or "gate_up_proj" will be present as substrings in target_modules. As a result, targets_mlp_proj will evaluate to False, and MoE expert LoRA will be silently disabled.
To make this check robust and fully support regex patterns, we can compile target_modules as a regular expression and test it against dummy full module names of MoE experts (e.g., model.layers.0.mlp.experts.0.gate_proj). This mirrors how PEFT matches target modules and ensures both exact substrings and complex regex patterns are correctly handled.
targets_mlp_proj = False
if is_regex:
try:
pattern = re.compile(target_modules)
dummy_names = [
f"model.layers.0.mlp.experts.0.{name}"
for name in _MOE_BROAD_MLP_TARGETS
] + [
f"model.layers.0.experts.0.{name}"
for name in _MOE_BROAD_MLP_TARGETS
]
targets_mlp_proj = any(pattern.search(dummy_name) is not None for dummy_name in dummy_names)
except re.error:
pass
if is_regex and targets_mlp_proj:The auto-generated target regex always lists every projection leaf (q/k/v/o and gate/up/down), so keying detection on a proj name mis-fired: it enabled expert LoRA for attention-only regexes and dropped the mlp/ffn path regexes. Key on the mlp/ffn/feed_forward/experts path segment instead, which is present only when the MLP/experts are actually targeted. Add a regression test for the attention-only case.
|
Reworked the detection to key on the mlp / ffn / feed_forward / experts path segment rather than the projection names. The auto target regex always lists every projection leaf (q/k/v/o and gate/up/down), so a proj-name check both mis-fired on an attention-only regex and dropped |
|
@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: c6a611a286
ℹ️ 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".
| targets_mlp_path = any( | ||
| tag in target_modules for tag in ("mlp", "ffn", "feed_forward", "experts") | ||
| ) | ||
| if is_regex and targets_mlp_path and "proj" in target_modules: |
There was a problem hiding this comment.
When a scoped regex contains MLP path alternatives but its leaf group is restricted to attention projections such as q_proj|k_proj|v_proj|o_proj, this condition still passes because the generic "proj" substring comes from those attention leaves. get_moe_target_parameters then returns all MoE expert parameters, so attention-only LoRA runs on MoE models still train gate_up_proj/down_proj and pay the extra expert LoRA cost. Please require one of the actual MoE leaves (gate_proj, up_proj, down_proj, or gate_up_proj) in the regex instead of any proj substring.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. Detection now scopes to the leaves the regex actually names, and an mlp path alternative with only attention leaves no longer triggers. The auto regex mlp tag block still enables the broad set for fused-expert models whose leaves are plain Parameters.
| targets_mlp_path = any( | ||
| tag in target_modules for tag in ("mlp", "ffn", "feed_forward", "experts") | ||
| ) | ||
| if is_regex and targets_mlp_path and "proj" in target_modules: |
There was a problem hiding this comment.
When a regex is scoped directly to feed_forward or experts and names only one expert leaf, for example .*experts.*down_proj on the alternate Gemma-style layout, this branch now returns the entire broad target set rather than the leaf that appeared in the regex. That makes get_moe_target_parameters add both gate_up_proj and down_proj, so a user trying to LoRA only the down projection will unexpectedly train the fused gate/up parameter as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed. A regex naming a single expert leaf like .*experts.*down_proj now returns only that projection; generic .*mlp.*proj keeps the broad set.
An mlp path alternative with attention-only leaves, for example (mlp|self_attn).(q_proj|o_proj), no longer enables expert LoRA, and a regex naming a single expert leaf such as .*experts.*down_proj now targets only that projection instead of the whole broad set. Generic mlp projections (.*mlp.*proj) and the auto regex mlp tag block keep the broad set for fused-expert models whose leaves are plain Parameters.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cdbb3dcaa4
ℹ️ 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".
| if "mlp|feed_forward|ffn|dense" in target_modules: | ||
| return set(_MOE_BROAD_MLP_TARGETS) |
There was a problem hiding this comment.
When get_peft_regex emits the attention-only pattern described by this change, its path group can still contain mlp|feed_forward|ffn|dense while the leaf group is only q_proj|k_proj|v_proj|o_proj; the earlier named and generic-proj checks correctly do not treat that as expert intent, but this fallback then returns all MoE targets solely from the path alias substring. Because get_moe_target_parameters only sees the regex string and not the original finetune_mlp_modules flag, attention-only runs on fused MoE models still silently attach LoRA to gate_up_proj and down_proj whenever that canonical alias block is present.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
get_peft_regex only appends the mlp tags (mlp|feed_forward|ffn|dense) to the path group when finetune_mlp_modules=True (peft_utils.py: regex_components += mlp_tags is guarded by that flag). An attention-only regex emits self_attn|attention|attn|mixer as its path group, so the tag block is never present. Confirmed empirically for both a fused-expert MoE and a per-expert-Linear MoE: the attention-only regex contains no feed_forward / tag block and get_moe_target_parameters returns None. Branch 3 only fires when mlp finetuning is on (fused-expert case), which is intended. No code change.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cdbb3dcaa4
ℹ️ 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".
| if "mlp|feed_forward|ffn|dense" in target_modules: | ||
| return set(_MOE_BROAD_MLP_TARGETS) |
There was a problem hiding this comment.
When an explicit attention-only list is routed through get_peft_regex because a family scope is active (for example FastVisionModel.get_peft_model's _scoping path), the generated regex can contain the full component block mlp|feed_forward|ffn|dense while its leaf group is still only q_proj|k_proj|v_proj|o_proj. This fallback therefore returns all MoE expert parameters even though no MLP projection leaf was selected, so language-only/attention-only MoE VLM finetunes still train the experts and pay the extra expert LoRA cost.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. Confirmed: an attention-only explicit list under family scoping (vision off) routes through get_peft_regex and the emitted regex carries the full mlp|feed_forward|ffn|dense component block while its leaf group is only q/k/v/o_proj, so the fallback trained the experts for a language-only request. The string alone cannot tell this apart from the fused-expert auto regex, so I now feed the caller's original leaf list into expert detection; only the auto (None/all-linear) path still relies on the regex. Attention-only lists no longer target experts, and fused-expert auto detection is unchanged. Added a regression test.
An attention-only explicit target_modules list routed through get_peft_regex for family scoping (e.g. FastVisionModel with vision layers off) yields a regex carrying the full mlp|feed_forward|ffn|dense component block even though its leaf group only names q/k/v/o_proj. Keying expert detection on that regex trained the experts for a language-only/attention-only request. Use the caller's original leaf list for detection; only the auto path uses the regex, where the mlp block is the sole MLP-intent signal on fused-expert models.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9d98b36ca9
ℹ️ 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".
| _moe_targets = _moe_detect_target if _moe_detect_target is not None else target_modules | ||
| target_parameters = get_moe_target_parameters(model, _moe_targets) |
There was a problem hiding this comment.
When an explicit projection list is routed through the family filters above, this switches MoE detection back to the unscoped original list. For example, with target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"] and finetune_mlp_modules=False, lines 1664-1679 correctly remove MLP targets from target_modules, but this still passes the original list to get_moe_target_parameters, which adds mlp.experts.* via target_parameters; the run ends up training MoE MLP experts despite an attention-only/frozen-MLP request. Please use the scoped result whenever the layer-family filters exclude MLP/language modules, while preserving the original-list behavior only for the attention-only leaf-list case it was meant to fix.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, this was a real gap. Fixed in 787d3b9: MoE expert detection now prefers the original leaf list only when both MLP and language families are still in scope. When finetune_mlp_modules or finetune_language_layers is False, the scoped regex (which already drops the experts) is honored instead of re-adding the original list's gate/up/down leaves. Factored into _select_moe_detection_targets with unit tests over the full matrix (frozen-mlp, vision-only, in-scope-mlp, attention-only).
…oE expert detection When an explicit leaf list that names MLP projections (gate_proj/up_proj/down_proj) is routed through get_peft_regex under finetune_mlp_modules=False, the scoped regex correctly drops the MLP leaves, but MoE expert detection was still keyed on the original list and re-added mlp.experts.* via target_parameters, training the experts the caller had frozen. Same gap for finetune_language_layers=False on vision-only runs. Prefer the original list only when MLP and language families are both in scope (preserving the attention-only fix); otherwise honor the scoped result so the frozen family is respected. Factored the choice into _select_moe_detection_targets with unit tests over the full selection matrix.
for more information, see https://pre-commit.ci
|
@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". |
1 participant
Summary
_moe_target_set_from_stringtreated any regex containing the substringmlporffnas targeting the expert MLP projections. Unsloth's auto-generated attention-only regex listsmlp,ffnandfeed_forwardas allowed intermediate path segments while its final group matches only q/k/v/o projections, so attention-only finetuning on MoE models silently enabled expert LoRA as well: the experts were trained and every MoE layer paid the extra expert LoRA grouped matmuls.What this does
Detects expert intent from the projection names themselves (
gate_proj/up_proj/down_proj/gate_up_proj) instead of themlpsubstring.