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

Shared expert

Open inside →×

learnedself.shared_experts

A SwiGLU of 2,048 hidden features applied to every position.

Why

Tokens routed to different experts may still need some common knowledge, and several routed experts may then each learn it. This expert runs for every token and is meant to hold that common part, which the authors of the design say reduces redundancy among the routed experts and leaves them more specialised. DeepSeekMoE §3.2

h=a⊙uh=a\odot u
yshared=hW2y^{\mathrm{shared}}=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
yshared=hW2y^{\mathrm{shared}}=hW_{2}
in
xffnx_{\mathrm{ffn}} [B batch, T positions, 6144 d]
out
ysharedy^{\mathrm{shared}} [B batch, T positions, 6144 d]
parameters
37,748,736 · × 75 MoE blocks = 2,831,155,200
Parameters

37,748,736 parameters per instance, 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 75 independent MoE blocks: 37,748,736×75=2,831,155,20037{,}748{,}736 \times 75 = 2{,}831{,}155{,}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.

Explanation

Its 2,048 hidden features are moe_intermediate_size × n_shared_experts (modeling_glm_moe_dsa.py 572–574).

Codemodeling_glm_moe_dsa.py 465–583
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
modeling_glm_moe_dsa.py 465–478
class GlmMoeDsaMLP(nn.Module):
    def __init__(self, config, intermediate_size=None):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj
Sources3

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

Direct evidence

The dense feed-forward network is a SwiGLU MLP, down(silu(gate x) ⊙ up x), computed in the input’s precision.

mlp · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 465–478

Read code excerpt
class GlmMoeDsaMLP(nn.Module):
    def __init__(self, config, intermediate_size=None):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj