mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
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:
parent
c520662c12
commit
c7b8666ce4
3 changed files with 44 additions and 0 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue