DL S015 · `nn.Module`, `nn.Sequential`, and the Layer Abstraction

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 15 / 80 · Module M03 · ~2 hours

The story — the day PyTorch stopped feeling like a script

Adam Paszke, the lead author of the original PyTorch paper, tells the story of the moment nn.Module clicked in his head. He was porting a research codebase from Torch7 (the Lua predecessor). In Torch7, you built a network with nn.Sequential():add(nn.Linear(...)). It worked, but anything non-sequential — a residual connection, a two-headed model, a GAN with two networks sharing weights — required manual bookkeeping of which tensors were parameters. When Paszke rewrote it in Python and realised he could just override __setattr__ to auto-detect parameters as attributes, he wrote it in one evening. The next morning he tested it on a two-headed model. It worked without any bookkeeping. He knew they had the right abstraction. Every model in every PyTorch codebase since — nanoGPT, Stable Diffusion, Llama, Whisper — is a subclass of that same nn.Module.

Yesterday we wrote a working MNIST classifier with raw tensors. It works. And it's already ugly. Four parameter tensors floating around at module scope, a forward() function that references them by name, an update loop that lists them by hand. Add a layer? Touch six places. Save the model? Manually pickle four tensors. Load it in a different script? Hope you spelled the names right.

Every serious PyTorch codebase — from the smallest research script to GPT-4's training rig — is organized around one class: nn.Module. It's the thing that makes a model feel like an object. Parameters live inside it. Forward is a method. Iteration, saving, moving to GPU, switching between train and eval mode — all one-liners because the module knows what's inside it.

The name is unassuming. The impact is enormous. Every layer you'll ever use (Linear, Conv2d, LayerNorm, MultiheadAttention, TransformerBlock) is an nn.Module. Every model composed of them is an nn.Module. Once you can write one, you can write anything.

Prerequisites

• Session 014 — tensors, autograd, .backward().

1 · The smallest possible nn.Module

Three parts, in order:

1. super().__init__() — calls nn.Module.__init__(), which sets up the internal dicts (_parameters, _modules, _buffers). Forget this line and everything silently breaks.
2. self.W = nn.Parameter(...) — wrapping a tensor in nn.Parameter tells the module "this is trainable, register it". nn.Module.__setattr__ is overridden to detect Parameters and put them in self._parameters.
3. def forward(self, x) — the actual math. You never call this directly; you call model(x) which invokes __call__, which calls forward plus runs hooks.

Usage:

layer = Linear(784, 128)
x = torch.randn(32, 784)
y = layer(x)                        # NOT layer.forward(x)
print(y.shape)                      # torch.Size([32, 128])
print(list(layer.parameters()))     # [W, b] — auto-discovered

You could have written all this with raw tensors (we did, yesterday). The win is what comes next: composition.

2 · Composition — modules that contain modules

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = Linear(784, 128)      # sub-module
        self.fc2 = Linear(128, 10)       # sub-module
 
    def forward(self, x):
        h = torch.relu(self.fc1(x))
        return self.fc2(h)
 
model = MLP()
print(sum(p.numel() for p in model.parameters()))     # 101 770

Nothing special had to happen for .parameters() to find W, b inside fc1 and fc2. nn.Module.__setattr__ detected that self.fc1 is another nn.Module, put it in self._modules, and .parameters() recursively walks children. This is the entire mechanism.

3 · The three ways layers get moved / iterated / saved

3.1 .to(device) — moves everything

Because parameters are registered, .to finds them all. If you used a plain torch.Tensor for weights (forgot nn.Parameter), it stays on CPU. Silent bug. Your GPU forward pass will crash with expected all tensors to be on same device.

3.2 .parameters() — feeds the optimizer

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

That's the whole reason we're doing this. No more listing [W1, b1, W2, b2] by hand. The optimizer holds an iterator over model parameters and updates them in .step().

3.3 .state_dict() — saves everything

torch.save(model.state_dict(), "mlp.pt")
 
# ... later, in a different script ...
model = MLP()
model.load_state_dict(torch.load("mlp.pt"))

Never torch.save(model) — that pickles the whole object including class definitions, breaks when you refactor, is a security hole when loading untrusted files. state_dict is just a dict of name → tensor and it's forward-compatible across almost every code change that doesn't rename layers.

4 · .train() vs .eval() — the mode switch that matters

Two layers behave differently in train vs eval mode: Dropout (S020) and BatchNorm (S021). Both check self.training. .train() sets it True on the module and all children; .eval() sets it False.

model.train()               # dropout on, batchnorm updates running stats
loss = criterion(model(xb), yb)
loss.backward()
optimizer.step()
 
model.eval()
with torch.no_grad():
    acc = (model(x_val).argmax(-1) == y_val).float().mean()

Forgetting .eval() before validation is one of the top-5 PyTorch bugs. Your validation numbers get randomly worse because dropout is still zeroing 20% of activations, and your batchnorm stats are being computed from the tiny val batch instead of the running average. Sometimes you spend a week thinking your model is broken.

5 · nn.Sequential — when order is all you need

Half the models you'll ever write are "apply these layers in order". PyTorch has a shorthand:

model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)

That's the entire MLP. It's an nn.Module under the hood — same .parameters(), .to, .state_dict. Forward is "call each child in order". Use it when your model is a straight pipe. Use a real class when you have branches, residuals, or conditional logic.

5.1 The OrderedDict version — named layers

I use this pattern a lot. The names show up in state_dict keys and TensorBoard graphs and debugging — worth the extra typing.

6 · The full MNIST training loop, PyTorch-idiomatic

import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
 
# (Datasets/loaders — full treatment in Session 016)
train_ds = datasets.MNIST("./data", train=True, download=True,
                          transform=transforms.ToTensor())
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
 
# Model
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)
 
# Loss + optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
 
# Train
model.train()
for epoch in range(3):
    for xb, yb in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(xb), yb)
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch}: loss {loss.item():.3f}")

That is a complete, idiomatic MNIST training loop. This is what every PyTorch training loop looks like, from tutorials to Meta's production. The 5-line body — zero grad → forward → loss → backward → step — is the same for a 3-layer MLP and a 175B-parameter transformer. That invariance is the beauty of the abstraction.

7 · Under the hood — the metaclass magic (in one paragraph)

nn.Module.__setattr__ is overridden. When you write self.fc1 = nn.Linear(...), Python calls __setattr__("fc1", linear_instance). The overridden version checks: "Is this an nn.Parameter? Store in ._parameters. Is it an nn.Module? Store in ._modules. Otherwise, store as a normal attribute." Then .parameters() walks ._parameters and recursively .parameters() of each entry in ._modules. That's the whole magic. Once you know this, no nn.Module behavior is mysterious.

8 · War stories

AttributeError: cannot assign module before Module.__init__() call

self.layers = nn.ModuleList([nn.Linear(sizes[i], sizes[i+1])
                             for i in range(len(sizes)-1)])

model.train()
# ... training ...
model.eval()
with torch.no_grad():
    # ... validation ...

9 · Diagram — the module tree

.parameters() does a depth-first walk of that tree and yields every yellow node.

9.5 · The seven container types you'll meet

Once you see composition, you notice PyTorch ships seven ways to hold sub-modules or parameters. Pick the right one and your model is elegant; pick wrong and you get the list-of-layers bug from §8.

The buffer row is easy to forget and shows up constantly in real code. RoPE precomputes a (max_seq_len, head_dim/2) cos/sin table — that's a buffer. Causal attention masks are buffers. BatchNorm's running_mean and running_var are buffers. If you find yourself storing a non-trainable tensor as self.mask = ... and then wondering why it stayed on CPU after .to('cuda'), replace with self.register_buffer('mask', ...).

9.6 · Hooks — the debugger you didn't know you had

Every nn.Module supports forward hooks and backward hooks — callbacks that run whenever the module is called or whenever gradients flow through it. They're the most-underused debugging tool in PyTorch.

def print_shape(module, inputs, output):
    print(f"{module.__class__.__name__:12s} in={inputs[0].shape}  out={output.shape}")
 
for layer in model.modules():
    layer.register_forward_hook(print_shape)
 
model(torch.randn(4, 784))
# Flatten     in=torch.Size([4, 784])  out=torch.Size([4, 784])
# Linear      in=torch.Size([4, 784])  out=torch.Size([4, 128])
# ReLU        in=torch.Size([4, 128])  out=torch.Size([4, 128])
# Linear      in=torch.Size([4, 128])  out=torch.Size([4, 10])

This is how you inspect activations in a foreign codebase without touching its code. Backward hooks let you catch NaN gradients at the exact layer they appear (Session 019). Grad-CAM, feature visualisation, layer-wise gradient statistics — all built on hooks.

2025 note: hooks compose with torch.compile as of 2.5 (previously they caused graph breaks). You can now leave hooks installed during compiled training.

9.7 · nn.Module meets torch.compile and FSDP2

A quick preview of two features you'll use heavily in later modules, because they're both built directly on nn.Module's tree structure.

torch.compile(model) — traces model.forward and every child module's forward. Because the module tree is introspectable, TorchDynamo can build one fused graph across dozens of layers. Rule of thumb: compile at the outermost module, not each child.

FSDP2 (fully_shard, PyTorch 2.4+) — the successor to FSDP1 that we cover in Module M09. It walks the module tree, wraps each sub-module in a DTensor-sharded parameter, and unshards on demand during forward/backward. The unit of sharding is the sub-module. This is why you name your transformer blocks as separate modules — FSDP2 shards per block automatically.

The abstraction we're building today is what makes distributed training possible without touching your training loop. When you eventually train a 7B parameter model across 8 GPUs, the only line that changes is model = fully_shard(model). Everything else — the forward, the optimizer.step(), the checkpointing — is unchanged. That is the ROI on learning nn.Module well.

10 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (smallest Module) and §9.5 (containers table) in 24 hours. Full session on day 7.

Next session (S016): DataLoader — shuffling, batching, worker processes (and why they deadlock), collate_fn for variable-length NLP inputs, pin_memory and the CPU→GPU pipeline.

Sticky note (keep on your desk): super().__init__() → declare Parameters/Modules → def forward. Save state_dict, never the object. .train()/.eval() sandwich — always both.

Previous: ← DL S014 · PyTorch Tensors and Autograd · Next: DL S016 · Datasets, DataLoaders, and Collate →