Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danielhanchen wants to merge 6 commits into
base: main
Choose a base branch
from
+35 −0
Open

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

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f9d010d
Auto-enable grouped MoE on loaded / PEFT'd models via loader hook
danielhanchen Jun 28, 2026
c0bc6ef
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 28, 2026
51f07a6
Install grouped-MoE loader wrappers before PatchFastRL
danielhanchen Jun 28, 2026
0085ca4
Re-evaluate grouped MoE after loading a PEFT adapter
danielhanchen Jul 1, 2026
ef8d27b
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 1, 2026
ec6fe59
Trim comments in the grouped MoE loader hooks
danielhanchen Jul 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions unsloth/models/llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3736,3 +3736,17 @@ def _for_training(m):
from .rl import PatchFastRL

PatchFastRL(FastLanguageModel = FastLlamaModel)


# Auto-enable grouped-GEMM MoE (transformers<5 ModuleList experts) on built / PEFT'd
# models. Wraps the loader leaves once; guarded so it never breaks model loading.
try:
from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Point the grouped-MoE import at a shipped module

In normal installs using the declared unsloth_zoo>=2026.6.7 dependency, this module path does not resolve, and the broad except silently skips the whole wrapping block. That means transformers<5 ModuleList MoE models loaded through this path (and the matching vision.py block) keep using the old expert-loop implementation instead of the intended grouped-GEMM patch, so the new auto-enable behavior is effectively a no-op unless users happen to have an unpublished zoo build.

Useful? React with 👍 / 👎.

FastLlamaModel.from_pretrained = staticmethod(
wrap_loader_for_grouped_moe(FastLlamaModel.from_pretrained)
)
Comment on lines +3742 to +3744

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Patch PEFT adapter loads after adapters attach

When a user loads an existing LoRA/PEFT adapter repo through FastLanguageModel.from_pretrained, this wrapped leaf only returns the base model; loader.py attaches the adapter afterward via PeftModel.from_pretrained and dispatch_model.patch_peft_model at lines 865-878. The grouped-MoE wrapper therefore never sees the final PEFT model for that common adapter-loading path, while wrapping get_peft_model only covers newly-created adapters, so loaded MoE adapter checkpoints miss the intended PEFT-aware grouped-GEMM conversion.

Useful? React with 👍 / 👎.

FastLlamaModel.get_peft_model = staticmethod(

@gemini-code-assist gemini-code-assist Bot Jun 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To ensure that the grouped-GEMM MoE loader wrapper is always active and cannot be bypassed, it is highly recommended to apply this patch before calling PatchFastRL(FastLanguageModel = FastLlamaModel). If PatchFastRL or any other downstream registration copies or wraps from_pretrained or get_peft_model from FastLlamaModel, applying this patch afterwards might result in the RL-patched entry points using the unwrapped/original methods, silently missing the MoE optimization.

wrap_loader_for_grouped_moe(FastLlamaModel.get_peft_model)
)
except Exception:
pass
13 changes: 13 additions & 0 deletions unsloth/models/vision.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2119,3 +2119,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)
)
Comment on lines +2127 to +2129

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 28, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Move grouped MoE enable after adapter attachment

When FastModel/FastVisionModel loads a PEFT adapter repo (is_peft), loader.py first calls FastBaseModel.from_pretrained and only attaches the adapter later via PeftModel.from_pretrained. Wrapping this leaf enables grouped MoE on the base model before any LoRA modules exist, so the no-LoRA-on-experts eligibility check can pass and there is no later recheck after adapter attachment; expert-targeted MoE adapters can then run through the grouped forward built from base weights and silently ignore those adapter weights. Please move or rerun the grouped-MoE hook after the is_peft adapter attach path as well.

Useful? React with 👍 / 👎.

FastBaseModel.get_peft_model = staticmethod(
wrap_loader_for_grouped_moe(FastBaseModel.get_peft_model)
)
except Exception:
pass
Loading