GLM-5 · Final RMSNorm · OpenModelAtlas

Final RMSNorm

learnedself.norm(hidden_states)

Rescale each position’s features using their root mean square, then apply learned gain.

\begin{aligned}m_2(x)&=\frac{1}{d}\sum_{j=1}^d x_j^2\\y&=\gamma_{\mathrm{final}}\odot\frac{x}{\sqrt{m_2(x)+\varepsilon}}\\\varepsilon&=10^{-5}\end{aligned}

in: $x_{78}$ [B batch, T positions, 6144 d]
out: $z$ [B batch, T positions, 6144 d]
parameters: 6,144

Normalized in float32, cast back to the input type, then multiplied by the gain (modeling_glm_moe_dsa.py 58–62).

Parameters

6,144 parameters per instance.

Exact per-instance accounting
Component	Calculation	Parameters
Learned gain	$6{,}144$	6,144
Per-instance total		6,144

Each normalization instance owns its learned gain vector.

Derived from the configured feature width; the model total matches the checkpoint’s tensors.

Codemodeling_glm_moe_dsa.py 48–802

modeling_glm_moe_dsa.py 48–62

class GlmMoeDsaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        """
        GlmMoeDsaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

modeling_glm_moe_dsa.py 648

    _keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"]

modeling_glm_moe_dsa.py 696–739

        if inputs_embeds is None:
            inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        if position_ids is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
            position_ids = position_ids.unsqueeze(0)

        # It may already have been prepared by e.g. `generate`
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            mask_kwargs = {
                "config": self.config,
                "inputs_embeds": inputs_embeds,
                "attention_mask": attention_mask,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
                "allow_is_causal_skip": False,  # Always force creation to account for causality in the indexer
            }
            causal_mask_mapping = {"indexed_attention": create_causal_mask(**mask_kwargs)}

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)

        topk_indices = None  # MAIN DIFF with DSV3.2
        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
            hidden_states, topk_indices = decoder_layer(
                hidden_states,
                attention_mask=causal_mask_mapping["indexed_attention"],
                position_embeddings=position_embeddings,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                prev_topk_indices=topk_indices,  # MAIN DIFF with DSV3.2
                **kwargs,
            )

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values,
        )

modeling_glm_moe_dsa.py 749–753

    def __init__(self, config):
        super().__init__(config)
        self.model = GlmMoeDsaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

modeling_glm_moe_dsa.py 799–802

        hidden_states = outputs.last_hidden_state
        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

Sources2

Direct evidence

RMSNorm owns one gain per feature; it normalizes in float32, casts back to the input precision and multiplies by the gain.

rmsnorm · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 48–62

Read code excerpt

class GlmMoeDsaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        """
        GlmMoeDsaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

Direct evidence

The model embeds tokens, numbers positions from the cached length, always builds the causal mask (the indexer needs it), applies the 78 blocks, normalizes and projects to vocabulary logits with an untied lm_head without upcasting them. The checkpoint’s layer 78, the multi-token-prediction layer, is ignored on load.

transformer · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 648–802

Read code excerpt

    _keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"]

… lines 649–695 omitted …

        if inputs_embeds is None:
            inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        if position_ids is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
            position_ids = position_ids.unsqueeze(0)

        # It may already have been prepared by e.g. `generate`
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            mask_kwargs = {
                "config": self.config,
                "inputs_embeds": inputs_embeds,
                "attention_mask": attention_mask,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
                "allow_is_causal_skip": False,  # Always force creation to account for causality in the indexer
            }
            causal_mask_mapping = {"indexed_attention": create_causal_mask(**mask_kwargs)}

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)

        topk_indices = None  # MAIN DIFF with DSV3.2
        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
            hidden_states, topk_indices = decoder_layer(
                hidden_states,
                attention_mask=causal_mask_mapping["indexed_attention"],
                position_embeddings=position_embeddings,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                prev_topk_indices=topk_indices,  # MAIN DIFF with DSV3.2
                **kwargs,
            )

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values,
        )

… lines 740–748 omitted …

    def __init__(self, config):
        super().__init__(config)
        self.model = GlmMoeDsaModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

… lines 754–798 omitted …

        hidden_states = outputs.last_hidden_state
        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])