fix glm4 moe

[ch-wan]

Motivation

GLM4 inherits deepseek but #8539's update is not applied to GLM4. Some logits are broken.

Modifications

Accuracy Test

Benchmark & Profiling

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Summary of Changes

Hello @ch-wan, 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 aims to resolve an issue where GLM4 Mixture of Experts (MoE) models exhibited broken logits due to a missing update from a previous pull request (#8539). The core of the fix involves refining the communication strategy within the MoE layers by introducing and enabling a 'reduce-scatter' mechanism, which should correctly handle tensor parallel operations and restore proper model behavior.

Highlights

• MoE Communication Logic Adjustment: Modified the forward_normal_dual_stream and forward_normal methods in glm4_moe.py to introduce a new boolean parameter, use_reduce_scatter. This parameter is now incorporated into the conditional logic that determines whether tensor_model_parallel_all_reduce is performed, allowing for more flexible or optimized communication strategies within the MoE layer.
• Enable Reduce-Scatter for MoE Layer: Updated the __init__ method to pass allow_reduce_scatter=True when initializing the MoELayer. This explicitly enables a reduce-scatter communication pattern for the GLM4 MoE implementation, which is crucial for addressing the reported issue with broken logits and aligning with previous updates.

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 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 or fill out our survey 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

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 fixes an issue in GLM-4 MoE models by incorporating a reduce_scatter related fix from DeepSeek V2. The changes add the use_reduce_scatter flag to avoid redundant all_reduce operations when reduce_scatter is used.

However, the implementation is incomplete. The use_reduce_scatter flag is not passed to the MLP/MoE layers, rendering the changes ineffective. I've provided comments on these issues and a suggestion to refactor duplicated code for better maintainability. Addressing these points will ensure the fix is complete.

[gemini-code-assist]

The use_reduce_scatter parameter is added to the signature of forward_normal, but it's not used in the method's body. This prevents the reduce_scatter optimization from being applied.

Update the conditions for calling tensor_model_parallel_all_reduce to include and not use_reduce_scatter.

[gemini-code-assist]

While enabling allow_reduce_scatter here is correct, the forward method of Glm4MoeDecoderLayer does not pass the use_reduce_scatter flag to self.mlp. This makes the changes ineffective, as use_reduce_scatter will always be its default value (False) inside the MLP/MoE layers.

To fix this, Glm4MoeDecoderLayer.forward needs to:

1. Calculate use_reduce_scatter by calling self.layer_communicator.should_use_reduce_scatter(forward_batch).
2. Pass this flag to self.mlp(...).

Additionally, Glm4MoeMLP.forward needs to be updated to accept and use the use_reduce_scatter parameter.

[gemini-code-assist]

To improve readability and maintainability, extract the condition into a local variable. This makes the logic clearer and avoids repeating the same complex condition.

            should_all_reduce = (
                self.tp_size > 1
                and not can_fuse_mlp_allreduce
                and not use_reduce_scatter
            )
            if should_all_reduce: