mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
Scope MoE expert LoRA detection to actual MLP projection targets (#6849)
* Scope MoE expert LoRA detection to actual MLP projection targets _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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments * Detect MoE expert LoRA via mlp path segment, not proj names 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. * Scope expert LoRA targets to the leaves a regex names 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. * Route explicit leaf list into MoE expert detection 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Respect finetune_mlp_modules and finetune_language_layers scope for MoE 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
22bd86ecb7
commit
7cc1752a64
3 changed files with 257 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ __all__ = [
|
|||
"maybe_prefetch_hf_snapshot",
|
||||
"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",
|
||||
|
|
@ -3913,8 +3914,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)
|
||||
|
||||
return set()
|
||||
|
|
@ -4000,6 +4018,31 @@ def get_moe_target_parameters(model, target_modules = None) -> 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
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from ._utils import (
|
|||
_get_text_only_config,
|
||||
_is_family_text_decoder,
|
||||
_apply_text_only_key_mapping,
|
||||
_select_moe_detection_targets,
|
||||
set_task_config_attr,
|
||||
)
|
||||
from ._utils import *
|
||||
|
|
@ -1703,6 +1704,16 @@ class FastBaseModel:
|
|||
)
|
||||
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,
|
||||
|
|
@ -1780,9 +1791,21 @@ class FastBaseModel:
|
|||
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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue