GLM-5 · MoE transformer block

78 transformer blocks

MoE transformer block

learnedGlmMoeDsaDecoderLayer(config, layer_idx)

MLA and a mixture of experts, each added to the residual stream; blocks 4–78.

x'_n=x_n+\Delta_{\mathrm{attn}}

x_{n+1}=x'_n+\Delta_{\mathrm{ffn}}

All 43 steps

x_{\mathrm{attn}}=\operatorname{RMSNorm}(x_n)

\text{Multi-head latent attention:}

c^{{\htmlClass{vl-stream-q}{Q}}}=xW_{DQ}

\hat c^{{\htmlClass{vl-stream-q}{Q}}}=\operatorname{RMSNorm}(c^{{\htmlClass{vl-stream-q}{Q}}})

c^{KV}=xW_{DKV}

\hat c^{KV}=\operatorname{RMSNorm}(c^{KV})

{\htmlClass{vl-stream-q}{q^{C}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{UQ}

{\htmlClass{vl-stream-q}{Q^{C}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{C}}})

{\htmlClass{vl-stream-q}{q^{\mathrm{pe}}}}=\hat c^{{\htmlClass{vl-stream-q}{Q}}}W_{QR}

{\htmlClass{vl-stream-q}{Q^{\mathrm{pe}}}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q^{\mathrm{pe}}}})

{\htmlClass{vl-stream-q}{Q^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-q}{Q^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})

{\htmlClass{vl-stream-q}{Q}}=[{\htmlClass{vl-stream-q}{Q^{C}}};\,{\htmlClass{vl-stream-q}{Q^{R}}}]

{\htmlClass{vl-stream-k}{k^{C}}}=\hat c^{KV}W_{UK}

{\htmlClass{vl-stream-k}{K^{C}}}=\operatorname{reshape}({\htmlClass{vl-stream-k}{k^{C}}})

{\htmlClass{vl-stream-k}{k^{\mathrm{pe}}}}=xW_{KR}

{\htmlClass{vl-stream-k}{k^{R}_{t}}}=\operatorname{RoPE}({\htmlClass{vl-stream-k}{k^{\mathrm{pe}}_{t}}},{\htmlClass{vl-stream-context}{R_t}})

{\htmlClass{vl-stream-k}{K_i}}=[{\htmlClass{vl-stream-k}{K^{C}_i}};\,{\htmlClass{vl-stream-k}{k^{R}}}]

{\htmlClass{vl-stream-v}{v}}=\hat c^{KV}W_{UV}

{\htmlClass{vl-stream-v}{V}}=\operatorname{reshape}({\htmlClass{vl-stream-v}{v}})

w=xW_{w}

S^{I}_{t,s}=\frac{1}{\sqrt{32}\sqrt{128}}\sum_{h=1}^{32}w_{t,h}\,\operatorname{ReLU}\!\big(\tilde {\htmlClass{vl-stream-q}{Q^{I}_{h,t}}}\cdot\tilde {\htmlClass{vl-stream-k}{k^{I}_{s}}}\big)

\tilde S^{I}=S^{I}+{\htmlClass{vl-stream-context}{M}}

{\htmlClass{vl-stream-context}{M^{\mathrm{sel}}_{t,s}}}=\begin{cases}{\htmlClass{vl-stream-context}{M_{t,s}}}&s\in\operatorname{top}_{2,048}(\tilde S^{I}_{t})\\-\infty&\text{otherwise}\end{cases}

x'_n=x_n+\Delta_{\mathrm{attn}}

x_{\mathrm{ffn}}=\operatorname{RMSNorm}(x'_n)

\text{Mixture of experts:}

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}

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

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}}

x_{n+1}=x'_n+\Delta_{\mathrm{ffn}}

in: $x_{3}$ [B batch, T positions, 6144 d] $R$ [T positions, 32 pairs] $M$ [T queries, T keys]
out: $x_{n+1}$ [B batch, T positions, 6144 d]
parameters: 9,877,404,672 · × 75 MoE blocks = 740,805,350,400

Parameters

9,877,404,672 parameters per MoE block, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Attention RMSNorm	$6{,}144$	6,144
Multi-head latent attention	174,394,112 across its components	174,394,112
MoE RMSNorm	$6{,}144$	6,144
Mixture of experts	9,702,998,272 across its components	9,702,998,272
Per-instance total		9,877,404,672

Across 75 independent MoE blocks: $9{,}877{,}404{,}672 \times 75 = 740{,}805{,}350{,}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.

Codemodeling_glm_moe_dsa.py 586–627

modeling_glm_moe_dsa.py 586–627

class GlmMoeDsaDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = GlmMoeDsaAttention(config, layer_idx)

        self.mlp = GlmMoeDsaMoE(config) if config.mlp_layer_types[layer_idx] == "sparse" else GlmMoeDsaMLP(config)

        self.input_layernorm = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = False,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        prev_topk_indices: torch.Tensor | None = None,  # MAIN DIFF with DSV3.2
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        # Self Attention
        hidden_states, _, topk_indices = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            position_embeddings=position_embeddings,
            prev_topk_indices=prev_topk_indices,  # MAIN DIFF with DSV3.2
            **kwargs,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states, topk_indices

Sources1

Direct evidence

A block applies input RMSNorm, attention and a residual addition, then post-attention RMSNorm, the feed-forward network and a second residual addition; the first first_k_dense_replace blocks use the dense MLP, the rest MoE.

block · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 586–627

Read code excerpt

class GlmMoeDsaDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = GlmMoeDsaAttention(config, layer_idx)

        self.mlp = GlmMoeDsaMoE(config) if config.mlp_layer_types[layer_idx] == "sparse" else GlmMoeDsaMLP(config)

        self.input_layernorm = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = False,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        prev_topk_indices: torch.Tensor | None = None,  # MAIN DIFF with DSV3.2
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        # Self Attention
        hidden_states, _, topk_indices = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            position_embeddings=position_embeddings,
            prev_topk_indices=prev_topk_indices,  # MAIN DIFF with DSV3.2
            **kwargs,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states, topk_indices