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

Routed expert

Open inside →×

learnedself.experts

One of 256 independent SwiGLU experts; each position runs only the 8 its router selected.

h=a⊙uh=a\odot u
Ei=hW2E_i=hW_{2}
All 5 steps
g=xffnW1g=x_{\mathrm{ffn}}W_{1}
u=xffnW3u=x_{\mathrm{ffn}}W_{3}
a=g σ(g)a=g\,\sigma(g)
h=a⊙uh=a\odot u
Ei=hW2E_i=hW_{2}
in
xffnx_{\mathrm{ffn}} [B batch, T positions, 6144 d]
out
EiE_i [B batch, T positions, 6144 d]
parameters
37,748,736 · × 75 MoE blocks × 256 experts = 724,775,731,200
Parameters

37,748,736 parameters per expert, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Gate projection6,144×2,0486{,}144 \times 2{,}04812,582,912
Up projection6,144×2,0486{,}144 \times 2{,}04812,582,912
Down projection2,048×6,1442{,}048 \times 6{,}14412,582,912
Per-instance total37,748,736

Across 256 independent experts: 37,748,736×256=9,663,676,41637{,}748{,}736 \times 256 = 9{,}663{,}676{,}416 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 runs expert ii only on the positions routed to it and adds its weighted output back with index_add_ (modeling_glm_moe_dsa.py 541–557).

Codemodeling_glm_moe_dsa.py 522–583
modeling_glm_moe_dsa.py 522–559
@use_experts_implementation
class GlmMoeDsaExperts(nn.Module):
    """Collection of expert weights stored as 3D tensors."""

    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_local_experts
        self.hidden_dim = config.hidden_size
        self.intermediate_dim = config.moe_intermediate_size
        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(
        self,
        hidden_states: torch.Tensor,
        top_k_index: torch.Tensor,
        top_k_weights: torch.Tensor,
    ) -> torch.Tensor:
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1)
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            if expert_idx == self.num_experts:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
            current_hidden_states = self.act_fn(gate) * up
            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

        return final_hidden_states
modeling_glm_moe_dsa.py 562–574
class GlmMoeDsaMoE(nn.Module):
    """
    A mixed expert module containing shared experts.
    """

    def __init__(self, config: GlmMoeDsaConfig):
        super().__init__()
        self.config = config
        self.experts = GlmMoeDsaExperts(config)
        self.gate = GlmMoeDsaTopkRouter(config)
        self.shared_experts = GlmMoeDsaMLP(
            config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
        )
modeling_glm_moe_dsa.py 535–559
    def forward(
        self,
        hidden_states: torch.Tensor,
        top_k_index: torch.Tensor,
        top_k_weights: torch.Tensor,
    ) -> torch.Tensor:
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1)
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            if expert_idx == self.num_experts:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
            current_hidden_states = self.act_fn(gate) * up
            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

        return final_hidden_states
modeling_glm_moe_dsa.py 576–583
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residuals = hidden_states
        orig_shape = hidden_states.shape
        _, topk_weights, topk_indices = self.gate(hidden_states)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape)
        hidden_states = hidden_states + self.shared_experts(residuals)
        return hidden_states
Sources3

Direct evidence

The implementation holds the routed experts as stacked tensors, the gate and up projections fused into one; each expert is a SwiGLU of moe_intermediate_size hidden features applied to the tokens routed to it and weighted by its routing weight.

expert · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 522–559

Read code excerpt
@use_experts_implementation
class GlmMoeDsaExperts(nn.Module):
    """Collection of expert weights stored as 3D tensors."""

    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_local_experts
        self.hidden_dim = config.hidden_size
        self.intermediate_dim = config.moe_intermediate_size
        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(
        self,
        hidden_states: torch.Tensor,
        top_k_index: torch.Tensor,
        top_k_weights: torch.Tensor,
    ) -> torch.Tensor:
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1)
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            if expert_idx == self.num_experts:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
            current_hidden_states = self.act_fn(gate) * up
            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

        return final_hidden_states

Direct evidence

The MoE layer holds the routed experts, the router and one shared expert MLP of n_shared_experts × moe_intermediate_size features.

moe-init · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 562–574

Read code excerpt
class GlmMoeDsaMoE(nn.Module):
    """
    A mixed expert module containing shared experts.
    """

    def __init__(self, config: GlmMoeDsaConfig):
        super().__init__()
        self.config = config
        self.experts = GlmMoeDsaExperts(config)
        self.gate = GlmMoeDsaTopkRouter(config)
        self.shared_experts = GlmMoeDsaMLP(
            config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
        )

Direct evidence

Each token’s output is the router-weighted sum of its selected experts plus the shared expert applied to every token.

moe-forward · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 535–583

Read code excerpt
    def forward(
        self,
        hidden_states: torch.Tensor,
        top_k_index: torch.Tensor,
        top_k_weights: torch.Tensor,
    ) -> torch.Tensor:
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1)
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        for expert_idx in expert_hit:
            expert_idx = expert_idx[0]
            if expert_idx == self.num_experts:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
            current_hidden_states = self.act_fn(gate) * up
            current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

        return final_hidden_states

… lines 560–575 omitted …

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residuals = hidden_states
        orig_shape = hidden_states.shape
        _, topk_weights, topk_indices = self.gate(hidden_states)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape)
        hidden_states = hidden_states + self.shared_experts(residuals)
        return hidden_states