fix: minor fix cutlass moe #7442
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @mickqian, 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!
I've addressed a critical bug in the Cutlass Mixture of Experts (MoE) implementation that was causing NaN values in the output. This was primarily due to incorrect handling of expert outputs and their associated weights, specifically an issue with truncation. The fix involves adjusting the expected tensor layouts for MoE weights and directly implementing the weighted sum logic, bypassing a previously problematic helper function. This ensures numerical stability and correctness for FP8 MoE operations.
Highlights
- MoE Weight Tensor Layout Correction: I've updated the
cutlass_fused_experts_fp8function to expectw1_qandw2_qweight tensors in a transposed layout, specifically(E, n*2, k)forw1_qand(E, k, n)forw2_q. This change is reflected in docstrings and internal assertions to ensure correct shape validation. - Direct Weighted Sum Calculation for Expert Outputs: I've replaced the call to
apply_shuffle_mul_sumwith explicit PyTorch operations (shuffle, reshape, element-wise multiplication withtopk_weights, and sum) withincutlass_fused_experts_fp8. This directly addresses the bug causingNaNoutputs by ensuring correct aggregation of expert results without truncation issues. - Optimized Weight Handling in FP8 Quantization: I've removed redundant
transpose(1, 2)operations when passing weights tocutlass_fused_experts_fp8from the FP8 quantization layer. This aligns the weight preparation with the new expected tensor layout of the MoE kernel, potentially improving efficiency by avoiding unnecessary transpositions.
Using Gemini Code Assist
The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and 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 to provide feedback.
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
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
3a9567c to
4e688c0
Compare
There was a problem hiding this comment.
Code Review
The pull request addresses a bug in the calculation of weighted hidden states from experts' output in the apply_shuffle_mul_sum function, which was causing incorrect truncation and leading to NaNs in the outputs. The fix involves modifying the shape assertions and the final calculation of the result in cutlass_fused_experts_fp8. Additionally, the code removes transpose operations in fp8.py before calling cutlass_fused_experts_fp8. The changes seem reasonable, but it's crucial to ensure that the weight shapes are correct after removing the transpose operations.
3927e91 to
e7c9d35
Compare
Motivation
Some bugs in calculating weighted hidden_states from experts output, in
apply_shuffle_mul_sum, where experts' output were wrongly truncated, leading to nans in outputs.Will fix in the future, if the kernel above shows an unneglectable speed difference with
shuffle_sum+ mul approachModifications
Checklist