GLM-5 · Select experts · OpenModelAtlas

78 transformer blocks › 75 MoE blocks › Mixture of experts › Router

Select experts

buffer · not trainedtorch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)

Keep the top 8 of 256 experts, ranked with a correction bias.

{\htmlData{light=selected select}{E^{\mathrm{sel}}_i}}=\big[\,i\in\operatorname{top}_{8}({\htmlData{light=affinities sigmoid}{s}}+b)\,\big]

in: $s$ [B batch, T positions, 256 experts]
out: $E^{\mathrm{sel}}$ [B batch, T positions, 256 experts]
parameters: 256 · × 75 MoE blocks = 19,200

Parameters

256 parameters owned by this component

Exact per-instance accounting
Component	Calculation	Parameters
learned vector	$256$	256
Per-instance total		256

Across 75 independent MoE blocks: $256 \times 75 = 19{,}200$ 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.

These are the correction bias, a float32 buffer that gradients do not train. The atlas counts its values with the parameters because the checkpoint stores them.

Explanation

The configuration sets one expert group, so the implementation’s group stage keeps all 256 experts.

The bias $b$ (256 correction values) is added only when choosing experts; the weights use the unbiased affinities (modeling_glm_moe_dsa.py 498–514).

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