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

Merged
danielhanchen merged 8 commits into from
Jul 6, 2026
+44 −0
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 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
bd7d58d
Merge remote-tracking branch 'origin/main' into moe-grouped-autohook
danielhanchen Jul 6, 2026
25a03e2
Re-evaluate grouped MoE after loading a PEFT adapter on the vision path
danielhanchen Jul 6, 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
13 changes: 13 additions & 0 deletions unsloth/models/llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3743,4 +3743,17 @@ def _for_training(m):

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

@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 👍 / 👎.

@danielhanchen danielhanchen Jul 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Intentional. The import is a guarded no-op until unsloth-zoo#837 lands, which ships this exact module (temporary_patches/moe_grouped_modulelist.py). This loader change is meant to land independently and activate once the zoo side merges, so the path already points at the shipped module.

FastLlamaModel.from_pretrained = staticmethod(
wrap_loader_for_grouped_moe(FastLlamaModel.from_pretrained)
)
Comment on lines +3750 to +3752

@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 👍 / 👎.

@danielhanchen danielhanchen Jul 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Covered by the post-attach auto_enable_grouped_moe re-eval in the FastLanguageModel loader path, which runs after PeftModel.from_pretrained on the loaded-adapter path. Same re-eval now added to the vision path in 25a03e2.

FastLlamaModel.get_peft_model = staticmethod(
wrap_loader_for_grouped_moe(FastLlamaModel.get_peft_model)
)
except Exception:
pass

PatchFastRL(FastLanguageModel = FastLlamaModel)
18 changes: 18 additions & 0 deletions unsloth/models/loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -876,6 +876,15 @@ def from_pretrained(
)
# 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}""
Expand Down Expand Up @@ -1809,6 +1818,15 @@ def _patched_car(
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:
Expand Down
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 @@ -2207,3 +2207,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 +2215 to +2217

@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 👍 / 👎.

@danielhanchen danielhanchen Jul 6, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Fixed in 25a03e2. The FastModel / FastVisionModel is_peft path now re-runs auto_enable_grouped_moe after PeftModel.from_pretrained, so an expert-LoRA block restores the original loop while an attention-only adapter keeps the grouped path, matching the FastLanguageModel loader path.

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