Skip to content
Merged
41 changes: 41 additions & 0 deletions tests/test_moe_lora_targets.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,3 +49,44 @@ def test_explicit_dotted_module_target_does_not_discover_moe_parameters():
)
is None
)


@pytest.mark.parametrize(
"target_modules",
[
# Attention-only auto-regex lists every projection leaf (incl. gate/up/down)
# but its path segment is attention-only, so experts must NOT be targeted.
r"(?:\bmodel\.layers\.[\d]{1,}\.(?:self_attn|attention|attn|mixer)\.(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))",
".*self_attn.*proj",
# An mlp path alternative with attention-only leaves is still attention-only.
r"model\.layers\.\d+\.(?:mlp|self_attn)\.(?:q_proj|k_proj|v_proj|o_proj)",
],
)
def test_attention_only_regex_does_not_discover_moe_parameters(target_modules):
from unsloth.models._utils import get_moe_target_parameters
assert get_moe_target_parameters(_FakeMoeModel(), target_modules) is None


def test_single_leaf_regex_targets_only_that_projection():
from unsloth.models._utils import get_moe_target_parameters
assert get_moe_target_parameters(_FakeMoeModel(), ".*experts.*down_proj") == [
"mlp.experts.down_proj",
]
assert get_moe_target_parameters(_FakeMoeModel(), ".*mlp.*gate_proj") == [
"mlp.experts.gate_up_proj",
]


def test_auto_regex_mlp_tag_block_discovers_moe_on_fused_models():
# get_peft_regex on a fused-expert model lists only attention Linears as
# leaves; the mlp tag block is the remaining signal of MLP finetune intent.
from unsloth.models._utils import get_moe_target_parameters
both_auto = (
r"(?:\bmodel\.layers\.[\d]{1,}\."
r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer)\."
r"(?:(?:q_proj|k_proj|v_proj|o_proj)))"
)
assert get_moe_target_parameters(_FakeMoeModel(), both_auto) == [
"mlp.experts.gate_up_proj",
"mlp.experts.down_proj",
]
21 changes: 19 additions & 2 deletions unsloth/models/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3507,8 +3507,25 @@ def _moe_target_set_from_string(target_modules: str) -> set[str]:
return {target_modules}

is_regex = re.search(r"[*+?()[\]{}|\\^$]", target_modules) is not None
targets_mlp = "mlp" in target_modules or "ffn" in target_modules
if is_regex and "proj" in target_modules and targets_mlp:
# Key detection on the mlp/ffn/experts path segment (absent from an
# attention-only regex), never on q/k/v/o leaves alone.
targets_mlp_path = any(
tag in target_modules for tag in ("mlp", "ffn", "feed_forward", "experts")
)
if not is_regex or not targets_mlp_path:
return set()
# Explicit expert leaves scope the target set to exactly those leaves.
named = {name for name in _MOE_BROAD_MLP_TARGETS if name in target_modules}
if named:
return named
# A generic projection under an mlp path (e.g. ".*mlp.*proj"): any proj
# occurrence that is not an attention leaf name.
if re.search(r"(?<![qkvo]_)(?<!out_)(?<!in_)proj", target_modules):
return set(_MOE_BROAD_MLP_TARGETS)
# The auto regex on fused-expert models lists only attention Linears as
# leaves; its mlp tag block is the remaining MLP-intent signal. A regex
# like "(mlp|self_attn).(q_proj|o_proj)" has neither and stays attention-only.
if "mlp|feed_forward|ffn|dense" in target_modules:
return set(_MOE_BROAD_MLP_TARGETS)
Comment on lines +3529 to 3530

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Avoid re-enabling experts for attention-only regexes

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 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

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.

Comment on lines +3529 to 3530

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Do not infer experts from the auto component block

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 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

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.


return set()
Expand Down
Loading