diff --git a/tests/test_moe_lora_targets.py b/tests/test_moe_lora_targets.py index 994d39f261..7f9b9a0485 100644 --- a/tests/test_moe_lora_targets.py +++ b/tests/test_moe_lora_targets.py @@ -49,3 +49,190 @@ 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", + ] + + +def test_explicit_attention_only_list_does_not_discover_moe_parameters(): + # An explicit attention-only leaf list names no MLP projection, so experts + # must never be targeted. get_peft_model routes this ORIGINAL list (not the + # scoped regex) into detection precisely because family scoping makes + # get_peft_regex emit its full "mlp|feed_forward|ffn|dense" component block + # even for an attention-only request (see the regex below), which the + # string fallback cannot distinguish from the fused-expert auto regex. + from unsloth.models._utils import get_moe_target_parameters + + attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"] + assert get_moe_target_parameters(_FakeMoeModel(), attn_only_list) is None + assert get_moe_target_parameters(_FakeMoeModel(), tuple(attn_only_list)) is None + + # The regex get_peft_regex emits for that same attention-only list under a + # vision-off family scope carries the mlp component block, so the string + # path would wrongly enable experts -- hence detection must use the list. + scoped_regex = ( + r"(?:.*?(?:language|text).*?" + 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(), scoped_regex) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_frozen_mlp_full_list_does_not_discover_moe_parameters(): + # Regression: an explicit list that names MLP leaves together with + # finetune_mlp_modules=False must NOT train experts. get_peft_regex scopes + # the MLP leaves out (its emitted regex carries no mlp tag block), so + # detection has to key on that SCOPED regex -- keying on the original list + # would let its gate/up/down leaves silently re-enable the frozen experts. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + # Representative of what get_peft_regex emits for that list under + # finetune_mlp_modules=False: attention-only path, no mlp component block. + scoped_regex = ( + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = False, + finetune_language_layers = True, + ) + assert selected is scoped_regex + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None + + +def test_frozen_language_full_list_does_not_discover_moe_parameters(): + # Vision-only request (finetune_language_layers=False) with a full leaf list + # must not reach the language-model experts either. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = ["q_proj", "gate_proj", "up_proj", "down_proj"] + scoped_regex = ( + r"(?:.*?(?:vision|visual|image).*?" + r"(?:self_attn|attention|attn|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = False, + ) + assert selected is scoped_regex + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None + + +def test_in_scope_mlp_full_list_still_discovers_moe_parameters(): + # With MLP and language both in scope, an explicit list that names MLP + # leaves SHOULD enable the experts (unchanged behavior): the original list + # is preferred and carries the gate/up/down intent. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + scoped_regex = r".*self_attn.*proj" # unused: original list is preferred + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = True, + ) + assert selected is original_list + assert get_moe_target_parameters(_FakeMoeModel(), selected) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_attention_only_list_prefers_original_when_in_scope(): + # The case the PR originally fixed: an attention-only list routed through + # get_peft_regex under a family scope (e.g. vision-off) still keeps experts + # off, because with MLP+language in scope detection uses the original + # attention-only list rather than the regex's spurious mlp component block. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"] + scoped_regex = ( # carries the spurious mlp block get_peft_regex always adds + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + attn_only_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = True, + ) + assert selected is attn_only_list + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 599a5c0262..c8b85fd44b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -85,6 +85,7 @@ "hf_login", "is_moe_model", "get_moe_target_parameters", + "_select_moe_detection_targets", "make_fast_generate_wrapper", "_mark_unsloth_disable_data_parallel", "_patch_transformers_trainer_data_parallel", @@ -3507,8 +3508,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"(? Optional[List[str return None +def _select_moe_detection_targets( + original_target_modules, + scoped_target_modules, + finetune_mlp_modules = True, + finetune_language_layers = True, +): + """Pick what get_moe_target_parameters keys expert detection on. + + Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an + attention-only request is not pushed into the experts by get_peft_regex's + ``mlp|feed_forward|ffn|dense`` component block (which the string fallback + cannot tell apart from a fused-expert auto regex). + + But only when the MLP and language families are BOTH still in scope. If the + caller scoped MLP or language OFF (``finetune_mlp_modules=False`` or + ``finetune_language_layers=False``) the scoped regex already drops the MoE + experts, and reusing the original list -- which may still name gate/up/down + leaves -- would wrongly re-introduce them. In that case honor the scoped + result so the frozen-MLP / vision-only request is respected. + """ + if original_target_modules is not None and finetune_mlp_modules and finetune_language_layers: + return original_target_modules + return scoped_target_modules + + def make_fast_generate_wrapper(original_generate): """ Creates a wrapper around model.generate that checks for incorrect diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 5ab55152db..f1fa6fe188 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -37,6 +37,7 @@ _get_text_only_config, _is_family_text_decoder, _apply_text_only_key_mapping, + _select_moe_detection_targets, set_task_config_attr, ) from ._utils import * @@ -1629,6 +1630,16 @@ def get_peft_model( ) else: _audio_kwargs = {} + # Remember the caller's ORIGINAL explicit leaf list for MoE expert + # detection. When an explicit list is routed through get_peft_regex for + # family scoping below, the generated regex carries get_peft_regex's full + # "mlp|feed_forward|ffn|dense" component block even when the caller named + # only attention leaves (q/k/v/o_proj). Keying expert detection on that + # regex would train the experts for an attention-only request. The + # original list carries the true leaf intent, so use it for MoE detection; + # only the auto (None / "all-linear") path relies on the regex, whose mlp + # block is the sole remaining MLP-intent signal on fused-expert models. + _moe_detect_target = target_modules if type(target_modules) in (list, tuple) else None if target_modules is None or target_modules == "all-linear": target_modules = get_peft_regex( model, @@ -1706,9 +1717,21 @@ def get_peft_model( loftq_config, lora_dropout, bias, init_lora_weights, model ) - # Auto-detect MoE models and populate target_parameters for expert layers + # Auto-detect MoE models and populate target_parameters for expert layers. + # Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an + # attention-only request does not train experts via get_peft_regex's mlp block, + # but only when MLP and language families are both still in scope. If the caller + # scoped MLP or language OFF (finetune_mlp_modules / finetune_language_layers + # False), the scoped regex already dropped the experts, so honor it instead of + # re-introducing the original list's gate/up/down leaves. if target_parameters is None: - target_parameters = get_moe_target_parameters(model, target_modules) + _moe_targets = _select_moe_detection_targets( + _moe_detect_target, + target_modules, + finetune_mlp_modules = finetune_mlp_modules, + finetune_language_layers = finetune_language_layers, + ) + target_parameters = get_moe_target_parameters(model, _moe_targets) if finetune_last_n_layers is not None and layers_to_transform is None: _total_layers = _get_total_transformer_layers(model)