DeepSeek V4.1 · Final RMSNorm

Final RMSNorm

learnedself.norm(h) · 1269

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^{-20}\end{aligned}

in: $h_{\mathrm{out}}$ [B batch, 1 position, 5120 d]
out: $z$ [B batch, 1 position, 5120 d]
parameters: 5,120

Parameters

5,120 parameters per instance.

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

Each normalization instance owns its learned gain vector.

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

Codemodel.py 281–293, 1241–1272

model.py 281–293

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor):
        dtype = x.dtype
        x = x.float()
        var = x.square().mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return (self.weight * x).to(dtype)

model.py 1241–1272

    @torch.inference_mode()
    def forward(
        self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
    ):
        """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
        has already filled, so the padding it uses internally never reaches here. token_types /
        images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
        inside the first (start_pos 0) chunk."""
        image_mask = None if token_types is None else token_types >= 0  # TEXT is -1
        # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
        engram_mask = None if image_mask is None else ~image_mask
        engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
        h = self.embed(input_ids)
        if images is not None:
            assert start_pos == 0, "image spans must be prefilled in a single chunk"
            self.merge_image_embeddings(images, h)
        # Expand to hc_mult copies for Hyper-Connections
        h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
        main_hiddens = []
        pre_mix = make_identity_pre_mix(h, self.hc_mult)
        for i, layer in enumerate(self.layers):
            if layer.engram is not None:
                h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
            # the MTP head reads the attention input of its target layers, not their output
            if i in self.target_layer_ids:
                main_hiddens.append(h.mean(dim=2))
            h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
        h = layer.hc_pre(h, pre_mix)
        logits = self.head(self.norm(h))
        output_ids = sample(logits, self.temperature)
        main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
        return output_ids, logits, main_hidden

Sources2

Direct evidence

RMSNorm divides by the root mean square over the features, computed in float32, and multiplies by a learned gain.

rmsnorm · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 281–293

Read code excerpt

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor):
        dtype = x.dtype
        x = x.float()
        var = x.square().mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return (self.weight * x).to(dtype)

Direct evidence

Tokens are embedded and repeated into four streams; each layer, after an Engram step where it has one, reads and writes the streams; the last reading mix collapses them, then the final norm and the head.

transformer-forward · deepseek-ai/DeepSeek-V4.1-Flash / inference/model.py
Pinned revision 2cba9e42aa02 ↗ · lines 1241–1272

Read code excerpt

    @torch.inference_mode()
    def forward(
        self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None
    ):
        """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it
        has already filled, so the padding it uses internally never reaches here. token_types /
        images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie
        inside the first (start_pos 0) chunk."""
        image_mask = None if token_types is None else token_types >= 0  # TEXT is -1
        # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask
        engram_mask = None if image_mask is None else ~image_mask
        engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None
        h = self.embed(input_ids)
        if images is not None:
            assert start_pos == 0, "image spans must be prefilled in a single chunk"
            self.merge_image_embeddings(images, h)
        # Expand to hc_mult copies for Hyper-Connections
        h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
        main_hiddens = []
        pre_mix = make_identity_pre_mix(h, self.hc_mult)
        for i, layer in enumerate(self.layers):
            if layer.engram is not None:
                h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask)
            # the MTP head reads the attention input of its target layers, not their output
            if i in self.target_layer_ids:
                main_hiddens.append(h.mean(dim=2))
            h, pre_mix = layer(h, start_pos, pre_mix, image_mask)
        h = layer.hc_pre(h, pre_mix)
        logits = self.head(self.norm(h))
        output_ids = sample(logits, self.temperature)
        main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None
        return output_ids, logits, main_hidden