[Model] feat: support BailingMoe V2

[AlpinDale]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @AlpinDale, 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 introduces comprehensive support for the BailingMoe V2 model architecture. It refactors the existing BailingMoe implementation to incorporate advanced features such as configurable query-key normalization in attention, a more flexible Mixture-of-Experts (MoE) routing mechanism with various scoring functions and expert grouping options, and a hybrid MLP structure that allows for dense MLPs in early layers and MoE MLPs in later layers. These changes aim to enhance the model's performance and configurability, ensuring compatibility with the new V2 specification.

Highlights

• BailingMoe V2 Support: Introduces a new model class, BailingMoeV2ForCausalLM, and registers it within the system, enabling the use of the BailingMoe V2 architecture.
• Enhanced MoE Routing: Implements more sophisticated Mixture-of-Experts (MoE) routing with configurable score functions (softmax/sigmoid), optional expert biases, grouped top-k selection, and a routed scaling factor for improved expert selection and output scaling.
• Query-Key Normalization: Adds support for optional query-key normalization within the attention mechanism, allowing for the use of either RMSNorm or standard LayerNorm based on configuration.
• Hybrid MLP Architecture: Enables a hybrid approach where initial decoder layers can utilize standard Dense MLPs, while subsequent layers employ the Mixture-of-Experts (MoE) MLP, configurable via first_k_dense_replace.
• Improved Weight Loading: Refines the weight loading process for BailingMoeForCausalLM to be more robust, gracefully handling tied word embeddings and potential missing parameters during model initialization.

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

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 adds support for BailingMoe V2. The changes include adding new features like QK normalization, a more flexible MoE router, and conditional use of dense MLP layers. Several bug fixes and robustness improvements have also been made, such as correcting weight loading for tied embeddings and handling of None residuals in normalization layers. The code quality is good. I have one suggestion to refactor a part of the MoE forward pass for better readability and efficiency.

[gemini-code-assist]

The shared_output is computed at the beginning of the forward pass but used only at the end. This can be inefficient if self.shared_experts(hidden_states) is a costly operation. Additionally, the if self.shared_experts: check is performed twice.

For better readability and efficiency, it's better to compute shared_output just before it's used and combine the logic within a single conditional block.

Suggested change

	if self.shared_experts:
	shared_output = self.shared_experts(hidden_states)
	# router_logits: (num_tokens, n_experts)
	router_logits, _ = self.gate(hidden_states)
	router_logits = self.gate(hidden_states.to(self.router_dtype))
	router_logits = router_logits.to(hidden_states.dtype)
	final_hidden_states = self.experts(hidden_states=hidden_states,
	router_logits=router_logits)
	if self.num_shared_experts > 0:
	final_hidden_states *= self.routed_scaling_factor
	if self.shared_experts:
	final_hidden_states = final_hidden_states + shared_output
	# router_logits: (num_tokens, n_experts)
	router_logits = self.gate(hidden_states.to(self.router_dtype))
	router_logits = router_logits.to(hidden_states.dtype)
	final_hidden_states = self.experts(hidden_states=hidden_states,
	router_logits=router_logits)
	final_hidden_states *= self.routed_scaling_factor
	if self.shared_experts:
	shared_output = self.shared_experts(hidden_states)
	final_hidden_states = final_hidden_states + shared_output