[MoE] Improve moe kernels for unsloth fine tuning#3812
Conversation
Summary of ChangesHello @Datta0, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request focuses on significantly improving the performance and usability of Mixture-of-Experts (MoE) kernels, particularly for Qwen3 models, within the Unsloth framework. The core changes involve integrating Triton-optimized grouped GEMM operations, implementing a robust caching system for kernel configurations to reduce overhead, and ensuring better compatibility with "torch.compile" by handling tracing and dynamic kernel parameters. These enhancements aim to provide a more efficient and streamlined fine-tuning experience for MoE architectures. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| N = W.shape[0] // num_experts | ||
| assert K == W.shape[1], f"K ({K}) must match W.shape[1] ({W.shape[1]})" | ||
|
|
||
| if fuse_mul_post: |
There was a problem hiding this comment.
The fused-mul path now references _FUSED_MUL_WARN but the module no longer defines the flag after the TMA refactor, so any call with fuse_mul_post=True hits a NameError before executing the kernel. This is a regression from the previous version where the guard was initialized and will crash inference paths that rely on post-mul fusion.
Useful? React with 👍 / 👎.
| expert_mask = torch.nn.functional.one_hot( | ||
| selected_experts, num_classes = num_experts | ||
| ).permute(2, 1, 0) | ||
|
|
||
| token_counts_by_expert = expert_mask.sum(dim = 1).int() |
There was a problem hiding this comment.
In the Triton forward for Qwen3 MoE, token_counts_by_expert is computed as expert_mask.sum(dim=1) which yields a [num_experts, num_tokens] matrix instead of the expected 1‑D counts per expert. The grouped GEMM kernels read m_sizes as a length‑num_experts vector, so this flattened matrix feeds arbitrary per‑token entries into the kernel, producing incorrect routing and outputs whenever the Triton MoE path is enabled.
Useful? React with 👍 / 👎.
![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 introduces a valuable auto-tuning cache system for MoE Triton kernels, which will prevent re-tuning on each run and significantly improve startup times. The changes also enhance torch.compile compatibility and add a new Triton-optimized path for Qwen3-MoE models. While the overall direction is excellent, I've found a critical bug in the token routing logic within the new qwen3_moe_triton.py file that needs to be addressed. Additionally, there are a couple of smaller issues regarding TMA support detection and error message handling that I've flagged for improvement.
| def _check_tma_support(): | ||
| import triton.language as tl | ||
|
|
||
| gpu_supports_tma = torch.cuda.get_device_capability()[0] >= 9 | ||
| # Check for both old experimental and new stable API names | ||
| triton_has_tma_api = hasattr(tl, "make_tensor_descriptor") or hasattr( |
There was a problem hiding this comment.
The current implementation of _check_tma_support can lead to runtime errors. It checks for the existence of either make_tensor_descriptor or _experimental_make_tensor_descriptor, but the kernels now exclusively use make_tensor_descriptor. If a user has a Triton version with only the experimental API, _SUPPORTS_TMA will be True, but the kernel launch will fail with an AttributeError. The check should be made stricter to only hasattr(tl, 'make_tensor_descriptor') to match the kernel's expectation, or the kernels should be updated to handle both API versions.
| def _check_tma_support(): | |
| import triton.language as tl | |
| gpu_supports_tma = torch.cuda.get_device_capability()[0] >= 9 | |
| # Check for both old experimental and new stable API names | |
| triton_has_tma_api = hasattr(tl, "make_tensor_descriptor") or hasattr( | |
| def _check_tma_support(): | |
| import triton.language as tl | |
| gpu_supports_tma = torch.cuda.get_device_capability()[0] >= 9 | |
| # Kernels now use the stable `make_tensor_descriptor` API. | |
| triton_has_tma_api = hasattr(tl, 'make_tensor_descriptor') | |
| return gpu_supports_tma and triton_has_tma_api |
This reverts commit da9e2c4.
This reverts commit 9bb0db2.
* remote unused hybridcache * cleanup
for more information, see https://pre-commit.ci
* Improve MoE performance * small changes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix imports * disable autotune * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * LoRA for MoE * Make autotune default * make dy contiguous * use non lora model as base for RL * Revert "use non lora model as base for RL" This reverts commit bc8f156. * fixup derp * non TMA [T4] * Revert "non TMA [T4]" This reverts commit 3530456. * Fixes for VL MoE and v5 transformers * [transformers] [v5] remove unused hybridcache (unslothai#3910) * remote unused hybridcache * cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * No double compile for qwen3moe * Fix top_k on trl GRPO * Recognise GLM as MoE * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix missing RotaryEmbeddingConfigMixin * Licensing for autotuning cache * Cleanup --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Erland366 <erland.pg366@gmail.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
3 participants
This might come in handy for unslothai/unsloth-zoo#396
We do prefer using grouped_mm there but this is a fallback for that before going pure pytorch mode
Needs transformers v5 there