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.<i> and mlp.experts.down_projs.<i> 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.<i> / down_projs.<i>). 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

@@ -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,12 +4062,16 @@ 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:
+        if _moe_parameter_exists(model, gate_up_name):
             moe_params.append(gate_up_name)
 
     if "down_proj" in target_set:
+        if _moe_parameter_exists(model, down_name):
             moe_params.append(down_name)
 
     if moe_params:
@@ -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.<i> /
+    down_projs.<i> 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.<i>`` 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.<i> / down_projs.<i> 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,

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)

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: