LLM Primer

Modules and Forward

inputmoduleforward rulenamed parametersoutputa module packages computation with trainable statethe optimizer updates parameters, not arbitrary temporary tensors
Modules make larger models easier to name, inspect, save, and train.

A module packages computation.

In PyTorch, most model pieces are written as subclasses of torch.nn.Module. The module owns parameters and defines a forward method.

import torch
from torch import nn

class TinyLinear(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(3, 2)

    def forward(self, x):
        return self.layer(x)

The module is not magic. It is a named container around a function and its trainable state.

Calling the module runs the forward computation:

model = TinyLinear()
x = torch.randn(4, 3)
scores = model(x)

print(scores.shape)  # torch.Size([4, 2])

Read the example by shape. The input has 4 examples with 3 features each. The linear layer produces 2 scores per example, so the output shape is (4, 2).

Modules become important when models are built from many pieces. Instead of passing every weight tensor around by hand, the model carries its own named parts.

Answer it first, then check.

The module boundary is an engineering boundary. The mathematics inside it is still a forward function.