PYTORCH · BUILD IT YOURSELF

Lab: build and test every transformer block

Full build path and setup · Download transformer parts · Download training runner · Download correctness tests

One token update has two jobs

Attention shares information across allowed positions. A feed-forward network transforms each position separately. Residual connections keep an unchanged path around both operations. Normalization controls the scale of each token's features. These parts have different jobs. A transformer block is their composition.

Our code uses a pre-normalization layout. One block computes x = x + attention(norm1(x)), followed by x = x + feed_forward(norm2(x)). Each branch returns the same shape as x. For a batch of 2 sequences, 5 positions, and width 32, that shape is (2,5,32).

Implement normalization as arithmetic

For each token, calculate the mean of its features. Subtract it. Divide by the square root of the feature variance plus a small positive number. That number prevents division by zero. Finally apply a learned scale and shift per feature.

import torch

def layer_norm(x, scale, shift, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    variance = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
    return (x - mean) * torch.rsqrt(variance + eps) * scale + shift

x = torch.tensor([[1., 3.]])
y = layer_norm(x, torch.ones(2), torch.zeros(2))
print(y)  # approximately [[-1, 1]]

The variance here divides by the feature count. It is not the unbiased sample-variance estimate. keepdim=True retains a final axis of length one. That lets subtraction and division broadcast across features. Normalizing across the batch would change the operation and mix unrelated examples.

RMSNorm is a related alternative. It divides by the root mean square without subtracting a mean. This course's transformer uses LayerNorm. Do not treat the names as interchangeable when reading another model's code.

Build a feature transformation

A basic feed-forward branch expands width C to 4C, applies GELU, then projects back to C. With C=32, the intermediate width is 128. GELU is a smooth nonlinear activation. Without a nonlinearity, two consecutive linear maps could collapse into one linear map.

The expansion factor is a choice, not a requirement. The small complete runner uses width 64 and hidden width 128, so its expansion is 2C. The example below uses 4C to make the two configurable sizes explicit.

For (B,T,C) input, a linear layer acts on the last axis. The same learned weights apply at every position. It does not mix positions. Attention supplies that mixing elsewhere.

from torch import nn
ffn = nn.Sequential(
    nn.Linear(32, 128),
    nn.GELU(),
    nn.Linear(128, 32),
)
assert ffn(torch.randn(2, 5, 32)).shape == (2, 5, 32)

nn.Sequential calls these functions in order. The object also registers their parameters. model.parameters() can therefore return the weights to the optimizer. A function that constructs new layers on every call would reset its learned weights. Create layers once, then reuse them.

Keep the identity path

Calculate one feature: if the input feature is 2.5 and the attention branch returns -0.4, the residual result is 2.1. The branch modifies the input. It does not replace the input completely.

A residual connection adds the branch output to its input. If the branch outputs zero, the result is exactly the input. This gives gradients a direct path through the addition. It does not make an unstable optimizer safe, and it does not remove the need to choose a learning rate.

Test the shape before adding. Broadcasting can hide an incorrect branch shape. For example, a branch shaped (B,1,C) may be added to every position without an error. The operation runs, but it is not the block you intended.

Training mode is part of the model behavior

Dropout randomly removes some activations during training. It rescales retained values to preserve their expected size. Evaluation disables that randomness. model.eval() changes layer behavior. torch.no_grad() disables gradient recording. They do different jobs; use both for ordinary evaluation.

Start correctness experiments with dropout zero. Once shapes, gradients, and masks pass, add regularization as a separate experiment. The tiny synthetic tasks are not evidence that one dropout rate works for language modeling.

From one block to a stack

An encoder block uses unrestricted self-attention except for padding. A decoder-only block uses causal self-attention. A sequence-to-sequence decoder adds a third branch: cross-attention over encoder outputs. It needs its own normalization and residual connection.

Use separate normalization modules for separate branches. Use separate block instances for separate layers. Repeating the same object in a Python list shares weights accidentally. Store the instances in nn.ModuleList. The final model normally applies one more normalization before its output projection in this pre-normalized design.

Inspect gradients rather than guessing

Predict that removing a residual path will change both forward values and gradient routes. Then compare gradient norms under the same seed and input. One noisy run is a debugging clue, not a general performance claim.

After one backward pass, check every intended trainable parameter. Its gradient should exist and contain finite numbers. A missing gradient can reveal a disconnected branch. A zero gradient is not automatically a bug, but persistent zeros deserve inspection.

Read the executable source

These source listings are copied from the runnable files by scripts/sync-build-code.mjs. The training runner and module belong in the same folder. Read a section, run the tests, then make one small change.

class ManualLayerNorm(nn.Module):
    """Layer normalization written from mean and variance for inspection."""

    def __init__(self, width: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.bias = nn.Parameter(torch.zeros(width))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        variance = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
        normalized = (x - mean) * torch.rsqrt(variance + self.eps)
        return normalized * self.weight + self.bias

class FeedForward(nn.Module):
    """Position-wise ``D -> hidden -> D`` network used by each block."""

    def __init__(self, d_model: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.input_projection = nn.Linear(d_model, hidden_dim)
        self.output_projection = nn.Linear(hidden_dim, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.output_projection(self.dropout(F.gelu(self.input_projection(x)))))

class EncoderBlock(nn.Module):
    """Pre-norm encoder block: non-causal self-attention, then FFN."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.attention_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(self, x: torch.Tensor, valid_tokens: torch.Tensor) -> torch.Tensor:
        attention_input = self.attention_norm(x)
        x = x + self.self_attention(
            attention_input,
            key_padding_mask=valid_tokens,
            query_padding_mask=valid_tokens,
        )
        x = x + self.feed_forward(self.ffn_norm(x))
        return x * valid_tokens.unsqueeze(-1).to(x.dtype)

class DecoderBlock(nn.Module):
    """Pre-norm decoder block with causal self-attention and cross-attention."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.self_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.cross_norm = ManualLayerNorm(d_model)
        self.cross_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid_tokens: torch.Tensor,
        memory_valid_tokens: torch.Tensor,
    ) -> torch.Tensor:
        normalized = self.self_norm(x)
        x = x + self.self_attention(
            normalized,
            key_padding_mask=target_valid_tokens,
            query_padding_mask=target_valid_tokens,
            causal=True,
        )
        normalized = self.cross_norm(x)
        x = x + self.cross_attention(
            normalized,
            memory,
            key_padding_mask=memory_valid_tokens,
            query_padding_mask=target_valid_tokens,
        )
        x = x + self.feed_forward(self.ffn_norm(x))
        return x * target_valid_tokens.unsqueeze(-1).to(x.dtype)

class DecoderOnlyBlock(nn.Module):
    """Pre-norm causal block that can append to a per-layer KV cache."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.attention_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(
        self,
        x: torch.Tensor,
        valid_tokens: torch.Tensor,
        *,
        cache: KVCache | None = None,
        use_cache: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache]:
        attention_result = self.self_attention(
            self.attention_norm(x),
            key_padding_mask=valid_tokens,
            query_padding_mask=valid_tokens,
            causal=True,
            cache=cache,
            use_cache=use_cache,
        )
        if use_cache:
            attention_output, new_cache = attention_result
        else:
            attention_output = attention_result
        x = x + attention_output
        x = x + self.feed_forward(self.ffn_norm(x))
        x = x * valid_tokens.unsqueeze(-1).to(x.dtype)
        if use_cache:
            return x, new_cache
        return x

Further reading

Primary references: Attention Is All You Need, RoFormer, and PyTorch attention API and mask conventions. This implementation is a small teaching model, not a reproduction of a pretrained model.