Final RMSNorm
×learnedself.norm(h) · 1269
Rescale each position’s features using their root mean square, then apply learned gain.
- in
- [B batch, T positions, 5120 d]
- out
- [B batch, T positions, 5120 d]
- parameters
- 5,120
Parameters
5,120 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Learned gain | 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
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) @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_hiddenSources2
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