Transformer Block Assembly › Assembling the Transformer Block

class TransformerBlock:
    """
    Complete transformer block with pre-normalization.

    Architecture:
        x -> RMSNorm -> Attention -> +x -> RMSNorm -> FFN -> +

    This is the standard block structure used in GPT-2/3, LLaMA, etc.
    """

    def __init__(self, d_model, n_heads, d_ff, use_swiglu=True):
        """
        Args:
            d_model: Model dimension
            n_heads: Number of attention heads
            d_ff: Feed-forward hidden dimension
            use_swiglu: Whether to use SwiGLU (True) or standard FFN (False)
        """
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_ff = d_ff

        # Normalization layers (one before each sublayer)
        self.norm1 = RMSNorm(d_model)
        self.norm2 = RMSNorm(d_model)

        # Attention sublayer
        self.attention = MultiHeadSelfAttention(d_model, n_heads)

        # Feed-forward sublayer
        if use_swiglu:
            self.ffn = SwiGLUFFN(d_model, d_ff)
        else:
            self.ffn = StandardFFN(d_model, d_ff)

    def __call__(self, x, mask=None):
        """
        Process input through the transformer block.

        Args:
            x: Input tensor of shape (seq_len, d_model)
            mask: Optional causal mask for attention

        Returns:
            Output tensor of shape (seq_len, d_model)
        """
        # Attention sublayer with pre-norm and residual
        normed = self.norm1(x)
        attn_out = self.attention(normed, mask)
        h = x + attn_out  # Residual connection

        # FFN sublayer with pre-norm and residual
        normed = self.norm2(h)
        ffn_out = self.ffn(normed)
        y = h + ffn_out  # Residual connection

        return y

    def num_parameters(self):
        """Count total parameters in the block."""
        # Normalization: 2 * d_model (two RMSNorm layers)
        norm_params = 2 * self.d_model

        # Attention: 4 * d_model^2 (Q, K, V, O projections)
        attn_params = 4 * self.d_model * self.d_model

        # FFN: depends on variant
        if isinstance(self.ffn, SwiGLUFFN):
            # SwiGLU: 3 * d_model * d_ff (gate, up, down)
            ffn_params = 3 * self.d_model * self.d_ff
        else:
            # Standard: 2 * d_model * d_ff (W1, W2)
            ffn_params = 2 * self.d_model * self.d_ff

        return norm_params + attn_params + ffn_params

# Create a transformer block with typical hyperparameters
d_model = 512
n_heads = 8
d_ff = 1376  # ~2.7x for SwiGLU (equivalent capacity to 4x standard)

block = TransformerBlock(d_model, n_heads, d_ff, use_swiglu=True)

# Create sample input: 16 tokens with 512-dimensional representations
seq_len = 16
x = np.random.randn(seq_len, d_model) * 0.02  # Small init like embeddings

# Process through the block
y = block(x)