MoE: Cast routing_weights dtype correctly

[mmathew23]

Update compiled to cast routing_weight tensors to router_logits.dtype instead of hidden_states.dtype.

This will patch the forward method, and any other methods that cast routing_weights to hidden_states.dtype.

Qwen3 Moe Test:
https://colab.research.google.com/drive/1z9KopzSJylyM3g8KJ3bnvvpo97LISHTm?usp=sharing

[gemini-code-assist]

Summary of Changes

Hello @mmathew23, I'm Gemini Code Assist¹! 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 implements a crucial fix for Mixture of Experts (MoE) models by correcting a dtype mismatch during the compilation process. Previously, routing_weights were erroneously cast to the hidden_states.dtype, which could lead to incorrect behavior. The updated compiler now ensures these weights are correctly cast to router_logits.dtype, improving model stability and accuracy. This change enhances the compiler's patching capabilities and automates the application of this specific fix to relevant modules.

Highlights

• Corrected MoE Routing Weight Dtype Casting: The pull request resolves an issue where routing_weights were incorrectly cast to hidden_states.dtype in Mixture of Experts (MoE) models. They are now correctly cast to router_logits.dtype.
• Enhanced Compiler Flexibility: The create_standalone_class function has been extended with a new_methods parameter, allowing for more granular replacement of specific methods within a compiled class.
• New Patching Mechanism for MoE: A dedicated function, patch_moe_routing_weights_cast, has been introduced to specifically identify and apply the correct dtype casting for MoE routing weights using regular expressions.
• Automated Module Identification: The compiler now automatically detects modules that contain the routing_weights.to(hidden_states.dtype) pattern and applies the new MoE routing weights cast patch during the compilation process.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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

1. 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. ↩

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to patch methods in compiled classes, which is then used to fix an issue with routing_weights dtype casting in MoE models. The changes are logical and address the described problem. I've identified a few areas for improvement, mainly concerning error handling and type hint correctness. Specifically, using more specific exception handling instead of broad except Exception would improve robustness, and correcting a return type hint will enhance code clarity and maintainability.

[gemini-code-assist]

Using a broad except Exception can hide unexpected errors and make debugging more difficult. It's better to catch more specific exceptions that you expect to occur here. For instance, getattr can raise AttributeError, and inspect.getsource can raise TypeError or OSError. Catching these specific exceptions would provide more precise error handling.

Suggested change

	except Exception as e:
	if os.environ.get("UNSLOTH_LOGGING_ENABLED", "0") == "1":
	print(f"Unsloth: Failed to replace method {method_name} in {module} with error = {str(e)}")
	except (AttributeError, TypeError, OSError) as e:
	if os.environ.get("UNSLOTH_LOGGING_ENABLED", "0") == "1":
	print(f"Unsloth: Failed to replace method {method_name} in {module} with error = {str(e)}")

[gemini-code-assist]

Similar to other parts of the code, this try...except block uses a broad except Exception. This can make debugging difficult by catching and silencing unexpected errors. Consider catching more specific exceptions that might be raised by create_standalone_class. If the goal is to catch any failure, it would be beneficial to log the full traceback when logging is enabled to get more context on the failure.