Modules, Parameters, Devices, and Reproducibility

PyTorch models are ordinary Python objects with one important convention: learnable tensors must be registered so PyTorch can find them. This lesson teaches how nn.Module, nn.Parameter, device placement, and random seeds fit together. The outcome is practical: you should be able to build a small model, inspect exactly what will be optimized, move every relevant tensor to the same compute device, save and reload weights, and make experiments repeatable enough to debug.

Purpose and Outcome

In earlier PyTorch foundation work, tensors and autograd explained how a scalar loss produces gradients. A trainable neural network adds durable state. A layer such as nn.Linear owns weight and bias tensors, exposes them to optimizers, participates in nested model structure, and knows how to serialize its state. Those behaviors come from nn.Module, not from magic in the training loop.

By the end of this chapter, you will distinguish a plain tensor attribute from a registered parameter, explain why model.to(device) does not move every tensor you create, use state_dict() as the canonical checkpoint interface, and set reproducibility controls without expecting impossible bit-for-bit guarantees across every hardware and kernel combination.

How Module State Works Internally

nn.Module overrides attribute assignment. When you assign a child module, PyTorch stores it in the module’s internal _modules registry. When you assign an nn.Parameter, PyTorch stores it in _parameters. When you call register_buffer, PyTorch stores a non-optimized tensor in _buffers. The recursive APIs, including parameters(), named_parameters(), state_dict(), train(), eval(), and to(), walk these registries.

A Parameter is a tensor subclass intended to be optimized. It defaults to requires_grad=True. A tensor can have gradients without being registered, but an optimizer built from model.parameters() will not see it unless it is a registered parameter. This is the most common hidden bug in hand-written modules: computation uses a tensor, gradients may even exist, but the training step never updates it.

Buffers solve a different problem. Batch normalization running statistics, positional masks, and normalization constants often belong to the model and should move with the model and appear in checkpoints, but should not receive gradient updates. Register those tensors as buffers. Use persistent=False for temporary buffers that should move devices but should not be saved.

Device placement is also registry-driven. Calling model.to("cuda") or model.to(device) moves registered parameters and buffers recursively. It does not move local tensors created inside forward unless you create them on the right device. Prefer deriving new tensors from an existing input, parameter, or buffer with APIs such as x.new_zeros(...), torch.zeros_like(...), or torch.arange(..., device=x.device).

API Anatomy

A minimal custom module inherits from nn.Module, calls super().__init__(), creates submodules or parameters in __init__, and defines forward. You normally call the module as model(x), which invokes hooks and then forward. Calling forward directly bypasses that machinery.

Worked Example 1: Registered Versus Unregistered Weights

The first example contrasts a broken model with a correct one. Both perform a matrix multiplication, but only the second model registers its weight as an nn.Parameter. The deterministic output is the number of visible parameters.

import torch
from torch import nn

class HiddenWeight(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.weight = torch.randn(3, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight

class RegisteredWeight(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.randn(3, 2))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight

bad = HiddenWeight()
good = RegisteredWeight()
print(sum(p.numel() for p in bad.parameters()))
print(sum(p.numel() for p in good.parameters()))

The expected output is 0 and then 6. The bad model can run forward, but an optimizer created with bad.parameters() has nothing to update. The good model exposes six scalar parameters, matching the 3 x 2 weight matrix.

Worked Example 2: Parameters, Buffers, and Device-Safe Forward

This example builds a small standardization layer. The scale is learnable, while the mean is a buffer. The forward method avoids creating CPU tensors accidentally by using registered state that already follows the input device.

import torch
from torch import nn

class StandardizeThenScale(nn.Module):
    def __init__(self, feature_mean: torch.Tensor) -> None:
        super().__init__()
        self.register_buffer("feature_mean", feature_mean.float())
        self.scale = nn.Parameter(torch.ones_like(feature_mean, dtype=torch.float32))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return (x - self.feature_mean) * self.scale

layer = StandardizeThenScale(torch.tensor([1.0, 2.0, 3.0]))
x = torch.tensor([[2.0, 4.0, 6.0]])
y = layer(x)
print(y.tolist())
print(list(layer.state_dict().keys()))

The expected output is [[1.0, 2.0, 3.0]], followed by keys containing scale and feature_mean. The optimizer should update scale, not feature_mean. A checkpoint should still save both because the mean is required to reproduce inference.

Worked Example 3: Repeatable Initialization and Checkpoint Round Trip

Reproducibility begins with controlling random number generators and recording enough state to repeat an experiment. The next script seeds PyTorch, initializes a model, saves its state_dict to memory, loads it into another instance, and verifies identical predictions for the same input.

import io
import torch
from torch import nn

torch.manual_seed(7)
model_a = nn.Sequential(nn.Linear(4, 5), nn.ReLU(), nn.Linear(5, 1))
x = torch.ones(2, 4)
out_a = model_a(x)

buffer = io.BytesIO()
torch.save(model_a.state_dict(), buffer)
buffer.seek(0)

model_b = nn.Sequential(nn.Linear(4, 5), nn.ReLU(), nn.Linear(5, 1))
model_b.load_state_dict(torch.load(buffer, weights_only=True))
out_b = model_b(x)

print(torch.allclose(out_a, out_b))

The expected output is True. Notice that model_b‘s random initialization does not matter after loading the checkpoint. The names and shapes in the saved state_dict must match the receiving architecture, which is why changing layer names or dimensions breaks strict loading.

Design Choices and Trade-offs

Use built-in modules when possible. They are tested, integrate with initialization utilities, expose expected state names, and usually have optimized kernels. Write custom modules when you need a new computation, a reusable composition of layers, or explicit state such as a learned temperature, mask, or calibration constant.

Prefer child modules over manually managed lists. A plain Python list containing layers is invisible to module traversal. Use nn.ModuleList when the number of layers is dynamic and you need to iterate in forward. Use nn.Sequential when the computation is a simple chain. Use nn.ParameterList or nn.ParameterDict for dynamic sets of learnable tensors.

Device choice trades speed, memory, determinism, and availability. CPU execution is often easiest to debug and may be deterministic for small examples. GPUs accelerate large tensor operations but introduce device mismatch errors, asynchronous execution, and kernels whose deterministic alternatives can be slower. Mixed precision improves throughput and memory usage but requires attention to numeric range and loss scaling in training.

Reproducibility has levels. For teaching and debugging, seed random generators and use fixed data order. For experiment comparison, record code revision, data version, hyperparameters, device type, library configuration, and checkpoint. For strict audits, enable deterministic algorithms where practical and accept that unsupported operations may fail or slow down.

Failure Modes and Troubleshooting

Symptom: the loss changes during backpropagation but a custom weight never updates. Cause: the tensor is a plain attribute instead of an nn.Parameter, or it is stored inside a plain list. Diagnose: print dict(model.named_parameters()).keys() and check whether the name appears. Correct: wrap learnable tensors in nn.Parameter and use ModuleList, ParameterList, or named child modules for collections.

Symptom: PyTorch raises an error saying tensors are on different devices. Cause: an input is on GPU while a mask, index, or constant created in forward defaults to CPU. Diagnose: print the .device of the input, parameters, buffers, and newly created tensor. Correct: create new tensors with device=x.device or derive them from existing tensors.

Symptom: loading a checkpoint reports missing or unexpected keys. Cause: the class definition changed, a layer was renamed, or a wrapper such as data parallelism changed key prefixes. Diagnose: compare checkpoint.keys() with model.state_dict().keys(). Correct: restore the matching architecture, write an explicit key migration, or use non-strict loading only when you intentionally accept uninitialized or ignored state.

Symptom: repeated runs produce different metrics. Cause: random initialization, shuffled data, nondeterministic kernels, dropout left in training mode, or data preprocessing outside the seed path. Diagnose: fix seeds, run on a tiny dataset, compare the first batch, parameter norms, and first prediction. Correct: seed all relevant generators, call model.eval() for evaluation, persist splits, and record the environment.

Reliability, Performance, and Security Implications

Reliability depends on treating model state as a named interface. Stable state names make checkpoints, fine-tuning, and partial loading understandable. Device-safe tensor creation prevents failures that appear only when moving from a laptop CPU to an accelerator. Reproducible scripts reduce debugging time because you can bisect changes against a repeatable baseline.

Performance depends on minimizing unnecessary device transfers. Moving data or model state between CPU and GPU inside each forward pass can dominate runtime. Put the model on the target device once, move each batch once, and keep temporary tensors on the same device. Use buffers for reusable constants so they move with the model instead of being rebuilt repeatedly.

Checkpoint security matters because model files may be exchanged between systems. Prefer loading weight dictionaries rather than arbitrary Python objects, keep checkpoints from untrusted sources isolated, and verify architecture, shapes, and expected keys before using loaded weights in a larger pipeline.

Hands-on Lab: Build, Move, Save, and Verify

Prerequisites: Python with PyTorch installed, a terminal or notebook, and optional CUDA or other accelerator support. The lab works on CPU. If an accelerator is available, it also verifies device movement.

1. Create a file or notebook cell with a custom TinyRegressor module containing two nn.Linear layers and a registered buffer named input_mean.
2. Choose device = torch.device("cuda" if torch.cuda.is_available() else "cpu"), instantiate the model, and call model.to(device).
3. Create a batch with torch.randn(8, 3, device=device), run a forward pass, and assert that the output shape is (8, 1).
4. Print parameter names, buffer names, and the device of each tensor. Verify that every registered tensor is on the chosen device.
5. Save model.state_dict() to a temporary path, create a fresh model, move it to the same device, load the state dict, and compare predictions with torch.allclose.
6. Cleanup by deleting the temporary checkpoint. In a notebook, also clear large tensors or restart the runtime if you used a memory-constrained accelerator.

import os
import tempfile
import torch
from torch import nn

class TinyRegressor(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.register_buffer("input_mean", torch.tensor([0.5, 0.0, -0.5]))
        self.net = nn.Sequential(nn.Linear(3, 4), nn.ReLU(), nn.Linear(4, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x - self.input_mean)

torch.manual_seed(123)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TinyRegressor().to(device)
batch = torch.randn(8, 3, device=device)
out = model(batch)
assert out.shape == (8, 1)
assert all(p.device == device for p in model.parameters())
assert all(b.device == device for b in model.buffers())

with tempfile.NamedTemporaryFile(delete=False) as tmp:
    path = tmp.name
try:
    torch.save(model.state_dict(), path)
    clone = TinyRegressor().to(device)
    clone.load_state_dict(torch.load(path, map_location=device, weights_only=True))
    assert torch.allclose(model(batch), clone(batch))
    print("verified", device.type, round(float(out.mean()), 6))
finally:
    os.remove(path)

Verification is successful when the script prints a line beginning with verified and no assertion fails. The rounded mean may differ by device and PyTorch build, so the reliable check is the shape, device assertions, and checkpoint round trip.

Assessment Exercises

1. A custom module stores three nn.Linear layers in a plain Python list. Explain why training may fail and rewrite the storage using the right PyTorch container.
2. You need a fixed vocabulary mask that must move to GPU and be saved with the model, but it must not be optimized. Should it be a parameter, buffer, or local tensor? Justify the choice.
3. A checkpoint fails with one missing key and one unexpected key after you renamed encoder to backbone. Describe a careful migration procedure.
4. Design a minimal reproducibility checklist for comparing two learning-rate schedules. Include state, data, randomness, and evaluation mode.
5. Find and correct the bug in a forward method that calls torch.zeros(batch_size, hidden_size) while the input batch may be on GPU.

Summary

nn.Module gives PyTorch models structure by registering child modules, learnable parameters, and non-learnable buffers. Optimizers update registered parameters; device moves and checkpoints traverse registered state. Correct model code therefore depends on putting state in the right registry, creating temporary tensors on the right device, and saving weights through state_dict. Reproducibility is a disciplined experiment practice: control randomness, persist model and data state, record configuration, and verify behavior with small deterministic checks before trusting large training runs.