feat(transformers): add Glm4Moe (v4.54.1) #1409
Conversation
Summary of ChangesHello @alien-0119, 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 introduces the Glm4Moe model, a significant addition to the Highlights
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
|
![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 the Glm4Moe model, including its implementation, auto-class registrations, and tests. The changes are mostly correct, but there are a few issues to address. I've found an incorrect module import, some incorrect type hints in the model implementation, an inefficient loop that could be a performance bottleneck, a placeholder model name in a docstring example, and an incorrect test case. Addressing these points will improve the correctness, maintainability, and performance of the new model.
| def moe(self, hidden_states: Tensor, topk_indices: Tensor, topk_weights: Tensor): | ||
| r""" | ||
| CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused | ||
| to not have to do a loop here (deepseek has 256 experts soooo yeah). | ||
| """ | ||
| final_hidden_states = mint.zeros_like(hidden_states, dtype=topk_weights.dtype) | ||
| expert_mask = mint.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) | ||
| expert_mask = expert_mask.permute(2, 0, 1) | ||
|
|
||
| for expert_idx in range(len(self.experts)): | ||
| expert = self.experts[expert_idx] | ||
| mask = expert_mask[expert_idx] | ||
| token_indices, weight_indices = mint.where(mask) | ||
|
|
||
| if token_indices.numel() > 0: | ||
| expert_weights = topk_weights[token_indices, weight_indices] | ||
| expert_input = hidden_states[token_indices] | ||
| expert_output = expert(expert_input) | ||
| weighted_output = expert_output * expert_weights.unsqueeze(-1) | ||
| final_hidden_states.index_add_(0, token_indices, weighted_output) | ||
|
|
||
| # in original deepseek, the output of the experts are gathered once we leave this module | ||
| # thus the moe module is itelsf an IsolatedParallel module | ||
| # and all expert are "local" meaning we shard but we don't gather | ||
| return final_hidden_states.type(hidden_states.dtype) |
There was a problem hiding this comment.
The current implementation of the Mixture of Experts (MoE) layer iterates over experts in a for loop. As noted in the docstring, this is inefficient and can be a significant performance bottleneck, especially with a large number of experts. For better performance, consider vectorizing this operation or using fused expert kernels if available in MindSpore. This would involve techniques like creating a single large weight matrix for all experts and using bmm or similar batched operations.
| past_key_value: Optional[Cache] = None, | ||
| cache_position: Optional[Tensor] = None, | ||
| **kwargs: Unpack[FlashAttentionKwargs], | ||
| ) -> tuple[Tensor, Optional[Tensor], Optional[tuple[Tensor]]]: |
There was a problem hiding this comment.
The return type hint for this method is -> tuple[Tensor, Optional[Tensor], Optional[tuple[Tensor]]], which suggests three return values. However, the method only returns two values: attn_output and attn_weights. The past_key_value is updated in-place and not returned. Please correct the type hint to match the actual return values.
| ) -> tuple[Tensor, Optional[Tensor], Optional[tuple[Tensor]]]: | |
| ) -> tuple[Tensor, Optional[Tensor]]: |
| cache_position: Optional[Tensor] = None, | ||
| position_embeddings: Optional[tuple[Tensor, Tensor]] = None, # necessary, but kept here for BC | ||
| **kwargs: Unpack[TransformersKwargs], | ||
| ) -> tuple[Tensor]: |
There was a problem hiding this comment.
The return type hint for this method is -> tuple[Tensor], but it returns a single Tensor (hidden_states), not a tuple. Please correct the type hint to -> Tensor.
| ) -> tuple[Tensor]: | |
| ) -> Tensor: |
| >>> model = Glm4MoeForCausalLM.from_pretrained("meta-glm4_moe/Glm4Moe-2-7b-hf") | ||
| >>> tokenizer = AutoTokenizer.from_pretrained("meta-glm4_moe/Glm4Moe-2-7b-hf") |
There was a problem hiding this comment.
The model names used in the example, meta-glm4_moe/Glm4Moe-2-7b-hf, appear to be placeholders and do not correspond to a valid model on the Hugging Face Hub. Please update this to a correct, loadable model name, or use a more explicit placeholder like your-namespace/your-model-name to make the example runnable and more helpful for users.
| >>> model = Glm4MoeForCausalLM.from_pretrained("meta-glm4_moe/Glm4Moe-2-7b-hf") | |
| >>> tokenizer = AutoTokenizer.from_pretrained("meta-glm4_moe/Glm4Moe-2-7b-hf") | |
| >>> model = Glm4MoeForCausalLM.from_pretrained("your-namespace/your-glm4-moe-model") | |
| >>> tokenizer = AutoTokenizer.from_pretrained("your-namespace/your-glm4-moe-model") |
alien-0119
force-pushed
the
glm4_moe_master
branch
from
3c08d74 to
b95138e
Compare
alien-0119
force-pushed
the
glm4_moe_master
branch
from
b95138e to
dc77245
Compare
1 participant
What does this PR do?
Adds # (feature)
Add Glm4Moe model and fast ut.
Usage Example:
Cannot load GLM-4.5/4.6 due to too many parameters of weights (exceeding 100B).
Performance:
Experiments were tested on Ascend Atlas 800T A2 machines with mindspore 2.7.0 pynative mode.
Before submitting
What's New. Here are thedocumentation guidelines
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@xxx