Transformer Block Assembly › Stacking Blocks: Building a Complete Model

class TransformerStack:
    """
    Stack of transformer blocks forming a complete encoder or decoder.
    """

    def __init__(
        self,
        n_layers,
        d_model,
        n_heads,
        d_ff,
        use_swiglu=True,
        is_decoder=False,
    ):
        """
        Args:
            n_layers: Number of transformer blocks
            d_model: Model dimension
            n_heads: Number of attention heads
            d_ff: Feed-forward hidden dimension
            use_swiglu: Whether to use SwiGLU FFN
            is_decoder: Whether to apply causal masking
        """
        self.n_layers = n_layers
        self.is_decoder = is_decoder

        # Create blocks
        if is_decoder:
            self.blocks = [
                DecoderBlock(d_model, n_heads, d_ff, use_swiglu)
                for _ in range(n_layers)
            ]
        else:
            self.blocks = [
                TransformerBlock(d_model, n_heads, d_ff, use_swiglu)
                for _ in range(n_layers)
            ]

        # Final normalization (after all blocks)
        self.final_norm = RMSNorm(d_model)

    def __call__(self, x):
        """
        Process input through all transformer blocks.

        Args:
            x: Input tensor of shape (seq_len, d_model)

        Returns:
            Output tensor of shape (seq_len, d_model)
        """
        # Pass through each block sequentially
        for block in self.blocks:
            x = block(x)

        # Final normalization
        x = self.final_norm(x)

        return x

    def num_parameters(self):
        """Total parameters in the stack."""
        block_params = sum(block.num_parameters() for block in self.blocks)
        # Final norm adds d_model parameters
        return block_params + self.blocks[0].d_model


# Create a small transformer stack
transformer = TransformerStack(
    n_layers=6,
    d_model=256,
    n_heads=4,
    d_ff=688,
    use_swiglu=True,
    is_decoder=True,
)

# Track representations through layers
def trace_through_stack(stack, x):
    """Collect representations after each layer."""
    representations = [x.copy()]

    current = x
    for block in stack.blocks:
        current = block(current)
        representations.append(current.copy())

    # After final norm
    final = stack.final_norm(current)
    representations.append(final.copy())

    return representations


# Trace through the stack
x_trace = np.random.randn(8, 256) * 0.02
layer_outputs = trace_through_stack(transformer, x_trace)