nn.Module: Registration, Composition and Checkpoint Contracts

The story — ownership matters beyond forward arithmetic

Four global tensors work for a toy network. Once there are layers, state, checkpointing and mode changes, a model needs explicit ownership. nn.Module provides parameter/buffer/child registration and recursive operations over that structure. It does not compute gradients itself: autograd follows executed tensor operations; an optimizer updates the objects it was given.

The original story about an overnight implementation and proprietary GPT-4 internals was unsupported. The source-verifiable mechanism is ordinary Python attribute assignment and module registries, not “metaclass magic.” A model may share modules/parameters, so the structure is tree-like but need not be a strict mathematical tree.

Prerequisites and outcomes

Bring raw tensors/autograd. We use PyTorch2.9.0+cpu. Learn to initialize Module, register parameters and buffers, compose containers, inspect names/counts, save and restore compatible state, distinguish eval from no-grad, remove debugging hooks, and understand the module-level choices behind compilation and FSDP2. Optional references: official neural-net tutorial, Module API, Patrick Loeber, Aladdin Persson.

1 · The smallest possible nn.Module

A parameter is a Tensor subclass used as a registration marker. Its default requires_grad=True is convenient, but registration and differentiation are separate. A plain requiring-grad tensor can receive an autograd gradient and still be absent from model.parameters().

import torch
import torch.nn as nn
import torch.nn.functional as F
import io
import copy
from collections import OrderedDict
 
torch.set_num_threads(1)
torch.manual_seed(15)
 
class Linear(nn.Module):
    def __init__(self,in_features,out_features):
        super().__init__()
        self.W=nn.Parameter(torch.randn(in_features,out_features)*.01)
        self.b=nn.Parameter(torch.zeros(out_features))
    def forward(self,x):return x@self.W+self.b
 
layer=Linear(784,128)
x=torch.randn(32,784)
assert layer(x).shape==(32,128)
assert sum(p.numel() for p in layer.parameters())==100480
assert set(dict(layer.named_parameters()))=={'W','b'}

Call layer(x) so Module's call machinery can invoke hooks; directly calling forward bypasses that machinery. Our custom W is (in,out) while nn.Linear.weight is (out,in). Forgetting super().__init__() before assigning a Parameter/Module normally raises immediately, rather than silently succeeding.

Registration experiment — keep the bias accounting correct

If only W loses its wrapper, b remains registered: the count falls from100480 to128, not zero. A plain tensor can still be given explicitly to an optimizer, but this lesson uses the model registry deliberately.

class PlainWeight(nn.Module):
    def __init__(self):
        super().__init__()
        self.W=torch.randn(784,128,requires_grad=True)
        self.b=nn.Parameter(torch.zeros(128))
    def forward(self,x):return x@self.W+self.b
bad=PlainWeight();bad(torch.ones(1,784)).sum().backward()
assert bad.W.grad is not None
assert sum(p.numel() for p in bad.parameters())==128
assert 'W' not in bad.state_dict()
try:
    class Broken(nn.Module):
        def __init__(self):self.fc=nn.Linear(2,2)
    Broken()
except AttributeError:pass
else:raise AssertionError('missing Module initialization must reject')

2 · Composition — modules that contain modules

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1=Linear(784,128)
        self.fc2=Linear(128,10)
    def forward(self,x):return self.fc2(torch.relu(self.fc1(x)))
model=MLP()
assert sum(p.numel() for p in model.parameters())==101770
assert set(model.state_dict())=={'fc1.W','fc1.b','fc2.W','fc2.b'}

Assigned Parameters enter the parameter registry; assigned Modules enter the child registry. Ordinary tensors do not automatically become parameters or buffers. Register persistent non-parameter state with register_buffer; persistent=False buffers still move with the module but do not appear in its state dictionary. A buffer is not automatically optimizer-owned, although one could explicitly supply a requiring-grad tensor to an optimizer.

The Lego metaphor is recursive ownership: inspecting/moving/saving the top module reaches registered children. It does not remove the need to understand control flow, shape constraints or shared objects.

3 · Moving, iterating and saving

3.1 to moves registered tensors

model.to(device=...,dtype=...) changes registered parameters and buffers according to dtype/device rules and returns the module. It does not move arbitrary tensor attributes or a separately created input batch. The CPU test uses a dtype conversion to expose that same registry boundary without requiring CUDA. Integer buffers keep integer dtype under a floating dtype conversion.

class StateProbe(nn.Module):
    def __init__(self):
        super().__init__()
        self.w=nn.Parameter(torch.ones(2))
        self.register_buffer('mean',torch.zeros(2))
        self.register_buffer('cache',torch.ones(2),persistent=False)
        self.loose=torch.zeros(2)
s=StateProbe().to(dtype=torch.float64)
assert s.w.dtype==s.mean.dtype==s.cache.dtype==torch.float64
assert s.loose.dtype==torch.float32
assert set(s.state_dict())=={'w','mean'}
assert list(dict(s.named_parameters()))==['w']

3.2 Parameters feed the optimizer

Construct the optimizer after final placement/wrapping where objects can change. The optimizer materializes references to supplied parameter objects in parameter groups; it is not a permanently live query of the module. Adding a layer later does not silently add it to an existing optimizer. Shared parameters are normally deduplicated by named_parameters/parameters traversal.

3.3 A state dictionary is state, not a whole experiment

It contains parameters and persistent buffers (and modules may define extra state), not a universal architecture specification, training mode, optimizer trajectory or loader cursor. Rename/reshape layers and strict loading should identify incompatibility. strict=False is not permission to ignore incompatible tensor shapes.

torch.save uses serialization including pickle metadata even for state dictionaries. Saving the whole model couples loading to Python class/module names; classes are not simply embedded as source code. Load only trusted artifacts and use weights_only=True and an explicit map location. Restricted loading reduces code-execution exposure but is not a complete denial-of-service or all-file-safety guarantee.

The in-memory test below saves a compatible state/config, reloads it and compares outputs. For a “best checkpoint” held in memory, clone/deepcopy state because state_dict() is a shallow dictionary of references to current tensors.

4 · train versus eval is orthogonal to grad mode

Dropout is stochastic in training and identity at evaluation; BatchNorm usually updates running statistics and uses batch statistics in training, then uses running statistics in eval when configured to track them. These are prominent examples, not the only modules that can branch on self.training. Custom modules and other built-ins can also care. BatchNorm with track_running_stats=False does not acquire frozen running statistics in eval.

m=nn.Sequential(nn.Linear(4,4),nn.Dropout(.5),nn.BatchNorm1d(4))
a=torch.randn(8,4)
m.train();train_out=m(a)
assert m[2].num_batches_tracked.item()==1
m.eval();eval_out=m(a)
assert eval_out.requires_grad # eval does not disable autograd
assert m[2].num_batches_tracked.item()==1
assert not torch.allclose(train_out,eval_out)
with torch.no_grad():assert not m(a).requires_grad
assert {'2.running_mean','2.running_var','2.num_batches_tracked'} <= set(m.state_dict())

Restore model.train() after validation before the next training phase. Eval plus no-grad is a common inference combination; gradient-based analysis may intentionally use eval with autograd. A validation gap does not prove missing eval, and no fixed5% penalty is implied.

5 · Sequential when order is the architecture

seq=nn.Sequential(OrderedDict([
    ('fc1',nn.Linear(784,128)),('relu',nn.ReLU()),('fc2',nn.Linear(128,10))]))
assert seq.fc1.weight.shape==(128,784)
assert sum(p.numel() for p in seq.parameters())==101770
assert seq(torch.randn(4,784)).shape==(4,10)

Sequential forwards one object through its children. That object can itself be a tuple/dict and a child can implement a residual, so “strictly incapable of multiple values or residuals” is too strong. A custom Module with explicit return x+self.block(x) is usually clearer when branching is a defining feature. OrderedDict names become state keys; replacing numeric positions with names is a checkpoint migration, not just presentation.

6 · Complete configurable MLP and training contract

This replaces the old dataset-downloading example with an executable synthetic default and an explicit local-data adapter. The network accepts flattened features or image batches via flatten(1), owns all layers in ModuleList, exposes optional dropout, and retains a non-trainable input offset as a buffer.

class ConfigMLP(nn.Module):
    def __init__(self,in_features=784,hidden_sizes=(128,),classes=10,dropout=0.):
        super().__init__()
        sizes=[in_features]+list(hidden_sizes)+[classes]
        self.layers=nn.ModuleList([nn.Linear(a,b) for a,b in zip(sizes,sizes[1:])])
        self.drop=nn.Dropout(dropout)
        self.register_buffer('offset',torch.zeros(in_features))
    def forward(self,x):
        x=x.flatten(1)-self.offset
        for layer in self.layers[:-1]:x=self.drop(F.relu(layer(x)))
        return self.layers[-1](x)
 
def train_epoch(model,loader,optimizer):
    model.train();total=0.;n=0
    for xb,yb in loader:
        optimizer.zero_grad(set_to_none=True)
        loss=F.cross_entropy(model(xb),yb)
        loss.backward();optimizer.step()
        total+=loss.item()*len(yb);n+=len(yb)
    if not n:raise ValueError('empty training loader')
    return total/n
 
def evaluate(model,loader):
    model.eval();total=0.;n=0;correct=0
    with torch.no_grad():
        for xb,yb in loader:
            logits=model(xb)
            total+=F.cross_entropy(logits,yb,reduction='sum').item()
            correct+=(logits.argmax(-1)==yb).sum().item();n+=len(yb)
    if not n:raise ValueError('empty evaluation loader')
    return total/n,correct/n
 
from torch.utils.data import TensorDataset,DataLoader
cfg=dict(in_features=4,hidden_sizes=(8,6),classes=2,dropout=.1)
net=ConfigMLP(**cfg)
sx=torch.randn(40,4);sy=(sx[:,0]+sx[:,1]>0).long()
loader=DataLoader(TensorDataset(sx,sy),batch_size=7,shuffle=True,
                  generator=torch.Generator().manual_seed(2),num_workers=0)
opt=torch.optim.SGD(net.parameters(),lr=.1)
start=evaluate(net,loader)[0]
for _ in range(25):train_epoch(net,loader,opt)
finish=evaluate(net,loader)[0]
assert finish<start
print('SYNTHETIC Module fit only',start,finish)

For already-local MNIST from DL011, wrap normalized float32 (N,784) or (N,1,28,28) arrays and int64 labels in TensorDataset; use55k training/5k validation and untouched10k test. Create ConfigMLP(784,(128,),10), iterate the training loader with train_epoch, deepcopy state when validation loss improves, restore it, and call evaluate once on test. No torchvision download is needed. The PyTorch loader lesson also shows download=False for existing torchvision-format files. Synthetic loss improvement here is not digit accuracy.

def fit_local_arrays(X,y,Xv,yv,Xt,yt,epochs=3):
    def wrap(a,b,shuffle):
        ds=TensorDataset(torch.as_tensor(a,dtype=torch.float32),torch.as_tensor(b,dtype=torch.int64))
        return DataLoader(ds,batch_size=128,shuffle=shuffle,num_workers=0,
                          generator=torch.Generator().manual_seed(0))
    tr,va,te=wrap(X,y,True),wrap(Xv,yv,False),wrap(Xt,yt,False)
    model=ConfigMLP();optimizer=torch.optim.SGD(model.parameters(),lr=.1)
    best=None;bestloss=float('inf')
    for _ in range(epochs):
        train_epoch(model,tr,optimizer);vl,_=evaluate(model,va)
        if vl<bestloss:bestloss=vl;best=copy.deepcopy(model.state_dict())
    if best is None:raise ValueError('positive epochs required')
    model.load_state_dict(best)
    return model,evaluate(model,te)

7 · The registration mechanism and container failures

Module overrides __setattr__ and offers explicit registration methods. It validates special names/types and maintains registries; the actual implementation is more than “otherwise ordinary attribute.” Assigning a plain list of Modules does not recursively register its elements. ModuleList/ModuleDict do. ParameterList/ParameterDict do likewise for parameters. Existing attribute names can impose type restrictions, so demonstrate a new module rather than replacing an already registered Parameter with an arbitrary Tensor in place.

class BadList(nn.Module):
    def __init__(self):
        super().__init__();self.layers=[nn.Linear(2,3),nn.Linear(3,1)]
class GoodList(nn.Module):
    def __init__(self):
        super().__init__();self.layers=nn.ModuleList([nn.Linear(2,3),nn.Linear(3,1)])
assert sum(p.numel() for p in BadList().parameters())==0
assert sum(p.numel() for p in GoodList().parameters())==13
try:torch.optim.SGD(BadList().parameters(),lr=.1)
except ValueError:pass
else:raise AssertionError('empty optimizer parameter list should reject')

If other registered parameters exist, they may update while the unregistered part stays fixed. If there are none, an optimizer commonly raises on the empty list; “silently trains nothing” is not the only failure mode.

8 · Checkpoint round-trip and debugging discipline

net.eval();probe=sx[:8]
with torch.no_grad():before=net(probe).clone()
buf=io.BytesIO()
torch.save({'config':cfg,'model':copy.deepcopy(net.state_dict())},buf)
buf.seek(0)
checkpoint=torch.load(buf,map_location='cpu',weights_only=True)
restored=ConfigMLP(**checkpoint['config'])
restored.load_state_dict(checkpoint['model'],strict=True);restored.eval()
with torch.no_grad():after=restored(probe)
assert torch.equal(before,after)
# Renaming/missing keys is detected rather than silently called a successful load.
broken=dict(checkpoint['model']);broken.pop('layers.0.weight')
try:restored.load_state_dict(broken,strict=True)
except RuntimeError:pass
else:raise AssertionError('missing state key should reject')

A same-process bit-equal CPU test is not a promise across devices, versions or nondeterministic kernels. The review also tests a fresh Python process with the same pinned runtime and saved input/config. Exact training resume needs optimizer/scheduler/scaler states when used, RNG states, sampler/loader position and environment details in addition to model state. Do not load an untrusted whole-model pickle to rescue a failed state load.

9.5 · Seven registration choices

These are common choices, not an exhaustive theorem that PyTorch has exactly seven container types. A causal mask or rotary table can be a nonpersistent buffer when deterministically reconstructible; BatchNorm running statistics are typically persistent. Plain tensors remain invisible to recursive movement/checkpointing unless explicitly registered.

9.6 · Hooks — observe, then remove

Register forward hooks on leaf modules if that is what you want to inspect; hooking model.modules() also includes the root container, so its output appears too. Save the handles and remove them in finally. Hooks can observe or deliberately modify values depending on the API; careless retention of outputs can keep graphs alive. Use full backward hooks for their documented semantics and avoid in-place changes to hook views.

seen=[];handles=[]
def record_shape(module,inputs,output):
    seen.append((module.__class__.__name__,tuple(inputs[0].shape),tuple(output.shape)))
try:
    for child in seq.modules():
        if not list(child.children()):handles.append(child.register_forward_hook(record_shape))
    seq(torch.randn(4,784))
finally:
    for handle in handles:handle.remove()
assert [x[0] for x in seen]==['Linear','ReLU','Linear']
assert seen[-1][2]==(4,10)
old=len(seen);seq(torch.randn(4,784));assert len(seen)==old

Hooks and torch.compile have version-specific capture, guard and graph-break behavior. Printing from a hook or changing hooks during a compiled run can matter; the old blanket “all hooks compose as of2.5, leave them installed” claim was not warranted. Diagnose eagerly, remove instrumentation, then validate any compiled path independently.

9.7 · Module boundaries, compilation and FSDP2

A module tree provides useful units for capture and sharding, but does not force one fused graph. Choose compile boundaries according to guards, unsupported code, reuse and measured recompilation/latency. Keep an eager numerical reference; no compiler benchmark is run here.

PyTorch2.9's FSDP2 fully_shard is a composable in-place transformation, with parameters represented using DTensor and communication scheduled around computation. Applying it to chosen children bottom-up and then the root establishes communication groups; it does not automatically select the best per-block strategy merely because blocks have names. Create the optimizer after sharding as required by the API. Process-group/device-mesh setup, distributed input sampling, checkpoint conversion and resumption are substantial parts of a real run. It is not “change one line and everything including checkpoints is unchanged.” This is source-reviewed architecture guidance, not an8-GPU reproduction.

10 · Retention scaffold — six worked answers

1. Initialization: super initializes registries and hook/mode machinery. Assignment of a child before initialization raises in the demonstrated case.
2. Plain lists: elements are not registered recursively. ModuleList/ModuleDict fix module ownership; ParameterList/ParameterDict cover raw parameters. An empty optimizer can raise.
3. Modes: Dropout changes stochastic masking; BatchNorm changes batch/running-stat behavior according to configuration. Eval does not disable autograd, and other modules may also branch on mode.
4. Persistent non-parameter state: self.register_buffer('name',tensor); use persistent=False only if its omission from checkpoints is intended.
5. Saving: state_dict decouples tensor state from Python object reconstruction, but torch serialization still uses pickle metadata. Recreate compatible config/class and load trusted restricted state; it is not universally compatible with refactors.
6. Configurable MLP: §6 implements ModuleList, arbitrary hidden widths and dropout; §8 verifies state/config round-trip and strict missing-key detection. Exact same-runtime output parity is tested separately from cross-platform determinism or exact training resume.

Use custom Module for explicit branching, Sequential for honest chains, and functional parameter arguments or torch.func.functional_call when weights should be treated as data. Functional style can coexist with Module tools; it need not abandon state_dict. Reconstruct the three registries tomorrow before moving to data pipelines.

Sources and execution boundary

PyTorch2.9 Module, serialization, Linear source, compiler, FSDP2 fully_shard. Exact passages and hashes support registration, mode, save/load and sharding claims. Executed tests are CPU synthetic models and trusted self-produced checkpoints; no MNIST download/training, CUDA, compiler or distributed run is claimed.