GLM-5 · Mixture of experts

78 transformer blocks › 75 MoE blocks

Mixture of experts

learnedself.mlp

Each position runs 8 of 256 routed experts, weighted by the router, plus one shared expert that every position runs.

Why

Many small experts. With few large experts, the tokens sent to one expert are more likely to cover very different kinds of knowledge, and one expert has to hold all of it. The authors of DeepSeekMoE split each expert into smaller ones and activate more of them at the same parameters and compute, so that knowledge can be divided among experts; the number of possible combinations of active experts grows enormously. DeepSeekMoE §3.1

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

y^{\mathrm{routed}}=\sum_{i=1}^{256}G_i\,E_i

\Delta_{\mathrm{ffn}}=y^{\mathrm{routed}}+y^{\mathrm{shared}}

All 19 steps

\text{Router:}

z=x_{\mathrm{ffn}}W_{r}

s=\sigma(z)

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

G_{i}=2.5\,\frac{s_i\,E^{\mathrm{sel}}_i}{\sum_{j}s_j\,E^{\mathrm{sel}}_j}

\text{Routed expert:}

g=x_{\mathrm{ffn}}W_{1}

u=x_{\mathrm{ffn}}W_{3}

a=g\,\sigma(g)

h=a\odot u

E_i=hW_{2}

y^{\mathrm{routed}}=\sum_{i=1}^{256}G_i\,E_i

\text{Shared expert:}

g=x_{\mathrm{ffn}}W_{1}

u=x_{\mathrm{ffn}}W_{3}

a=g\,\sigma(g)

h=a\odot u

y^{\mathrm{shared}}=hW_{2}

\Delta_{\mathrm{ffn}}=y^{\mathrm{routed}}+y^{\mathrm{shared}}

in: $x_{\mathrm{ffn}}$ [B batch, T positions, 6144 d]
out: $\Delta_{\mathrm{ffn}}$ [B batch, T positions, 6144 d]
parameters: 9,702,998,272 · × 75 MoE blocks = 727,724,870,400

Parameters

9,702,998,272 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Router	1,573,120 across its components	1,573,120
Routed expert	$37{,}748{,}736 \times 256$	9,663,676,416
Shared expert	37,748,736 across its components	37,748,736
Per-instance total		9,702,998,272

Across 75 independent MoE blocks: $9{,}702{,}998{,}272 \times 75 = 727{,}724{,}870{,}400$ 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 flattens batch and positions into one token axis for the experts and restores the shape (modeling_glm_moe_dsa.py 580–581).

Codemodeling_glm_moe_dsa.py 535–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

Sources2

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