From 62a6eb2a3df395e2e0e94218cfc963318f5c7392 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 05:57:44 -0700 Subject: [PATCH] MoE LoRA: auto-target per-expert Linear experts (gpt-oss 4bit) instead of leaving them frozen (#6936) * models: auto-target per-expert Linear MoE experts for LoRA (gpt-oss 4bit) 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. and mlp.experts.down_projs. 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. / down_projs.). 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. * models: scope per-expert MoE targets, keep repeat get_peft_model idempotent, 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. * [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> --- unsloth/models/_utils.py | 119 +++++++++++++++++++++++++++++++++++++-- unsloth/models/llama.py | 18 ++++++ unsloth/models/vision.py | 29 ++++++++++ 3 files changed, 162 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1c75f8ce6..b68cb702b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -86,6 +86,8 @@ __all__ = [ "maybe_prefetch_hf_snapshot", "is_moe_model", "get_moe_target_parameters", + "get_moe_target_modules", + "warn_if_zoo_cannot_merge_moe_experts", "_select_moe_detection_targets", "make_fast_generate_wrapper", "_mark_unsloth_disable_data_parallel", @@ -4060,13 +4062,17 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str alternate_name = "experts.down_proj", ) - # gate_up_proj combines both gate_proj and up_proj in MoE - # Also match "gate_up_proj" directly since users may specify the fused name + # gate_up_proj combines gate_proj and up_proj; also match the fused name directly. + # Only target a fused expert Parameter that exists: per-expert Linear layouts + # (e.g. gpt-oss bnb-4bit) have no fused Parameter and are handled by + # get_moe_target_modules, so skip them rather than pass PEFT a dead path. if "gate_proj" in target_set or "up_proj" in target_set or "gate_up_proj" in target_set: - moe_params.append(gate_up_name) + if _moe_parameter_exists(model, gate_up_name): + moe_params.append(gate_up_name) if "down_proj" in target_set: - moe_params.append(down_name) + if _moe_parameter_exists(model, down_name): + moe_params.append(down_name) if moe_params: print( @@ -4077,6 +4083,111 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str return None +def _moe_parameter_exists(model, name: str) -> bool: + """True if ``name`` is an exact suffix of some parameter path on the model.""" + if not hasattr(model, "named_parameters"): + return False + try: + for parameter_name, _ in model.named_parameters(): + if parameter_name == name or parameter_name.endswith("." + name): + return True + except Exception: + return False + return False + + +def get_moe_target_modules(model, target_modules = None) -> List[str]: + """Per-expert ``target_modules`` suffixes for MoE models whose experts are stored + as per-expert ``nn.Linear`` ModuleLists rather than fused nn.Parameters. + + gpt-oss bnb-4bit is the canonical case (mlp.experts.gate_up_projs. / + down_projs. as Linear4bit): no fused Parameter, and the plain + gate/up/down_proj leaves do not match, so LoRA skips them. Returning the + per-expert suffixes makes PEFT attach via ordinary suffix matching (the + module-LoRA counterpart of get_moe_target_parameters). Returns [] for non-MoE, + fused-parameter MoEs, an absent per-expert layout, or a request that omits the + MLP experts (so an attention-only run does not train experts). + """ + if not is_moe_model(model): + return [] + if target_modules is None: + return [] + if isinstance(target_modules, str): + target_set = _moe_target_set_from_string(target_modules) + else: + target_set = { + target + for target in target_modules or () + if (isinstance(target, str) and "." not in target and target in _MOE_BROAD_MLP_TARGETS) + } + if not (target_set & _MOE_BROAD_MLP_TARGETS): + return [] + + if not hasattr(model, "named_modules"): + return [] + + # Scope the returned suffixes to the requested projection leaves, matching + # get_moe_target_parameters: gate_proj/up_proj/gate_up_proj map to the fused + # gate_up ModuleList (e.g. gate_up_projs); down_proj maps to the down ModuleList + # (e.g. down_projs). A down-only (or gate/up-only) request must not pull in the + # other projection. + want_gate_up = bool(target_set & {"gate_proj", "up_proj", "gate_up_proj"}) + want_down = "down_proj" in target_set + + targets = set() + for name, module in model.named_modules(): + if not isinstance(module, torch.nn.ModuleList) or len(module) == 0: + continue + parent, _, leaf = name.rpartition(".") + # ModuleList directly under an ``experts`` container, holding only Linear + # leaves (bnb Linear4bit / Linear8bitLt subclass nn.Linear). After PEFT has + # wrapped the experts the child is a LoRA layer whose ``base_layer`` is the + # Linear, so accept that too (keeps this idempotent across a re-wrapped model). + if not parent.endswith("experts"): + continue + if not all( + isinstance(child, torch.nn.Linear) + or isinstance(getattr(child, "base_layer", None), torch.nn.Linear) + for child in module + ): + continue + # Honor the requested subset: classify the ModuleList by projection role. + leaf_lower = leaf.lower() + is_down = "down" in leaf_lower + is_gate_up = (not is_down) and ("gate" in leaf_lower or "up" in leaf_lower) + if is_down and not want_down: + continue + if is_gate_up and not want_gate_up: + continue + # One entry per expert index; ``leaf.`` matches expert i in every layer. + for expert_index in range(len(module)): + targets.add(f"{leaf}.{expert_index}") + + return sorted(targets) + + +def warn_if_zoo_cannot_merge_moe_experts(): + """Warn once when the installed unsloth_zoo cannot fold per-expert Linear MoE LoRA + into a merged_16bit checkpoint. Older zoo releases keep the fused gate_up_proj / + down_proj tensors and drop the per-expert gate_up_projs. / down_projs. deltas, + so save_pretrained_merged("merged_16bit") would silently lose the expert training + (the LoRA adapter itself still saves and reloads correctly).""" + try: + from unsloth_zoo import saving_utils as _saving_utils + + # _fold_perexpert_lora_into_fused is the helper that folds these experts. + if hasattr(_saving_utils, "_fold_perexpert_lora_into_fused"): + return + except Exception: + return # cannot introspect zoo -> stay quiet rather than false-alarm + logger.warning_once( + "Unsloth: the installed unsloth_zoo will not fold these per-expert experts into " + "a merged_16bit checkpoint, so save_pretrained_merged('merged_16bit') would drop " + "the expert LoRA. Upgrade unsloth_zoo to merge them; saving the LoRA adapter is " + "unaffected." + ) + + def _select_moe_detection_targets( original_target_modules, scoped_target_modules, diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a1da09975..1f43f6144 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3105,6 +3105,11 @@ class FastLlamaModel: new_target_modules = list(target_modules) + list( modules_to_save if modules_to_save is not None else [] ) + # Per-expert Linear MoE experts (e.g. gpt-oss bnb-4bit) were auto-added to the + # saved target_modules when the adapter was first created. Recompute them so a + # repeat get_peft_model call with the same args stays idempotent instead of + # tripping the mismatch below. No-op for non per-expert-Linear models. + new_target_modules += get_moe_target_modules(model, target_modules) # Now check! new_target_modules = set(new_target_modules) @@ -3331,6 +3336,19 @@ class FastLlamaModel: if target_parameters is None: target_parameters = get_moe_target_parameters(model, target_modules) + # Per-expert Linear expert layouts (e.g. gpt-oss bnb-4bit) are Linear modules, + # not fused Parameters, so target them via target_modules. No-op otherwise. + _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) + if _added: + print( + f"Unsloth: Detected MoE model with per-expert Linear experts. " + f"Enabling LoRA on {len(_added)} expert projection modules." + ) + warn_if_zoo_cannot_merge_moe_experts() + if finetune_last_n_layers is not None and layers_to_transform is None: from .vision import _get_total_transformer_layers _total_layers = _get_total_transformer_layers(model) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 0a68a49fe..0235d80e9 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1807,6 +1807,35 @@ class FastBaseModel: ) target_parameters = get_moe_target_parameters(model, _moe_targets) + # Per-expert Linear expert layouts (e.g. gpt-oss bnb-4bit) target experts via + # target_modules, not fused Parameters. Extend either form PEFT accepts: a leaf + # list (explicit) or a regex string (auto / all-linear / scoped). No-op otherwise. + _moe_module_detect = _select_moe_detection_targets( + _moe_detect_target, + target_modules, + finetune_mlp_modules = finetune_mlp_modules, + finetune_language_layers = finetune_language_layers, + ) + _moe_module_targets = get_moe_target_modules(model, _moe_module_detect) + if _moe_module_targets: + if isinstance(target_modules, (list, tuple)): + target_modules = list(target_modules) + [ + target for target in _moe_module_targets if target not in target_modules + ] + elif isinstance(target_modules, str): + _expert_leaves = sorted({t.rsplit(".", 1)[0] for t in _moe_module_targets}) + _expert_alt = ( + r".*\.experts\.(?:" + + "|".join(re.escape(leaf) for leaf in _expert_leaves) + + r")\.\d+" + ) + target_modules = f"(?:{target_modules})|(?:{_expert_alt})" + print( + f"Unsloth: Detected MoE model with per-expert Linear experts. " + f"Enabling LoRA on {len(_moe_module_targets)} expert projection modules." + ) + warn_if_zoo_cannot_merge_moe_experts() + if finetune_last_n_layers is not None and layers_to_transform is None: _total_layers = _get_total_transformer_layers(model) if _total_layers is not None and _total_layers > 0: