From c7b8666ce43d223fb3ffb1ccce89d73e26291e4a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 6 Jul 2026 05:46:30 -0700
Subject: [PATCH] Auto-enable grouped MoE on loaded / PEFT'd models via loader
 hook (#6727)

* Auto-enable grouped MoE on loaded / PEFT'd models via loader hook

Wraps the FastLlamaModel and FastBaseModel from_pretrained / get_peft_model leaves with wrap_loader_for_grouped_moe so the grouped-GEMM MoE forward is installed on the live instance after the model and its compiled module are built. Gated by UNSLOTH_MOE_GROUPED and wrapped in try/except, so it is a no-op when the unsloth_zoo module is absent or no eligible MoE block exists.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Install grouped-MoE loader wrappers before PatchFastRL

* Re-evaluate grouped MoE after loading a PEFT adapter

When loading an existing adapter through FastLanguageModel.from_pretrained,
the base model is evaluated for grouped MoE when the wrapped from_pretrained
leaf returns, but the adapter is attached afterwards via PeftModel and
patch_peft_model. Re-run auto_enable_grouped_moe on the final model so
blocks whose experts gained LoRA are restored to the original loop,
attention-only adapters keep the grouped path on their frozen experts, and
recompute is re-derived from the final gradient-checkpointing state. Guarded
so it never blocks adapter loading.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim comments in the grouped MoE loader hooks

Shorten the loader re-eval and llama.py wrapper comments; code is unchanged
(verified comment-only).

* Re-evaluate grouped MoE after loading a PEFT adapter on the vision path

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 unsloth/models/llama.py  | 13 +++++++++++++
 unsloth/models/loader.py | 18 ++++++++++++++++++
 unsloth/models/vision.py | 13 +++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 417a88f480..564be09578 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -3824,4 +3824,17 @@ class FastLlamaModel:
 
 from .rl import PatchFastRL
 
+# Auto-enable grouped-GEMM MoE (tf<5 ModuleList experts) on built / PEFT'd models. Wrap the
+# loader leaves before PatchFastRL so downstream patchers see the wrapped versions. Guarded.
+try:
+    from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe
+    FastLlamaModel.from_pretrained = staticmethod(
+        wrap_loader_for_grouped_moe(FastLlamaModel.from_pretrained)
+    )
+    FastLlamaModel.get_peft_model = staticmethod(
+        wrap_loader_for_grouped_moe(FastLlamaModel.get_peft_model)
+    )
+except Exception:
+    pass
+
 PatchFastRL(FastLanguageModel = FastLlamaModel)
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 22cb65dc4a..2818b6ee80 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -896,6 +896,15 @@ class FastLanguageModel(FastLlamaModel):
             )
             # Patch it as well!
             model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing)
+            # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back
+            # to the original loop, an attention-only adapter keeps the grouped path. Guarded.
+            try:
+                from unsloth_zoo.temporary_patches.moe_grouped_modulelist import (
+                    auto_enable_grouped_moe,
+                )
+                auto_enable_grouped_moe(model)
+            except Exception:
+                pass  # optional speedup; never block model loading
 
         # Patch Tiled MLP
         # to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}""
@@ -1852,6 +1861,15 @@ class FastModel(FastBaseModel):
             model = FastBaseModel.post_patch_model(
                 model, use_gradient_checkpointing, trust_remote_code = trust_remote_code
             )
+            # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back
+            # to the original loop, an attention-only adapter keeps the grouped path. Guarded.
+            try:
+                from unsloth_zoo.temporary_patches.moe_grouped_modulelist import (
+                    auto_enable_grouped_moe,
+                )
+                auto_enable_grouped_moe(model)
+            except Exception:
+                pass  # optional speedup; never block model loading
 
         # Apply QAT if specified
         if qat_scheme is not None:
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index e80ed8917b..0a68a49fee 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -2304,3 +2304,16 @@ def check_dataset_for_missing_videos(
         warnings.warn(error_msg, stacklevel = 2)
 
     return missing
+
+
+# Auto-enable grouped-GEMM MoE (transformers<5 ModuleList experts); see llama.py.
+try:
+    from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe
+    FastBaseModel.from_pretrained = staticmethod(
+        wrap_loader_for_grouped_moe(FastBaseModel.from_pretrained)
+    )
+    FastBaseModel.get_peft_model = staticmethod(
+        wrap_loader_for_grouped_moe(FastBaseModel.get_peft_model)
+    )
+except Exception:
+    pass