PyTorch Foundations

A PyTorch tensor is a multi-dimensional array with a few extra fields: dtype (data type), device (CPU or which GPU), requires_grad (whether autograd tracks it), and grad (the accumulated gradient, populated by .backward). Everything else in the framework — modules, optimizers, distributed training, torch.compile — is bookkeeping around tensor operations.

The four ways to make a tensor:

import torch

# 1. From data
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

# 2. Zeros / ones / random / arange / linspace
b = torch.zeros(3, 4)
c = torch.randn(2, 3, 4)            # standard normal
d = torch.arange(0, 10, step=2)     # like np.arange

# 3. Like an existing tensor (matched shape and dtype)
e = torch.zeros_like(a)

# 4. From numpy (shares memory!)
import numpy as np
f = torch.from_numpy(np.array([1.0, 2.0, 3.0]))

The dtype defaults to float32 for floating-point tensors and int64 for integers. The device defaults to CPU. Move a tensor across devices with .to(device):

device = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.randn(1000, 1000).to(device)

A tensor's GPU memory (VRAM) is just its element count times its byte size: x.numel * x.element_size, which for float32 (4 bytes) means a 100M-parameter model needs ~0.4 GB just for the weights. CS336 lecture 2's resource accounting (referenced under "What you'll be able to do") extends this to a full training step: counting parameters, gradients, and optimizer state, AdamW in float32 costs roughly 4+4+(4+4) = 16 bytes per parameter, before activations.

requires_grad=True is the flag that turns a tensor into an autograd leaf. Parameters (the things an optimizer updates) have requires_grad=True. Activations and intermediate values do not need it; PyTorch infers requires_grad=True for any output of an operation where at least one input has requires_grad=True.

x = torch.randn(5, requires_grad=True)   # leaf, gradient will be populated
y = x ** 2                                # not a leaf, but requires_grad inferred
y.sum().backward()
print(x.grad)   # 2*x

Tensor operations follow the same broadcasting rules as numpy. Two tensors are broadcastable if, when you align their shapes right-justified, every dimension is either equal or one of them is 1.

a = torch.randn(3, 1, 5)
b = torch.randn(   4, 5)
(a + b).shape   # torch.Size([3, 4, 5])

The most useful operations to internalize early:

view vs reshape is the subtle one. view requires the tensor to be contiguous in memory; it's a zero-copy reinterpretation. reshape will copy if it has to. After transpose or permute, the tensor is usually non-contiguous; calling .view on it raises RuntimeError. The fix is either .reshape (which handles it transparently) or .contiguous.view (which makes an explicit copy first).

x = torch.randn(2, 3, 4)
y = x.permute(0, 2, 1)         # (2, 4, 3) but non-contiguous
# y.view(2, 12)                 # RuntimeError
y.contiguous().view(2, 12)     # ok
y.reshape(2, 12)               # also ok

einsum is the most expressive operation in PyTorch — it is generalized matrix multiplication with explicit index bookkeeping. Read the subscript string as named axes: a letter is summed over (contracted) iff it does not appear to the right of ->, while every letter listed on the right is kept in the output — so a letter shared across inputs but still present on the right (like the batch b) is a batched/parallel axis, run in parallel rather than contracted. So "bij,bjk->bik" is batched matmul (sum over j, keep b,i,k), and "bhid,bhjd->bhij" is the attention-score computation from a transformer. Learn einsum once and you write tensor code that reads like math.

When you operate on tensors with requires_grad=True, PyTorch builds a directed acyclic graph behind the scenes. Each operation records a grad_fn (the gradient function for its operation) and the input tensors it depended on. Calling .backward on a scalar output walks this graph in reverse, calling each grad_fn to accumulate gradients into the .grad field of every leaf tensor.

The graph is dynamic. It is rebuilt every forward pass. This is unlike TensorFlow 1.x, which had a static graph. The dynamism is what lets you use Python control flow (if, for, while) inside a model and have it differentiate correctly.

x = torch.tensor([2.0], requires_grad=True)
y = x ** 3 + 4 * x
print(y.grad_fn)   # <AddBackward0 object> — the most recent op
y.backward()
print(x.grad)      # tensor([16.]) — analytically 3x^2 + 4 at x=2

Three things to know:

Gradients accumulate. If you call .backward twice on the same graph (without retain_graph=True), you get an error. If you call .backward on two different outputs that share leaves, gradients accumulate (sum). This is why every training loop calls optimizer.zero_grad before each loss.backward — to clear the previous step's gradients.

torch.no_grad for inference. Inside with torch.no_grad:, PyTorch skips graph construction. This saves memory and time. Always wrap evaluation loops in it.

.detach cuts the graph. A detached tensor shares data with the original but is treated as a leaf with no grad_fn. Useful when you want to use a value as a constant without it contributing to gradients (e.g., for a target in a teacher-student loss).

import torch
import torch.nn as nn

model = nn.Linear(10, 1)
x = torch.randn(32, 10)
y_target = torch.randn(32, 1)

with torch.no_grad():
    y_pred_frozen = model(x)            # no graph built; no gradients
    
y_pred = model(x)
loss = ((y_pred - y_target) ** 2).mean()
loss.backward()
# model.weight.grad and model.bias.grad now populated

nn.Module is the way you organize parameters in PyTorch. It is a Python class with two responsibilities: store learnable parameters, and define a forward method that operates on them.

The minimal pattern:

import torch.nn as nn

class MyLayer(nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.01)
        self.bias = nn.Parameter(torch.zeros(out_features))

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

Three rules. First, always call super.__init__ in your __init__, before assigning anything else. PyTorch's machinery depends on it. Second, assign learnable tensors as nn.Parameter(...) — this is what gets them registered for optimizer.parameters calls. A plain torch.Tensor attribute is not a parameter and will be silently ignored. Third, define forward, not __call__. nn.Module.__call__ wraps forward with hooks and gradient tracking; if you override __call__ you bypass all of it.

nn.Module composes. A module that contains submodules registers their parameters automatically through Python attribute discovery:

class MLP(nn.Module):
    def __init__(self, d_in: int, d_hidden: int, d_out: int):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.fc2 = nn.Linear(d_hidden, d_out)
        self.activation = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(self.activation(self.fc1(x)))

m = MLP(784, 128, 10)
print(sum(p.numel() for p in m.parameters()))   # 101770

nn.Sequential is a convenience for the common "stack modules" pattern. Use it when there is no branching:

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

For anything with skip connections, conditional branches, or multiple outputs, write forward by hand.

Buffers. Sometimes you need to store a tensor in a module that is not a learnable parameter but should still move with .to(device) and serialize with the model. The classic example is the running mean and variance in BatchNorm, or a causal mask in attention. Use self.register_buffer("name", tensor).

class CausalAttention(nn.Module):
    def __init__(self, max_seq_len: int):
        super().__init__()
        # Upper-triangular mask; not learned, but should move with the module
        mask = torch.triu(torch.ones(max_seq_len, max_seq_len), diagonal=1).bool()
        self.register_buffer("mask", mask)

To make nn.Module concrete, here is what nn.Linear actually does, written by hand. Read it next to the PyTorch source (it is one of the cleaner ones).

import math

class Linear(nn.Module):
    """nn.Linear, re-implemented."""
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter("bias", None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        # Kaiming uniform init: same as nn.Linear's default.
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in = self.in_features
            bound = 1 / math.sqrt(fan_in)
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = x @ self.weight.T
        if self.bias is not None:
            out = out + self.bias
        return out

Three things to notice. First, weight is shape (out_features, in_features), not (in_features, out_features). PyTorch convention. Second, reset_parameters is a separate method that the default __init__ calls; this lets users re-initialize without recreating the module. Third, when bias=False, the convention is register_parameter("bias", None) rather than self.bias = None. This puts a recognized None-valued slot in the module's internal _parameters dict, reserving the name so a real nn.Parameter can later be assigned to it cleanly. A None slot is filtered out of named_parameters, parameters, and the state_dict alike, while self.bias still resolves to None as an attribute — exactly what the if self.bias is not None checks above rely on.

You will write custom nn.Module classes for anything you can't express with nn.Sequential. Transformer blocks, ResNet blocks, attention layers, mixture-of-experts routing — all of them are custom modules. The pattern in this section is what every one of them looks like.

A Dataset is anything with __len__ and __getitem__. A DataLoader wraps a Dataset to produce batches, shuffle, and parallelize via worker processes.

The two Dataset patterns:

from torch.utils.data import Dataset, DataLoader

class TabularDataset(Dataset):
    def __init__(self, X: np.ndarray, y: np.ndarray):
        self.X = torch.from_numpy(X).float()
        self.y = torch.from_numpy(y).long()

    def __len__(self) -> int:
        return len(self.X)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.X[idx], self.y[idx]

That's a "map-style" dataset. For streamed or infinite data, subclass IterableDataset and override __iter__.

For the DataLoader, the defaults are usually wrong for production. Set:

loader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,          # crucial during training; defaults to False
    num_workers=4,         # parallel data loading; 0 means main process only
    pin_memory=True,       # speeds up CPU→GPU transfers
    drop_last=True,        # drop the last partial batch during training
    persistent_workers=True,  # avoid worker startup cost each epoch
)

num_workers > 0 is what makes large data pipelines fast. Each worker is a separate process that prefetches batches in parallel with the main process running the model. Set it to the number of CPU cores you have, then tune down if you see CPU saturation.

pin_memory=True allocates the batch in page-locked memory on the host, which lets CUDA's async transfer engine move it to the GPU without an extra copy. Always set it for GPU training.

An optimizer is a small object that knows how to update parameters given their gradients. PyTorch's optimizers all follow the same interface:

import torch.optim as optim

opt = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# In the training loop:
opt.zero_grad()           # clear last step's gradients
loss.backward()           # accumulate this step's gradients
opt.step()                # update parameters in-place

The four optimizers you will actually use:

SGD with momentum. The classical optimizer. optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=1e-4). Best for convnets when tuned well; the LR has a narrow sweet spot.

Adam (Kingma & Ba, 2014). Adaptive per-parameter learning rates, with exponential moving averages of both the gradient (momentum) and the squared gradient (variance). The update is $\theta_t \leftarrow \theta_{t-1} - \eta \hat{v}_t / (\sqrt{\hat{s}_t} + \epsilon)$ where $\hat{v}_t$ is the bias-corrected first moment and $\hat{s}_t$ is the bias-corrected second moment. optim.Adam(params, lr=3e-4, betas=(0.9, 0.999)). Forgiving with respect to LR. Karpathy's "Adam at 3e-4 is safe" defaults are real.

AdamW (Loshchilov & Hutter, 2019). Adam with decoupled weight decay. The original Adam combines L2 regularization with the adaptive rescaling, which has the side effect of making weight decay weaker on parameters with large second moments. AdamW separates the two. For LLM training, AdamW is the default. optim.AdamW(params, lr=3e-4, weight_decay=0.1).

LBFGS for tiny problems where you can fit the whole dataset in one batch. Almost never used for deep networks.

For LLM-scale training, the modern frontier is Lion and Muon, but you should not reach for them until you have AdamW working. They are covered briefly in another chapter.

Parameter groups let you give different hyperparameters to different parts of the model:

opt = optim.AdamW([
    {"params": model.embed.parameters(), "lr": 1e-5},
    {"params": model.head.parameters(),  "lr": 1e-3},
], weight_decay=0.1)

Standard pattern for fine-tuning: small LR for the pretrained backbone, large LR for the newly-initialized head.

A scheduler adjusts the learning rate during training. PyTorch wraps these in torch.optim.lr_scheduler. The three patterns you will see most often:

StepLR / MultiStepLR: decay LR by a factor every N epochs. Classic for ImageNet-style training.

sched = optim.lr_scheduler.StepLR(opt, step_size=30, gamma=0.1)
for epoch in range(100):
    train_one_epoch(...)
    sched.step()    # call once per epoch, after the optimizer steps

CosineAnnealingLR: smoothly decay LR from lr to eta_min over T_max steps following a cosine curve. The standard for transformer training, often combined with linear warmup at the start.

sched = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10_000, eta_min=1e-5)

OneCycleLR: warm up then anneal in a single cycle. Smith's "1cycle" policy from 13-fastbook/16_accel_sgd. Often gives best results for fixed-budget training.

sched = optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-3,
                                       total_steps=10_000,
                                       pct_start=0.1)

The canonical PyTorch training step is four lines:

def train_step(model, opt, x, y, loss_fn):
    opt.zero_grad()                  # 1. clear last step's gradients
    logits = model(x)                # 2. forward
    loss = loss_fn(logits, y)
    loss.backward()                  # 3. backward — populates .grad on every param
    opt.step()                       # 4. update parameters in-place
    return loss.item()

Around those four lines wraps the rest: data loading, device transfer, evaluation, logging, checkpointing. Here is the full pattern, batteries included:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

def train(model, train_dl, val_dl, opt, sched, loss_fn, n_epochs, device):
    model.to(device)
    history = {"train_loss": [], "val_loss": [], "val_acc": []}

    for epoch in range(n_epochs):
        model.train()
        running_loss = 0.0
        for x, y in train_dl:
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            logits = model(x)
            loss = loss_fn(logits, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            opt.step()
            if sched is not None:
                sched.step()
            running_loss += loss.item() * x.size(0)

        train_loss = running_loss / len(train_dl.dataset)

        # Eval
        model.eval()
        val_loss, correct = 0.0, 0
        with torch.no_grad():
            for x, y in val_dl:
                x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
                logits = model(x)
                val_loss += loss_fn(logits, y).item() * x.size(0)
                correct += (logits.argmax(1) == y).sum().item()

        val_loss /= len(val_dl.dataset)
        val_acc = correct / len(val_dl.dataset)
        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["val_acc"].append(val_acc)
        print(f"epoch {epoch}: train {train_loss:.4f}  val {val_loss:.4f}  acc {val_acc:.4f}")

    return history

Five details that matter.

set_to_none=True in zero_grad is faster than setting gradients to a zero tensor, because the optimizer can detect None and skip the update for unused parameters. As of PyToranother chapter.x this is the default but it is good to be explicit.

non_blocking=True in .to(device) works in concert with pin_memory=True in the DataLoader. The data transfer happens asynchronously while the previous batch is still computing. Without both, you serialize the transfer and waste GPU time.

clip_grad_norm_(..., max_norm=1.0) clips the gradient to a maximum L2 norm. This prevents the occasional huge gradient from destabilizing training. Standard for transformers; nice-to-have for CNNs. Note the trailing underscore: this is in-place.

model.train and model.eval switch dropout and batchnorm modes. Train uses stochastic behavior; eval uses fixed statistics. Forget to switch back and your eval numbers are wrong; forget to switch to train at the start of an epoch and you train without dropout.

with torch.no_grad: around the eval loop saves a few percent of memory and time. Cheap insurance.

Two paths. The simple version saves only the parameters:

# Save
torch.save(model.state_dict(), "model.pt")

# Load
model = MLP(784, 128, 10)            # must construct with same architecture
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()

state_dict is a Python dict of {layer_name: tensor}. It is the canonical serialization format because it does not depend on your class definitions surviving (whereas torch.save(model, "x.pt") pickles the whole module, which breaks across refactors).

For resuming a training run, you also want the optimizer state, the scheduler state, the epoch count, and the random number generator state:

def save_checkpoint(path, model, opt, sched, epoch, rng_state):
    torch.save({
        "model": model.state_dict(),
        "optimizer": opt.state_dict(),
        "scheduler": sched.state_dict() if sched else None,
        "epoch": epoch,
        "rng_state": torch.get_rng_state(),
        "cuda_rng_state": torch.cuda.get_rng_state() if torch.cuda.is_available() else None,
    }, path)

def load_checkpoint(path, model, opt, sched, device):
    ckpt = torch.load(path, map_location=device)
    model.load_state_dict(ckpt["model"])
    opt.load_state_dict(ckpt["optimizer"])
    if sched and ckpt["scheduler"]:
        sched.load_state_dict(ckpt["scheduler"])
    torch.set_rng_state(ckpt["rng_state"])
    if torch.cuda.is_available() and ckpt["cuda_rng_state"] is not None:
        torch.cuda.set_rng_state(ckpt["cuda_rng_state"])
    return ckpt["epoch"]

Modern GPUs (V100 onwards) compute float16 / bfloat16 arithmetic 2-8× faster than float32. Mixed precision training runs most operations in float16 / bfloat16 while keeping a master copy of the weights in float32 for the optimizer update.

The two APIs:

from torch.amp import autocast, GradScaler

scaler = GradScaler()

for x, y in train_dl:
    x, y = x.to(device), y.to(device)
    opt.zero_grad(set_to_none=True)
    with autocast(device_type="cuda", dtype=torch.float16):
        logits = model(x)
        loss = loss_fn(logits, y)
    scaler.scale(loss).backward()
    scaler.unscale_(opt)                                   # un-scale grads before clip
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(opt)
    scaler.update()

autocast runs operations inside the block in float16 (or bfloat16) where it is safe and faster, falling back to float32 for operations where precision matters (softmax, layer norm). You do not have to manage this; PyTorch knows which is which.

GradScaler exists because float16 has a much smaller dynamic range than float32. Gradients of typical losses underflow to zero in float16. The scaler multiplies the loss by a large constant before backward, then divides the gradients back down before the optimizer step. The scale is adaptive: it grows when no overflows happen, shrinks when they do.

bfloat16 does not need GradScaler because its range matches float32 (only the precision is reduced). If your GPU supports bf16 (A100, H100, RTX 30/40), use it: autocast(device_type="cuda", dtype=torch.bfloat16) and drop the scaler entirely. Llama, Mistral, and most modern LLMs are trained in bf16.

Speedup: 1.5-3× on memory-bound layers, ~2× on transformer attention, less on small models where you are kernel-launch-bound.

torch.compile is PyToranother chapter.x's JIT compiler. One line gets you anywhere from 5% to 3× speedup with no code changes:

model = MLP(784, 128, 10).to(device)
model = torch.compile(model)   # ← one line; rest of the code is unchanged

Under the hood, torch.compile traces the model's forward pass, fuses operations (collapses several elementwise ops into one GPU pass instead of one launch each), generates kernels — the actual GPU programs that run an op, here emitted via Triton — and re-uses them on subsequent calls. The first call is slow (compilation overhead, often 30 seconds or more); subsequent calls are fast.

Three things to know:

The default mode is "default" which trades compile time for runtime. "reduce-overhead" minimizes Python overhead and helps small models. "max-autotune" searches over kernel implementations and is best for large models with stable shapes.

torch.compile recompiles when shapes change. If your batch size is sometimes 128 and sometimes 32, you pay the compile cost twice. Use dynamic=True to compile once with symbolic shapes, at the cost of some speedup.

Some operations break compilation. Custom CUDA kernels, certain Python control flow, and a few exotic ops. The error messages are usually clear; the fallback is to mark just the problematic submodule with @torch.compiler.disable.

You will eventually want to train on multiple GPUs. The two patterns:

DataParallel (legacy, single-process, multi-GPU). One process spawns threads; each thread runs the model on one GPU; gradients are averaged. Easy to set up, slow because of GIL contention and the master-GPU bottleneck. Use only for prototyping.

DistributedDataParallel (DDP, the real one). One process per GPU. Each process runs the full model on its GPU, computes gradients on its batch, and all-reduces gradients across processes after backward — an all-reduce is the collective operation that sums each process's gradient and hands every process the same averaged result, run over NVIDIA's nccl communication backend on GPUs. Scales linearly to ~64 GPUs.

The minimal DDP setup:

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")

    model = MLP(784, 128, 10).to(device)
    model = DDP(model, device_ids=[local_rank])

    # Use DistributedSampler so each rank sees a different slice of data
    sampler = torch.utils.data.DistributedSampler(train_ds)
    train_dl = DataLoader(train_ds, batch_size=128, sampler=sampler, ...)

    for epoch in range(n_epochs):
        sampler.set_epoch(epoch)   # required for shuffling to work
        ...   # rest is the standard training loop

Launch with torchrun --nproc_per_node=8 train.py. The launcher sets LOCAL_RANK, WORLD_SIZE, and the other env vars.

This is the depth another chapter needs. another chapter (deep training) covers FSDP, ZeRO, tensor parallelism, and the rest of the distributed-training landscape. For now, knowing DDP exists and how to use it on a small cluster is enough.