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>
This commit is contained in:
Daniel Han 2026-07-06 05:46:30 -07:00 committed by GitHub
parent c520662c12
commit c7b8666ce4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 44 additions and 0 deletions

View file

@ -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)

View file

@ -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:

View file

@ -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