Skip to content

[GLM-V and GLM-4.7] Cast to FP32 before gate projection for GLM model. #21660

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hnyls2002 merged 3 commits into from
Mar 30, 2026
+6 −1
Merged

[GLM-V and GLM-4.7] Cast to FP32 before gate projection for GLM model. #21660

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d26f7e0
Update glm4_moe.py
zRzRzRzRzRzRzR Mar 30, 2026
1718fa4
cache
zRzRzRzRzRzRzR Mar 30, 2026
582eddd
add comment
zRzRzRzRzRzRzR Mar 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion python/sglang/srt/models/glm4_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -327,9 +327,14 @@ def __init__(
self.e_score_correction_bias = nn.Parameter(
torch.empty((config.n_routed_experts), dtype=torch.float32)
)
# GLM requires FP32 gate projection; cache to avoid per-forward cast.
# FIXME: if gate weight is updated at runtime (e.g. expert rebalancing), _weight_fp32 must be invalidated.
self.register_buffer("_weight_fp32", None, persistent=False)

def forward(self, hidden_states):
logits = F.linear(hidden_states, self.weight, None)
if self._weight_fp32 is None:
self._weight_fp32 = self.weight.data.to(torch.float32)
logits = F.linear(hidden_states.to(torch.float32), self._weight_fp32, None)
return logits
Comment on lines 334 to 338
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Casting self.weight to float32 on every forward pass is inefficient as it creates a copy of the weight tensor each time. This can impact performance, especially for large models.

A better approach is to cache the float32 version of the weight. This can be done lazily on the first forward pass. The cached version should also be updated if the model is moved to a different device.

Here is a suggested implementation:

    def forward(self, hidden_states):
        # Cast to FP32 before gate projection for GLM-V model.
        if self.weight.dtype == torch.float32:
            weight_fp32 = self.weight
        else:
            if not hasattr(self, "_weight_fp32") or self._weight_fp32.device != self.weight.device:
                self._weight_fp32 = self.weight.to(torch.float32)
            weight_fp32 = self._weight_fp32

        logits = F.linear(hidden_states.to(torch.float32), weight_fp32, None)
        return logits



Expand Down
Loading