File size: 2,893 Bytes
b7858ac a05efc6 b7858ac a05efc6 1384d83 a05efc6 b7858ac a05efc6 1384d83 b7858ac a05efc6 b7858ac 1384d83 b7858ac a05efc6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """Decoder layer — a single transformer block.
Each block applies pre-norm attention followed by pre-norm MLP, with residual
connections around both sublayers:
normed = RMSNorm(x)
h = x + Attention(normed, ...)
normed = RMSNorm(h)
out = h + MLP(normed)
Pre-norm keeps the residual stream unnormalised. Gradients flow more cleanly
through unnormalised residuals at depth, and each sublayer receives a stable,
normalised view of the signal.
Two independent RMSNorm instances are used — one before attention, one before MLP.
They learn different scalings because they precede layers with different dynamic
ranges. Sharing them would be wrong.
torch.nn.RMSNorm is used directly (available from PyTorch 2.4+). It omits mean
subtraction, is faster than LayerNorm, and proved more stable at scale.
"""
import torch
import torch.nn as nn
from transformers import PretrainedConfig
from transformers.cache_utils import Cache
from .attention import GroupedQueryAttention
from .mlp import SwiGLUMLP
class DecoderLayer(nn.Module):
"""A single pre-norm transformer decoder block.
Composes GroupedQueryAttention and SwiGLUMLP with residual connections and
independent RMSNorm instances on each sublayer input.
Args:
config: Model config passed through to attention and MLP. Must also expose
``hidden_size`` and ``rms_norm_eps``.
"""
def __init__(self, config: PretrainedConfig) -> None:
super().__init__()
self.attn_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.mlp_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.attention = GroupedQueryAttention(config)
self.mlp = SwiGLUMLP(config)
def forward(
self,
x: torch.Tensor,
position_ids: torch.Tensor,
cache: Cache | None = None,
layer_idx: int = 0,
causal_mask: torch.Tensor | None = None,
) -> torch.Tensor:
"""Apply one decoder block to the input.
Args:
x: Input of shape (batch, seq_len, hidden_size).
position_ids: Absolute positions of shape (batch, seq_len).
cache: HuggingFace Cache object for KV accumulation, or None when
caching is disabled. Passed through to attention unchanged.
layer_idx: Cache slot index for this layer. Each layer has its own
index so they accumulate independently within the shared cache.
causal_mask: Optional boolean attention mask of shape
(1, 1, seq_len, kv_len). Passed through to attention unchanged.
Returns:
Output tensor of shape (batch, seq_len, hidden_size).
"""
attn_out = self.attention(self.attn_norm(x), position_ids, cache, layer_idx, causal_mask)
h = x + attn_out
return h + self.mlp(self.mlp_norm(h))
|