Mixture of experts

Checkpoint configuration and implementation zai-org/GLM-5 @ c183ef8, huggingface/transformers @ 5e4d630 · parameter count checked against the checkpoint’s tensor headers · How to read

78 transformer blocks › 75 MoE blocks › Mixture of experts

learnedself.gate

Score all 256 experts for each position and keep 8; their weights are normalized, then scaled by 2.5.

Eisel=[ i∈top⁡8(s+b) ]E^{\mathrm{sel}}_i=\big[\,i\in\operatorname{top}_{8}(s+b)\,\big]
Gi=2.5 si Eisel∑jsj EjselG_{i}=2.5\,\frac{s_i\,E^{\mathrm{sel}}_i}{\sum_{j}s_j\,E^{\mathrm{sel}}_j}
All 4 steps
z=xffnWrz=x_{\mathrm{ffn}}W_{r}
s=σ(z)s=\sigma(z)
Eisel=[ i∈top⁡8(s+b) ]E^{\mathrm{sel}}_i=\big[\,i\in\operatorname{top}_{8}(s+b)\,\big]
Gi=2.5 si Eisel∑jsj EjselG_{i}=2.5\,\frac{s_i\,E^{\mathrm{sel}}_i}{\sum_{j}s_j\,E^{\mathrm{sel}}_j}
in
xffnx_{\mathrm{ffn}} [B batch, T positions, 6144 d]
out
GG [B batch, T positions, 256 experts]
parameters
1,573,120 · × 75 MoE blocks = 117,984,000
Parameters

1,573,120 parameters per instance, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Expert scores6,144×2566{,}144 \times 2561,572,864
Select experts256256256
Per-instance total1,573,120

Across 75 independent MoE blocks: 1,573,120×75=117,984,0001{,}573{,}120 \times 75 = 117{,}984{,}000 parameters.

Derived from the pinned configuration and transformers implementation, and equal to the count of the checkpoint’s tensors read from its safetensors headers, less the multi-token-prediction layer the implementation does not load. The total includes the routers’ 19,200 correction-bias values, which the implementation holds as buffers.

Explanation

The implementation returns each position’s 8 weights with their expert indices (modeling_glm_moe_dsa.py 519); GG spreads them over all 256 experts, zero for the other 248.

Codemodeling_glm_moe_dsa.py 481–647
modeling_glm_moe_dsa.py 481–492
class GlmMoeDsaTopkRouter(nn.Module):
    def __init__(self, config: GlmMoeDsaConfig):
        super().__init__()
        self.top_k = config.num_experts_per_tok
        self.num_experts = config.num_local_experts
        self.hidden_dim = config.hidden_size
        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
        self.routed_scaling_factor = config.routed_scaling_factor
        self.num_group = config.n_group
        self.topk_group = config.topk_group
        self.norm_topk_prob = config.norm_topk_prob
        self.e_score_correction_bias = nn.Buffer(torch.zeros((self.num_experts), dtype=torch.float32))
modeling_glm_moe_dsa.py 647
    _keep_in_fp32_modules_strict = ["e_score_correction_bias"]
modeling_glm_moe_dsa.py 494–519
    def forward(self, hidden_states):
        hidden_states = hidden_states.view(-1, self.hidden_dim)
        router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
        scores = router_logits.sigmoid()
        scores_for_choice = scores + self.e_score_correction_bias
        group_scores = (
            scores_for_choice.view(-1, self.num_group, self.num_experts // self.num_group)
            .topk(2, dim=-1)[0]
            .sum(dim=-1)
        )
        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
        group_mask = torch.zeros_like(group_scores)
        group_mask.scatter_(1, group_idx, 1)
        score_mask = (
            group_mask.unsqueeze(-1)
            .expand(-1, self.num_group, self.num_experts // self.num_group)
            .reshape(-1, self.num_experts)
        )
        scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
        topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
        topk_weights = scores.gather(1, topk_indices)
        if self.norm_topk_prob:
            denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
            topk_weights /= denominator
        topk_weights = topk_weights * self.routed_scaling_factor
        return router_logits, topk_weights, topk_indices
Sources2

Direct evidence

The router owns one score vector per routed expert and a per-expert correction bias, kept in float32 and stored in the checkpoint as a buffer, used only for selection.

gate-init · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 481–647

Read code excerpt
class GlmMoeDsaTopkRouter(nn.Module):
    def __init__(self, config: GlmMoeDsaConfig):
        super().__init__()
        self.top_k = config.num_experts_per_tok
        self.num_experts = config.num_local_experts
        self.hidden_dim = config.hidden_size
        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
        self.routed_scaling_factor = config.routed_scaling_factor
        self.num_group = config.n_group
        self.topk_group = config.topk_group
        self.norm_topk_prob = config.norm_topk_prob
        self.e_score_correction_bias = nn.Buffer(torch.zeros((self.num_experts), dtype=torch.float32))

… lines 493–646 omitted …

    _keep_in_fp32_modules_strict = ["e_score_correction_bias"]

Direct evidence

Routing: float32 sigmoid scores, the bias added for selection, the best groups kept by their top-two scores (one group here), the top-8 experts chosen, their unbiased scores renormalized and scaled by the route scale.

gate-forward · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 494–519

Read code excerpt
    def forward(self, hidden_states):
        hidden_states = hidden_states.view(-1, self.hidden_dim)
        router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
        scores = router_logits.sigmoid()
        scores_for_choice = scores + self.e_score_correction_bias
        group_scores = (
            scores_for_choice.view(-1, self.num_group, self.num_experts // self.num_group)
            .topk(2, dim=-1)[0]
            .sum(dim=-1)
        )
        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
        group_mask = torch.zeros_like(group_scores)
        group_mask.scatter_(1, group_idx, 1)
        score_mask = (
            group_mask.unsqueeze(-1)
            .expand(-1, self.num_group, self.num_experts // self.num_group)
            .reshape(-1, self.num_experts)
        )
        scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
        topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
        topk_weights = scores.gather(1, topk_indices)
        if self.norm_topk_prob:
            denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
            topk_weights /= denominator
        topk_weights = topk_weights * self.routed_scaling_factor
        return router_logits, topk_weights, topk_indices