[MoE] Updates to MoE, fixes to qwen static method#574
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3e54e64f00
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for Gemma4 MoE, including specialized LoRA extraction and optimized grouped GEMM backends. It improves compatibility with Transformers v5 by handling variations in function signatures (such as the presence or absence of cache_position) across Gemma3 and CSM models. Furthermore, it updates GRPO training logic to support additional token type IDs and refines MoE merging to handle diverse 3D tensor formats. Review feedback highlights that commenting out @torch.compiler.disable decorators may inadvertently re-enable compilation for incompatible modules, potentially leading to runtime errors.
I am having trouble creating individual review comments. Click here to see my feedback.
unsloth_zoo/compiler.py (1252)
Commenting out the @torch.compiler.disable decorator effectively re-enables compilation for modules that were previously explicitly marked as incompatible with torch.compile. If these modules still contain operations that cause graph breaks or compilation errors, this change will lead to runtime failures when disable is set to True (e.g., when UNSLOTH_COMPILE_DISABLE=partial). If the intention is to allow compilation, it would be cleaner to remove the logic or use @torch.compile explicitly rather than leaving commented-out code in the generated source.
unsloth_zoo/compiler.py (4022-4026)
Similar to the change in create_standalone_class, commenting out the @torch.compiler.disable decorator in the function source code will cause torch.compile to attempt to compile these functions if they are called from a compiled context. This bypasses the explicit disabling mechanism for functions in disable_compile_functions. If these functions are now compatible with compilation, they should be removed from the list; otherwise, commenting out the decorator may re-introduce previously avoided compilation errors.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9133bc567f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0d0d10ba67
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Brings PR unslothai#527's moeFix branch up to current main (HEAD 57bbdc0). Conflict resolution: - unsloth_zoo/temporary_patches/qwen3_moe.py: took main's version (PR unslothai#574 already added staticmethod() + a refactored extractor that uses extract_moe_lora_weights_for_grouped_mm; my local layout-aware rewrite is now redundant). - unsloth_zoo/temporary_patches/qwen3_vl_moe.py: took main's version (same reason; PR unslothai#574 already wrapped the extractor with staticmethod). Auto-merged cleanly: - unsloth_zoo/temporary_patches/glm4_moe.py: keeps the new patch_glm4_moe_standard registration alongside main's helper-based refactor of patch_glm4_moe. - unsloth_zoo/saving_utils.py: PR unslothai#647 (saving fixes) is already in main as faee224, so my three saving cherry-picks are subsumed. Stashed-then-dropped local shims: - The local _active_merge_device backport in saving_utils.py and the _unsloth_get_mm_token_id / _unsloth_fix_mm_token_type_ids compat shims in rl_replacements.py were never committed; both symbols now exist in main, so the shims were dropped. Verified post-merge: - python -c "import unsloth; import unsloth_zoo.temporary_patches.{glm4_moe,moe_bnb_transformers,qwen3_moe,misc}" succeeds - patch_glm4_moe_standard, _ParamShapeProxy, patch_peft_param_wrapper_4bit_expert_shape, patch_peft_param_wrapper_merge_4bit are all reachable.
2 participants
Also fixes for transformers v5.3