Module 01: Foundations — PyTorch and the Mathematics of Deep Learning

Table of Contents

• What is PyTorch?
• Installation
• Core Philosophy
• Architecture Stack
• Mathematical Prerequisites
• Linear Algebra
• Calculus for Deep Learning
• Probability and Statistics
• Optimization Theory
• Information Theory

---

What is PyTorch?

PyTorch is an open-source machine learning framework developed primarily by Meta AI (formerly Facebook AI Research). It provides two core capabilities:

• N-dimensional tensor computation — similar to NumPy but with GPU acceleration
• Automatic differentiation — computes gradients of arbitrary computational graphs

A Brief History

• 2002: Torch was created in Lua at NYU by Ronan Collobert and others.
• 2016: PyTorch 0.1 was released by Facebook AI Research (FAIR), bringing the

Torch tensor library to Python with automatic differentiation built in.

• 2018: PyTorch 1.0 merged the research-focused PyTorch with the production-focused

Caffe2, adding TorchScript for model export.

• 2022: PyTorch 2.0 introduced torch.compile(), a compiler-based approach that

can dramatically speed up models with a single line of code.

• 2023: PyTorch moved to the Linux Foundation, becoming a truly community-governed project.
• 2024-2025: Continued evolution with FlexAttention, torch.export improvements,

and expanded hardware support (Intel XPU, Apple MPS, AMD ROCm).

PyTorch vs TensorFlow

Why PyTorch won in research: The key reason is debuggability. When your model produces NaN values or wrong outputs, you can insert a breakpoint() anywhere in your forward pass, inspect tensors, and understand what's happening. In graph-based frameworks, you're debugging a compiled representation, not your original code.

The PyTorch Ecosystem

PyTorch is not just one library — it's a constellation:

• torchvision: Computer vision (datasets, models, transforms)
• torchaudio: Audio processing
• torchtext: NLP utilities
• PyTorch Lightning / Fabric: Training framework that reduces boilerplate
• HuggingFace Transformers: Built on PyTorch, the dominant NLP/LLM library
• TorchServe: Model serving for production
• ONNX Runtime: Export PyTorch models to a portable format
• torch.export / ExecuTorch: On-device deployment (mobile, embedded)

---

Installation

CPU-Only Installation (Recommended for Learning)

# Using pip (simplest)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

# Using conda
conda install pytorch torchvision torchaudio cpuonly -c pytorch

With CUDA (for NVIDIA GPUs)

# CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# CUDA 12.1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# CUDA 12.4
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

Verify Installation

import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Number of GPUs: {torch.cuda.device_count()}")

# Quick test
x = torch.rand(3, 3)
print(f"Random tensor:\n{x}")

---

Core Philosophy

1. Eager Execution (Define-by-Run)

In PyTorch, operations execute immediately as Python runs them. There is no separate "compilation" step before you can see results:

import torch

x = torch.tensor([1.0, 2.0, 3.0])
y = x * 2          # This executes RIGHT NOW — not later in a session
print(y)            # tensor([2., 4., 6.])

Why this matters: You can use standard Python control flow (if, for, while) inside your models. The computation graph is built dynamically as your code runs, which means different inputs can follow different code paths.

2. Dynamic Computation Graphs

Unlike static graph frameworks, PyTorch rebuilds the computation graph every time you run a forward pass. This means:

def dynamic_model(x):
    if x.sum() > 0:         # This condition is evaluated at runtime
        return x * 2
    else:
        return x * 3

Each call to dynamic_model can follow a different path through the code. The autograd graph is constructed on the fly and torn down after .backward().

3. Python-First Design

PyTorch is designed to feel like an extension of Python and NumPy, not a separate language embedded in Python. You think in Python, you debug in Python, you profile in Python. The C++ backend handles performance; you rarely need to think about it.

4. torch.compile — The Best of Both Worlds

Starting with PyTorch 2.0, you can optionally compile your models for speed while keeping the eager-mode development experience:

model = MyModel()
compiled_model = torch.compile(model)  # One line for significant speedup
output = compiled_model(input_data)

This uses TorchDynamo to trace your Python code, TorchInductor to generate optimized kernels, and falls back to eager mode for unsupported patterns.

---

Architecture Stack

Understanding PyTorch's layered architecture helps you know where to look when debugging or optimizing. From top to bottom:

┌──────────────────────────────────────────────────┐
│  Python Frontend (torch, torch.nn, etc.)         │  ← You write code here
├──────────────────────────────────────────────────┤
│  torch.compile (TorchDynamo + TorchInductor)     │  ← Optional compilation
├──────────────────────────────────────────────────┤
│  Autograd Engine                                 │  ← Automatic differentiation
├──────────────────────────────────────────────────┤
│  ATen (A Tensor Library)                         │  ← Core tensor operations
├──────────────────────────────────────────────────┤
│  C10 (Caffe2 + PyTorch Core)                     │  ← Dispatcher, memory, dtypes
├──────────────────────────────────────────────────┤
│  Hardware Backends (CPU/CUDA/MPS/XPU/ROCm)       │  ← Actual computation
└──────────────────────────────────────────────────┘

Python Frontend: The torch module, torch.nn, torch.optim, etc. — the high-level API you interact with daily.

torch.compile: TorchDynamo captures your Python code as an FX graph. TorchInductor generates optimized C++/CUDA/Triton kernels from that graph.

Autograd: The engine that tracks operations on tensors with requires_grad=True and computes gradients via reverse-mode automatic differentiation.

ATen: "A Tensor Library" — over 2,000 operators written in C++ that implement the actual math (add, matmul, conv2d, etc.). When you call torch.add(a, b), this is where the computation happens.

C10: The core library providing the dispatcher (routes operations to the right backend), memory allocators, dtype system, and device abstraction.

Hardware Backends: BLAS libraries (MKL for CPU, cuBLAS for GPU), cuDNN for convolutions, and other hardware-specific optimized libraries.

---

Mathematical Prerequisites

Deep learning sits at the intersection of linear algebra, calculus, probability, and optimization. You don't need a PhD in mathematics, but you need working familiarity with these concepts. This section teaches them through PyTorch code.

---

Linear Algebra

Linear algebra is the language of deep learning. Neural networks are fundamentally sequences of matrix multiplications interspersed with non-linear functions.

Scalars, Vectors, Matrices, and Tensors

import torch

scalar = torch.tensor(3.14)             # 0-D tensor (scalar)
vector = torch.tensor([1.0, 2.0, 3.0])  # 1-D tensor (vector)
matrix = torch.tensor([[1, 2], [3, 4]])  # 2-D tensor (matrix)
tensor_3d = torch.randn(2, 3, 4)        # 3-D tensor

print(f"Scalar shape: {scalar.shape}")   # torch.Size([])
print(f"Vector shape: {vector.shape}")   # torch.Size([3])
print(f"Matrix shape: {matrix.shape}")   # torch.Size([2, 2])
print(f"3D shape: {tensor_3d.shape}")    # torch.Size([2, 3, 4])

Why tensors? In deep learning, data naturally has multiple dimensions. An image is a 3D tensor (channels × height × width). A batch of images is 4D (batch × channels × height × width). A batch of sequences of word embeddings is 3D (batch × sequence_length × embedding_dim).

Norms — Measuring Vector Magnitude

A norm measures the "size" of a vector. Different norms emphasize different properties:

• L1 norm (Manhattan): Sum of absolute values. Encourages sparsity in optimization.

\( \|x\|_1 = \sum_i |x_i| \)

• L2 norm (Euclidean): Square root of sum of squares. The "ordinary" distance.

\( \|x\|_2 = \sqrt{\sum_i x_i^2} \)

• L∞ norm (Max): Largest absolute value. Used in adversarial robustness.

\( \|x\|_\infty = \max_i |x_i| \)

v = torch.tensor([3.0, -4.0])
print(f"L1 norm: {torch.norm(v, p=1)}")      # 7.0
print(f"L2 norm: {torch.norm(v, p=2)}")      # 5.0
print(f"L∞ norm: {torch.norm(v, p=float('inf'))}")  # 4.0

Why norms matter in deep learning: L2 regularization (weight decay) penalizes large L2 norms of weight vectors, preventing overfitting. L1 regularization drives weights to exactly zero, performing feature selection. Gradient clipping uses norms to prevent exploding gradients.

Dot Product — Measuring Similarity

The dot product of two vectors measures their alignment:

\( a \cdot b = \sum_i a_i b_i = \|a\| \|b\| \cos\theta \)

a = torch.tensor([1.0, 0.0])
b = torch.tensor([0.0, 1.0])
c = torch.tensor([1.0, 1.0])

print(f"a·b (perpendicular): {torch.dot(a, b)}")    # 0.0
print(f"a·c (45 degrees): {torch.dot(a, c)}")       # 1.0
print(f"a·a (parallel): {torch.dot(a, a)}")          # 1.0

Why dot products matter: Attention mechanisms in Transformers compute dot products between query and key vectors to measure relevance. The output of a linear layer y = Wx + b is a batch of dot products between weight rows and the input vector.

Matrix Multiplication — The Core Operation

Matrix multiplication is the single most important operation in deep learning. Every linear layer, every attention head, every convolution can be expressed as matrix multiplications.

For matrices A (m×n) and B (n×p), the result C = AB is (m×p):

\( C_{ij} = \sum_k A_{ik} B_{kj} \)

A = torch.tensor([[1., 2.], [3., 4.]])  # 2×2
B = torch.tensor([[5., 6.], [7., 8.]])  # 2×2

C = A @ B  # or torch.matmul(A, B) or torch.mm(A, B)
print(f"A @ B =\n{C}")
# tensor([[19., 22.],
#         [43., 50.]])

The n dimensions are "consumed" by the multiplication.

Eigendecomposition

A square matrix A can be decomposed as A = QΛQ⁻¹, where Q contains eigenvectors and Λ is a diagonal matrix of eigenvalues. An eigenvector v satisfies Av = λv — the matrix only scales it, doesn't change its direction.

A = torch.tensor([[2., 1.], [1., 2.]], dtype=torch.float)
eigenvalues, eigenvectors = torch.linalg.eig(A)
print(f"Eigenvalues: {eigenvalues}")
print(f"Eigenvectors:\n{eigenvectors}")

Why eigendecomposition matters: Principal Component Analysis (PCA) uses eigendecomposition to find the directions of maximum variance in data. The condition number (ratio of largest to smallest eigenvalue) tells you how numerically stable a problem is.

Singular Value Decomposition (SVD)

SVD generalizes eigendecomposition to non-square matrices: A = UΣVᵀ.

• U: left singular vectors (column space basis)
• Σ: singular values (diagonal, non-negative, sorted)
• Vᵀ: right singular vectors (row space basis)

A = torch.tensor([[1., 2., 3.], [4., 5., 6.]], dtype=torch.float)
U, S, Vh = torch.linalg.svd(A)
print(f"U shape: {U.shape}, S shape: {S.shape}, Vh shape: {Vh.shape}")

Why SVD matters: Low-rank approximation via SVD is used in LoRA (Low-Rank Adaptation) for efficient fine-tuning of large language models. It's also the mathematical foundation of matrix factorization in recommender systems.

---

Calculus for Deep Learning

Derivatives and the Chain Rule

A derivative measures how a function's output changes when its input changes:

\( f'(x) = \lim_{h \to 0} \frac{f(x+h) - f(x)}{h} \)

The chain rule is the single most important calculus concept for deep learning. If y = f(g(x)), then:

\( \frac{dy}{dx} = \frac{dy}{dg} \cdot \frac{dg}{dx} \)

Neural networks are compositions of functions: output = f_n(f_{n-1}(...f_1(x)...)). Backpropagation is just the chain rule applied repeatedly.

x = torch.tensor(2.0, requires_grad=True)
y = x**3 + 2*x**2 + x  # y = x³ + 2x² + x
y.backward()             # dy/dx = 3x² + 4x + 1
print(f"dy/dx at x=2: {x.grad}")  # 3(4) + 4(2) + 1 = 21

Gradients — Multivariable Derivatives

When a function has multiple inputs, the gradient is the vector of all partial derivatives:

\( \nabla f = \left[\frac{\partial f}{\partial x_1}, \frac{\partial f}{\partial x_2}, \ldots\right] \)

The gradient points in the direction of steepest ascent. To minimize a loss, we move in the opposite direction: gradient descent.

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
f = (x**2).sum()  # f = x₁² + x₂² + x₃²
f.backward()
print(f"Gradient: {x.grad}")  # [2, 4, 6] — the gradient ∇f = 2x

Jacobians and Hessians

The Jacobian generalizes the gradient to vector-valued functions. If f: ℝⁿ → ℝᵐ, the Jacobian J is an m×n matrix where J_ij = ∂f_i/∂x_j.

The Hessian is the matrix of second derivatives: H_ij = ∂²f/∂x_i∂x_j. It tells you about the curvature of the loss surface — whether you're at a minimum, maximum, or saddle point.

from torch.autograd.functional import jacobian, hessian

def f(x):
    return torch.stack([x[0]**2 + x[1], x[0] * x[1]**2])

x = torch.tensor([1.0, 2.0])
J = jacobian(f, x)
print(f"Jacobian:\n{J}")
# [[2*x0, 1   ],    = [[2, 1],
#  [x1^2, 2*x0*x1]]    [4, 4]]

---

Probability and Statistics

Probability Distributions

PyTorch's torch.distributions module provides a rich set of probability distributions. These are essential for:

• Initializing weights (normal, uniform)
• Variational autoencoders (reparameterization trick)
• Reinforcement learning (sampling actions)
• Bayesian neural networks

from torch.distributions import Normal, Bernoulli, Categorical

normal = Normal(loc=0.0, scale=1.0)  # mean=0, std=1
sample = normal.sample((5,))
log_prob = normal.log_prob(torch.tensor(0.0))
print(f"Samples: {sample}")
print(f"Log prob of 0: {log_prob}")  # log(1/√(2π)) ≈ -0.9189

Expectation and Variance

• Expectation E[X]: The average value of a random variable.
• Variance Var(X) = E[(X - E[X])²]: How spread out the values are.

samples = torch.randn(100000)  # Standard normal
print(f"Mean ≈ {samples.mean():.4f}")  # ≈ 0
print(f"Var ≈ {samples.var():.4f}")    # ≈ 1

Cross-Entropy — The Loss Function of Classification

Cross-entropy measures the difference between two probability distributions p (true) and q (predicted):

\( H(p, q) = -\sum_i p_i \log(q_i) \)

When p is a one-hot vector (classification), this simplifies to: \( H(p, q) = -\log(q_{\text{true class}}) \)

This is why we use log-softmax + negative log likelihood, which PyTorch combines into nn.CrossEntropyLoss:

import torch.nn as nn

logits = torch.tensor([[2.0, 1.0, 0.1]])  # Raw model output
target = torch.tensor([0])                  # True class is 0

loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, target)
print(f"Cross-entropy loss: {loss.item():.4f}")

KL Divergence

KL divergence measures how one probability distribution differs from a reference:

\( D_{KL}(P \| Q) = \sum_i P(i) \log\frac{P(i)}{Q(i)} \)

It's asymmetric: D_KL(P||Q) ≠ D_KL(Q||P). Used in VAEs (variational autoencoders) to keep the learned latent distribution close to a prior (usually standard normal).

import torch.nn.functional as F

p = torch.tensor([0.4, 0.3, 0.3])  # True distribution
q = torch.tensor([0.33, 0.33, 0.34])  # Predicted distribution

kl = F.kl_div(q.log(), p, reduction='sum')
print(f"KL divergence: {kl.item():.4f}")

---

Optimization Theory

Gradient Descent — The Foundation

Gradient descent minimizes a function by repeatedly taking steps opposite to the gradient. The update rule:

\( \theta_{t+1} = \theta_t - \alpha \nabla L(\theta_t) \)

where α is the learning rate — the most important hyperparameter in deep learning.

• Too large: Overshoots the minimum, loss oscillates or diverges
• Too small: Converges extremely slowly, may get stuck in local minima
• Just right: Smooth convergence to a good minimum

# Minimizing f(x) = (x - 3)² from scratch
x = torch.tensor(0.0, requires_grad=True)
lr = 0.1

for step in range(50):
    loss = (x - 3) ** 2
    loss.backward()
    with torch.no_grad():
        x -= lr * x.grad
    x.grad.zero_()

print(f"Final x: {x.item():.6f}")  # ≈ 3.0

Stochastic Gradient Descent (SGD)

In practice, computing the gradient over the entire dataset is expensive. SGD estimates the gradient using a random mini-batch:

\( \theta_{t+1} = \theta_t - \alpha \nabla L_{\text{batch}}(\theta_t) \)

The noise from mini-batch sampling actually helps — it can escape local minima and leads to solutions that generalize better.

SGD with Momentum

Plain SGD can oscillate in narrow valleys. Momentum adds a "velocity" term that accumulates past gradients, smoothing the trajectory:

\( v_t = \beta v_{t-1} + \nabla L(\theta_t) \) \( \theta_{t+1} = \theta_t - \alpha v_t \)

β (typically 0.9) controls how much history to keep. Think of it like a ball rolling downhill — it builds up speed in consistent directions.

Adam — Adaptive Moment Estimation

Adam combines momentum with per-parameter adaptive learning rates. It maintains two running averages:

• m (first moment): exponential moving average of gradients (like momentum)
• v (second moment): exponential moving average of squared gradients

\( m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t \) \( v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \) \( \hat{m}_t = m_t / (1 - \beta_1^t) \) (bias correction) \( \hat{v}_t = v_t / (1 - \beta_2^t) \) (bias correction) \( \theta_{t+1} = \theta_t - \alpha \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon) \)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))

Why Adam works well: Parameters that receive sparse, infrequent gradients get larger effective learning rates (because v is small). Parameters with frequent, large gradients get smaller effective learning rates (because v is large). This adaptive behavior is crucial for training models where different parameters need different learning rates.

When to use what:

• SGD + Momentum: Often achieves better final performance but requires careful

learning rate tuning and scheduling. Preferred for vision models (ResNet, etc.).

• Adam/AdamW: Faster convergence, less sensitive to learning rate. Preferred

for Transformers, LLMs, and when you want quick results.

• AdamW: Adam with decoupled weight decay — almost always preferred over plain Adam.

---

Information Theory

Entropy — Measuring Uncertainty

Entropy quantifies the uncertainty in a probability distribution:

\( H(p) = -\sum_i p_i \log p_i \)

• Maximum entropy: uniform distribution (maximum uncertainty)
• Minimum entropy (0): all probability mass on one outcome (certainty)

uniform = torch.tensor([0.25, 0.25, 0.25, 0.25])
certain = torch.tensor([1.0, 0.0, 0.0, 0.0])

H_uniform = -(uniform * uniform.log()).sum()
H_certain = -(certain * (certain + 1e-8).log()).sum()

print(f"Entropy of uniform: {H_uniform:.4f}")  # 1.3863 (= ln(4))
print(f"Entropy of certain: {H_certain:.4f}")   # ≈ 0.0

Cross-Entropy Loss Derivation

Why do we use cross-entropy as a loss function for classification?

• Maximum likelihood: We want to find model parameters θ that maximize the

probability of the observed data: argmax_θ P(data|θ).

• Log transformation: Taking the log converts the product of probabilities

into a sum (numerically stable, easier to optimize): argmax_θ Σ log P(y_i | x_i, θ)

• Negation: Maximizing log-likelihood = minimizing negative log-likelihood:

argmin_θ -Σ log P(y_i | x_i, θ)

• Softmax output: For classification, P(y=k|x) = softmax(logits)_k, so:

Loss = -log(softmax(logits)_{true_class})

This is exactly cross-entropy between the one-hot true distribution and the softmax predicted distribution. The connection is not coincidental — cross-entropy loss is the information-theoretically optimal loss for classification.

logits = torch.tensor([2.0, 1.0, 0.1])
probs = torch.softmax(logits, dim=0)
true_class = 0

nll_loss = -torch.log(probs[true_class])
print(f"NLL loss: {nll_loss.item():.4f}")

ce_loss = nn.CrossEntropyLoss()(logits.unsqueeze(0), torch.tensor([true_class]))
print(f"CE loss:  {ce_loss.item():.4f}")  # Same value

---

What's Next?

With these mathematical foundations, you're ready to dive into PyTorch's tensor system in Module 02. The linear algebra you learned here will appear every time you work with layers (matrix multiplication), losses (norms, cross-entropy), and optimization (gradients, Adam). The key insight is that deep learning is applied linear algebra + calculus, automated by PyTorch's autograd system.

Run math_with_pytorch.py in this directory to see all these concepts in action with runnable code.

---

📓 Open Notebook — Interactive version of this module

Module 02: Tensors — The Complete Guide

Table of Contents

• What is a Tensor?
• Tensor Creation
• Data Types
• Device Management
• Tensor Properties
• Element-wise Operations
• Reduction Operations
• Matrix Operations
• Tensor Manipulation
• Indexing Deep Dive
• Broadcasting
• Views vs Copies
• In-place Operations
• Strides Explained
• NumPy Interop

---

What is a Tensor?

A tensor is a multi-dimensional array — the fundamental data structure of PyTorch. If you know NumPy, a tensor is like np.ndarray but with two superpowers: GPU acceleration and automatic differentiation.

The word "tensor" comes from mathematics, where it refers to objects that obey certain transformation rules. In deep learning, we use the word more loosely to mean "an n-dimensional array of numbers."

Dimensionality taxonomy:

import torch

scalar = torch.tensor(3.14)                  # 0-D
vector = torch.tensor([1, 2, 3])             # 1-D
matrix = torch.tensor([[1, 2], [3, 4]])      # 2-D
cube = torch.randn(2, 3, 4)                  # 3-D

print(scalar.ndim, vector.ndim, matrix.ndim, cube.ndim)
# 0, 1, 2, 3

Why "tensor" instead of "array"? The name emphasizes that these objects carry metadata (dtype, device, gradient tracking) and participate in PyTorch's autograd system. A NumPy array can't run on a GPU or compute gradients.

---

Tensor Creation

PyTorch offers many ways to create tensors. Each serves a specific purpose:

From Python data

# From a list
t = torch.tensor([1, 2, 3])          # int64 by default
t = torch.tensor([1.0, 2.0, 3.0])    # float32 by default

# From nested lists (matrix)
t = torch.tensor([[1, 2], [3, 4]])

# With explicit dtype
t = torch.tensor([1, 2, 3], dtype=torch.float32)

Constant-fill tensors

# Zeros and ones
torch.zeros(3, 4)         # 3x4 matrix of zeros
torch.ones(2, 3, 5)       # 2x3x5 tensor of ones
torch.full((2, 3), 7.0)   # 2x3 matrix filled with 7.0

# Identity matrix
torch.eye(4)               # 4x4 identity matrix

Why these exist: torch.zeros is used to initialize bias terms, accumulation buffers, and masks. torch.ones creates multiplicative identities. torch.eye creates identity matrices for residual connections and initialization schemes.

Random tensors

# Uniform [0, 1)
torch.rand(3, 4)

# Standard normal (mean=0, std=1)
torch.randn(3, 4)

# Random integers
torch.randint(low=0, high=10, size=(3, 4))

# Random permutation
torch.randperm(10)  # A random shuffling of [0, 1, ..., 9]

Why random tensors matter: Weight initialization is crucial for training. torch.randn is the basis of Gaussian initialization. torch.randperm is used for shuffling dataset indices. Setting torch.manual_seed(42) makes random operations reproducible.

Sequences

# Integer sequence [0, 1, 2, ..., 9]
torch.arange(10)
torch.arange(2, 10)       # [2, 3, ..., 9]
torch.arange(0, 1, 0.1)   # [0.0, 0.1, ..., 0.9]

# Evenly spaced (specify count, not step)
torch.linspace(0, 1, steps=5)    # [0.0, 0.25, 0.5, 0.75, 1.0]
torch.logspace(0, 2, steps=3)    # [10^0, 10^1, 10^2] = [1, 10, 100]

Why these exist: torch.arange creates position indices. torch.linspace is used for evaluating functions over a range (plotting, interpolation). torch.logspace creates logarithmic learning rate schedules.

Uninitialized tensors

# Uninitialized — contains whatever was in memory!
torch.empty(3, 4)

Warning: torch.empty does NOT fill with zeros. It allocates memory and returns whatever garbage was there. Use it only when you'll immediately overwrite all values, as it avoids the cost of zero-filling. This matters in performance-critical code.

Like-functions (matching shape/dtype/device of an existing tensor)

x = torch.randn(3, 4, dtype=torch.float32)
torch.zeros_like(x)      # Same shape, dtype, device — filled with zeros
torch.ones_like(x)       # Same but with ones
torch.randn_like(x)      # Same but with random normal values
torch.empty_like(x)      # Same but uninitialized
torch.full_like(x, 5.0)  # Same but filled with 5.0

Why these exist: When writing generic code (custom layers, loss functions), you often need to create a tensor with the same properties as an input. These functions handle dtype and device automatically, avoiding common bugs.

---

Data Types

Every tensor has a dtype (data type). Choosing the right one affects memory, speed, and numerical precision:

float32 vs float16 vs bfloat16

This choice is one of the most important practical decisions in training:

• float32: Full precision. Use for the optimizer state and any numerically

sensitive operations (loss computation, normalization).

• float16: Half precision. 2x memory savings, faster on GPUs with tensor cores.

BUT has a limited range (max ~65504), which means gradients can overflow. Requires loss scaling for stable training.

• bfloat16: Same memory as float16 but with the same exponent range as float32.

This means it almost never overflows. Preferred for LLM training because you get memory savings without the numerical headaches of float16.

x = torch.randn(1000, 1000)
print(f"float32: {x.element_size()} bytes per element, {x.nelement() * x.element_size() / 1e6:.1f} MB")

x16 = x.to(torch.float16)
print(f"float16: {x16.element_size()} bytes per element, {x16.nelement() * x16.element_size() / 1e6:.1f} MB")

xbf = x.to(torch.bfloat16)
print(f"bfloat16: {xbf.element_size()} bytes per element")

Casting

x = torch.tensor([1, 2, 3])          # int64
x_float = x.float()                   # → float32
x_half = x.half()                     # → float16
x_double = x.double()                 # → float64
x_int = x_float.int()                 # → int32
x_long = x_float.long()               # → int64
x_bool = x.bool()                     # → bool (0=False, nonzero=True)
x_cast = x.to(torch.float32)          # General casting

---

Device Management

Tensors live on a specific device. Operations between tensors require them to be on the same device — this is a common source of runtime errors.

# CPU (default)
x_cpu = torch.randn(3, 4)
print(x_cpu.device)  # cpu

# Check for GPU availability
if torch.cuda.is_available():
    x_gpu = torch.randn(3, 4, device='cuda')
    x_gpu = x_cpu.to('cuda')      # Move to GPU
    x_gpu = x_cpu.cuda()          # Shorthand
    x_back = x_gpu.cpu()          # Move back to CPU

# Apple Silicon
if torch.backends.mps.is_available():
    x_mps = x_cpu.to('mps')

# Context manager for default device
with torch.device('cpu'):
    x = torch.randn(3, 4)  # Created on CPU

Common error: RuntimeError: Expected all tensors to be on the same device. This happens when you mix CPU and GPU tensors in an operation. Always check .device when debugging.

---

Tensor Properties

Every tensor carries metadata that you can inspect:

x = torch.randn(2, 3, 4, dtype=torch.float32)

print(f"Shape: {x.shape}")              # torch.Size([2, 3, 4])
print(f"Size (same): {x.size()}")       # torch.Size([2, 3, 4])
print(f"Dimensions: {x.ndim}")          # 3
print(f"Data type: {x.dtype}")          # torch.float32
print(f"Device: {x.device}")            # cpu
print(f"Total elements: {x.numel()}")   # 24 (= 2*3*4)
print(f"Element size: {x.element_size()} bytes")  # 4 (float32)
print(f"Total memory: {x.nelement() * x.element_size()} bytes")  # 96
print(f"Strides: {x.stride()}")         # (12, 4, 1)
print(f"Is contiguous: {x.is_contiguous()}")  # True
print(f"Requires grad: {x.requires_grad}")    # False

Understanding shape vs size: They're identical. x.shape is a property, x.size() is a method. Most people use .shape (following NumPy convention). You can index into shape: x.shape[0] gives the first dimension.

---

Element-wise Operations

Element-wise operations apply a function independently to each element. The shapes of input tensors must be compatible (same shape or broadcastable).

Arithmetic

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])

a + b        # tensor([5., 7., 9.])     — or torch.add(a, b)
a - b        # tensor([-3., -3., -3.])   — or torch.sub(a, b)
a * b        # tensor([4., 10., 18.])    — or torch.mul(a, b)
a / b        # tensor([0.25, 0.4, 0.5])  — or torch.div(a, b)
a ** 2       # tensor([1., 4., 9.])      — or torch.pow(a, 2)
a // b       # Floor division
a % b        # Modulo

Mathematical functions

x = torch.tensor([0.0, 1.0, 2.0])

torch.exp(x)      # e^x:  [1.0, 2.718, 7.389]
torch.log(x + 1)  # ln(x+1) — add 1 to avoid log(0)
torch.sqrt(x)     # [0.0, 1.0, 1.414]
torch.abs(x - 1)  # |x-1|: [1.0, 0.0, 1.0]
torch.sin(x)      # Sine
torch.cos(x)      # Cosine
torch.tanh(x)     # Hyperbolic tangent (activation function)
torch.sigmoid(x)  # 1 / (1 + e^(-x)) (activation function)

Clamping

x = torch.tensor([-3.0, -1.0, 0.5, 2.0, 5.0])
torch.clamp(x, min=0.0)             # ReLU! [0, 0, 0.5, 2, 5]
torch.clamp(x, min=-1.0, max=1.0)   # Clip to [-1, 1]

Why clamping matters: torch.clamp(x, min=0) is literally the ReLU activation function. Gradient clipping uses torch.clamp on gradient norms. Value clipping prevents numerical instability.

---

Reduction Operations

Reductions collapse one or more dimensions by aggregating values. Understanding the dim parameter is critical.

The dim parameter

When you specify dim=k, that dimension is "collapsed" (removed from the output):

x = torch.tensor([[1., 2., 3.],
                   [4., 5., 6.]])  # Shape: (2, 3)

x.sum()            # 21.0 — all elements (scalar output)
x.sum(dim=0)       # [5., 7., 9.] — collapse rows → shape (3,)
x.sum(dim=1)       # [6., 15.] — collapse columns → shape (2,)
x.sum(dim=1, keepdim=True)  # [[6.], [15.]] — shape (2, 1) — keeps the dim

Mental model for dim: Think "I'm reducing ALONG this axis." dim=0 means "go down the rows" (collapse them). dim=1 means "go across the columns." keepdim=True keeps the reduced dimension as size 1, which is essential for broadcasting the result back.

Common reductions

x = torch.tensor([[1., 2., 3.],
                   [4., 5., 6.]])

x.mean(dim=1)           # [2., 5.] — row means
x.std(dim=1)            # Standard deviation per row
x.var(dim=1)            # Variance per row
x.max(dim=1)            # Returns (values, indices) — both the max and where
x.min(dim=0)            # Returns (values, indices) along dim 0
x.argmax(dim=1)         # Index of max in each row: [2, 2]
x.argmin(dim=0)         # Index of min in each column: [0, 0, 0]
x.prod(dim=1)           # Product: [6., 120.]
torch.norm(x, dim=1)    # L2 norm per row

The max/min return type

max and min return a named tuple with .values and .indices:

vals, idxs = x.max(dim=1)
print(f"Max values: {vals}")    # [3., 6.]
print(f"Max indices: {idxs}")   # [2, 2]

---

Matrix Operations

Matrix multiplication varieties

# 2D @ 2D: standard matrix multiply
A = torch.randn(3, 4)
B = torch.randn(4, 5)
C = A @ B                # Shape: (3, 5)
C = torch.mm(A, B)       # Equivalent, only for 2D
C = torch.matmul(A, B)   # Most general

# Batched matrix multiply: each batch is independent
A = torch.randn(10, 3, 4)   # 10 matrices of shape 3x4
B = torch.randn(10, 4, 5)   # 10 matrices of shape 4x5
C = A @ B                    # Shape: (10, 3, 5) — 10 independent matmuls
C = torch.bmm(A, B)          # Equivalent, only for 3D

# Vector-matrix products
v = torch.randn(4)
M = torch.randn(4, 5)
result = v @ M               # Shape: (5,) — vector × matrix
result = M.T @ v              # Shape: (5,) — equivalent

Einstein summation (einsum)

einsum is the Swiss army knife of tensor operations. It uses subscript notation to specify arbitrary contractions:

# Matrix multiply: "ik,kj->ij" means sum over k
A = torch.randn(3, 4)
B = torch.randn(4, 5)
C = torch.einsum('ik,kj->ij', A, B)  # Same as A @ B

# Batch matrix multiply
A = torch.randn(10, 3, 4)
B = torch.randn(10, 4, 5)
C = torch.einsum('bik,bkj->bij', A, B)  # Same as torch.bmm

# Dot product
a = torch.randn(5)
b = torch.randn(5)
d = torch.einsum('i,i->', a, b)  # Same as torch.dot

# Outer product
outer = torch.einsum('i,j->ij', a, b)  # Shape: (5, 5)

# Trace
M = torch.randn(4, 4)
tr = torch.einsum('ii->', M)  # Same as torch.trace(M)

# Batch diagonal
B_diag = torch.einsum('bii->bi', torch.randn(3, 4, 4))  # Diagonal of each batch

Why einsum? It makes complex tensor operations readable. Instead of chains of transpose, reshape, and matmul, one einsum string describes the operation declaratively. Attention mechanisms are often clearest in einsum notation.

---

Tensor Manipulation

Reshaping

x = torch.arange(12)  # [0, 1, 2, ..., 11]

# view: returns a view (shares memory) — requires contiguous input
x.view(3, 4)       # Shape: (3, 4)
x.view(2, 2, 3)    # Shape: (2, 2, 3)
x.view(-1, 4)      # -1 is inferred: (3, 4)
x.view(-1)         # Flatten: (12,)

# reshape: like view but works on non-contiguous tensors (may copy)
x.reshape(3, 4)
x.reshape(-1, 6)   # (2, 6)

# flatten: collapse dimensions
y = torch.randn(2, 3, 4)
y.flatten()              # Shape: (24,) — all dims
y.flatten(1)             # Shape: (2, 12) — flatten from dim 1 onward
y.flatten(start_dim=1, end_dim=2)  # Shape: (2, 12)

view vs reshape: view always returns a view (no memory copy), but requires the tensor to be contiguous in memory. reshape returns a view if possible, but will copy data if necessary. Rule of thumb: use reshape unless you specifically need to guarantee no copy (then use view and handle the contiguity yourself).

Transposing and permuting

x = torch.randn(2, 3, 4)

# transpose: swap exactly two dimensions
x.transpose(0, 1)   # Shape: (3, 2, 4)
x.transpose(1, 2)   # Shape: (2, 4, 3)

# For 2D matrices, .T is shorthand
m = torch.randn(3, 4)
m.T                  # Shape: (4, 3)

# permute: reorder ALL dimensions at once
x.permute(2, 0, 1)  # Shape: (4, 2, 3) — moved dim 2 to front

Common use case: Images come as (batch, H, W, C) from some libraries but PyTorch expects (batch, C, H, W). Fix with: img.permute(0, 3, 1, 2).

Squeezing and unsqueezing

x = torch.randn(1, 3, 1, 4)

x.squeeze()       # Remove ALL size-1 dims → shape (3, 4)
x.squeeze(0)      # Remove dim 0 if size 1 → shape (3, 1, 4)
x.squeeze(2)      # Remove dim 2 if size 1 → shape (1, 3, 4)

y = torch.randn(3, 4)
y.unsqueeze(0)    # Add dim at position 0 → shape (1, 3, 4)
y.unsqueeze(1)    # Add dim at position 1 → shape (3, 1, 4)
y.unsqueeze(-1)   # Add dim at end → shape (3, 4, 1)

Why these matter: Many PyTorch operations expect specific numbers of dimensions. A single image (3, H, W) needs unsqueeze(0) to become a batch of 1 (1, 3, H, W) for a model. After processing, squeeze(0) removes the batch dimension.

Concatenation and stacking

a = torch.randn(2, 3)
b = torch.randn(2, 3)

# cat: join along EXISTING dimension
torch.cat([a, b], dim=0)   # Shape: (4, 3) — stack vertically
torch.cat([a, b], dim=1)   # Shape: (2, 6) — join horizontally

# stack: join along NEW dimension
torch.stack([a, b], dim=0)  # Shape: (2, 2, 3) — new dim 0 indexes a vs b
torch.stack([a, b], dim=1)  # Shape: (2, 2, 3) — new dim inserted at 1

Key difference: cat glues tensors along an existing dimension (dimensions must match elsewhere). stack creates a new dimension and places tensors along it (all tensors must have exactly the same shape).

Splitting and chunking

x = torch.arange(12).reshape(4, 3)

# split: split into pieces of given size
pieces = torch.split(x, 2, dim=0)  # Two pieces of size 2 each
# pieces[0] shape: (2, 3), pieces[1] shape: (2, 3)

# chunk: split into N roughly equal pieces
chunks = torch.chunk(x, 3, dim=0)  # Three chunks (sizes 2, 1, 1)

Expanding and repeating

x = torch.tensor([[1], [2], [3]])  # Shape: (3, 1)

# expand: view-based (no memory copy), only expands size-1 dims
x.expand(3, 4)    # Shape: (3, 4) — [1,1,1,1], [2,2,2,2], [3,3,3,3]
x.expand(-1, 4)   # -1 means "keep this dim's size"

# repeat: actually copies data
x.repeat(1, 4)    # Shape: (3, 4) — same result but data is copied
x.repeat(2, 3)    # Shape: (6, 3) — repeat 2x along dim 0, 3x along dim 1

expand vs repeat: expand is free (just changes strides, no memory allocation). repeat copies data. Always prefer expand when possible. But be careful: modifying an expanded tensor affects all "copies" since they share memory.

---

Indexing Deep Dive

Basic indexing (returns views)

x = torch.arange(20).reshape(4, 5)
# tensor([[ 0,  1,  2,  3,  4],
#         [ 5,  6,  7,  8,  9],
#         [10, 11, 12, 13, 14],
#         [15, 16, 17, 18, 19]])

x[0]        # Row 0: [0, 1, 2, 3, 4]
x[0, 2]     # Element at row 0, col 2: 2
x[-1]       # Last row: [15, 16, 17, 18, 19]
x[-1, -1]   # Last element: 19

Slicing (returns views)

x[1:3]        # Rows 1 and 2
x[:, 2:4]     # All rows, columns 2 and 3
x[::2]        # Every other row (stride 2)
x[:, ::-1]    # Reverse columns (or use torch.flip)
x[1:3, 2:5]   # Submatrix: rows 1-2, cols 2-4

Boolean (mask) indexing (returns copies)

x = torch.randn(4, 4)
mask = x > 0
positives = x[mask]       # 1-D tensor of all positive values
x[mask] = 0               # Set all positive values to 0
x[x < -1] = -1            # Clamp from below

# Common pattern: conditional replacement
scores = torch.randn(5)
scores[scores < 0] = 0    # ReLU by hand

Why boolean indexing returns copies: The selected elements aren't contiguous in memory, so they can't form a view. Any modification to the result won't affect the original tensor (but assigning back with x[mask] = val does work).

Fancy (advanced) indexing

x = torch.arange(20).reshape(4, 5)
rows = torch.tensor([0, 2, 3])
cols = torch.tensor([1, 3, 4])

x[rows]              # Select rows 0, 2, 3 → shape (3, 5)
x[rows, cols]         # Select (0,1), (2,3), (3,4) → [1, 13, 19]
x[:, [0, 2, 4]]      # Select columns 0, 2, 4 → shape (4, 3)

gather and scatter

gather collects values from a tensor according to an index tensor. Think of it as "for each position in the index, fetch the value at that index from the source."

# gather: collect elements
src = torch.tensor([[1, 2, 3],
                     [4, 5, 6]])
index = torch.tensor([[0, 2],
                       [1, 0]])
result = torch.gather(src, dim=1, index=index)
# result[0,0] = src[0, index[0,0]] = src[0, 0] = 1
# result[0,1] = src[0, index[0,1]] = src[0, 2] = 3
# result = tensor([[1, 3], [5, 4]])

# scatter: the inverse of gather — place values at index positions
dst = torch.zeros(2, 3, dtype=torch.long)
dst.scatter_(1, index, src=torch.tensor([[10, 20], [30, 40]]))

Where gather is used: Gathering token embeddings by index, selecting log probabilities at target positions (for computing cross-entropy manually), implementing top-k selection.

index_select and masked_select

x = torch.randn(4, 5)
torch.index_select(x, dim=0, index=torch.tensor([0, 3]))  # Select rows 0 and 3
torch.index_select(x, dim=1, index=torch.tensor([1, 4]))  # Select cols 1 and 4

mask = x > 0
selected = torch.masked_select(x, mask)  # 1-D tensor of all True-mask elements

---

Broadcasting

Broadcasting lets you operate on tensors with different shapes without explicit copying. PyTorch follows NumPy's broadcasting rules.

The Rules (applied right-to-left)

• If tensors have different numbers of dimensions, prepend 1s to the smaller

tensor's shape until they match.

• For each dimension, sizes must either be equal or one of them must be 1.
• A dimension of size 1 is "stretched" to match the other.

# Example: (3, 4) + (4,) → (3, 4) + (1, 4) → (3, 4)
A = torch.randn(3, 4)
b = torch.randn(4)        # Automatically becomes (1, 4), then broadcast to (3, 4)
C = A + b                  # Works! Each row of A gets b added

# Example: (3, 1) + (1, 4) → (3, 4)
col = torch.tensor([[1.], [2.], [3.]])   # Shape: (3, 1)
row = torch.tensor([[10., 20., 30., 40.]])  # Shape: (1, 4)
result = col + row  # Shape: (3, 4) — addition table!
# [[11, 21, 31, 41],
#  [12, 22, 32, 42],
#  [13, 23, 33, 43]]

Common broadcasting patterns

# Subtract the mean from each row (feature normalization)
x = torch.randn(100, 10)       # 100 samples, 10 features
mean = x.mean(dim=0)            # Shape: (10,)
x_centered = x - mean           # Broadcasting: (100, 10) - (10,) → (100, 10)

# Subtract mean from each column
col_mean = x.mean(dim=1, keepdim=True)  # Shape: (100, 1)
x_col_centered = x - col_mean            # (100, 10) - (100, 1) → (100, 10)

Why keepdim=True matters: Without it, x.mean(dim=1) has shape (100,). You can't subtract shape (100,) from shape (100, 10) — it's ambiguous. With keepdim=True, shape is (100, 1), which broadcasts unambiguously.

Broadcasting failure

# This FAILS:
a = torch.randn(3, 4)
b = torch.randn(5)
# a + b → Error! 4 ≠ 5 and neither is 1

---

Views vs Copies

Understanding when PyTorch shares memory vs copies it is crucial for both performance and correctness.

Operations that return views (share memory)

x = torch.arange(12).reshape(3, 4)

# These all share memory with x:
y = x.view(4, 3)         # Reshape (requires contiguous)
y = x.reshape(4, 3)      # Reshape (may return view)
y = x.T                   # Transpose
y = x.transpose(0, 1)     # Transpose
y = x[0]                  # Basic indexing
y = x[:2]                 # Slicing
y = x.unsqueeze(0)        # Add dimension
y = x.squeeze()           # Remove size-1 dims
y = x.expand(2, 3, 4)     # Expand size-1 dims
y = x.narrow(0, 0, 2)     # Narrowing

# Modifying y also modifies x!
y = x.view(4, 3)
y[0, 0] = 999
print(x[0, 0])  # 999 — shared memory!

Operations that return copies

x = torch.arange(12).reshape(3, 4)

# These create new tensors (independent memory):
y = x.clone()             # Explicit copy
y = x.contiguous()        # Copy if not already contiguous
y = x[x > 5]              # Boolean indexing
y = x[[0, 2]]             # Fancy indexing
y = x.repeat(2, 1)        # Repeat (not expand)

How to check

x = torch.arange(6).reshape(2, 3)
y = x.view(3, 2)
z = x.clone()

print(x.data_ptr() == y.data_ptr())  # True — same memory
print(x.data_ptr() == z.data_ptr())  # False — different memory
print(x.storage().data_ptr() == y.storage().data_ptr())  # True

---

In-place Operations

In-place operations modify a tensor's data directly without allocating new memory. They are indicated by a trailing underscore:

x = torch.tensor([1.0, 2.0, 3.0])

x.add_(1)          # x is now [2, 3, 4] — no new tensor created
x.mul_(2)          # x is now [4, 6, 8]
x.zero_()          # x is now [0, 0, 0]
x.fill_(5)         # x is now [5, 5, 5]
x.clamp_(min=0)    # ReLU in-place
x.uniform_()       # Fill with uniform random numbers
x.normal_()        # Fill with normal random numbers

Autograd implications

In-place operations can break autograd because they modify data that the computation graph may need for backward:

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x * 2
y.add_(1)     # This MAY cause an error during backward!
# PyTorch tracks in-place modifications via a version counter.
# If backward() needs y's original value, this will fail.

Rule of thumb: Avoid in-place operations on tensors that require gradients, unless you're explicitly operating on gradient buffers (like zeroing them) or on detached tensors. The one exception is param.data manipulation, which bypasses autograd entirely.

---

Strides Explained

Strides are the mechanism that makes views, transposes, and slices work without copying data. A stride tells PyTorch how many elements to skip in physical memory to advance one position along each dimension.

How strides work

x = torch.arange(12).reshape(3, 4)
# Memory layout: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
# Shape: (3, 4), Strides: (4, 1)
#
# To go from x[0,0] to x[1,0]: skip 4 elements (stride of dim 0)
# To go from x[0,0] to x[0,1]: skip 1 element (stride of dim 1)

print(f"Strides: {x.stride()}")  # (4, 1)

Transpose changes strides, not data

x = torch.arange(6).reshape(2, 3)
print(f"x strides: {x.stride()}")     # (3, 1) — row-major

xt = x.T
print(f"x.T strides: {xt.stride()}")  # (1, 3) — column-major
print(f"x.T is contiguous: {xt.is_contiguous()}")  # False!

# The data in memory hasn't changed! Only the interpretation (strides) changed.
# x:   [0, 1, 2, 3, 4, 5] read as 2x3 with strides (3, 1)
# x.T: [0, 1, 2, 3, 4, 5] read as 3x2 with strides (1, 3)

Why contiguity matters: Some operations (like view) require contiguous memory. If a tensor isn't contiguous, call .contiguous() first (which copies data to a contiguous layout) or use .reshape() which handles this automatically.

Strides enable zero-copy slicing

x = torch.arange(20).reshape(4, 5)
y = x[::2, ::2]  # Every other row, every other column
print(f"y strides: {y.stride()}")  # (10, 2)
# Stride of 10 means skip 10 elements (2 rows of 5) to get next row
# Stride of 2 means skip 2 elements to get next column
# No data was copied!

---

NumPy Interop

PyTorch and NumPy can share memory, enabling seamless interoperation.

Conversion

import numpy as np

# NumPy → PyTorch (shares memory by default!)
np_array = np.array([1.0, 2.0, 3.0])
tensor = torch.from_numpy(np_array)
np_array[0] = 999
print(tensor[0])  # 999.0 — shared memory!

# PyTorch → NumPy (shares memory by default!)
tensor = torch.tensor([1.0, 2.0, 3.0])
np_array = tensor.numpy()
tensor[0] = 999
print(np_array[0])  # 999.0 — shared memory!

The shared memory warning

This is the single most common source of subtle bugs when mixing PyTorch and NumPy. If you modify one, the other changes too. To get an independent copy:

# Safe conversion (independent copy)
np_array = tensor.clone().numpy()           # PyTorch → NumPy (safe)
tensor = torch.from_numpy(np_array.copy())  # NumPy → PyTorch (safe)

GPU tensors and NumPy

NumPy only works on CPU. You must move GPU tensors to CPU first:

# If x is on GPU:
# x.numpy()          # ERROR! Can't convert CUDA tensor to numpy
# x.cpu().numpy()    # OK — move to CPU first
# x.detach().cpu().numpy()  # OK — also detaches from autograd

dtype compatibility

# NumPy float64 → PyTorch float64 (NOT the default float32!)
np_f64 = np.array([1.0, 2.0])  # numpy default is float64
t = torch.from_numpy(np_f64)
print(t.dtype)  # torch.float64

# If you want float32, cast explicitly
t32 = torch.from_numpy(np_f64).float()

---

What's Next?

With tensors mastered, Module 03 covers autograd — PyTorch's automatic differentiation engine that makes all of deep learning possible. You'll learn how PyTorch tracks operations on tensors and computes gradients automatically.

Run the example files in this directory to practice:

• creation_and_properties.py — tensor creation and inspection
• operations.py — element-wise and reduction operations
• indexing_and_slicing.py — all forms of indexing
• broadcasting.py — broadcasting rules with examples
• views_strides_memory.py — views, strides, and memory layout

---

📓 Open Notebook — Interactive version of this module

Module 03: Autograd — Automatic Differentiation

Table of Contents

• What is Automatic Differentiation?
• Forward Mode vs Reverse Mode
• The Computation Graph
• Leaf Tensors vs Intermediate Tensors
• requires_grad, grad_fn, grad
• The backward() Function
• Gradient Accumulation
• torch.no_grad() vs torch.inference_mode()
• detach()
• Custom Autograd Functions
• gradcheck and gradgradcheck
• Higher-Order Gradients
• torch.autograd.grad()
• Jacobian and Hessian Computation
• Common Pitfalls
• Autograd Hooks
• Compiled Autograd

---

What is Automatic Differentiation?

Automatic differentiation (AD) is a technique for computing exact derivatives of functions expressed as computer programs. It is NOT:

• Symbolic differentiation (like Mathematica/Sympy): These manipulate

mathematical expressions symbolically, which can lead to expression swell (the derivative expression becomes exponentially larger than the original).

• Numerical differentiation (finite differences): Computing

(f(x+h) - f(x)) / h is simple but introduces truncation and rounding errors, and scales poorly with the number of parameters (requires one function evaluation per parameter).

Instead, AD exploits the fact that every computer program, no matter how complex, is composed of elementary operations (+, *, sin, exp, etc.) whose derivatives are known. By applying the chain rule systematically, AD computes exact derivatives at machine precision.

import torch

x = torch.tensor(3.0, requires_grad=True)
y = torch.sin(x) * torch.exp(x)

# PyTorch knows the derivative of sin, exp, and *.
# It chains them together to get dy/dx exactly.
y.backward()
print(f"dy/dx at x=3: {x.grad.item():.6f}")

# Verify: d/dx[sin(x)*exp(x)] = cos(x)*exp(x) + sin(x)*exp(x)
manual = (torch.cos(torch.tensor(3.0)) + torch.sin(torch.tensor(3.0))) * torch.exp(torch.tensor(3.0))
print(f"Manual:       {manual.item():.6f}")

---

Forward Mode vs Reverse Mode

There are two ways to apply the chain rule through a computation:

Forward Mode (Tangent Mode)

Propagates derivatives forward through the computation, alongside the primal (original) computation. For f: R^n → R^m:

• Computes one column of the Jacobian per forward pass
• Cost: O(n) forward passes for full Jacobian
• Efficient when n << m (few inputs, many outputs)

Reverse Mode (Adjoint Mode) = Backpropagation

• Computes one row of the Jacobian per backward pass
• Cost: O(m) backward passes for full Jacobian
• Efficient when m << n (many inputs, few outputs)

Why neural networks use reverse mode: A neural network has millions of parameters (inputs to the loss function) but produces a single scalar loss (one output). Reverse mode computes ALL gradients (∂loss/∂param for every param) in a single backward pass — regardless of how many parameters there are. Forward mode would require one pass per parameter, which is millions of times slower.

This asymmetry is why:

• Training uses reverse mode (backpropagation): 1 scalar output, millions of inputs
• Forward-mode AD is used for Jacobian-vector products in some optimization methods
• PyTorch supports both: backward() for reverse mode, torch.autograd.forward_ad

for forward mode

---

The Computation Graph

When you perform operations on tensors with requires_grad=True, PyTorch builds a directed acyclic graph (DAG) that records every operation. This graph is essential for computing gradients.

How the graph is built

x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)

# Each operation creates a node in the graph
a = x * y        # MulBackward node
b = a + x        # AddBackward node
c = b.sin()      # SinBackward node

print(f"c.grad_fn: {c.grad_fn}")  # SinBackward
print(f"b.grad_fn: {b.grad_fn}")  # AddBackward
print(f"a.grad_fn: {a.grad_fn}")  # MulBackward
print(f"x.grad_fn: {x.grad_fn}")  # None (leaf tensor)

The graph looks like:

x ──→ MulBackward(a=x*y) ──→ AddBackward(b=a+x) ──→ SinBackward(c=sin(b))
y ──↗                    x ──↗

Graph lifecycle

The graph is built dynamically during the forward pass and consumed (destroyed) during backward. This is the "dynamic graph" feature of PyTorch:

• Forward pass: Operations are recorded as nodes in the graph.
• backward(): The graph is traversed in reverse order, computing gradients.
• Graph destroyed: After backward, the graph is freed (by default).
• Next forward: A new graph is built from scratch.

This means every iteration can follow a different code path — if statements, for loops, and variable-length sequences all just work.

---

Leaf Tensors vs Intermediate Tensors

Leaf tensors

A leaf tensor is one that was created directly by the user (not by an operation). Leaf tensors are the "starting points" of the computation graph.

x = torch.tensor(1.0, requires_grad=True)   # Leaf
w = torch.randn(3, 4, requires_grad=True)    # Leaf
b = torch.zeros(4, requires_grad=True)        # Leaf

print(f"x.is_leaf: {x.is_leaf}")  # True
print(f"w.is_leaf: {w.is_leaf}")  # True

Intermediate tensors

Tensors created by operations on other tensors are intermediate. They have a grad_fn that records how they were created.

y = x * 2        # Intermediate
z = w @ b        # Intermediate

print(f"y.is_leaf: {y.is_leaf}")      # False
print(f"y.grad_fn: {y.grad_fn}")      # MulBackward0

Key difference: where gradients are stored

By default, PyTorch only stores gradients for leaf tensors. This is a memory optimization — intermediate gradients are computed during backward but discarded immediately after use.

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
z = y * 3

z.backward()
print(f"x.grad: {x.grad}")   # 12.0 — stored because x is a leaf
print(f"y.grad: {y.grad}")   # None — not stored because y is intermediate

If you need gradients for intermediate tensors, use retain_grad():

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
y.retain_grad()  # Tell PyTorch to keep y's gradient
z = y * 3
z.backward()
print(f"y.grad: {y.grad}")   # 3.0 — now it's stored

---

requires_grad, grad_fn, grad

These three attributes form the autograd trinity:

requires_grad

A boolean flag indicating whether this tensor participates in gradient computation.

# Tensors that need gradients
w = torch.randn(3, 3, requires_grad=True)   # Set at creation
b = torch.zeros(3)
b.requires_grad_(True)                       # Set after creation (in-place)

# Tensors that don't
data = torch.randn(100, 3)                   # Data doesn't need gradients
labels = torch.randint(0, 2, (100,))          # Labels don't need gradients

Rule: If ANY input to an operation has requires_grad=True, the output will also have requires_grad=True (gradient tracking propagates forward).

grad_fn

A reference to the backward function (the operation that created this tensor). Leaf tensors have grad_fn=None.

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2      # y.grad_fn = PowBackward0
z = y.sum()      # z.grad_fn = SumBackward0

You can traverse the graph by following grad_fn links:

print(z.grad_fn)                     # SumBackward0
print(z.grad_fn.next_functions)      # Links to y's grad_fn

grad

The accumulated gradient tensor. Only populated after backward() is called. Only leaf tensors have .grad by default.

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
loss = (x ** 2).sum()
loss.backward()
print(f"x.grad: {x.grad}")  # [2., 4., 6.] — the gradient ∂loss/∂x = 2x

---

The backward() Function

backward() is the workhorse of training. It computes gradients by traversing the computation graph in reverse order (topological sort).

How it works

x = torch.tensor(3.0, requires_grad=True)
y = x ** 3 + 2 * x ** 2 - 5 * x
y.backward()
# Computes dy/dx = 3x² + 4x - 5 = 27 + 12 - 5 = 34
print(f"dy/dx at x=3: {x.grad.item()}")  # 34.0

The gradient argument

When the output is not a scalar, you must provide a gradient argument (also called the "upstream gradient" or "grad_output"):

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2  # y is a vector, not a scalar

# Must provide gradient (Jacobian-vector product)
y.backward(gradient=torch.tensor([1.0, 1.0, 1.0]))
print(f"x.grad: {x.grad}")  # [2., 4., 6.]

# The gradient argument acts as weights in a weighted sum:
# effectively computing d(sum(gradient * y))/dx

Why? PyTorch's backward pass always computes vector-Jacobian products (VJPs). For a scalar output, the "vector" is implicitly 1.0. For vector outputs, you must supply it. In practice, this gradient comes from the downstream loss.

Graph destruction

By default, backward() destroys the computation graph after use:

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
y.backward()
# y.backward()  # ERROR: graph already freed

# To keep the graph, use retain_graph=True
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
y.backward(retain_graph=True)
y.backward()  # OK — graph was retained

---

Gradient Accumulation

Gradients accumulate by default. This is the single most common source of bugs for PyTorch beginners.

x = torch.tensor(1.0, requires_grad=True)

# First backward
y = x * 2
y.backward()
print(f"After first backward: x.grad = {x.grad}")  # 2.0

# Second backward WITHOUT zeroing
y = x * 3
y.backward()
print(f"After second backward: x.grad = {x.grad}")  # 5.0 (= 2.0 + 3.0)
# The gradient ACCUMULATED!

Why accumulation exists

Gradient accumulation is actually a feature, not a bug. It's useful for:

• Large effective batch sizes: Process mini-batches one at a time, accumulate

gradients, then take one optimizer step. This lets you simulate large batches on limited GPU memory.

• Multiple losses: If your model has multiple loss terms, you can backward

each one separately and the gradients add up correctly.

How to zero gradients

# Method 1: Manual zeroing
x.grad.zero_()

# Method 2: Optimizer (most common)
optimizer.zero_grad()           # Zeros all parameter gradients

# Method 3: Set to None (slightly faster)
optimizer.zero_grad(set_to_none=True)
# or
x.grad = None

The standard training loop pattern

# for batch in dataloader:
#     optimizer.zero_grad()       # 1. Zero gradients
#     output = model(batch)       # 2. Forward pass
#     loss = criterion(output)    # 3. Compute loss
#     loss.backward()             # 4. Backward pass (compute gradients)
#     optimizer.step()            # 5. Update parameters

---

torch.no_grad() vs torch.inference_mode()

Both disable gradient tracking, but for different purposes:

torch.no_grad()

Disables gradient computation. Tensors created inside still have requires_grad based on their inputs, but no graph is built.

x = torch.tensor(1.0, requires_grad=True)

with torch.no_grad():
    y = x * 2
    print(f"y.requires_grad: {y.requires_grad}")  # False
    print(f"y.grad_fn: {y.grad_fn}")               # None

Use for: Validation loops, parameter updates, evaluation metrics.

torch.inference_mode()

A stricter, more optimized version of no_grad(). Disables autograd entirely and enables additional optimizations.

x = torch.tensor(1.0, requires_grad=True)

with torch.inference_mode():
    y = x * 2
    # Even stricter: y is an InferenceTensor, cannot be used with autograd at all

Use for: Production inference, deployment, any time you're certain you won't need gradients.

When to use which

---

detach()

detach() creates a new tensor that shares data but is disconnected from the computation graph.

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
z = y.detach()  # z shares data with y but has no grad_fn

print(f"y.requires_grad: {y.requires_grad}")  # True
print(f"z.requires_grad: {z.requires_grad}")  # False
print(f"z.grad_fn: {z.grad_fn}")               # None
print(f"Same data: {y.data_ptr() == z.data_ptr()}")  # True

Common uses

• Stopping gradient flow: In GANs, you detach the generator output when

training the discriminator, so gradients don't flow back to the generator.

• Converting to NumPy: tensor.detach().numpy() — you must detach first

if the tensor requires gradients.

• Target networks: In reinforcement learning, the target network's output

is detached so it's treated as a constant.

# GAN-style gradient stopping
# fake = generator(noise)
# d_loss = discriminator(fake.detach())  # Stops gradient to generator

---

Custom Autograd Functions

When PyTorch's built-in operations aren't sufficient, you can define custom forward and backward passes.

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

# Usage
x = torch.randn(5, requires_grad=True)
y = MyReLU.apply(x)
y.sum().backward()
print(f"x: {x}")
print(f"x.grad: {x.grad}")

The ctx object

ctx (context) is used to pass information from forward to backward:

• ctx.save_for_backward(tensor1, tensor2, ...): Save tensors needed for backward.

These are stored efficiently and checked for version consistency.

• ctx.saved_tensors: Retrieve saved tensors in backward.
• ctx.needs_input_grad: Tuple of booleans indicating which inputs need gradients.
• ctx.mark_dirty(tensor): Mark tensors modified in-place.
• ctx.mark_non_differentiable(tensor): Mark outputs that don't need gradients.

Rules for custom functions

• forward receives the actual tensor values and returns output tensors.
• backward receives the upstream gradient (grad_output) for each output

and must return one gradient per input (or None if that input doesn't need gradients).

• The number of tensors returned by backward must match the number of

inputs to forward (excluding ctx).

---

gradcheck and gradgradcheck

These utilities verify that your custom autograd function computes correct gradients by comparing against numerical finite differences.

from torch.autograd import gradcheck, gradgradcheck

class MySigmoid(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        result = 1 / (1 + torch.exp(-x))
        ctx.save_for_backward(result)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        result, = ctx.saved_tensors
        return grad_output * result * (1 - result)

# gradcheck requires float64 for numerical precision
x = torch.randn(5, dtype=torch.float64, requires_grad=True)
test = gradcheck(MySigmoid.apply, (x,), eps=1e-6, atol=1e-4)
print(f"Gradient check passed: {test}")

# gradgradcheck checks second derivatives
test2 = gradgradcheck(MySigmoid.apply, (x,), eps=1e-6, atol=1e-4)
print(f"Double gradient check passed: {test2}")

---

Higher-Order Gradients

By default, backward() only computes first-order gradients. To compute higher-order gradients (gradients of gradients), use create_graph=True.

x = torch.tensor(2.0, requires_grad=True)
y = x ** 4  # y = x^4

# First derivative: dy/dx = 4x³
grad1 = torch.autograd.grad(y, x, create_graph=True)[0]
print(f"dy/dx = {grad1.item()}")  # 32.0

# Second derivative: d²y/dx² = 12x²
grad2 = torch.autograd.grad(grad1, x, create_graph=True)[0]
print(f"d²y/dx² = {grad2.item()}")  # 48.0

# Third derivative: d³y/dx³ = 24x
grad3 = torch.autograd.grad(grad2, x)[0]
print(f"d³y/dx³ = {grad3.item()}")  # 48.0

Why create_graph=True

Without create_graph=True, the backward pass doesn't build a graph for itself. This means you can't differentiate through the backward pass. With it enabled, the gradient computation is itself differentiable.

Use cases:

• Gradient penalty (WGAN-GP): Penalize the norm of gradients, which requires

computing the gradient of the gradient norm.

• MAML (meta-learning): Differentiate through the inner optimization loop.
• Physics-informed neural networks: Loss functions involving derivatives of

the network output.

---

torch.autograd.grad()

torch.autograd.grad() computes gradients without storing them in .grad attributes. This is useful for:

• Computing gradients of specific outputs w.r.t. specific inputs
• Higher-order gradients
• Functional-style gradient computation

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()

# Compute gradient without storing in x.grad
grad = torch.autograd.grad(y, x)[0]
print(f"Gradient: {grad}")        # [2., 4., 6.]
print(f"x.grad: {x.grad}")        # None — not stored

Multiple outputs and inputs

x = torch.tensor(1.0, requires_grad=True)
y = torch.tensor(2.0, requires_grad=True)

out1 = x * y
out2 = x ** 2 + y ** 2

# Gradients of both outputs w.r.t. both inputs
grads = torch.autograd.grad(
    outputs=[out1, out2],
    inputs=[x, y],
    grad_outputs=[torch.tensor(1.0), torch.tensor(1.0)]
)
print(f"d(out1+out2)/dx: {grads[0].item()}")  # y + 2x = 2 + 2 = 4
print(f"d(out1+out2)/dy: {grads[1].item()}")  # x + 2y = 1 + 4 = 5

---

Jacobian and Hessian Computation

Jacobian

The Jacobian matrix contains all partial derivatives of a vector-valued function. For f: R^n → R^m, J is m × n where J_ij = ∂f_i/∂x_j.

from torch.autograd.functional import jacobian

def f(x):
    return torch.stack([x[0]**2 + x[1], x[0] * x[1]**2])

x = torch.tensor([1.0, 2.0])
J = jacobian(f, x)
print(f"Jacobian:\n{J}")
# [[2*x0, 1   ],     = [[2, 1],
#  [x1^2, 2*x0*x1]]     [4, 4]]

Hessian

H is n × n where H_ij = ∂²f/∂x_i∂x_j.

from torch.autograd.functional import hessian

def g(x):
    return x[0]**3 + x[0]*x[1]**2 + x[1]**3

x = torch.tensor([1.0, 2.0])
H = hessian(g, x)
print(f"Hessian:\n{H}")
# [[6*x0,    2*x1 ],     = [[6,  4],
#  [2*x1, 2*x0+6*x1]]      [4, 14]]

Efficient Jacobian-vector products (JVPs) and vector-Jacobian products (VJPs)

Computing the full Jacobian is expensive (O(n) backward passes). Often you only need the product of the Jacobian with a specific vector:

from torch.autograd.functional import jvp, vjp

def f(x):
    return torch.stack([x[0]**2, x[0]*x[1]])

x = torch.tensor([1.0, 2.0])
v = torch.tensor([1.0, 0.0])  # Direction vector

# JVP: J @ v (forward mode — one forward pass)
_, jvp_result = jvp(f, (x,), (v,))
print(f"JVP (J @ v): {jvp_result}")

# VJP: v^T @ J (reverse mode — one backward pass)
_, vjp_fn = vjp(f, x)
vjp_result = vjp_fn(torch.tensor([1.0, 0.0]))
print(f"VJP (v^T @ J): {vjp_result}")

---

Common Pitfalls

1. Forgetting to zero gradients

x = torch.tensor(1.0, requires_grad=True)
for i in range(3):
    loss = x * (i + 1)
    loss.backward()
    print(f"Step {i}: x.grad = {x.grad.item()}")
    # Without zeroing: 1.0, 3.0, 6.0 (accumulated!)
    # x.grad.zero_()  # Uncomment to fix

2. In-place operations on tensors that need gradients

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x * 2

# This WILL cause problems:
# y.add_(1)  # In-place modification of a tensor needed for backward
# y.backward(torch.tensor([1.0, 1.0]))  # RuntimeError!

3. Gradient not flowing through integer operations

x = torch.tensor(3.0, requires_grad=True)
y = x.int()  # Casting to int breaks gradient flow
# y has no grad_fn — gradient is lost!

4. NaN gradients

Common causes:

• log(0): Use log(x + epsilon) or torch.clamp(x, min=1e-8)
• sqrt(0): derivative of sqrt at 0 is infinity. Use sqrt(x + epsilon)
• 0/0: Can occur in normalization layers when variance is zero
• Division by a very small number: use torch.clamp on denominators

x = torch.tensor(0.0, requires_grad=True)
# y = torch.log(x)   # Will give -inf, grad is inf
y = torch.log(x + 1e-8)  # Safe
y.backward()
print(f"Safe log grad: {x.grad}")

5. Modifying parameters without torch.no_grad()

x = torch.tensor(1.0, requires_grad=True)
# x = x - 0.1 * x.grad  # This creates a NEW tensor, breaking the leaf status
# Instead:
# with torch.no_grad():
#     x -= 0.1 * x.grad

---

Autograd Hooks

Hooks let you inspect or modify gradients during the backward pass without changing the model code.

Tensor hooks

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

def print_grad(grad):
    print(f"  Hook received gradient: {grad}")

x.register_hook(print_grad)
y = (x ** 2).sum()
y.backward()  # Hook fires during backward

Gradient modification hooks

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

def clip_grad(grad):
    return torch.clamp(grad, -1.0, 1.0)

x.register_hook(clip_grad)
y = (x ** 3).sum()
y.backward()
print(f"Clipped gradient: {x.grad}")  # Gradients are clipped to [-1, 1]

Module hooks

For nn.Module, you can register hooks on the entire module:

import torch.nn as nn

model = nn.Linear(5, 3)

def forward_hook(module, input, output):
    print(f"Forward: input shape={input[0].shape}, output shape={output.shape}")

def backward_hook(module, grad_input, grad_output):
    print(f"Backward: grad_output shape={grad_output[0].shape}")

model.register_forward_hook(forward_hook)
model.register_full_backward_hook(backward_hook)

x = torch.randn(2, 5)
y = model(x)
y.sum().backward()

---

Compiled Autograd

PyTorch 2.x introduced compiled autograd, which applies torch.compile to the backward pass. This can significantly speed up training by:

• Fusing backward operations into optimized kernels
• Eliminating Python overhead in the backward pass
• Enabling whole-graph optimization across forward and backward

# Basic usage (requires torch >= 2.0)
model = torch.nn.Linear(100, 10)
compiled_model = torch.compile(model)

# Both forward AND backward are compiled
x = torch.randn(32, 100)
y = compiled_model(x)
loss = y.sum()
loss.backward()  # This backward is also compiled

The compilation happens lazily — the first iteration traces the graph, subsequent iterations use the compiled version. This means the first iteration is slower (compilation overhead) but all subsequent iterations are faster.

---

Summary

Autograd is the engine that makes deep learning practical. Key takeaways:

• Reverse-mode AD is efficient for neural networks (one backward pass for

all parameters).

• The computation graph is built dynamically during forward and consumed

during backward.

• Always zero gradients before each optimization step.
• Use torch.no_grad() for validation and inference_mode() for deployment.
• Custom Functions let you define operations with hand-written gradients.
• Hooks let you inspect and modify gradients without changing model code.

Run the example files to see these concepts in action:

• gradient_basics.py — basic gradient computation and the training loop
• computation_graph.py — visualizing and understanding the computation graph
• custom_functions.py — writing custom autograd functions
• higher_order_gradients.py — second derivatives, Jacobians, and Hessians

---

📓 Open Notebook — Interactive version of this module

Module 04: Neural Networks in PyTorch

Table of Contents

• What is nn.Module?
• Module Lifecycle
• Parameters vs Buffers
• Container Modules
• Linear Layers
• Convolution Layers
• Pooling Layers
• Normalization Layers
• Activation Functions
• Dropout
• Recurrent Layers
• Transformer Layers
• Embedding Layers
• Loss Functions
• Functional API
• Weight Initialization
• Hooks
• State Dict

---

What is nn.Module?

torch.nn.Module is the base class for ALL neural network components in PyTorch. Every layer, every model, every building block inherits from this class. When you build a neural network in PyTorch, you create a class that inherits from nn.Module.

Think of nn.Module as a container that:

• Holds learnable parameters (weights and biases)
• Defines how data flows through the network (the forward method)
• Provides utilities for moving to GPU, saving/loading, switching train/eval modes
• Builds a tree of sub-modules (layers inside your model)

import torch
import torch.nn as nn

class MyNetwork(nn.Module):
    def __init__(self):
        super().__init__()  # MUST call parent __init__
        self.linear1 = nn.Linear(784, 256)
        self.linear2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = torch.relu(self.linear1(x))
        x = self.linear2(x)
        return x

Key insight: When you assign an nn.Module or nn.Parameter as an attribute of your module (in __init__), PyTorch automatically registers it. This means it will show up in .parameters(), be moved when you call .to(device), and be saved in .state_dict().

---

Module Lifecycle

__init__: Construction

In __init__, you define all the layers and parameters your network needs. You MUST call super().__init__() first. Any nn.Module or nn.Parameter assigned as an attribute gets automatically registered.

class Model(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        # These are automatically registered as sub-modules
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, output_dim)
        # This is NOT registered (plain Python attribute)
        self.activation_name = "relu"

forward: The Forward Pass

The forward method defines how input data flows through your network. You NEVER call forward() directly — instead, you call the module as a function: model(x). This is because __call__ does extra work (hooks, checks) before calling forward.

def forward(self, x):
    x = torch.relu(self.layer1(x))
    x = self.layer2(x)
    return x

# Correct usage:
output = model(input_tensor)  # calls __call__ which calls forward

# WRONG — never do this:
# output = model.forward(input_tensor)

Train vs Eval Mode

Modules have two modes that affect behavior of certain layers (Dropout, BatchNorm):

model.train()   # Sets training mode (dropout active, batchnorm uses batch stats)
model.eval()    # Sets evaluation mode (dropout disabled, batchnorm uses running stats)

# Check current mode
print(model.training)  # True or False

This is critical: forgetting model.eval() during inference causes incorrect results because Dropout still drops neurons and BatchNorm uses batch statistics instead of learned running statistics.

---

Parameters vs Buffers

Parameters

Parameters are tensors that require gradients and are updated by the optimizer. They represent the learnable weights of your model.

class CustomLayer(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        # Manual parameter creation
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
    
    def forward(self, x):
        return x @ self.weight.t() + self.bias

Buffers

Buffers are tensors that are part of the module's state but do NOT require gradients. They are saved in state_dict() and moved with .to(device), but the optimizer ignores them.

Common use cases:

• Running mean/variance in BatchNorm
• Fixed positional encodings
• Binary masks that don't change during training

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        # Buffer: saved in state_dict, moved with .to(), but NOT optimized
        self.register_buffer('running_mean', torch.zeros(10))
        # Non-persistent buffer: moved with .to() but NOT saved
        self.register_buffer('temp_mask', torch.ones(10), persistent=False)

register_parameter and register_buffer

class ExplicitRegistration(nn.Module):
    def __init__(self):
        super().__init__()
        # Explicit parameter registration (equivalent to self.weight = nn.Parameter(...))
        self.register_parameter('weight', nn.Parameter(torch.randn(5, 3)))
        # Can register None — useful for optional parameters
        self.register_parameter('optional_bias', None)
        # Buffer registration
        self.register_buffer('counter', torch.tensor(0))

Traversing the Module Tree

model = MyNetwork()

# All parameters (recursively)
for name, param in model.named_parameters():
    print(f"{name}: shape={param.shape}, requires_grad={param.requires_grad}")

# All sub-modules (recursively)
for name, module in model.named_modules():
    print(f"{name}: {type(module).__name__}")

# Direct children only
for name, module in model.named_children():
    print(f"{name}: {type(module).__name__}")

# All buffers
for name, buf in model.named_buffers():
    print(f"{name}: shape={buf.shape}")

---

Container Modules

nn.Sequential

Chains modules in order. Input flows through each module sequentially.

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)
# Equivalent to calling each in order: output = layer3(relu(layer2(relu(layer1(x)))))
output = model(input_tensor)

Use when: Your network is a simple chain of operations with no branching.

nn.ModuleList

A list of modules. Does NOT define a forward pass — you iterate manually.

class MultiHeadModel(nn.Module):
    def __init__(self, num_heads):
        super().__init__()
        self.heads = nn.ModuleList([nn.Linear(256, 10) for _ in range(num_heads)])
    
    def forward(self, x):
        return [head(x) for head in self.heads]

Use when: You need a variable number of layers that you'll iterate over yourself. WARNING: A plain Python list [] will NOT register the modules!

nn.ModuleDict

A dictionary of modules, accessed by string keys.

class MultiTaskModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Linear(784, 256)
        self.heads = nn.ModuleDict({
            'classification': nn.Linear(256, 10),
            'regression': nn.Linear(256, 1),
        })
    
    def forward(self, x, task):
        features = torch.relu(self.backbone(x))
        return self.heads[task](features)

Use when: You need named access to different sub-modules (multi-task, configurable architectures).

nn.ParameterList and nn.ParameterDict

Same idea but for raw parameters instead of modules:

class CustomModel(nn.Module):
    def __init__(self, num_experts):
        super().__init__()
        self.expert_weights = nn.ParameterList(
            [nn.Parameter(torch.randn(256, 256)) for _ in range(num_experts)]
        )
        self.config_params = nn.ParameterDict({
            'scale': nn.Parameter(torch.ones(1)),
            'shift': nn.Parameter(torch.zeros(1)),
        })

---

Linear Layers

nn.Linear

Applies a linear transformation: y = xW^T + b

# Input:  (batch_size, in_features)   e.g., (32, 784)
# Output: (batch_size, out_features)  e.g., (32, 256)
linear = nn.Linear(in_features=784, out_features=256, bias=True)

# Weight shape: (out_features, in_features) = (256, 784)
# Bias shape:   (out_features,) = (256,)

The math: For input x of shape (, in_features), output y of shape (, out_features): y_i = sum_j(x_j * W_ij) + b_i

nn.Bilinear

Applies a bilinear transformation: y = x1^T A x2 + b

# Two inputs of potentially different sizes
bilinear = nn.Bilinear(in1_features=20, in2_features=30, out_features=40)
input1 = torch.randn(128, 20)
input2 = torch.randn(128, 30)
output = bilinear(input1, input2)  # shape: (128, 40)

nn.LazyLinear

Infers in_features from the first input — useful for prototyping:

lazy = nn.LazyLinear(out_features=256)
# in_features is determined on first forward pass
output = lazy(torch.randn(32, 784))  # Now it knows in_features=784

---

Convolution Layers

Core Concepts

Convolutions slide a kernel (filter) across the input, computing dot products at each position.

Key parameters:

• kernel_size: Size of the sliding window (e.g., 3 means 3x3 for Conv2d)
• stride: How far the kernel moves each step (default=1)
• padding: Zero-padding added to input borders
• dilation: Spacing between kernel elements (dilated/atrous convolution)
• groups: Split input channels into groups for grouped convolution

Output size formula (for each spatial dimension):

output_size = floor((input_size + 2*padding - dilation*(kernel_size-1) - 1) / stride + 1)

nn.Conv1d

For sequential/temporal data (text, audio, time series).

# Input:  (batch, in_channels, length)      e.g., (32, 1, 100)
# Output: (batch, out_channels, new_length) e.g., (32, 16, 98)
conv1d = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3)

nn.Conv2d

For image data. The most commonly used convolution.

# Input:  (batch, in_channels, height, width)       e.g., (32, 3, 224, 224)
# Output: (batch, out_channels, new_height, new_width) e.g., (32, 64, 112, 112)
conv2d = nn.Conv2d(
    in_channels=3,       # RGB input
    out_channels=64,     # 64 filters
    kernel_size=3,       # 3x3 kernel
    stride=2,            # Downsample by 2
    padding=1            # Same padding for stride=1
)

# Weight shape: (out_channels, in_channels/groups, kernel_h, kernel_w)
# = (64, 3, 3, 3)

nn.Conv3d

For volumetric data (video, 3D medical images).

# Input:  (batch, channels, depth, height, width)
conv3d = nn.Conv3d(in_channels=3, out_channels=64, kernel_size=3, padding=1)

ConvTranspose2d (Transposed Convolution)

Used for upsampling — goes from smaller to larger spatial dimensions. Often called "deconvolution" (technically incorrect name).

# Input:  (batch, in_channels, H, W)
# Output: (batch, out_channels, H*2, W*2) with stride=2
upsample = nn.ConvTranspose2d(
    in_channels=64, out_channels=32,
    kernel_size=4, stride=2, padding=1
)

Depthwise Separable Convolution

A two-step convolution that's much more efficient:

• Depthwise: Apply one filter per input channel (groups=in_channels)
• Pointwise: 1x1 convolution to mix channels

class DepthwiseSeparable(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size=3, padding=1):
        super().__init__()
        # Depthwise: each input channel gets its own filter
        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size,
                                   padding=padding, groups=in_ch)
        # Pointwise: 1x1 conv to combine channels
        self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1)
    
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

---

Pooling Layers

Pooling reduces spatial dimensions while retaining important information.

MaxPool2d

Takes the maximum value in each pooling window:

# Input:  (batch, channels, 224, 224)
# Output: (batch, channels, 112, 112)  — halves spatial dims
pool = nn.MaxPool2d(kernel_size=2, stride=2)

AvgPool2d

Takes the average value in each pooling window:

pool = nn.AvgPool2d(kernel_size=2, stride=2)

AdaptiveAvgPool2d — Global Average Pooling

Outputs a fixed spatial size regardless of input size. Setting output to (1,1) gives "global average pooling" — commonly used before the final classifier.

# No matter what spatial size comes in, output is (batch, channels, 1, 1)
gap = nn.AdaptiveAvgPool2d(output_size=(1, 1))
# Then flatten: (batch, channels, 1, 1) -> (batch, channels)

This is the modern replacement for large fully-connected layers at the end of CNNs.

---

Normalization Layers

BatchNorm (nn.BatchNorm1d, BatchNorm2d)

Normalizes across the batch dimension. For each feature/channel:

Formula:

y = (x - E[x]) / sqrt(Var[x] + eps) * gamma + beta

Where gamma (weight) and beta (bias) are learnable parameters.

Training behavior: Uses batch mean and variance, updates running statistics. Eval behavior: Uses stored running mean and variance (fixed).

bn = nn.BatchNorm2d(num_features=64)  # 64 channels
# Maintains: running_mean, running_var (buffers), weight, bias (parameters)

When to use: CNNs with large batch sizes. Not suitable for batch_size=1 or variable batch sizes (use LayerNorm or GroupNorm instead).

LayerNorm

Normalizes across the feature dimensions (not the batch). Each sample is normalized independently.

# For a transformer with hidden_size=512
ln = nn.LayerNorm(normalized_shape=512)
# For image data: normalize over (C, H, W)
ln_image = nn.LayerNorm([64, 32, 32])

When to use: Transformers, RNNs, any case where batch stats are unreliable.

GroupNorm

Splits channels into groups and normalizes within each group. A middle ground between BatchNorm (all channels) and InstanceNorm (each channel separately).

gn = nn.GroupNorm(num_groups=32, num_channels=256)

When to use: When batch size is small, or when you want BatchNorm-like behavior without batch dependency.

InstanceNorm

Normalizes each channel of each sample independently. Equivalent to GroupNorm with num_groups = num_channels.

inst_norm = nn.InstanceNorm2d(num_features=64)

When to use: Style transfer, image generation tasks.

RMSNorm

Root Mean Square Layer Normalization — simpler than LayerNorm (no mean subtraction):

rms_norm = nn.RMSNorm(normalized_shape=512)

Formula: y = x / RMS(x) * gamma where RMS(x) = sqrt(mean(x^2) + eps)

When to use: Modern LLMs (LLaMA, etc.) — slightly faster than LayerNorm.

---

Activation Functions

Activations introduce non-linearity. Without them, stacking linear layers is equivalent to a single linear layer.

ReLU: f(x) = max(0, x)

The default choice. Simple, fast, but can "die" (output 0 for all inputs).

nn.ReLU(inplace=False)  # inplace=True saves memory but can cause issues

LeakyReLU: f(x) = max(alpha*x, x) (default alpha=0.01)

Prevents dying ReLU by allowing small negative slope.

nn.LeakyReLU(negative_slope=0.01)

PReLU: f(x) = max(alpha*x, x) where alpha is LEARNED

nn.PReLU(num_parameters=1)  # One alpha per channel if num_parameters=num_channels

GELU: f(x) = x * Phi(x) where Phi is the CDF of standard normal

Used in Transformers (BERT, GPT). Smooth approximation of ReLU.

nn.GELU(approximate='none')  # 'tanh' for faster approximation

SiLU/Swish: f(x) = x * sigmoid(x)

Smooth, non-monotonic. Used in EfficientNet, many modern architectures.

nn.SiLU()

Mish: f(x) = x * tanh(softplus(x))

Similar to SiLU but slightly different properties.

nn.Mish()

Sigmoid: f(x) = 1 / (1 + exp(-x))

Squashes to [0, 1]. Used for binary classification output, gates.

nn.Sigmoid()

Tanh: f(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))

Squashes to [-1, 1]. Used in RNN gates.

nn.Tanh()

Softmax: f(x_i) = exp(x_i) / sum(exp(x_j))

Outputs a probability distribution (sums to 1). Used for multi-class classification.

nn.Softmax(dim=-1)  # Usually along the last dimension

---

Dropout

Dropout randomly zeros elements during training to prevent overfitting. During evaluation, dropout is disabled and outputs are unchanged.

Key insight: During training, remaining elements are scaled by 1/(1-p) so that expected values remain the same at test time (inverted dropout).

nn.Dropout

dropout = nn.Dropout(p=0.5)  # 50% of elements zeroed during training

nn.Dropout2d

Drops entire channels (feature maps) for Conv2d outputs:

dropout2d = nn.Dropout2d(p=0.1)  # Drops entire channels

nn.AlphaDropout

For use with SELU activation — maintains self-normalizing property:

alpha_dropout = nn.AlphaDropout(p=0.1)

---

Recurrent Layers

nn.RNN

Basic recurrent layer: h_t = tanh(x_t W_ih^T + h_{t-1} W_hh^T + b)

rnn = nn.RNN(input_size=128, hidden_size=256, num_layers=2,
             batch_first=True, bidirectional=False, dropout=0.1)
# Input:  (batch, seq_len, input_size)
# Output: (batch, seq_len, hidden_size * num_directions)
# Hidden: (num_layers * num_directions, batch, hidden_size)
output, h_n = rnn(input_seq, h_0)

nn.LSTM

Long Short-Term Memory — solves vanishing gradients with gates:

Gate equations:

f_t = sigmoid(W_f [h_{t-1}, x_t] + b_f)    # Forget gate
i_t = sigmoid(W_i [h_{t-1}, x_t] + b_i)    # Input gate
g_t = tanh(W_g [h_{t-1}, x_t] + b_g)       # Cell candidate
o_t = sigmoid(W_o [h_{t-1}, x_t] + b_o)    # Output gate
c_t = f_t * c_{t-1} + i_t * g_t            # Cell state
h_t = o_t * tanh(c_t)                       # Hidden state

lstm = nn.LSTM(input_size=128, hidden_size=256, num_layers=2,
               batch_first=True, bidirectional=True, dropout=0.1)
output, (h_n, c_n) = lstm(input_seq)
# output shape: (batch, seq_len, hidden_size * 2) for bidirectional

nn.GRU

Gated Recurrent Unit — simpler than LSTM with fewer parameters:

Gate equations:

r_t = sigmoid(W_r [h_{t-1}, x_t])  # Reset gate
z_t = sigmoid(W_z [h_{t-1}, x_t])  # Update gate
n_t = tanh(W_n [r_t * h_{t-1}, x_t])  # New gate
h_t = (1 - z_t) * n_t + z_t * h_{t-1}  # Hidden state

gru = nn.GRU(input_size=128, hidden_size=256, num_layers=2,
             batch_first=True, bidirectional=True)
output, h_n = gru(input_seq)

---

Transformer Layers

nn.MultiheadAttention

Computes scaled dot-product attention across multiple heads:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, dropout=0.1,
                            batch_first=True)
# Self-attention: query = key = value
attn_output, attn_weights = mha(query, key, value, key_padding_mask=mask)

nn.TransformerEncoderLayer

One layer of a Transformer encoder (self-attention + feedforward):

encoder_layer = nn.TransformerEncoderLayer(
    d_model=512, nhead=8, dim_feedforward=2048,
    dropout=0.1, activation='gelu', batch_first=True,
    norm_first=True  # Pre-norm (more stable training)
)

nn.TransformerEncoder

Stack of TransformerEncoderLayers:

encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
output = encoder(src, src_key_padding_mask=padding_mask)

---

Embedding Layers

nn.Embedding

Lookup table that maps integer indices to dense vectors:

# vocab_size=10000, embedding_dim=256
embed = nn.Embedding(num_embeddings=10000, embedding_dim=256, padding_idx=0)
# Input: (batch, seq_len) of integer indices
# Output: (batch, seq_len, embedding_dim)
token_ids = torch.tensor([[1, 45, 234, 0, 0]])  # 0 = padding
embeddings = embed(token_ids)  # shape: (1, 5, 256)

padding_idx: The embedding at this index is always zero and is not updated.

nn.EmbeddingBag

More efficient when you need to sum/mean/max embeddings (e.g., bag-of-words):

embed_bag = nn.EmbeddingBag(num_embeddings=10000, embedding_dim=256, mode='mean')
# Returns one vector per "bag" — no need to manually average

---

Loss Functions

nn.CrossEntropyLoss

For multi-class classification. Combines LogSoftmax + NLLLoss.

Formula: loss = -log(exp(x_y) / sum(exp(x_j))) where y is the true class.

criterion = nn.CrossEntropyLoss()
# logits: (batch, num_classes) — RAW scores, NOT softmax
# target: (batch,) — class indices (integers)
loss = criterion(logits, targets)

nn.BCEWithLogitsLoss

For binary or multi-label classification. Combines Sigmoid + BCELoss.

Formula: loss = -[ylog(sigmoid(x)) + (1-y)log(1-sigmoid(x))]

criterion = nn.BCEWithLogitsLoss()
# logits: (batch, num_labels) — RAW scores
# target: (batch, num_labels) — 0.0 or 1.0

nn.MSELoss

Mean Squared Error for regression.

Formula: loss = mean((y_pred - y_true)^2)

criterion = nn.MSELoss()

nn.L1Loss

Mean Absolute Error for regression.

Formula: loss = mean(|y_pred - y_true|)

nn.HuberLoss (Smooth L1)

Combination of L1 and L2 — less sensitive to outliers than MSE.

Formula: L2 for |error| < delta, L1 for |error| >= delta.

criterion = nn.HuberLoss(delta=1.0)

nn.KLDivLoss

Kullback-Leibler Divergence — measures how one distribution differs from another.

Formula: loss = y_true * (log(y_true) - x)

criterion = nn.KLDivLoss(reduction='batchmean', log_target=False)
# Input must be log-probabilities!

nn.TripletMarginLoss

For metric learning with (anchor, positive, negative) triplets.

Formula: loss = max(d(anchor, positive) - d(anchor, negative) + margin, 0)

criterion = nn.TripletMarginLoss(margin=1.0)
loss = criterion(anchor, positive, negative)

nn.CosineEmbeddingLoss

Measures cosine similarity between pairs.

criterion = nn.CosineEmbeddingLoss(margin=0.0)
# target: +1 (similar) or -1 (dissimilar)
loss = criterion(x1, x2, target)

---

Functional API

torch.nn.functional (commonly imported as F) provides the same operations as nn.Module layers but as pure functions without stored state.

import torch.nn.functional as F

# Module version (has stored parameters):
relu_module = nn.ReLU()
output = relu_module(x)

# Functional version (stateless):
output = F.relu(x)

When to use Module vs Functional:

• Use Module when the operation has learnable parameters (Linear, Conv, BatchNorm)
• Use Module when the operation has different train/eval behavior (Dropout, BatchNorm)
• Use Functional for stateless operations (relu, softmax in forward pass)
• Use Functional when you need the operation in a custom forward pass without

wanting to register it as a sub-module

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)  # Module: has parameters
        self.bn = nn.BatchNorm2d(64)     # Module: has state (running stats)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = F.relu(x)         # Functional: no state needed
        x = F.dropout(x, p=0.5, training=self.training)  # Must pass training flag!
        return x

---

Weight Initialization

Proper initialization prevents vanishing/exploding gradients at the start of training.

Xavier (Glorot) Initialization

Designed for sigmoid/tanh activations. Keeps variance constant across layers.

Formula: W ~ Uniform(-sqrt(6/(fan_in+fan_out)), sqrt(6/(fan_in+fan_out)))

nn.init.xavier_uniform_(layer.weight)
nn.init.xavier_normal_(layer.weight)

Kaiming (He) Initialization

Designed for ReLU activations. Accounts for the fact that ReLU zeros out half the inputs.

Formula: W ~ Normal(0, sqrt(2/fan_in))

nn.init.kaiming_uniform_(layer.weight, mode='fan_in', nonlinearity='relu')
nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')

When to use each:

• Xavier: sigmoid, tanh activations
• Kaiming: ReLU, LeakyReLU activations
• Normal/Uniform: When you want simple random initialization
• Zeros: For biases (common default)
• Ones: For normalization layer weights

def init_weights(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, nonlinearity='relu')
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')

model.apply(init_weights)  # Recursively apply to all modules

---

Hooks

Hooks let you inspect or modify intermediate values during forward/backward passes without changing the model code.

Forward Hook

Called after forward() completes. Receives (module, input, output).

def print_output_shape(module, input, output):
    print(f"{module.__class__.__name__}: output shape = {output.shape}")

hook_handle = model.layer1.register_forward_hook(print_output_shape)
# Later: hook_handle.remove()

Forward Pre-Hook

Called before forward(). Receives (module, input). Can modify the input.

def modify_input(module, args):
    # args is a tuple of inputs
    return (args[0] * 2,)  # Double the input

handle = model.layer1.register_forward_pre_hook(modify_input)

Backward Hook

Called during backward pass. Can inspect or modify gradients.

def print_grad(module, grad_input, grad_output):
    print(f"Grad output norm: {grad_output[0].norm()}")

handle = model.layer1.register_full_backward_hook(print_grad)

Common Use Cases:

• Feature extraction from intermediate layers
• Gradient visualization/debugging
• Gradient clipping per layer
• Activation statistics for debugging training

---

State Dict

The state dict is an OrderedDict mapping parameter/buffer names to tensors. It's the standard way to save and load models.

Saving and Loading

# Save
torch.save(model.state_dict(), 'model_weights.pth')

# Load
model = MyModel()  # Create model with same architecture
model.load_state_dict(torch.load('model_weights.pth', weights_only=True))

Partial Loading (strict=False)

When architectures don't match exactly:

# Load only matching keys, ignore missing/unexpected keys
state_dict = torch.load('pretrained.pth', weights_only=True)
model.load_state_dict(state_dict, strict=False)

Inspecting State Dict

state_dict = model.state_dict()
for key, tensor in state_dict.items():
    print(f"{key}: shape={tensor.shape}, dtype={tensor.dtype}")

Modifying Before Loading

# Remove prefix from keys (e.g., from DataParallel)
state_dict = torch.load('model.pth', weights_only=True)
new_state_dict = {}
for k, v in state_dict.items():
    new_key = k.replace('module.', '')  # Remove DataParallel prefix
    new_state_dict[new_key] = v
model.load_state_dict(new_state_dict)

---

Summary

---

📓 Open Notebook — Interactive version of this module

Module 05: Optimizers and Learning Rate Schedulers

Table of Contents

• Optimizer Fundamentals
• SGD — Stochastic Gradient Descent
• Adam — Adaptive Moment Estimation
• AdamW — Decoupled Weight Decay
• Other Optimizers
• Learning Rate Schedulers
• Practical Advice
• Gradient Clipping
• Compiled Optimizers

---

Optimizer Fundamentals

An optimizer updates model parameters to minimize the loss function. In PyTorch, all optimizers inherit from torch.optim.Optimizer and share this interface:

import torch.optim as optim

# Create optimizer — pass it the parameters to optimize
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
for batch in dataloader:
    optimizer.zero_grad()       # Clear old gradients
    output = model(batch)       # Forward pass
    loss = criterion(output, target)
    loss.backward()             # Compute gradients
    optimizer.step()            # Update parameters

Parameter Groups

Optimizers support different settings for different parameter groups:

optimizer = optim.SGD([
    {'params': model.backbone.parameters(), 'lr': 0.001},   # Lower LR for backbone
    {'params': model.head.parameters(), 'lr': 0.01},         # Higher LR for head
], momentum=0.9, weight_decay=1e-4)

Optimizer State Dict

# Save optimizer state (for training resumption)
torch.save(optimizer.state_dict(), 'optimizer.pth')

# Load optimizer state
optimizer.load_state_dict(torch.load('optimizer.pth'))

The state dict contains:

• state: per-parameter state (momentum buffers, Adam moments, step counts)
• param_groups: hyperparameters (lr, momentum, weight_decay, etc.)

---

SGD

Stochastic Gradient Descent is the simplest optimizer but with momentum is still competitive.

Vanilla SGD

Update rule: theta = theta - lr * gradient

SGD with Momentum

Momentum accumulates past gradients to smooth updates and escape shallow local minima.

Classical momentum:

v_t = momentum * v_{t-1} + gradient_t
theta_t = theta_{t-1} - lr * v_t

Nesterov momentum (look-ahead):

v_t = momentum * v_{t-1} + gradient(theta - lr * momentum * v_{t-1})
theta_t = theta_{t-1} - lr * v_t

Nesterov is generally better — it evaluates the gradient at the "look-ahead" position, which provides better correction.

SGD with Weight Decay (L2 Regularization)

gradient_with_wd = gradient + weight_decay * theta
theta = theta - lr * gradient_with_wd

Weight decay adds a penalty proportional to parameter magnitude, pushing weights toward zero.

optimizer = optim.SGD(
    model.parameters(),
    lr=0.1,
    momentum=0.9,
    weight_decay=1e-4,
    nesterov=True
)

When to use SGD:

• Computer vision tasks (often gives better generalization than Adam)
• When you have a good learning rate schedule
• Large batch training
• When you want maximum control

---

Adam

Adam (Adaptive Moment Estimation) adapts the learning rate for each parameter based on first and second moments of the gradients.

Algorithm Step-by-Step

Initialize: m_0 = 0, v_0 = 0, t = 0

For each step:
  t = t + 1
  g_t = gradient at step t
  
  # Update biased first moment estimate (mean of gradients)
  m_t = beta1 * m_{t-1} + (1 - beta1) * g_t
  
  # Update biased second moment estimate (mean of squared gradients)
  v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2
  
  # Bias correction (crucial in early steps)
  m_hat_t = m_t / (1 - beta1^t)
  v_hat_t = v_t / (1 - beta2^t)
  
  # Update parameters
  theta_t = theta_{t-1} - lr * m_hat_t / (sqrt(v_hat_t) + eps)

Why bias correction? Since m_0 = 0 and v_0 = 0, the estimates are biased toward zero in early training. Dividing by (1 - beta^t) corrects this — as t grows large, the correction approaches 1 and has no effect.

Default hyperparameters:

• lr = 0.001
• beta1 = 0.9 (momentum for first moment)
• beta2 = 0.999 (momentum for second moment)
• eps = 1e-8 (numerical stability)

optimizer = optim.Adam(
    model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0  # This is L2 regularization, NOT decoupled
)

When to use Adam:

• Default choice for most tasks
• NLP, transformers
• When you don't want to tune the learning rate carefully
• Faster convergence early in training

---

AdamW

AdamW fixes a subtle but important problem with Adam's weight decay implementation.

The Problem with Adam + Weight Decay

In standard Adam with weight_decay, the weight decay is applied to the gradient:

g_t = gradient + weight_decay * theta   # L2 regularization added to gradient

But then Adam's adaptive learning rate scales this differently per parameter, effectively applying DIFFERENT regularization strengths to different parameters. This breaks the intended uniform regularization.

Decoupled Weight Decay (AdamW)

AdamW applies weight decay directly to the parameters AFTER the Adam update:

theta_t = theta_{t-1} - lr * adam_update - lr * weight_decay * theta_{t-1}

This means every parameter gets the same relative shrinkage regardless of its gradient history.

Impact: AdamW generalizes better than Adam+L2, especially for transformers and large models. It's now the default for most modern training.

optimizer = optim.AdamW(
    model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01  # Decoupled! Common value: 0.01 to 0.1
)

When to use AdamW:

• Transformers (GPT, BERT, ViT, etc.)
• When using weight decay (almost always)
• Default recommendation for most tasks today

---

Other Optimizers

Adagrad

Adapts learning rate based on accumulated squared gradients. Good for sparse data but learning rate decays to zero over time (problematic for long training).

optimizer = optim.Adagrad(model.parameters(), lr=0.01)

RMSprop

Fixes Adagrad's decaying learning rate by using exponential moving average of squared gradients. Predecessor to Adam (Adam = RMSprop + momentum).

optimizer = optim.RMSprop(model.parameters(), lr=0.01, alpha=0.99)

Adadelta

Similar to RMSprop but eliminates the need to set an initial learning rate.

optimizer = optim.Adadelta(model.parameters(), lr=1.0, rho=0.9)

LBFGS

Limited-memory BFGS — a quasi-Newton method. Uses second-order information (Hessian approximation). Much more expensive per step but converges in fewer steps. Requires a closure.

optimizer = optim.LBFGS(model.parameters(), lr=1.0, max_iter=20)

def closure():
    optimizer.zero_grad()
    output = model(input)
    loss = criterion(output, target)
    loss.backward()
    return loss

optimizer.step(closure)

RAdam (Rectified Adam)

Adam with a variance-rectification term that provides an automatic warmup effect. Addresses the high variance of Adam in early training.

Muon

A newer optimizer designed for large language models. Uses momentum and sign-based updates for more stable training at scale.

Adafactor

Memory-efficient alternative to Adam. Factorizes the second moment matrix to reduce memory from O(mn) to O(m+n). Popular for training very large models.

---

Learning Rate Schedulers

Schedulers adjust the learning rate during training. The general pattern:

optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

for epoch in range(100):
    train(...)
    scheduler.step()  # Update LR after each epoch

StepLR

Decay by gamma every step_size epochs.

# LR: 0.1 -> 0.01 (at epoch 30) -> 0.001 (at epoch 60)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

MultiStepLR

Decay at specific milestones.

# LR: 0.1 -> 0.01 (at epoch 30) -> 0.001 (at epoch 80)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 80], gamma=0.1)

ExponentialLR

Multiply LR by gamma every epoch.

scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)

CosineAnnealingLR

Smoothly decays LR following a cosine curve from initial LR to eta_min.

# Decays from lr to eta_min over T_max epochs
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

OneCycleLR

Implements the 1-cycle policy: ramp up LR, then decay. Often achieves best results.

scheduler = optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=0.01, total_steps=1000,
    pct_start=0.3,  # 30% warmup, 70% decay
    anneal_strategy='cos'
)
# Note: step() after each BATCH, not each epoch!

ReduceLROnPlateau

Reduce LR when a metric plateaus (most practical for validation loss).

scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=10
)
# Must pass the metric value:
scheduler.step(val_loss)

LinearLR

Linearly scale LR from start_factor to end_factor over total_iters.

# Warmup: LR goes from 0.001*0.1 to 0.001 over 10 epochs
scheduler = optim.lr_scheduler.LinearLR(
    optimizer, start_factor=0.1, end_factor=1.0, total_iters=10
)

SequentialLR (Warmup + Cosine)

Chain multiple schedulers together:

warmup = optim.lr_scheduler.LinearLR(optimizer, start_factor=0.1, total_iters=10)
cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=90)
scheduler = optim.lr_scheduler.SequentialLR(
    optimizer, schedulers=[warmup, cosine], milestones=[10]
)

CosineAnnealingWarmRestarts

Cosine annealing with periodic restarts (warm restarts increase exploration).

scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer, T_0=10, T_mult=2  # First restart at 10, then 20, 40, ...
)

---

Practical Advice

Which Optimizer for Which Task?

Learning Rate Selection

• Learning Rate Finder: Start very small, increase exponentially, plot loss vs LR.

Choose LR where loss is decreasing steepest (typically 1/10 of the minimum).

• Rule of thumb: If training is unstable, reduce LR by 3-10x.
• Linear scaling rule: When increasing batch size by N, multiply LR by N (with warmup).

Warmup Strategies

Warmup is crucial for:

• Large learning rates
• Large batch sizes
• Transformer training

Common approach: Linear warmup for 5-10% of total training, then cosine decay.

---

Gradient Clipping

Prevents exploding gradients by limiting gradient magnitude.

clip_grad_norm_ (recommended)

Clips the total norm of all gradients:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

This scales all gradients uniformly to ensure the total L2 norm doesn't exceed max_norm. It preserves gradient direction.

clip_grad_value_

Clips each gradient element independently:

torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

Each element is clamped to [-clip_value, clip_value]. Can change gradient direction.

Usage in training loop:

for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch).sum()
    loss.backward()
    # Clip AFTER backward(), BEFORE step()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

---

Compiled Optimizers

PyTorch 2.0+ can compile optimizers with torch.compile for significant speedups:

optimizer = optim.AdamW(model.parameters(), lr=0.001)
# The optimizer step is fused and optimized
# This happens automatically when the model is compiled

@torch.compile
def train_step(model, x, y):
    output = model(x)
    loss = F.cross_entropy(output, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return loss

Benefits:

• Fused kernels: multiple optimizer ops become one kernel launch
• Reduced memory traffic
• Horizontal fusion across parameter groups
• Can give 10-20% speedup on optimizer step

---

Summary

---

📓 Open Notebook — Interactive version of this module

Module 06: Data Loading in PyTorch

Table of Contents

• Dataset Class
• IterableDataset
• Built-in Datasets
• DataLoader
• Worker Processes
• Custom Collate Functions
• Samplers
• Data Augmentation
• Memory-Mapped Datasets
• Best Practices

---

Dataset Class

The torch.utils.data.Dataset is the base class for all map-style datasets. You implement two methods:

• __len__(): Returns the total number of samples
• __getitem__(index): Returns one sample given an index

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

Key properties:

• Supports random access by index
• Length is known ahead of time
• DataLoader can shuffle efficiently
• Can be split into train/val subsets

When data doesn't fit in memory:

class LargeFileDataset(Dataset):
    def __init__(self, file_paths):
        self.file_paths = file_paths
    
    def __len__(self):
        return len(self.file_paths)
    
    def __getitem__(self, idx):
        # Load only one file at a time
        data = load_file(self.file_paths[idx])
        return process(data)

---

IterableDataset

For streaming data or when random access is not possible/efficient:

from torch.utils.data import IterableDataset

class StreamDataset(IterableDataset):
    def __init__(self, url):
        self.url = url
    
    def __iter__(self):
        # Yield samples one at a time
        for line in open_stream(self.url):
            yield process(line)

Use cases:

• Reading from network streams
• Very large files that can only be read sequentially
• Data generated on-the-fly
• Databases with cursor-based access

Key differences from map-style Dataset:

• No __len__() — size may be unknown
• No __getitem__() — no random access
• DataLoader cannot shuffle (must shuffle upstream)
• Multi-worker requires careful work splitting

Multi-worker IterableDataset:

class ShardedDataset(IterableDataset):
    def __init__(self, file_list):
        self.file_list = file_list
    
    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            # Single-process loading
            files = self.file_list
        else:
            # Split files among workers
            per_worker = len(self.file_list) // worker_info.num_workers
            start = worker_info.id * per_worker
            end = start + per_worker
            files = self.file_list[start:end]
        
        for f in files:
            for sample in read_file(f):
                yield sample

---

Built-in Datasets

TensorDataset

Wraps tensors — each sample is a tuple indexed along the first dimension:

from torch.utils.data import TensorDataset

X = torch.randn(1000, 20)
y = torch.randint(0, 5, (1000,))
dataset = TensorDataset(X, y)
# dataset[0] returns (X[0], y[0])

ConcatDataset

Concatenates multiple datasets end-to-end:

from torch.utils.data import ConcatDataset

combined = ConcatDataset([dataset_train, dataset_extra])
# len(combined) == len(dataset_train) + len(dataset_extra)

Subset

Selects a subset of a dataset by indices:

from torch.utils.data import Subset

# Manual train/val split
indices = torch.randperm(len(dataset))
train_set = Subset(dataset, indices[:800])
val_set = Subset(dataset, indices[800:])

random_split

Convenience function for splitting:

from torch.utils.data import random_split

train_set, val_set, test_set = random_split(
    dataset, [0.7, 0.15, 0.15],  # Fractions
    generator=torch.Generator().manual_seed(42)
)

---

DataLoader

The DataLoader is the workhorse that turns a Dataset into an iterable of batches:

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=32,          # Samples per batch
    shuffle=True,           # Randomize order each epoch
    num_workers=4,          # Parallel data loading processes
    pin_memory=True,        # Speed up CPU->GPU transfer
    drop_last=False,        # Drop incomplete final batch?
    persistent_workers=True,# Keep workers alive between epochs
    prefetch_factor=2,      # Batches to prefetch per worker
)

for batch_X, batch_y in loader:
    # batch_X: (32, ...), batch_y: (32, ...)
    output = model(batch_X)
    loss = criterion(output, batch_y)
    ...

Key Parameters:

batch_size: Number of samples per batch. Larger = faster training but more memory. Typical values: 16, 32, 64, 128, 256.

shuffle: Randomize sample order each epoch. Always True for training, False for validation/test (for reproducible evaluation).

num_workers: Number of subprocess workers for data loading. Set to 0 for debugging (main process only). Typical: 2-8 depending on CPU cores and data complexity.

pin_memory: Pre-allocates batch tensors in pinned (page-locked) memory. Makes CPU->GPU transfers faster. Always True when training on GPU.

drop_last: If True, drops the last batch if it's smaller than batch_size. Important for BatchNorm (needs consistent batch size) and distributed training.

persistent_workers: Keeps worker processes alive across epochs instead of respawning. Saves the cost of worker initialization. Use with num_workers > 0.

prefetch_factor: Number of batches each worker pre-loads. Higher values use more memory but can hide I/O latency better. Default is 2.

---

Worker Processes

How Workers Work

When num_workers > 0, DataLoader spawns separate processes (not threads!) that:

• Receive batch indices from the main process
• Call dataset.__getitem__() for each index
• Apply the collate function to form a batch
• Send the batch back to the main process via shared memory

Why num_workers > 0 is Faster

• Data loading (disk I/O, decompression, augmentation) happens in parallel
• While GPU trains on batch N, workers prepare batch N+1, N+2, ...
• Python's GIL doesn't affect separate processes

Common Pitfalls

• Too many workers: More workers != always faster. Each worker duplicates the

dataset object in memory. Start with num_workers=4 and tune.

• Fork vs Spawn: On Linux, workers are forked (fast, shares memory). On macOS/Windows,

workers are spawned (slow startup, separate memory). Set: ``python torch.multiprocessing.set_start_method('spawn') # If needed ``

• Random state in workers: Each worker gets the same random seed by default.

Use worker_init_fn to set different seeds: ```python def worker_init_fn(worker_id): seed = torch.initial_seed() % 2**32 numpy.random.seed(seed + worker_id)

loader = DataLoader(..., worker_init_fn=worker_init_fn) ```

• File handles: If your dataset opens files, each worker opens its own copy.

With many workers, you can run out of file descriptors.

• Shared memory limits: Workers send data via shared memory. Very large batches

can exhaust /dev/shm. Increase shared memory or reduce batch size.

---

Custom Collate Functions

The collate function converts a list of samples into a batch. The default collator stacks tensors, but custom collation handles variable-length sequences:

def custom_collate(batch):
    """Pad variable-length sequences to same length."""
    sequences, labels = zip(*batch)
    # Pad sequences to max length in this batch
    lengths = [len(s) for s in sequences]
    max_len = max(lengths)
    padded = torch.zeros(len(sequences), max_len)
    for i, (seq, length) in enumerate(zip(sequences, lengths)):
        padded[i, :length] = seq
    labels = torch.tensor(labels)
    lengths = torch.tensor(lengths)
    return padded, labels, lengths

loader = DataLoader(dataset, batch_size=32, collate_fn=custom_collate)

Common custom collate patterns:

• Padding variable-length sequences
• Creating attention masks
• Handling nested data structures (dicts, lists)
• Filtering out None values (corrupt samples)

---

Samplers

Samplers control the order in which indices are provided to the Dataset.

SequentialSampler

Iterates indices 0, 1, 2, ..., N-1. Used when shuffle=False.

RandomSampler

Random permutation of indices. Used when shuffle=True.

from torch.utils.data import RandomSampler

# With replacement (for bootstrap):
sampler = RandomSampler(dataset, replacement=True, num_samples=10000)

WeightedRandomSampler

For handling class imbalance — oversamples rare classes:

from torch.utils.data import WeightedRandomSampler

# Assign weight to each sample (higher weight = sampled more often)
class_counts = [1000, 100, 50]  # Imbalanced classes
class_weights = 1.0 / torch.tensor(class_counts, dtype=torch.float)
sample_weights = class_weights[labels]  # Weight per sample

sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(dataset),
    replacement=True
)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
# Note: cannot use shuffle=True with a sampler

SubsetRandomSampler

Random sampling from a fixed set of indices (for train/val splits):

from torch.utils.data import SubsetRandomSampler

indices = list(range(len(dataset)))
train_indices = indices[:800]
val_indices = indices[800:]

train_loader = DataLoader(dataset, batch_size=32,
                          sampler=SubsetRandomSampler(train_indices))
val_loader = DataLoader(dataset, batch_size=32,
                        sampler=SubsetRandomSampler(val_indices))

BatchSampler

Groups indices into batches (useful for custom batching logic):

from torch.utils.data import BatchSampler, SequentialSampler

# Create batches of similar-length sequences (for efficient padding)
sampler = BatchSampler(
    SequentialSampler(dataset),
    batch_size=32,
    drop_last=False
)

DistributedSampler

For multi-GPU training — ensures each GPU sees different data:

from torch.utils.data import DistributedSampler

sampler = DistributedSampler(
    dataset,
    num_replicas=world_size,  # Number of GPUs
    rank=rank,                # This GPU's index
    shuffle=True
)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)

# IMPORTANT: Must call set_epoch each epoch for proper shuffling
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)
    for batch in loader:
        ...

---

Data Augmentation

Data augmentation applies random transformations during training to increase data diversity and reduce overfitting.

torchvision.transforms.v2

The modern transforms API for images:

from torchvision.transforms import v2

train_transform = v2.Compose([
    v2.RandomResizedCrop(224),
    v2.RandomHorizontalFlip(p=0.5),
    v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    v2.RandomRotation(15),
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

val_transform = v2.Compose([
    v2.Resize(256),
    v2.CenterCrop(224),
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

MixUp

Linearly interpolates between pairs of samples:

mixed_input = lambda * input_i + (1-lambda) * input_j
mixed_target = lambda * target_i + (1-lambda) * target_j

CutMix

Cuts a patch from one image and pastes it onto another:

mixed_input = mask * input_i + (1-mask) * input_j
mixed_target = lambda * target_i + (1-lambda) * target_j

Where lambda is the area ratio.

Custom Augmentations

class CustomAugmentation:
    def __init__(self, p=0.5):
        self.p = p
    
    def __call__(self, image):
        if torch.rand(1) < self.p:
            return apply_augmentation(image)
        return image

---

Memory-Mapped Datasets

For datasets too large to fit in RAM, memory-mapping lets the OS page data in/out:

import numpy as np

class MemmapDataset(Dataset):
    def __init__(self, data_path, shape, dtype='float32'):
        # Memory-map the file (doesn't load into RAM)
        self.data = np.memmap(data_path, dtype=dtype, mode='r', shape=shape)
    
    def __len__(self):
        return self.data.shape[0]
    
    def __getitem__(self, idx):
        # Only this sample is loaded into RAM
        return torch.from_numpy(self.data[idx].copy())

Benefits:

• Works with datasets much larger than RAM
• OS handles caching transparently
• Sequential access patterns are very efficient
• Multiple workers can share the same memory map

---

Best Practices

Performance Checklist

• pin_memory=True when using GPU training
• persistent_workers=True with num_workers > 0 to avoid respawn cost
• prefetch_factor=2 (default) is usually good; increase if I/O-bound
• num_workers: Start with 4, increase until CPU is saturated or you run out of RAM
• drop_last=True for training with BatchNorm

Data Loading Bottleneck Detection

If your GPU utilization is low, data loading might be the bottleneck:

import time

for batch in loader:
    start = time.time()
    output = model(batch)
    loss.backward()
    optimizer.step()
    gpu_time = time.time() - start
    # If data loading time >> gpu_time, you're data-bound

Memory Efficiency

• Use uint8 for images until the augmentation step
• Convert to float32 only in the transform/collate
• Use memory-mapped files for very large datasets
• Consider chunked reading for CSV/parquet files

Reproducibility

# Seed everything for reproducible data loading
def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    import numpy as np
    np.random.seed(worker_seed)
    import random
    random.seed(worker_seed)

g = torch.Generator()
g.manual_seed(42)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    worker_init_fn=seed_worker,
    generator=g,
)

Common Patterns

Training with validation:

train_loader = DataLoader(train_set, batch_size=64, shuffle=True, 
                          num_workers=4, pin_memory=True)
val_loader = DataLoader(val_set, batch_size=128, shuffle=False,
                        num_workers=4, pin_memory=True)

Infinite data loader (for step-based training):

def infinite_loader(loader):
    while True:
        for batch in loader:
            yield batch

for step, batch in enumerate(infinite_loader(train_loader)):
    if step >= max_steps:
        break
    train_step(batch)

Progress bar with tqdm:

from tqdm import tqdm

for batch in tqdm(loader, desc="Training"):
    ...

---

📓 Open Notebook — Interactive version of this module

Module 07: The Complete Training Guide

Overview

Training a neural network is where theory meets practice. This module covers everything from the basic training loop to advanced techniques used by state-of-the-art models. By the end, you'll understand not just how to train models, but why each step exists and how to diagnose problems.

---

1. Basic Training Loop Anatomy

Every training loop in PyTorch follows the same fundamental pattern:

forward pass → compute loss → backward pass → optimizer step

Let's break each step down:

Step 1: Forward Pass

predictions = model(inputs)

Data flows through the model's layers. Each layer applies its transformation (matrix multiply, activation, normalization, etc.) and PyTorch records the operations in a computational graph for later backpropagation.

Step 2: Compute Loss

loss = loss_fn(predictions, targets)

The loss function measures how far the model's predictions are from the true targets. Common losses:

• nn.CrossEntropyLoss() — classification (combines LogSoftmax + NLLLoss)
• nn.MSELoss() — regression (mean squared error)
• nn.BCEWithLogitsLoss() — binary classification (numerically stable)

Step 3: Backward Pass

loss.backward()

PyTorch walks backward through the computational graph, computing the gradient of the loss with respect to every parameter that has requires_grad=True. These gradients accumulate in each parameter's .grad attribute.

Step 4: Optimizer Step

optimizer.step()

The optimizer uses the computed gradients to update the model's parameters. Different optimizers (SGD, Adam, AdamW) use different update rules, but all read from .grad and modify .data.

The Missing Step: zero_grad()

optimizer.zero_grad()

This clears old gradients before computing new ones. Without it, gradients accumulate across iterations (which is sometimes intentional — see gradient accumulation below).

Complete Minimal Loop

model.train()
for epoch in range(num_epochs):
    for batch_inputs, batch_targets in dataloader:
        optimizer.zero_grad()           # Clear old gradients
        predictions = model(batch_inputs)  # Forward pass
        loss = loss_fn(predictions, batch_targets)  # Compute loss
        loss.backward()                 # Compute gradients
        optimizer.step()                # Update parameters

---

2. train() vs eval() Mode

What model.train() Does

Sets the model to training mode. This affects layers that behave differently during training vs inference:

• Dropout: Randomly zeros elements during training. During eval, all

elements pass through (scaled appropriately).

• BatchNorm: During training, uses batch statistics (mean/var of current

batch) and updates running statistics. During eval, uses the accumulated running statistics.

What model.eval() Does

model.eval()

Sets the model to evaluation mode. Dropout is disabled, BatchNorm uses running statistics instead of batch statistics.

Common Mistake

# WRONG: Forgetting to switch modes
def evaluate(model, test_loader):
    # model is still in train() mode!
    # Dropout is randomly zeroing activations
    # BatchNorm is using (and updating!) batch statistics
    total_correct = 0
    for inputs, targets in test_loader:
        outputs = model(inputs)
        ...

# CORRECT:
def evaluate(model, test_loader):
    model.eval()
    with torch.no_grad():  # Also disable gradient computation
        total_correct = 0
        for inputs, targets in test_loader:
            outputs = model(inputs)
            ...
    model.train()  # Switch back after evaluation

torch.no_grad() vs model.eval()

These are different things:

• model.eval() — changes layer behavior (dropout, batchnorm)
• torch.no_grad() — disables gradient computation (saves memory/compute)

For evaluation, you typically want BOTH.

---

3. zero_grad() — Why and How

Why Gradients Accumulate

PyTorch accumulates gradients by default. After loss.backward(), the .grad attribute of each parameter ADDS to whatever was already there:

param.grad += new_gradient  # This is what happens internally

This design choice enables gradient accumulation (discussed later), but it means you must manually clear gradients each iteration.

set_to_none=True Optimization

optimizer.zero_grad(set_to_none=True)

Instead of setting gradients to zero tensors, this sets them to None. Benefits:

• Slightly less memory (no zero tensor allocated)
• Can be marginally faster
• The gradient will be lazily created on the next backward pass

This is now the default in modern PyTorch (>= 2.0). The only reason to use set_to_none=False is if your code explicitly checks if param.grad is not None.

---

4. Mixed Precision Training (AMP)

What It Is

Mixed precision training uses lower-precision floating point numbers (float16 or bfloat16) for most operations, while keeping critical operations in float32. This is faster because:

• Lower precision operations use less memory bandwidth
• Hardware (GPUs, modern CPUs) has specialized units for half-precision math
• Smaller tensors mean more data fits in cache

float16 vs float32

The torch.amp API (Modern, Device-Agnostic)

# The modern way (PyTorch 2.0+)
with torch.amp.autocast(device_type='cpu', dtype=torch.bfloat16):
    output = model(input)
    loss = loss_fn(output, target)

The autocast context manager automatically casts operations to the specified lower precision where safe, and keeps float32 where needed (e.g., loss computation, softmax, layer norm).

GradScaler (for float16 only)

float16 has a limited range. Small gradients can underflow to zero. GradScaler solves this by:

• Scaling the loss UP before backward (so gradients are larger)
• Unscaling gradients before optimizer step
• Skipping steps where gradients contain inf/nan (and reducing the scale)

scaler = torch.amp.GradScaler()

for inputs, targets in dataloader:
    optimizer.zero_grad()
    with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
        output = model(inputs)
        loss = loss_fn(output, targets)

    scaler.scale(loss).backward()     # Scaled backward
    scaler.step(optimizer)            # Unscale + step (or skip)
    scaler.update()                   # Adjust scale factor

BFloat16 Doesn't Need GradScaler

bfloat16 has the same exponent range as float32, so gradients don't underflow. You can use it without GradScaler:

with torch.amp.autocast(device_type='cpu', dtype=torch.bfloat16):
    output = model(input)
    loss = loss_fn(output, target)

loss.backward()  # No scaler needed
optimizer.step()

---

5. BFloat16 vs Float16 — When to Use Each

Float16 (FP16)

• Pros: Widely supported, maximum speed on older GPUs (V100)
• Cons: Limited range (max 65504), requires GradScaler, can overflow/underflow
• Use when: Training on older NVIDIA GPUs, inference where range is known

BFloat16 (BF16)

• Pros: Same range as float32, no GradScaler needed, more stable training
• Cons: Less precision (7 mantissa bits vs 10), requires Ampere+ GPU or modern CPU
• Use when: Training large models (LLMs), when numerical stability matters

Practical Recommendation

• For training: prefer bfloat16 if your hardware supports it
• For inference: either works, float16 slightly more precise per value
• For CPU: bfloat16 is supported on modern x86 (AMX) and ARM

---

6. Gradient Accumulation

The Problem

You want a batch size of 256 but your memory only fits 32 samples.

The Solution

Accumulate gradients over multiple mini-batches before stepping:

accumulation_steps = 8  # Effective batch = 32 * 8 = 256

for i, (inputs, targets) in enumerate(dataloader):
    # Forward + backward (gradients accumulate)
    output = model(inputs)
    loss = loss_fn(output, targets)
    loss = loss / accumulation_steps  # Scale loss!
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Why Divide the Loss?

Without division, accumulated gradients are accumulation_steps times larger than a true large-batch gradient. Dividing the loss by accumulation_steps makes the accumulated gradient equivalent to computing it over one large batch.

Mathematically: grad(L/N) summed N times = grad(sum(L_i)/N) which equals the gradient of the mean loss over all N mini-batches.

---

7. Gradient Checkpointing (Activation Checkpointing)

The Problem

During the forward pass, PyTorch saves all intermediate activations for use in the backward pass. For deep models, this uses enormous amounts of memory.

The Solution

Don't save activations — recompute them during the backward pass. This trades compute time (~33% more) for memory (~60-80% savings).

Usage

from torch.utils.checkpoint import checkpoint

class DeepModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.block1 = HeavyBlock()
        self.block2 = HeavyBlock()
        self.block3 = HeavyBlock()

    def forward(self, x):
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        x = checkpoint(self.block3, x, use_reentrant=False)
        return x

use_reentrant Parameter

• use_reentrant=False (recommended): Uses a newer, more robust implementation.

Supports all autograd features correctly.

• use_reentrant=True (legacy): The old implementation. Has subtle bugs with

certain autograd features. Being phased out.

Always use use_reentrant=False for new code.

---

8. Gradient Clipping

The Problem

Exploding gradients: when gradients become extremely large, the optimizer takes huge steps that destabilize training. Common in RNNs and deep networks.

clip_grad_norm_

Scales all gradients so their combined L2 norm doesn't exceed a threshold:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

This preserves the direction of gradients but limits their magnitude. The most common approach.

clip_grad_value_

Clamps each gradient element independently to [-value, +value]:

torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

More aggressive — changes gradient direction. Rarely used in practice.

Where to Place Clipping

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()  # Step with clipped gradients

Always clip AFTER backward, BEFORE step.

---

9. Transfer Learning

The Concept

Take a model pre-trained on a large dataset (e.g., ImageNet with 1M images) and adapt it to your smaller dataset. The pre-trained model already knows useful features (edges, textures, shapes).

Strategy 1: Feature Extraction (Freeze Everything)

model = torchvision.models.resnet18(weights='IMAGENET1K_V1')

# Freeze all parameters
for param in model.parameters():
    param.requires_grad = False

# Replace the final classifier
model.fc = nn.Linear(512, num_classes)
# Only model.fc parameters will be trained

Strategy 2: Fine-tuning with Different Learning Rates

# Lower LR for pretrained backbone, higher LR for new head
optimizer = torch.optim.Adam([
    {'params': model.features.parameters(), 'lr': 1e-5},
    {'params': model.classifier.parameters(), 'lr': 1e-3},
])

Strategy 3: Progressive Unfreezing

Start with everything frozen except the head. Gradually unfreeze layers from top to bottom:

# Epoch 1-5: Only train the head
# Epoch 5-10: Unfreeze last block + head
# Epoch 10+: Unfreeze everything with small LR

This prevents the pre-trained features from being destroyed by large early gradients.

---

10. Fine-Tuning Strategies

Full Fine-Tuning

Unfreeze everything, train with a small learning rate:

for param in model.parameters():
    param.requires_grad = True
optimizer = Adam(model.parameters(), lr=1e-5)

Best when you have enough data and the domain differs from pre-training.

Linear Probing

Only train a linear layer on top of frozen features:

for param in model.parameters():
    param.requires_grad = False
probe = nn.Linear(feature_dim, num_classes)

Good baseline to check how useful the features are.

Progressive Unfreezing

Unfreeze one layer group at a time, starting from the output end:

# Phase 1: Just the head
# Phase 2: Head + last block
# Phase 3: Head + last 2 blocks
# Phase N: Everything

Which to Choose?

---

11. Knowledge Distillation

The Concept

Train a small "student" model to mimic a large "teacher" model. The student learns from the teacher's soft predictions (probability distributions) which contain more information than hard labels.

Why Soft Targets Help

Hard label: [0, 0, 1, 0] — "this is a cat" Soft prediction: [0.01, 0.05, 0.85, 0.09] — "this is mostly cat, slightly dog"

The soft predictions encode relationships between classes that hard labels miss.

Temperature Scaling

Higher temperature makes the distribution softer (more informative):

soft_teacher = F.softmax(teacher_logits / temperature, dim=-1)
soft_student = F.log_softmax(student_logits / temperature, dim=-1)
distill_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
distill_loss = distill_loss * (temperature ** 2)  # Scale back

The temperature ** 2 factor compensates for the reduced gradient magnitude at higher temperatures.

Combined Loss

total_loss = alpha * distill_loss + (1 - alpha) * hard_loss

Typically alpha = 0.5-0.9 (emphasize soft targets).

---

12. EMA (Exponential Moving Average)

What It Is

Maintain a running average of model parameters that smooths out training noise:

ema_param = decay * ema_param + (1 - decay) * current_param

Typical decay: 0.999 or 0.9999 (very slow moving average).

Why It Helps

• Reduces variance in the final model
• Often achieves better generalization than the final checkpoint
• Used in many SOTA models (diffusion models, GANs, etc.)

Implementation

class EMA:
    def __init__(self, model, decay=0.999):
        self.decay = decay
        self.shadow = {name: p.clone().detach()
                      for name, p in model.named_parameters()}

    @torch.no_grad()
    def update(self, model):
        for name, param in model.named_parameters():
            self.shadow[name].mul_(self.decay).add_(
                param.data, alpha=1 - self.decay
            )

    def apply(self, model):
        for name, param in model.named_parameters():
            param.data.copy_(self.shadow[name])

---

13. SWA (Stochastic Weight Averaging)

The Concept

Average model weights from multiple points in training (typically from a cyclical or high constant LR schedule). This tends to find flatter minima which generalize better.

PyTorch Built-in Support

from torch.optim.swa_utils import AveragedModel, SWALR

swa_model = AveragedModel(model)
swa_scheduler = SWALR(optimizer, swa_lr=0.05)

for epoch in range(swa_start, total_epochs):
    train_one_epoch(model)
    swa_model.update_parameters(model)
    swa_scheduler.step()

# Update batch normalization statistics
torch.optim.swa_utils.update_bn(train_loader, swa_model)

SWA vs EMA

• EMA: Continuous exponential average, gives more weight to recent params
• SWA: Equal-weight average of checkpoints, typically from later training

---

14. Label Smoothing

What It Does

Instead of training against hard targets [0, 0, 1, 0], use soft targets [0.033, 0.033, 0.9, 0.033]. This prevents the model from becoming overconfident.

Formula

smooth_target = (1 - smoothing) * one_hot + smoothing / num_classes

With smoothing=0.1 and 4 classes:

• True class: 0.9 + 0.1/4 = 0.925
• Other classes: 0.1/4 = 0.025

PyTorch Implementation

# Built-in support in CrossEntropyLoss
loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)

When to Use

• Large models prone to overconfidence
• When calibrated probabilities matter (not just accuracy)
• Generally helps with 0.05-0.1 smoothing; higher can hurt

---

15. Early Stopping

The Pattern

Stop training when validation loss stops improving to prevent overfitting:

class EarlyStopping:
    def __init__(self, patience=10, min_delta=0.001):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = float('inf')

    def __call__(self, val_loss):
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            return False  # Continue training
        self.counter += 1
        return self.counter >= self.patience  # Stop if patience exceeded

Best Practice

Always save the model at the best validation loss, not the final epoch:

if val_loss < best_val_loss:
    best_val_loss = val_loss
    torch.save(model.state_dict(), 'best_model.pt')

---

16. Logging and Monitoring

What to Track

• Training loss (per batch and per epoch average)
• Validation loss (per epoch)
• Learning rate (especially with schedulers)
• Gradient norm (detect exploding/vanishing gradients)
• Parameter statistics (weight magnitudes per layer)

Simple Logging Pattern

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (inputs, targets) in enumerate(train_loader):
        loss = train_step(inputs, targets)
        running_loss += loss.item()

        if (i + 1) % log_every == 0:
            avg_loss = running_loss / log_every
            print(f"Epoch {epoch}, Step {i+1}, Loss: {avg_loss:.4f}")
            running_loss = 0.0

    val_loss = evaluate(model, val_loader)
    print(f"Epoch {epoch}, Val Loss: {val_loss:.4f}")

When to Save Checkpoints

• After each epoch (for resume capability)
• When validation metric improves (for best model)
• At fixed intervals for long training runs

---

17. Reproducibility

The Full Reproducibility Recipe

import torch
import numpy as np
import random

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed(42)

# For fully deterministic operations
torch.use_deterministic_algorithms(True)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

Caveats

• torch.use_deterministic_algorithms(True) may raise errors for ops without

deterministic implementations

• Setting benchmark = False can slow down training
• DataLoader with num_workers > 0 needs worker seeding:

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

dataloader = DataLoader(
    dataset,
    worker_init_fn=seed_worker,
    generator=torch.Generator().manual_seed(42),
)

---

18. Common Training Debugging

Loss Not Decreasing

• Learning rate too high: Loss oscillates wildly. Try 10x smaller.
• Learning rate too low: Loss decreases extremely slowly. Try 10x larger.
• Bug in data pipeline: Verify labels match inputs.
• Wrong loss function: Ensure loss matches the task (e.g., CrossEntropy for

classification needs raw logits, not softmax outputs).

• Model too small: May lack capacity to fit even training data.

NaN Loss

• Learning rate too high: Gradients explode. Lower LR or add clipping.
• Division by zero: Check for zero denominators in custom loss.
• Log of zero/negative: Ensure inputs to log() are positive.
• Overflow in float16: Use GradScaler or switch to bfloat16.

Overfitting (Train Loss Low, Val Loss High)

• Add regularization (dropout, weight decay)
• Reduce model size
• Add data augmentation
• Use early stopping
• Get more data

Underfitting (Both Losses High)

• Increase model capacity (more layers, wider layers)
• Train longer
• Reduce regularization
• Check for bugs in the model architecture
• Verify the task is learnable with this architecture

Gradient Debugging

# Check for vanishing/exploding gradients
for name, param in model.named_parameters():
    if param.grad is not None:
        grad_norm = param.grad.norm()
        if grad_norm == 0:
            print(f"WARNING: Zero gradient in {name}")
        elif grad_norm > 100:
            print(f"WARNING: Large gradient in {name}: {grad_norm}")

---

19. Putting It All Together

A production-ready training loop combines many of these techniques:

def train(model, train_loader, val_loader, config):
    optimizer = AdamW(model.parameters(), lr=config.lr, weight_decay=0.01)
    scheduler = CosineAnnealingLR(optimizer, T_max=config.epochs)
    early_stop = EarlyStopping(patience=10)
    ema = EMA(model, decay=0.999)

    for epoch in range(config.epochs):
        model.train()
        for i, (inputs, targets) in enumerate(train_loader):
            with torch.amp.autocast('cpu', dtype=torch.bfloat16):
                output = model(inputs)
                loss = loss_fn(output, targets)
                loss = loss / config.accumulation_steps

            loss.backward()

            if (i + 1) % config.accumulation_steps == 0:
                clip_grad_norm_(model.parameters(), max_norm=1.0)
                optimizer.step()
                optimizer.zero_grad()
                ema.update(model)

        scheduler.step()
        val_loss = evaluate(model, val_loader)

        if early_stop(val_loss):
            break

---

Summary

---

📓 Open Notebook — Interactive version of this module

Module 08: torch.compile — The Complete Guide

Overview

torch.compile is PyTorch's JIT compiler that makes your models run faster without changing your code. Introduced in PyTorch 2.0, it captures your Python model into an optimized graph and generates efficient machine code.

---

1. What is torch.compile?

The Problem

Standard PyTorch executes operations one-by-one ("eager mode"). Each operation:

• Dispatches from Python to C++
• Launches a kernel
• Reads/writes memory
• Returns to Python

This means:

• Python overhead on every operation
• No opportunity to fuse adjacent operations
• Suboptimal memory access patterns

The Solution

torch.compile captures a sequence of operations, optimizes them as a group, and generates fused kernels that minimize memory traffic and Python overhead.

Expected Speedup

• Typical: 20-50% faster on GPU, 10-30% on CPU
• Best case: 2-3x for memory-bound models (transformers)
• Worst case: No speedup if the model has many graph breaks

---

2. How It Works — The 3 Stages

Stage 1: TorchDynamo (Graph Capture)

TorchDynamo intercepts Python bytecode execution and traces the operations into an FX graph (a simple intermediate representation).

How it works (simplified):

• Your function runs normally
• Dynamo watches which torch operations are called
• It records them into a graph
• For subsequent calls with the same "shape signature," it replays the graph

Guards: Dynamo records assumptions about inputs (dtype, device, shape). If an assumption is violated on a later call, the graph is invalidated and Dynamo retraces (recompiles) the function.

Graph breaks: When Dynamo encounters something it can't trace (like print() with a tensor, or data-dependent control flow), it "breaks" the graph — splits execution into: compiled-graph-1 → Python → compiled-graph-2.

Stage 2: AOTAutograd (Ahead-of-Time Autograd)

After capturing the forward graph, AOTAutograd:

• Traces the backward pass as well (ahead of time)
• Partitions into a forward graph and backward graph
• Both are passed to the backend for optimization

This is important because the backward pass can also be optimized and fused.

Stage 3: TorchInductor (Code Generation)

Inductor takes the graph and generates optimized code:

• For GPU: Generates Triton kernels (a Python-based GPU programming language)
• For CPU: Generates C++/OpenMP code

Key optimizations:

• Operator fusion: Combine multiple ops into one kernel (e.g., linear + relu)
• Memory planning: Reuse memory buffers, minimize allocations
• Layout optimization: Choose best memory layout for the hardware

---

3. Basic Usage

Three ways to use torch.compile:

# Method 1: Compile a model
compiled_model = torch.compile(model)
output = compiled_model(input)

# Method 2: Decorator on a function
@torch.compile
def my_function(x, y):
    return torch.matmul(x, y) + x

# Method 3: Compile a specific function
compiled_fn = torch.compile(my_function)

First Call is Slow

The first call triggers compilation (tracing + code generation). Subsequent calls with the same input shapes are fast:

compiled_model = torch.compile(model)

# First call: SLOW (compiles)
output = compiled_model(input_batch_1)

# Second call: FAST (uses compiled code)
output = compiled_model(input_batch_2)

---

4. Compilation Modes

torch.compile(model, mode="default")         # Balanced
torch.compile(model, mode="reduce-overhead") # Minimize framework overhead
torch.compile(model, mode="max-autotune")    # Maximum optimization effort

"default"

• Balanced between compilation time and runtime performance
• Uses a reasonable set of optimizations
• Good starting point

"reduce-overhead"

• Uses CUDA graphs to eliminate kernel launch overhead
• Best for models with many small kernels
• May increase memory usage
• GPU only

"max-autotune"

• Tries many kernel implementations and picks the fastest
• Much longer compilation time
• Best runtime performance
• Good for production deployment after development

Comparison

---

5. Graph Breaks

What Causes Graph Breaks

A graph break splits the compiled region into multiple segments, with Python execution in between. This reduces optimization opportunities.

Common causes:

• print() with tensor values — requires Python execution
• Unsupported Python operations — certain builtins
• Data-dependent control flow — if tensor.item() > 0:
• Calling non-compilable functions — some third-party code
• In-place operations on views (in some cases)
• Python side-effects — logging, writing to files

How to Find Graph Breaks

# Method 1: explain()
explanation = torch._dynamo.explain(model)(sample_input)
print(explanation)

# Method 2: fullgraph=True raises an error on any break
compiled = torch.compile(model, fullgraph=True)
try:
    compiled(input)
except Exception as e:
    print(f"Graph break: {e}")

# Method 3: Logging
import logging
torch._logging.set_logs(graph_breaks=True)

How to Fix Graph Breaks

# BAD: print causes graph break
def forward(self, x):
    x = self.linear(x)
    print(f"Shape: {x.shape}")  # GRAPH BREAK!
    return self.relu(x)

# GOOD: remove print or use torch._dynamo.config.suppress_errors
def forward(self, x):
    x = self.linear(x)
    return self.relu(x)

# BAD: data-dependent control flow
def forward(self, x):
    if x.sum() > 0:  # GRAPH BREAK! (value depends on data)
        return x * 2
    return x

# GOOD: use torch.where for data-dependent logic
def forward(self, x):
    return torch.where(x.sum() > 0, x * 2, x)

---

6. fullgraph=True

Forces compilation of the entire function as a single graph. If there would be any graph break, compilation fails with an error instead of silently degrading.

@torch.compile(fullgraph=True)
def my_fn(x):
    # This MUST be fully traceable — no graph breaks allowed
    return x.sin() + x.cos()

Use this when:

• You want maximum performance (no graph breaks = fully optimized)
• You want to catch non-compilable code early
• In production code where you've already fixed all breaks

---

7. Dynamic Shapes

The Problem: Recompilation

By default, torch.compile captures the exact shapes of inputs. If shapes change, it must recompile:

compiled_fn = torch.compile(fn)
compiled_fn(torch.randn(32, 64))   # Compiles for shape [32, 64]
compiled_fn(torch.randn(16, 64))   # Recompiles for shape [16, 64]!
compiled_fn(torch.randn(8, 64))    # Recompiles again!

The Solution: dynamic=True

compiled_fn = torch.compile(fn, dynamic=True)
compiled_fn(torch.randn(32, 64))   # Compiles with symbolic shapes
compiled_fn(torch.randn(16, 64))   # Reuses compiled code!
compiled_fn(torch.randn(8, 64))    # Reuses again!

With dynamic=True, Dynamo uses symbolic shapes (e.g., s0 instead of 32), and the generated code handles any batch size.

Automatic Dynamic Shapes

PyTorch can also automatically detect that a dimension varies and switch to dynamic shapes after seeing multiple sizes:

# After 2 recompilations on the same dimension, Dynamo marks it dynamic
compiled_fn(torch.randn(32, 64))   # Compile for [32, 64]
compiled_fn(torch.randn(16, 64))   # Recompile, mark dim 0 as dynamic
compiled_fn(torch.randn(8, 64))    # Uses dynamic code — no recompile!

mark_dynamic

For fine-grained control:

x = torch.randn(32, 64)
torch._dynamo.mark_dynamic(x, 0)  # Mark dimension 0 as dynamic
compiled_fn(x)  # Compiles with dynamic dim 0, static dim 1 (64)

---

8. Compilation Cache

How Caching Works

Compiled code is cached so you don't recompile every time you restart:

• In-memory cache: Within a process, same function + same guards = reuse
• Persistent cache: Across process restarts (PyTorch 2.1+), compiled

artifacts are stored on disk

Persistent Cache

# Enable persistent cache (stored in ~/.cache/torch/inductor/)
import torch._inductor.config
torch._inductor.config.fx_graph_cache = True

This means:

• First run: full compilation (slow)
• Second run: loads from cache (fast startup)

---

9. Compiler Stances

Stances control how the compiler handles recompilation:

# Don't compile at all — run in eager mode
torch.compiler.set_stance("force_eager")

# Warn (log) on recompilation instead of silently recompiling
torch.compiler.set_stance("eager_on_recompile")

# Error if recompilation would occur
torch.compiler.set_stance("fail_on_recompile")

# Default behavior
torch.compiler.set_stance("default")

Use fail_on_recompile in production to catch unexpected dynamic behavior that would hurt performance.

---

10. Debugging torch.compile

torch._dynamo.explain()

Shows what happened during compilation:

explanation = torch._dynamo.explain(compiled_fn)(input)
print(explanation)
# Shows: number of graphs, graph breaks, break reasons

TORCH_LOGS Environment Variable

# See graph breaks
TORCH_LOGS="graph_breaks" python script.py

# See what Dynamo captured
TORCH_LOGS="dynamo" python script.py

# See generated code
TORCH_LOGS="output_code" python script.py

# See recompilation reasons
TORCH_LOGS="recompiles" python script.py

Common Errors

• "torch._dynamo.exc.Unsupported" — An operation can't be traced.

Fix: rewrite using supported operations.

• Recompilation spam — Model keeps recompiling.

Fix: Use dynamic=True or ensure consistent input shapes.

• Incorrect results — Compiled code gives different outputs.

Fix: Report as a bug. Workaround: mark the function with torch._dynamo.disable().

---

11. torch._dynamo.reset()

Clears all compiled graphs and cached state:

torch._dynamo.reset()

Useful when:

• Testing different compilation settings
• Debugging compilation issues
• Benchmarking (to force recompilation)

---

12. Custom Backends

You can write your own backend that receives the FX graph:

def my_backend(gm: torch.fx.GraphModule, example_inputs):
    """
    Custom backend that receives the graph and returns a callable.

    Args:
        gm: The captured FX graph module
        example_inputs: Example inputs used during tracing

    Returns:
        A callable that takes the same inputs and produces outputs
    """
    # Inspect the graph
    print(f"Graph has {len(list(gm.graph.nodes))} nodes")

    # You can transform the graph here, or just return it as-is
    # (gm is already callable)
    return gm

compiled_fn = torch.compile(fn, backend=my_backend)

This is useful for:

• Profiling what operations are captured
• Custom optimizations
• Research on graph transformations

---

13. Compiled Autograd

By default, torch.compile only compiles the forward pass. The backward pass still runs in eager mode. Compiled Autograd compiles the backward too:

with torch._dynamo.compiled_autograd.enable(torch.compile(backend="inductor")):
    loss.backward()

Benefits:

• Backward pass also gets operator fusion
• End-to-end compilation of training step

---

14. Performance Tips

When torch.compile Helps Most

• Transformer models (lots of small ops to fuse)
• Memory-bound operations (fusion reduces memory traffic)
• Models with many element-wise operations in sequence
• Standard architectures using torch.nn modules

When It Doesn't Help Much

• Already compute-bound (large matmuls with batch dim)
• Heavy graph breaks (too much falls back to Python)
• Very small models (compilation overhead > runtime savings)
• Highly dynamic models (frequent recompilation)

Best Practices

• Profile first: Know where time is spent before compiling
• Start simple: torch.compile(model) with defaults
• Check for graph breaks: Use explain() or fullgraph=True
• Use dynamic=True if batch sizes vary
• Use max-autotune for production deployments
• Cache compiled code for fast startup

---

15. FX Graph Basics

The intermediate representation (IR) used by torch.compile is an FX graph:

import torch.fx

def fn(x, y):
    z = x + y
    return z.relu()

# Trace into FX graph
traced = torch.fx.symbolic_trace(fn)
print(traced.graph)

Output:

graph():
    %x : [num_users=1] = placeholder[target=x]
    %y : [num_users=1] = placeholder[target=y]
    %add : [num_users=1] = call_function[target=operator.add](args = (%x, %y))
    %relu : [num_users=1] = call_method[target=relu](args = (%add,))
    return relu

Node types:

• placeholder — function inputs
• call_function — calls to functions (like torch.add)
• call_method — method calls on tensors (.relu(), .view(), etc.)
• call_module — calls to nn.Module submodules
• output — return value

Understanding FX graphs helps when:

• Writing custom backends
• Debugging compilation issues
• Understanding what optimizations are applied

---

Summary

Quick Start

# Step 1: Just compile it
model = torch.compile(model)

# Step 2: Check for issues
explanation = torch._dynamo.explain(model)(sample_input)

# Step 3: Optimize
model = torch.compile(model, mode="max-autotune", fullgraph=True)

---

📓 Open Notebook — Interactive version of this module

Module 09: Attention Mechanisms — From Scratch to FlexAttention

Overview

Attention is the core mechanism behind transformers — the architecture powering GPT, BERT, LLaMA, and virtually all modern AI models. This module builds attention from first principles and works up to PyTorch's most advanced APIs.

---

1. What is Attention? (Intuitive Explanation)

The Analogy

Imagine you're reading a sentence: "The cat sat on the mat because it was tired."

What does "it" refer to? You (unconsciously) attend to earlier words and determine "it" = "cat". Attention is this mechanism: for each position in a sequence, it looks at ALL other positions and decides which ones are relevant.

The Core Idea

Given a query ("what am I looking for?"), attention:

• Compares the query against all keys ("what information is available?")
• Produces attention weights (how relevant each key is)
• Uses those weights to aggregate values ("what to return?")

---

2. Scaled Dot-Product Attention

The Formula

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

Where:

• Q (Query): [batch, seq_len, d_k] — what each position is looking for
• K (Key): [batch, seq_len, d_k] — what each position offers for matching
• V (Value): [batch, seq_len, d_v] — what each position actually provides
• d_k: dimension of keys/queries

Why Scale by sqrt(d_k)?

Without scaling, dot products grow with d_k. Large dot products push softmax into regions where it has extremely small gradients (saturation). Dividing by sqrt(d_k) keeps variance ~1 regardless of dimension.

Example: if Q and K entries are independent with mean 0, variance 1:

• Q @ K^T has variance = d_k
• After scaling: variance = d_k / d_k = 1

Step-by-Step with Shapes

# Input: Q, K, V each with shape [batch, seq_len, d_model]
# For simplicity, assume d_k = d_v = d_model

# Step 1: Compute attention scores
scores = Q @ K.transpose(-2, -1)    # [batch, seq_len, seq_len]

# Step 2: Scale
scores = scores / math.sqrt(d_k)    # [batch, seq_len, seq_len]

# Step 3: (Optional) Apply mask
scores = scores.masked_fill(mask == 0, float('-inf'))

# Step 4: Softmax to get attention weights
weights = softmax(scores, dim=-1)   # [batch, seq_len, seq_len]

# Step 5: Weighted sum of values
output = weights @ V                # [batch, seq_len, d_v]

---

3. Causal (Autoregressive) Attention

What It Is

In language models that generate text left-to-right, position i should only attend to positions 0, 1, ..., i (not future positions). This is enforced with a causal mask:

mask = [[1, 0, 0, 0],
        [1, 1, 0, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 1]]

Implementation

# Create causal mask (lower triangular)
seq_len = Q.shape[1]
causal_mask = torch.tril(torch.ones(seq_len, seq_len))

# Apply: set future positions to -inf before softmax
scores = scores.masked_fill(causal_mask == 0, float('-inf'))
# After softmax, -inf becomes 0 (no attention to future)

---

4. Multi-Head Attention

Why Multiple Heads?

A single attention head can only focus on one pattern (e.g., subject-verb agreement). Multiple heads let the model attend to different patterns simultaneously:

• Head 1 might attend to syntactic relationships
• Head 2 might attend to semantic similarity
• Head 3 might attend to positional proximity

How It Works

Instead of one attention with d_model dimensions, use h heads each with d_k = d_model / h dimensions:

# Input: x with shape [batch, seq_len, d_model]
# Project to Q, K, V for each head
Q = W_q(x)  # [batch, seq_len, d_model]
K = W_k(x)  # [batch, seq_len, d_model]
V = W_v(x)  # [batch, seq_len, d_model]

# Reshape to [batch, num_heads, seq_len, head_dim]
Q = Q.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)
K = K.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)
V = V.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)

# Apply attention independently per head
# output: [batch, num_heads, seq_len, head_dim]
output = scaled_dot_product_attention(Q, K, V)

# Concatenate heads and project
output = output.transpose(1, 2).reshape(batch, seq_len, d_model)
output = W_o(output)  # Final linear projection

---

5. F.scaled_dot_product_attention (SDPA)

PyTorch provides an optimized implementation:

import torch.nn.functional as F

output = F.scaled_dot_product_attention(
    query,      # [batch, heads, seq_q, dim]
    key,        # [batch, heads, seq_kv, dim]
    value,      # [batch, heads, seq_kv, dim_v]
    attn_mask=None,     # Optional mask
    dropout_p=0.0,      # Dropout probability
    is_causal=False,    # If True, applies causal mask automatically
    scale=None,         # Custom scale factor (default: 1/sqrt(dim))
)

Benefits

• Automatically selects the fastest available backend
• Memory-efficient (doesn't materialize the full attention matrix)
• Fused implementation (fewer memory reads/writes)

---

6. SDPA Backends

PyTorch has multiple backends for SDPA, each optimized for different cases:

Flash Attention

• Memory: O(N) instead of O(N^2)
• Speed: Fastest for long sequences
• How: Tiles the computation, never materializes full attention matrix
• Requires: GPU with compute capability >= 8.0 (A100+)

Memory-Efficient Attention

• Memory: O(N) via chunked computation
• Speed: Good general-purpose
• How: Processes attention in chunks
• Requires: Any GPU

cuDNN Attention

• Speed: Can be fastest for standard shapes
• Requires: cuDNN 8.9+

Math (Fallback)

• Memory: O(N^2) — materializes full attention matrix
• Speed: Slowest for long sequences
• Works: Everywhere (CPU and GPU)
• Use: When other backends aren't available or for debugging

Backend Selection

from torch.nn.attention import sdpa_kernel, SDPBackend

# Force a specific backend
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output = F.scaled_dot_product_attention(q, k, v)

# Use math backend for debugging (can inspect attention weights)
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(q, k, v)

---

7. Flash Attention Explained Simply

The Problem

Standard attention computes:

S = Q @ K^T    # [N, N] — this is HUGE for long sequences
P = softmax(S) # [N, N] — stored in memory
O = P @ V      # output

For N = 16384 (16K tokens): S alone is 16384^2 * 4 bytes = 1 GB!

The Flash Attention Trick

Instead of computing the full N×N matrix:

• Tile Q, K, V into small blocks that fit in GPU fast memory (SRAM)
• Compute attention block by block
• Use the online softmax trick to accumulate correct softmax results

across blocks without needing the full row

Result: O(N) memory instead of O(N^2), and faster due to fewer memory accesses.

Online Softmax

The key insight: you can compute softmax incrementally. As you process each block of keys, maintain a running maximum and running sum, then correct at the end. This avoids needing all scores simultaneously.

---

8. FlexAttention

What It Is

FlexAttention (torch.nn.attention.flex_attention) lets you define custom attention patterns using simple Python functions, while still getting the performance benefits of Flash Attention-like fused kernels.

Why It Exists

Previously, custom attention patterns (sliding window, ALiBi, document masking) required:

• Writing custom CUDA kernels (hard)
• Materializing masks (memory expensive)
• Giving up on fused implementations (slow)

FlexAttention compiles your pattern into an efficient kernel.

Core API

from torch.nn.attention.flex_attention import (
    flex_attention,
    create_block_mask,
)

# score_mod: modifies attention scores before softmax
def causal_score_mod(score, b, h, q_idx, kv_idx):
    return torch.where(q_idx >= kv_idx, score, float('-inf'))

# mask_mod: defines which positions can attend to which (for BlockMask)
def causal_mask_mod(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

# Create a BlockMask for efficiency
block_mask = create_block_mask(causal_mask_mod, B=1, H=1, Q_LEN=seq_len, KV_LEN=seq_len)

# Apply flex attention
output = flex_attention(query, key, value, block_mask=block_mask)

score_mod vs mask_mod

• score_mod(score, b, h, q_idx, kv_idx): Modifies the attention score

for a specific (query, key) pair. Can add biases, apply causal masking, etc. Returns the modified score.

• mask_mod(b, h, q_idx, kv_idx): Returns True/False for whether this

(query, key) pair should be allowed. Used by BlockMask to skip entire blocks of computation.

---

9. FlexAttention Patterns

Causal Attention

def causal(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

Sliding Window

def sliding_window(b, h, q_idx, kv_idx):
    return (q_idx - kv_idx).abs() <= window_size

Causal + Sliding Window

def causal_sliding(b, h, q_idx, kv_idx):
    return (q_idx >= kv_idx) & (q_idx - kv_idx <= window_size)

ALiBi (Attention with Linear Biases)

def alibi_score_mod(score, b, h, q_idx, kv_idx):
    slope = 2 ** (-(h + 1) * 8 / num_heads)
    bias = -slope * (q_idx - kv_idx).abs()
    return score + bias

Document Masking (Multiple Documents in One Sequence)

# document_id[i] tells which document position i belongs to
def document_mask(b, h, q_idx, kv_idx):
    return document_id[q_idx] == document_id[kv_idx]

Prefix LM (Bidirectional prefix + Causal suffix)

def prefix_lm(b, h, q_idx, kv_idx):
    # Allow bidirectional attention within prefix
    # Causal attention after prefix
    return (kv_idx < prefix_length) | (q_idx >= kv_idx)

---

10. Building a Transformer Block

Pre-Norm vs Post-Norm

# Post-norm (original Transformer paper):
x = x + attention(x)
x = layer_norm(x)

# Pre-norm (GPT-2, modern models — more stable training):
x = x + attention(layer_norm(x))

RMSNorm vs LayerNorm

# LayerNorm: normalize, then scale + shift
# Subtracts mean AND divides by std
y = (x - mean(x)) / std(x) * gamma + beta

# RMSNorm: only divides by RMS (no mean subtraction)
# Faster, often works just as well
y = x / rms(x) * gamma

Activation Functions

• GELU: Smooth approximation of ReLU, used in BERT/GPT
• SiLU (Swish): x * sigmoid(x), used in LLaMA, PaLM
• SwiGLU: Gated variant, used in modern LLMs

Complete Block

class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads):
        self.norm1 = RMSNorm(dim)
        self.attn = MultiHeadAttention(dim, num_heads)
        self.norm2 = RMSNorm(dim)
        self.ffn = SwiGLU(dim)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

---

11. Positional Encoding

Attention is permutation-invariant — it doesn't know token order. We add positional information explicitly.

Sinusoidal (Original Transformer)

PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

Each dimension oscillates at a different frequency. The model can learn to attend to relative positions via linear combinations.

Learned Positional Embeddings

self.pos_embed = nn.Embedding(max_seq_len, d_model)
x = x + self.pos_embed(positions)

Simple but limited to max_seq_len seen during training.

RoPE (Rotary Position Embeddings)

Used in LLaMA, Mistral, and most modern LLMs.

Key idea: encode position by rotating Q and K vectors in 2D subspaces. The dot product Q·K then naturally depends on relative position.

# For each pair of dimensions (2i, 2i+1):
q_rotated[2i]   = q[2i] * cos(theta) - q[2i+1] * sin(theta)
q_rotated[2i+1] = q[2i] * sin(theta) + q[2i+1] * cos(theta)
# Where theta = position * base_freq^(-2i/d)

Benefits:

• Relative position awareness
• Can extrapolate to longer sequences than training
• No additional parameters

---

12. KV Cache for Inference

The Problem

During autoregressive generation, each new token needs to attend to ALL previous tokens. Without caching, you recompute K and V for all previous tokens at every step.

The Solution: KV Cache

Cache the K and V projections from previous steps:

# Step 1: Process prompt
K_cache = key_projection(prompt)    # [batch, num_heads, prompt_len, head_dim]
V_cache = value_projection(prompt)  # [batch, num_heads, prompt_len, head_dim]

# Step 2+: Generate each new token
for _ in range(max_new_tokens):
    # Only project the NEW token
    new_k = key_projection(new_token)   # [batch, num_heads, 1, head_dim]
    new_v = value_projection(new_token) # [batch, num_heads, 1, head_dim]

    # Append to cache
    K_cache = torch.cat([K_cache, new_k], dim=2)
    V_cache = torch.cat([V_cache, new_v], dim=2)

    # Query only needs the new token, but attends to full cache
    new_q = query_projection(new_token) # [batch, num_heads, 1, head_dim]
    output = attention(new_q, K_cache, V_cache)

Memory vs Speed Tradeoff

• Without cache: O(n) compute per step, O(1) extra memory
• With cache: O(1) amortized compute per step, O(n) extra memory

The KV cache makes generation ~10x faster for long sequences.

---

Summary

---

📓 Open Notebook — Interactive version of this module

Module 10: Distributed Training in PyTorch

Table of Contents

• Why Distributed Training?
• Key Concepts
• Launching with torchrun
• Collective Operations
• DistributedDataParallel (DDP)
• DeviceMesh
• DTensor (Distributed Tensor)
• FSDP1 vs FSDP2
• FSDP2 (fully_shard)
• Tensor Parallelism
• Pipeline Parallelism
• Combining Strategies: 3D Parallelism
• Distributed Checkpointing (DCP)
• SymmetricMemory
• Context Parallel
• Practical Advice

---

Why Distributed Training?

As models grow larger and datasets expand, a single GPU becomes insufficient. Distributed training addresses three fundamental bottlenecks:

1. Model too large for one GPU (Memory) A model like Llama 70B requires ~140 GB just for parameters in FP16. Even the largest GPUs (H100 with 80 GB) cannot hold this model, let alone the optimizer states and activations needed for training. Distributed strategies split the model across GPUs.

2. Training too slow (Compute) Even when a model fits on one GPU, training can take weeks or months. By distributing data across N GPUs, each GPU processes 1/N of the data per step, achieving near-linear speedup. Training that takes 30 days on 1 GPU takes ~4 days on 8 GPUs.

3. Data too large (I/O and throughput) With terabytes of training data, increasing throughput by processing more samples in parallel reduces wall-clock training time proportionally.

The Parallelism Taxonomy

---

Key Concepts

World Size, Rank, and Local Rank

When you launch distributed training, you create multiple processes, each driving one GPU. These processes form a process group and coordinate via collective communication.

Node 0 (Machine 0)          Node 1 (Machine 1)
┌──────────────────┐        ┌──────────────────┐
│ GPU0   GPU1      │        │ GPU0   GPU1      │
│ rank=0 rank=1    │        │ rank=2 rank=3    │
│ local_rank=0     │        │ local_rank=0     │
│        local_rank=1       │        local_rank=1
└──────────────────┘        └──────────────────┘
        world_size = 4

• World size: Total number of processes across all machines
• Rank: Unique global identifier for each process (0 to world_size-1)
• Local rank: Identifier within a single machine (0 to num_local_gpus-1)

import torch.distributed as dist

dist.init_process_group(backend="nccl")  # or "gloo" for CPU
rank = dist.get_rank()
world_size = dist.get_world_size()
local_rank = int(os.environ["LOCAL_RANK"])

Process Groups

A process group is a subset of all processes that can communicate. The default process group includes all processes and is created by init_process_group. You can create sub-groups for specialized communication:

# Create a group with only ranks 0 and 1
subgroup = dist.new_group(ranks=[0, 1])

# Only ranks in the group participate in collectives on this group
if rank in [0, 1]:
    dist.all_reduce(tensor, group=subgroup)

Backends

NCCL (pronounced "nickel") is the standard for GPU training and provides the best performance. Gloo is useful for CPU-based experiments and prototyping.

# GPU training (most common)
dist.init_process_group(backend="nccl")

# CPU training or prototyping
dist.init_process_group(backend="gloo")

# Use NCCL for GPU ops, Gloo for CPU ops (advanced)
dist.init_process_group(backend="nccl")
cpu_group = dist.new_group(backend="gloo")

---

Launching with torchrun

torchrun is PyTorch's built-in launcher for distributed training. It replaces the older torch.distributed.launch. It sets up environment variables and spawns processes for you.

Single-Node Launch

# 4 GPUs on one machine
torchrun --nproc_per_node=4 train.py --arg1 val1

# Can also use for CPU (with gloo backend)
torchrun --nproc_per_node=2 train.py

Multi-Node Launch

On each machine, run torchrun with the same --master_addr and --master_port:

# Machine 0 (master)
torchrun \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr=192.168.1.100 \
    --master_port=29500 \
    train.py

# Machine 1
torchrun \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=1 \
    --master_addr=192.168.1.100 \
    --master_port=29500 \
    train.py

Environment Variables Set by torchrun

Your training script reads these:

import os

rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])

Elastic Launch

torchrun supports elastic training where nodes can join or leave:

torchrun \
    --nproc_per_node=4 \
    --nnodes=2:8 \       # min 2, max 8 nodes
    --rdzv_backend=c10d \
    --rdzv_endpoint=master:29500 \
    train.py

---

Collective Operations

Collective operations are communication primitives where all processes in a group participate. Understanding these is essential because all distributed strategies are built on top of them.

All-Reduce

Every process starts with a tensor. After all-reduce, every process has the element-wise sum (or other reduction) of all tensors.

Before:                After all_reduce(SUM):
Rank 0: [1, 2]        Rank 0: [10, 20]
Rank 1: [3, 4]   →    Rank 1: [10, 20]
Rank 2: [6, 14]       Rank 2: [10, 20]

This is the core operation in DDP: after each backward pass, gradients are all-reduced so every replica has the same averaged gradients.

tensor = torch.tensor([rank * 2.0, rank * 3.0])
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
# Now tensor is the same on all ranks

All-Gather

Each process contributes a tensor, and every process receives the concatenation of all tensors.

Before:                After all_gather:
Rank 0: [A]            Rank 0: [A, B, C]
Rank 1: [B]       →    Rank 1: [A, B, C]
Rank 2: [C]            Rank 2: [A, B, C]

FSDP uses all-gather to reconstruct full parameter tensors before the forward pass.

local_tensor = torch.tensor([rank])
gathered = [torch.zeros(1) for _ in range(world_size)]
dist.all_gather(gathered, local_tensor)
# gathered = [tensor([0]), tensor([1]), tensor([2])]

Reduce-Scatter

The inverse of all-gather. First reduces (sums) all tensors element-wise, then scatters the result so each rank gets a different chunk.

Before:                  After reduce_scatter:
Rank 0: [1, 2, 3]       Rank 0: [6]   (sum of position 0: 1+2+3)
Rank 1: [2, 3, 4]  →    Rank 1: [9]   (sum of position 1: 2+3+4+... wait)
Rank 2: [3, 4, 5]       Rank 2: [12]  (sum of position 2: 3+4+5)

More precisely, with 3 ranks each holding a 3-element tensor:

• Element-wise sum: [1+2+3, 2+3+4, 3+4+5] = [6, 9, 12]
• Scatter: Rank 0 gets [6], Rank 1 gets [9], Rank 2 gets [12]

FSDP uses reduce-scatter after backward to reduce gradients and distribute shards back to their owners.

output = torch.zeros(2)
input_tensor = torch.arange(world_size * 2, dtype=torch.float) + rank
dist.reduce_scatter_tensor(output, input_tensor, op=dist.ReduceOp.SUM)

Broadcast

One process sends a tensor to all other processes.

Before:                After broadcast(src=0):
Rank 0: [42, 7]        Rank 0: [42, 7]
Rank 1: [0, 0]    →    Rank 1: [42, 7]
Rank 2: [0, 0]         Rank 2: [42, 7]

tensor = torch.tensor([42.0, 7.0]) if rank == 0 else torch.zeros(2)
dist.broadcast(tensor, src=0)

Barrier

Synchronizes all processes. Every process blocks until all processes have reached the barrier. No data is exchanged.

dist.barrier()  # All ranks wait here until everyone arrives

Use barriers sparingly; they are expensive and often unnecessary when collectives already imply synchronization.

Reduce

Like all-reduce, but the result only goes to one destination rank.

Before:                After reduce(dst=0, SUM):
Rank 0: [1, 2]         Rank 0: [6, 9]
Rank 1: [2, 3]    →    Rank 1: [2, 3]  (unchanged)
Rank 2: [3, 4]         Rank 2: [3, 4]  (unchanged)

Scatter

One process distributes different chunks to each process.

Before (rank 0 has all data):     After scatter(src=0):
Rank 0: [[A], [B], [C]]          Rank 0: [A]
Rank 1: []                  →    Rank 1: [B]
Rank 2: []                       Rank 2: [C]

Gather

The inverse of scatter. All processes send data to one destination.

Before:                After gather(dst=0):
Rank 0: [A]            Rank 0: [[A], [B], [C]]
Rank 1: [B]       →    Rank 1: [B]  (unchanged)
Rank 2: [C]            Rank 2: [C]  (unchanged)

---

DistributedDataParallel (DDP)

DDP is the simplest and most commonly used distributed training strategy. Each GPU holds a complete copy of the model. Training data is split across GPUs, and gradients are synchronized via all-reduce after each backward pass.

How DDP Works Internally

• Initialization: The model is replicated on each GPU. Parameters are

broadcast from rank 0 to ensure all replicas start identically.

• Forward pass: Each rank processes its own mini-batch independently.

No communication occurs during forward.

• Backward pass: As gradients are computed, DDP groups them into

buckets (default ~25 MB each). When a bucket is full, all-reduce starts immediately -- overlapping communication with computation for the remaining layers. This is called gradient bucketing.

• Optimizer step: After all-reduce completes, every rank has identical

averaged gradients. Each rank runs the optimizer independently, producing identical updated parameters.

                  Forward (independent)
                  ┌─────────┐
Rank 0: Data₀ →  │ Model₀  │ → Loss₀
                  └─────────┘
                  ┌─────────┐
Rank 1: Data₁ →  │ Model₁  │ → Loss₁
                  └─────────┘

                  Backward (all-reduce gradients)
                  ┌────────────────────────┐
                  │  All-Reduce Gradients  │
                  │  (bucketed, overlapped │
                  │   with backward comp)  │
                  └────────────────────────┘

                  Optimizer Step (independent, identical)

Complete DDP Setup

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup():
    dist.destroy_process_group()

def main():
    local_rank = setup()
    device = torch.device(f"cuda:{local_rank}")

    # Create model and move to GPU
    model = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
    ).to(device)

    # Wrap with DDP
    model = DDP(model, device_ids=[local_rank])

    # DistributedSampler ensures each rank sees different data
    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()

    for epoch in range(10):
        sampler.set_epoch(epoch)  # Shuffle differently each epoch
        for batch_x, batch_y in dataloader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            optimizer.zero_grad()
            output = model(batch_x)
            loss = loss_fn(output, batch_y)
            loss.backward()
            optimizer.step()

    cleanup()

if __name__ == "__main__":
    main()

DistributedSampler

The DistributedSampler partitions the dataset indices so each rank gets a different subset. It pads the dataset to make it evenly divisible by world_size.

Key: call sampler.set_epoch(epoch) each epoch to get different shuffling. Without this, every epoch uses the same data order per rank.

DDP Tips

• Access the underlying model via model.module (DDP wraps it).
• Save checkpoints only on rank 0 to avoid file conflicts.
• Use torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) before wrapping

with DDP if your model uses BatchNorm.

• The find_unused_parameters=True flag is needed if some parameters don't

receive gradients every iteration (e.g., conditional branches).

---

DeviceMesh

DeviceMesh is the foundation of modern distributed training in PyTorch. It provides a multi-dimensional abstraction over a group of devices, replacing the manual management of process groups.

What is a DeviceMesh?

A DeviceMesh represents a logical grid of devices. Each dimension of the mesh corresponds to a parallelism strategy:

from torch.distributed.device_mesh import init_device_mesh

# 1D mesh: all 8 GPUs in one dimension (simple data parallelism)
mesh_1d = init_device_mesh("cuda", (8,), mesh_dim_names=("dp",))

# 2D mesh: 4 data-parallel groups × 2 tensor-parallel groups
# GPUs arranged as:
#   TP=0  TP=1
#   ──────────
#   GPU0  GPU1    ← DP group 0
#   GPU2  GPU3    ← DP group 1
#   GPU4  GPU5    ← DP group 2
#   GPU6  GPU7    ← DP group 3
mesh_2d = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))

# 3D mesh: DP × TP × PP
mesh_3d = init_device_mesh(
    "cuda", (2, 2, 2), mesh_dim_names=("dp", "tp", "pp")
)

Why DeviceMesh?

Before DeviceMesh, you had to manually create process groups for each parallelism dimension and carefully track which ranks belonged to which groups. DeviceMesh automates this:

# Old way: manually creating groups
dp_groups = []
for i in range(0, 8, 2):
    dp_groups.append(dist.new_group([i, i+1]))

# New way: DeviceMesh handles it
mesh = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))
dp_mesh = mesh["dp"]  # Automatically creates the right groups
tp_mesh = mesh["tp"]

Accessing Sub-Meshes

You can slice a DeviceMesh to get a sub-mesh for a specific dimension:

mesh = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))

# Get the sub-mesh for data parallelism
dp_mesh = mesh["dp"]  # 1D mesh with 4 devices (for this rank's DP group)

# Get the sub-mesh for tensor parallelism
tp_mesh = mesh["tp"]  # 1D mesh with 2 devices (for this rank's TP group)

# These sub-meshes carry the correct process groups
# so you can pass them directly to FSDP, TP, etc.

DeviceMesh for 3D Parallelism

# 16 GPUs: 2 DP × 4 TP × 2 PP
mesh = init_device_mesh(
    "cuda", (2, 4, 2), mesh_dim_names=("dp", "tp", "pp")
)

# Each parallelism strategy gets its own sub-mesh
dp_mesh = mesh["dp"]
tp_mesh = mesh["tp"]
pp_mesh = mesh["pp"]

# Apply each strategy using its sub-mesh
# TP on tp_mesh, FSDP on dp_mesh, PP on pp_mesh

---

DTensor (Distributed Tensor)

DTensor is a tensor abstraction that represents a tensor distributed across multiple devices. It knows how the tensor is distributed and automatically handles the communication needed for operations.

Placement Types

DTensor uses placements to describe how a tensor's data is distributed across the devices in a DeviceMesh dimension:

Creating DTensors

from torch.distributed.tensor import DTensor, Shard, Replicate, distribute_tensor

mesh = init_device_mesh("cuda", (4,))

# Create a regular tensor and distribute it
big_tensor = torch.randn(16, 32)

# Shard along dim 0: each of 4 GPUs gets a [4, 32] chunk
sharded = distribute_tensor(big_tensor, mesh, placements=[Shard(0)])

# Replicate: each GPU gets the full [16, 32] tensor
replicated = distribute_tensor(big_tensor, mesh, placements=[Replicate()])

distribute_module

Instead of manually distributing each parameter, distribute_module distributes an entire module's parameters and handles input/output:

from torch.distributed.tensor import distribute_module, Shard, Replicate

def input_fn(mod, inputs, mesh):
    # How to distribute inputs
    return (distribute_tensor(inputs[0], mesh, [Shard(0)]),)

def output_fn(mod, outputs, mesh):
    # How to gather outputs
    return outputs.full_tensor()

model = nn.Linear(1024, 512)
distribute_module(
    model,
    device_mesh=mesh,
    input_fn=input_fn,
    output_fn=output_fn,
)

DTensor and Automatic Communication

The key insight: when you do operations on DTensors, PyTorch automatically inserts the right collectives. If you multiply a Shard(1) tensor by a Replicate() tensor, PyTorch knows it needs an all-reduce to get the correct result.

A: [Shard(1)]  ×  B: [Replicate()]  →  C: [Partial()]  →  all_reduce → C: [Replicate()]

This is how Tensor Parallelism works under the hood.

---

FSDP1 vs FSDP2

FSDP1 (FullyShardedDataParallel - Legacy)

FSDP1 was PyTorch's first fully sharded data parallelism implementation. It wraps the entire module:

# FSDP1 (legacy)
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model, auto_wrap_policy=size_based_auto_wrap_policy)

Problems with FSDP1:

• Module wrapping: FSDP1 wraps modules, creating a new FSDP module

hierarchy. This breaks model.layer1.weight access patterns.

• Non-composable: Hard to combine with other parallelism strategies.
• Complex API: Many constructor arguments, hard to reason about behavior.
• Flattened parameters: Parameters are flattened into a single 1D tensor,

making debugging and checkpointing harder.

FSDP2 (fully_shard - Current Standard)

FSDP2 is a ground-up rewrite with a composable, per-parameter design:

# FSDP2 (current)
from torch.distributed.fsdp import fully_shard

# Apply to submodules first, then root
for layer in model.layers:
    fully_shard(layer)
fully_shard(model)

Advantages of FSDP2:

• Composable: Works seamlessly with TP, PP, and other strategies.
• Per-parameter sharding: Each parameter is sharded independently (via

DTensor), preserving the original module structure.

• Simpler API: A single function call per module.
• Better debugging: Parameters remain accessible with their original names.
• DTensor-based: Built on DTensor, providing a clean abstraction.

Use FSDP2 for all new code. FSDP1 is maintained but not actively developed.

---

FSDP2 (fully_shard)

How Sharding Works

FSDP2 shards model parameters across data-parallel ranks. During training:

• Idle state: Each rank holds only its shard (1/N) of each parameter.

Memory usage is reduced by ~N× for parameters.

• Before forward: all-gather reconstructs the full parameters from all

shards. Each rank now temporarily has the full parameter for computation.

• Forward computation: Runs normally with the full parameters.

• After forward: Full parameters are freed (unless needed for backward).

Memory drops back to 1/N.

• Before backward: all-gather again to get full parameters for gradient

computation.

• After backward: reduce-scatter synchronizes gradients AND distributes

gradient shards. Each rank ends up with the gradient shard corresponding to its parameter shard.

• Optimizer step: Each rank updates only its parameter shard using its

gradient shard. No communication needed.

     ┌─────────────────────────────────────────────────┐
     │              FSDP2 Training Loop                │
     │                                                 │
     │  Idle: Each rank holds 1/N of params            │
     │         ↓                                       │
     │  all-gather → Full params → Forward             │
     │         ↓                                       │
     │  Free full params (keep shards)                 │
     │         ↓                                       │
     │  all-gather → Full params → Backward            │
     │         ↓                                       │
     │  reduce-scatter → Gradient shards               │
     │         ↓                                       │
     │  Optimizer step (on shards only)                │
     │         ↓                                       │
     │  Back to idle (1/N params + 1/N grads)          │
     └─────────────────────────────────────────────────┘

Basic FSDP2 Setup

import torch
import torch.nn as nn
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

class TransformerBlock(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, 8)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.ffn(self.norm2(x))
        return x

class TransformerModel(nn.Module):
    def __init__(self, d_model=512, n_layers=6):
        super().__init__()
        self.embed = nn.Embedding(10000, d_model)
        self.layers = nn.ModuleList(
            [TransformerBlock(d_model) for _ in range(n_layers)]
        )
        self.head = nn.Linear(d_model, 10000)

    def forward(self, x):
        x = self.embed(x)
        for layer in self.layers:
            x = layer(x)
        return self.head(x)

# Apply FSDP2: submodules first, then root
model = TransformerModel().cuda()
for layer in model.layers:
    fully_shard(layer)
fully_shard(model)  # Root module last

MixedPrecisionPolicy

Mixed precision reduces memory and increases throughput by using lower precision for computation while maintaining accuracy:

from torch.distributed.fsdp import MixedPrecisionPolicy

# BFloat16 for compute, FP32 for parameter storage
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,    # Cast params to bf16 for forward/backward
    reduce_dtype=torch.float32,     # Reduce gradients in fp32 for accuracy
)

for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)
fully_shard(model, mp_policy=mp_policy)

CPUOffloadPolicy

For extremely large models, offload parameters to CPU when not in use:

from torch.distributed.fsdp import CPUOffloadPolicy

offload_policy = CPUOffloadPolicy(pin_memory=True)

for layer in model.layers:
    fully_shard(layer, offload_policy=offload_policy)
fully_shard(model, offload_policy=offload_policy)

This trades compute speed for memory: parameters are moved to CPU after forward/backward, freeing GPU memory. pin_memory=True uses pinned (page- locked) CPU memory for faster CPU-GPU transfers.

FSDP2 + DeviceMesh

When combining FSDP with other parallelism, use DeviceMesh to specify which dimension is for data parallelism:

mesh = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))
dp_mesh = mesh["dp"]

# fully_shard on the DP dimension
for layer in model.layers:
    fully_shard(layer, mesh=dp_mesh)
fully_shard(model, mesh=dp_mesh)

---

Tensor Parallelism

Tensor Parallelism (TP) splits individual layers across GPUs. Unlike FSDP which shards parameters and reconstructs them for computation, TP keeps parameters split and performs distributed computation.

Why Tensor Parallelism?

• Reduces memory per GPU for very large individual layers
• Reduces latency per layer (each GPU does less work)
• Essential for models where single layers are too large for one GPU
• Works within a single node (requires fast interconnect like NVLink)

Column-wise and Row-wise Parallelism

For a linear layer Y = XW + b, there are two ways to split:

ColwiseParallel: Split W along columns (output dimension)

         GPU 0          GPU 1
W = [w₀ | w₁]   →   W₀ = w₀     W₁ = w₁
Y = X @ W       →   Y₀ = X @ W₀  Y₁ = X @ W₁
                     Y = [Y₀ | Y₁]  (gathered)

Each GPU computes a portion of the output. The input X is replicated.

RowwiseParallel: Split W along rows (input dimension)

          GPU 0          GPU 1
W = [w₀]          W₀ = w₀     W₁ = w₁
    [w₁]
X = [x₀ | x₁]    X₀ = x₀     X₁ = x₁
Y = X @ W    →    Y = X₀@W₀ + X₁@W₁  (all-reduce)

Input is split and each GPU computes a partial result that must be summed.

Using TP in PyTorch

from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
    SequenceParallel,
)

mesh = init_device_mesh("cuda", (world_size,))

# Parallelize specific layers
parallelize_plan = {
    "attn.qkv_proj": ColwiseParallel(),
    "attn.out_proj": RowwiseParallel(),
    "ffn.up_proj": ColwiseParallel(),
    "ffn.down_proj": RowwiseParallel(),
}

parallelize_module(model, mesh, parallelize_plan)

SequenceParallel

SequenceParallel splits the sequence dimension instead of replicating activations. Between ColwiseParallel and RowwiseParallel layers, activations can be kept split along the sequence dimension, reducing activation memory:

parallelize_plan = {
    "norm1": SequenceParallel(),
    "attn.qkv_proj": ColwiseParallel(input_layouts=Shard(0)),
    "attn.out_proj": RowwiseParallel(output_layouts=Shard(0)),
    "norm2": SequenceParallel(),
    "ffn.up_proj": ColwiseParallel(input_layouts=Shard(0)),
    "ffn.down_proj": RowwiseParallel(output_layouts=Shard(0)),
}

When to Use TP

• Large models with huge linear layers (e.g., LLMs with 8192+ hidden dim)
• When you have fast intra-node interconnect (NVLink)
• Typically TP degree of 2, 4, or 8 within a single node
• Beyond 8-way TP, the communication overhead usually outweighs benefits

---

Pipeline Parallelism

Pipeline Parallelism (PP) splits a model into sequential stages, each running on a different GPU (or set of GPUs). Data flows through stages sequentially, like an assembly line.

The Bubble Problem

Naive pipeline parallelism has a severe inefficiency: while stage 1 processes micro-batch 1, stages 2-N are idle. This idle time is called the pipeline bubble.

Naive (Sequential):
Stage 0: [F1][F2][F3][F4]
Stage 1:     [F1][F2][F3][F4]
Stage 2:         [F1][F2][F3][F4]
                              ↑ huge bubble, most GPUs idle most of the time

Micro-batching

The solution: split each mini-batch into multiple micro-batches and pipeline them:

With micro-batches:
Stage 0: [F1][F2][F3][F4][B4][B3][B2][B1]
Stage 1:     [F1][F2][F3][F4][B4][B3][B2][B1]
Stage 2:         [F1][F2][F3][F4][B4][B3][B2][B1]

F = forward, B = backward, number = micro-batch id

Pipeline Schedules

Different schedules offer different trade-offs between memory, bubble ratio, and implementation complexity:

Where p = number of pipeline stages, m = number of micro-batches, v = number of virtual stages (chunks).

GPipe: Simple but memory-hungry. All micro-batches do forward, then all do backward. Must store activations for all micro-batches simultaneously.

1F1B (One Forward One Backward): After a warm-up phase, alternates one forward and one backward. Steady-state memory is constant (only one micro- batch's activations at a time).

Interleaved 1F1B: Each rank handles multiple non-contiguous stages (e.g., rank 0 handles stages 0 and 4). This reduces the bubble because micro-batches cycle through stages faster.

Zero Bubble: Splits backward into two parts (input gradient and weight gradient) and overlaps weight gradient computation with the next forward pass. Nearly eliminates the bubble.

DualPipeV: A bidirectional schedule where micro-batches flow both forward and backward through the pipeline simultaneously, achieving near-zero bubble with better memory efficiency.

PP in PyTorch

from torch.distributed.pipelining import (
    pipeline,
    SplitPoint,
    ScheduleGPipe,
    Schedule1F1B,
    ScheduleInterleaved1F1B,
)

# Split model into stages
pipe = pipeline(
    model,
    mb_args=(torch.randn(batch_size, seq_len, d_model),),
    split_spec={
        "layers.3": SplitPoint.BEGINNING,  # Split before layer 3
        "layers.6": SplitPoint.BEGINNING,  # Split before layer 6
    },
)

# Get this rank's stage
stage = pipe.get_stage(rank, device)

# Create schedule
schedule = Schedule1F1B(stage, n_microbatches=8)

# Run
if rank == 0:
    schedule.step(input_data)
elif rank == num_stages - 1:
    losses = schedule.step()
else:
    schedule.step()

---

Combining Strategies: 3D Parallelism

For training the largest models (100B+ parameters), you combine multiple parallelism strategies. The standard combination is 3D parallelism: Data Parallel (FSDP) × Tensor Parallel × Pipeline Parallel.

DeviceMesh for 3D Parallelism

# 64 GPUs across 8 nodes, 8 GPUs per node
# 8 DP × 4 TP × 2 PP
mesh = init_device_mesh(
    "cuda", (8, 4, 2), mesh_dim_names=("dp", "tp", "pp")
)

Typical Assignment

• TP within a node: TP requires fast interconnect, so TP groups are

within a single node (connected by NVLink).

• PP across nodes: PP has less communication (only activations at stage

boundaries), so it can span nodes.

• FSDP across remaining GPUs: FSDP handles the data parallelism dimension.

Node 0:  [GPU0, GPU1, GPU2, GPU3]  ← TP group, PP stage 0
Node 1:  [GPU4, GPU5, GPU6, GPU7]  ← TP group, PP stage 1
...
Across nodes: FSDP groups

Code Pattern for 3D Parallelism

mesh = init_device_mesh(
    "cuda", (dp_size, tp_size, pp_size),
    mesh_dim_names=("dp", "tp", "pp"),
)

# 1. Apply Tensor Parallelism first (innermost)
tp_mesh = mesh["tp"]
for layer in model.layers:
    parallelize_module(layer, tp_mesh, tp_plan)

# 2. Apply FSDP (middle)
dp_mesh = mesh["dp"]
for layer in model.layers:
    fully_shard(layer, mesh=dp_mesh)
fully_shard(model, mesh=dp_mesh)

# 3. Apply Pipeline Parallelism (outermost)
pp_mesh = mesh["pp"]
# Split model into stages along pp_mesh

---

Distributed Checkpointing (DCP)

When training with FSDP, TP, or other distributed strategies, each rank holds only a shard of the model. Distributed Checkpointing (DCP) handles saving and loading these sharded states correctly.

Why Not torch.save?

torch.save requires gathering the full model to one rank, which:

• May not fit in memory for very large models
• Creates a bottleneck (one rank doing all the work)
• Produces a format tied to the original parallelism configuration

DCP saves each rank's shard independently and can reshard when loading with a different parallelism configuration.

Basic Save and Load

import torch.distributed.checkpoint as dcp

# Save
state_dict = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
dcp.save(state_dict, checkpoint_id="checkpoints/step_1000")

# Load
state_dict = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
dcp.load(state_dict, checkpoint_id="checkpoints/step_1000")
model.load_state_dict(state_dict["model"])
optimizer.load_state_dict(state_dict["optimizer"])

Async Save

For large models, checkpointing can take minutes. Async save moves the checkpoint writing to a background thread so training can continue:

# Async save returns a Future
future = dcp.async_save(state_dict, checkpoint_id="checkpoints/step_1000")

# Training continues immediately...
# Optionally wait for completion before the next checkpoint
future.result()

get_model_state_dict / set_model_state_dict

These utilities handle the complexity of getting/setting state dicts for models wrapped with FSDP, TP, etc.:

from torch.distributed.checkpoint.state_dict import (
    get_model_state_dict,
    set_model_state_dict,
    get_optimizer_state_dict,
    set_optimizer_state_dict,
    StateDictOptions,
)

# Get a "clean" state dict (handles FSDP/TP unwrapping)
model_state = get_model_state_dict(model)
optim_state = get_optimizer_state_dict(model, optimizer)

# Save
dcp.save({"model": model_state, "optim": optim_state}, checkpoint_id=path)

# Load
state = {"model": model_state, "optim": optim_state}
dcp.load(state, checkpoint_id=path)
set_model_state_dict(model, state["model"])
set_optimizer_state_dict(model, optimizer, state["optim"])

HuggingFace Format

DCP can save in HuggingFace-compatible format for interoperability:

from torch.distributed.checkpoint import HuggingFaceLoadPlanner

# Load a HuggingFace checkpoint
dcp.load(
    state_dict,
    checkpoint_id="path/to/hf_checkpoint",
    planner=HuggingFaceLoadPlanner(),
)

---

SymmetricMemory

SymmetricMemory is an intra-node optimization that provides direct GPU-to-GPU memory access using NVLink, bypassing the traditional collective communication libraries.

What is SymmetricMemory?

On a multi-GPU node with NVLink, GPUs can directly read/write each other's memory. SymmetricMemory allocates a shared memory region accessible by all GPUs in a group, enabling custom, low-latency communication patterns.

import torch.distributed.symmetric_memory as sm

# Allocate symmetric memory (same virtual address on all GPUs)
t = sm.empty_strided_p2p(
    size=(1024, 1024),
    stride=(1024, 1),
    dtype=torch.float32,
    device=torch.device(f"cuda:{local_rank}"),
)

# Direct GPU-to-GPU operations
sm.memcpy_p2p(dst=t_on_gpu1, src=t_on_gpu0)

When to Use

• Custom all-reduce implementations that exploit NVLink topology
• Fine-grained producer-consumer patterns between GPUs
• Overlapping communication with computation at a granularity finer than

what standard collectives offer

• Typically used in advanced performance optimization, not in everyday training

---

Context Parallel

Context Parallelism (CP) addresses the challenge of training with very long sequences. When the sequence length is so large that a single GPU cannot hold the activations for one sequence, CP splits the sequence across GPUs.

How It Works

CP distributes the sequence dimension across GPUs in a process group. For attention computation, this requires specialized communication because each position needs to attend to all other positions:

Sequence: [token_0, token_1, ..., token_8191]

GPU 0: [token_0 ... token_2047]
GPU 1: [token_2048 ... token_4095]
GPU 2: [token_4096 ... token_6143]
GPU 3: [token_6144 ... token_8191]

For attention: GPU 0 needs KV from all GPUs → ring attention

Ring Attention

Ring attention is a common CP implementation where KV pairs are passed around in a ring. Each GPU computes attention with its local Q and the received KV, then passes KV to the next GPU:

Step 1: GPU0 attends to KV₀, GPU1 to KV₁, ...
Step 2: GPUs pass KV to neighbor: GPU0 gets KV₃, GPU1 gets KV₀, ...
Step 3: Repeat until all KV seen by all GPUs

When to Use

• Sequence lengths > 8K-16K tokens (depends on model size and GPU memory)
• Long-document training, video models, genomics
• Often combined with TP and FSDP

---

Practical Advice

Choosing a Parallelism Strategy

Start: Does the model fit on 1 GPU with your batch size?
  │
  ├── YES → Use DDP
  │         Still too slow? → Increase DDP world size
  │
  └── NO  → Does the model fit on 1 GPU (batch_size=1)?
            │
            ├── YES → Use FSDP2
            │         Want even faster? → FSDP2 + TP
            │
            └── NO  → Use FSDP2 + TP
                      Still doesn't fit? → Add PP
                      Very long sequences? → Add CP

Rules of Thumb

• Start with DDP if your model fits on one GPU. It's the simplest and

most efficient.

• Move to FSDP2 when memory is the bottleneck. FSDP reduces per-GPU

memory at the cost of extra communication.

• Add TP when individual layers are very large (hidden_dim > 4096) or

when you need to reduce per-GPU memory further. Keep TP within a node.

• Add PP when you have many GPUs across nodes and want to reduce

cross-node communication. PP only communicates activations at stage boundaries.

• Use CP specifically for long-sequence training where sequence

activations dominate memory.

• Match TP to NVLink topology: Use TP degree 2, 4, or 8 matching the

NVLink connectivity within your node.

• More micro-batches reduce PP bubble: For pipeline parallelism, use

at least 2-4× as many micro-batches as pipeline stages.

Common Pitfalls

• Forgetting sampler.set_epoch(epoch): Causes same data order every

epoch with DDP/FSDP.

• Not sharding submodules before root in FSDP2: Always apply fully_shard

to submodules first, then the root module.

• TP across nodes: TP requires fast interconnect. Putting a TP group

across nodes with only InfiniBand (instead of NVLink) kills performance.

• Saving checkpoints on all ranks: Use rank 0 for simple saves, or DCP

for distributed saves.

• Not using no_sync() for gradient accumulation: When accumulating

gradients across multiple steps, wrap forward/backward in model.no_sync() to skip all-reduce on intermediate steps.

Gradient Accumulation with DDP/FSDP

accumulation_steps = 4
for i, (data, target) in enumerate(dataloader):
    # Use no_sync for intermediate steps to avoid wasteful all-reduce
    context = model.no_sync() if (i + 1) % accumulation_steps != 0 else nullcontext()
    with context:
        output = model(data)
        loss = loss_fn(output, target) / accumulation_steps
        loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

---

Files in This Module

---

📓 Open Notebook — Interactive version of this module

Module 11: Export and Deployment

Table of Contents

• Why Export?
• The PyTorch Deployment Landscape
• torch.export
• Static vs Dynamic Shapes
• draft_export
• Saving and Loading
• Graph Inspection
• AOTInductor
• NativeRT
• ONNX Export
• TorchServe
• ExecuTorch
• Quantization for Deployment
• Practical Workflow

---

Why Export?

During research, you run PyTorch in eager mode: Python executes operations one at a time, giving you full flexibility (print statements, breakpoints, dynamic control flow). But production deployment has different requirements:

• No Python dependency: Servers, mobile devices, and embedded systems may

not have (or want) a Python runtime.

• Performance: Compiler optimizations (operator fusion, memory planning,

kernel selection) require seeing the full computation graph ahead of time.

• Predictability: Production systems need deterministic latency and memory

usage. Eager mode's dynamism makes this hard to guarantee.

• Portability: The same model may run on server GPUs, edge TPUs, mobile

phones, or web browsers.

Export bridges the gap: it captures your Python model as a self-contained computation graph that can be optimized and deployed without Python.

Research (Python, eager)  →  Export (capture graph)  →  Deploy (optimized, no Python)

---

The PyTorch Deployment Landscape

PyTorch provides multiple paths from research to production:

                          ┌──────────────┐
                          │  Your Model  │
                          │  (nn.Module) │
                          └──────┬───────┘
                                 │
                          ┌──────┴───────┐
                          │ torch.export │
                          │ (capture     │
                          │  full graph) │
                          └──────┬───────┘
                                 │
              ┌──────────┬───────┼────────┬──────────┐
              ▼          ▼       ▼        ▼          ▼
        ┌──────────┐ ┌───────┐ ┌────┐ ┌───────┐ ┌─────────┐
        │AOTInductor│ │NativeRT│ │ONNX│ │Torch- │ │ExecuTorch│
        │(.so lib) │ │(C++   │ │    │ │Serve  │ │(mobile/ │
        │          │ │engine)│ │    │ │       │ │edge)    │
        └──────────┘ └───────┘ └────┘ └───────┘ └─────────┘
           Server      Server   Any    Server    Mobile/
           (C++/Py)    (C++)   runtime  (Py)     Embedded

---

torch.export

torch.export is the primary tool for capturing a PyTorch model as a complete, self-contained graph. It traces through your model and produces an ExportedProgram containing the full computation graph with no Python dependency.

Basic Export

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.linear(x))

model = MyModel()
example_input = (torch.randn(3, 10),)

# Export the model
exported = torch.export.export(model, example_input)

# The exported program can be called like the original
result = exported.module()(torch.randn(3, 10))

What ExportedProgram Contains

An ExportedProgram captures:

• Graph Module: An fx.GraphModule containing the computation graph as

a series of operations (ATen operators).

• Graph Signature: Maps graph inputs/outputs to parameters, buffers,

and user inputs.

• State Dict: The model's parameters and buffers.
• Range Constraints: Valid ranges for dynamic dimensions.
• Module Call graph: Preserves the module hierarchy for debugging.

exported = torch.export.export(model, example_input)

# Access the graph
print(exported.graph_module.graph)

# Access parameters
print(exported.state_dict.keys())

# The graph shows ATen-level operations
for node in exported.graph_module.graph.nodes:
    print(f"  {node.op}: {node.target}")

What Gets Captured

torch.export traces through your model's forward method and captures:

• All tensor operations as ATen ops
• Control flow (with restrictions — see below)
• Shape computations

It does NOT capture:

• Print statements or side effects
• Operations on non-tensor values that aren't related to shapes
• Dynamic Python control flow based on tensor values (use torch.cond instead)

Handling Control Flow

Standard if/else based on tensor values won't work:

# This will fail with torch.export:
def forward(self, x):
    if x.sum() > 0:  # Can't branch on tensor value!
        return x * 2
    return x * 3

# Use torch.cond instead:
def forward(self, x):
    return torch.cond(
        x.sum() > 0,
        lambda x: x * 2,
        lambda x: x * 3,
        (x,),
    )

Control flow based on tensor shapes (not values) is fine because shapes are known at export time (or constrained for dynamic shapes):

# This is fine:
def forward(self, x):
    if x.shape[0] > 5:  # Shape is known at export time
        return x[:5]
    return x

---

Static vs Dynamic Shapes

By default, torch.export traces with static shapes: the exported model only accepts inputs with the exact same shapes as the example inputs. For production, you usually want dynamic shapes so the model handles variable batch sizes, sequence lengths, etc.

The Dim API

Use torch.export.Dim to declare which dimensions are dynamic:

from torch.export import Dim, export

# Declare a dynamic dimension named "batch"
batch = Dim("batch", min=1, max=128)

# Export with dynamic first dimension
exported = export(
    model,
    (torch.randn(4, 10),),  # Example with batch=4
    dynamic_shapes={"x": {0: batch}},  # dim 0 is dynamic
)

# Now works with any batch size from 1 to 128:
exported.module()(torch.randn(1, 10))    # batch=1
exported.module()(torch.randn(64, 10))   # batch=64
exported.module()(torch.randn(128, 10))  # batch=128

Multiple Dynamic Dimensions

batch = Dim("batch", min=1, max=256)
seq_len = Dim("seq_len", min=1, max=2048)

exported = export(
    model,
    (torch.randn(4, 128, 512),),
    dynamic_shapes={"x": {0: batch, 1: seq_len}},
    # dim 0 = batch (dynamic), dim 1 = seq_len (dynamic), dim 2 = 512 (static)
)

Constraints Between Dimensions

When multiple inputs share a dimension (e.g., same batch size), use the same Dim object:

batch = Dim("batch", min=1, max=128)

def forward(self, x, y):
    return x + y  # x and y must have same batch size

exported = export(
    model,
    (torch.randn(4, 10), torch.randn(4, 20)),
    dynamic_shapes={
        "x": {0: batch},
        "y": {0: batch},  # Same Dim → enforces same batch size
    },
)

Automatic Dynamic Shapes

For convenience, you can let PyTorch infer dynamic shapes:

from torch.export import export

exported = export(
    model,
    (torch.randn(4, 10),),
    dynamic_shapes={"x": {0: Dim.AUTO}},
)

---

draft_export

When torch.export fails (due to unsupported Python constructs, dynamic control flow, etc.), draft_export helps you debug by producing a best- effort export with detailed error information.

from torch.export import draft_export

# If regular export fails:
try:
    exported = torch.export.export(model, example_input)
except Exception as e:
    print(f"Export failed: {e}")

# Use draft_export to get a partial graph + diagnostics
ep, report = draft_export(model, example_input)

# The report shows what went wrong and how to fix it
print(report)

draft_export returns:

• An ExportedProgram (possibly with graph breaks or approximations)
• A report detailing what couldn't be captured and suggesting fixes

Common fixes suggested by draft_export:

• Replace if tensor_val > 0 with torch.cond
• Replace for i in range(tensor.shape[0]) with bounded loop
• Mark data-dependent shapes with torch.export.Dim

---

Saving and Loading

PT2 Archive Format

The recommended format for saving exported programs is the PT2 Archive:

import torch

# Export
model = MyModel()
exported = torch.export.export(model, (torch.randn(1, 10),))

# Save as PT2 archive
torch.export.save(exported, "model.pt2")

# Load (no need for original model code!)
loaded = torch.export.load("model.pt2")
result = loaded.module()(torch.randn(1, 10))

The PT2 archive contains:

• The computation graph (serialized FX graph)
• Model weights (state dict)
• Metadata (dynamic shape constraints, etc.)

This is a self-contained format. You do not need the original Python model definition to load and run the model.

Saving for Different Backends

# Save for later AOTInductor compilation
torch.export.save(exported, "model.pt2")

# Save for ONNX (different path)
torch.onnx.export(model, example_input, "model.onnx", dynamo=True)

Versioning and Compatibility

PT2 archives are versioned. PyTorch maintains backward compatibility: a model saved with an older PyTorch version can be loaded with a newer version (within the same major version).

---

Graph Inspection

After exporting, you can inspect the computation graph to understand what was captured and verify correctness.

Viewing the Graph

exported = torch.export.export(model, example_input)

# Print the graph (human-readable)
print(exported.graph_module.graph)

# Print generated code (more readable)
print(exported.graph_module.code)

Walking the Graph Nodes

Each node in the graph represents an operation:

for node in exported.graph_module.graph.nodes:
    print(f"Op: {node.op:15s} | Target: {node.target} | Args: {node.args}")

Node types:

• placeholder: Input to the graph (parameters, buffers, user inputs)
• call_function: A function call (ATen operator)
• output: The graph's return value
• get_attr: Access a stored attribute

Listing All Operations

# Get the set of all ATen ops used in the graph
ops = set()
for node in exported.graph_module.graph.nodes:
    if node.op == "call_function":
        ops.add(str(node.target))

print(f"Operations used ({len(ops)}):")
for op in sorted(ops):
    print(f"  {op}")

Understanding the Graph Signature

sig = exported.graph_signature

# What are the inputs?
print("User inputs:", sig.input_specs)

# What are the outputs?
print("Outputs:", sig.output_specs)

# Parameters and buffers
print("Parameters:", [s for s in sig.input_specs if s.kind.name == "PARAMETER"])

---

AOTInductor

AOTInductor (Ahead-Of-Time Inductor) compiles an exported model into a native shared library (.so on Linux, .dylib on macOS). The result runs without Python, with maximum performance.

Compilation Flow

ExportedProgram → AOTInductor → .so shared library → C++ or Python inference

Python API

import torch

model = MyModel().eval()
example_input = (torch.randn(1, 3, 224, 224),)

# Export
exported = torch.export.export(model, example_input)

# Compile to .so
compiled_path = torch._inductor.aot_compile(
    exported.module(),
    example_input,
    options={"aot_inductor.output_path": "model.so"},
)

# Load and run in Python (for testing)
compiled_model = torch._inductor.aot_load(compiled_path)
result = compiled_model(torch.randn(1, 3, 224, 224))

Package API (Recommended)

The package API bundles the model, weights, and metadata:

# Compile and package
torch._inductor.aoti_compile_and_package(
    exported,
    package_path="model_package.pt2",
)

# Load the package
loaded = torch._inductor.aoti_load_package("model_package.pt2")
result = loaded(torch.randn(1, 3, 224, 224))

C++ Deployment

The primary use case for AOTInductor is deploying without Python:

#include <torch/csrc/inductor/aoti_runner/model_container_runner.h>

int main() {
    // Load the compiled model
    auto runner = torch::inductor::AOTIModelContainerRunner("model.so");

    // Create input tensor
    auto input = torch::randn({1, 3, 224, 224});
    std::vector<torch::Tensor> inputs = {input};

    // Run inference
    auto outputs = runner.run(inputs);
    auto result = outputs[0];

    return 0;
}

When to Use AOTInductor

• Maximum inference performance (compiled, optimized kernels)
• Deployment without Python
• Server-side GPU inference
• When you can compile ahead of time (shapes known in advance)

---

NativeRT

NativeRT is a C++ inference engine for running PyTorch models. While AOTInductor compiles to native code, NativeRT interprets the exported graph using PyTorch's C++ runtime.

NativeRT vs AOTInductor

When to Use NativeRT

• When AOTInductor compilation is too slow or complex
• When you need C++ inference without ahead-of-time compilation
• Quick prototyping of C++ deployment

Usage Pattern

# Export and save
exported = torch.export.export(model, example_input)
torch.export.save(exported, "model.pt2")

# In C++ with NativeRT:
# auto runner = torch::nativert::ModelRunner("model.pt2");
# auto output = runner.run(inputs);

---

ONNX Export

ONNX (Open Neural Network Exchange) is an open format for representing ML models. PyTorch can export to ONNX for running on ONNX Runtime, which supports multiple hardware backends.

Modern ONNX Export (Dynamo-based)

import torch

model = MyModel().eval()
example_input = (torch.randn(1, 3, 224, 224),)

# Export to ONNX using the new dynamo-based exporter
onnx_program = torch.onnx.export(
    model,
    example_input,
    dynamo=True,
)

# Save to file
onnx_program.save("model.onnx")

Running with ONNX Runtime

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("model.onnx")

# Get input name
input_name = session.get_inputs()[0].name

# Run inference
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = session.run(None, {input_name: input_data})

Dynamic Shapes with ONNX

from torch.export import Dim

batch = Dim("batch", min=1, max=128)

onnx_program = torch.onnx.export(
    model,
    (torch.randn(4, 3, 224, 224),),
    dynamo=True,
    dynamic_shapes={"x": {0: batch}},
)
onnx_program.save("model_dynamic.onnx")

When to Use ONNX

• Cross-framework deployment (model trained in PyTorch, deployed with

TensorFlow Serving, etc.)

• Hardware with ONNX Runtime support but not PyTorch C++ runtime
• When inference hardware vendor provides an ONNX Runtime backend

---

TorchServe

TorchServe is PyTorch's model serving framework. It handles the operational concerns of serving models in production: batching, scaling, monitoring, and A/B testing.

Overview

Client → TorchServe ┬→ Worker 1 → Model
   (HTTP/gRPC)       ├→ Worker 2 → Model
                     └→ Worker N → Model

Key Features

• Dynamic batching: Aggregates individual requests into batches for

efficient GPU utilization

• Multi-model serving: Serve multiple models from one instance
• Model versioning: A/B testing, canary deployments
• Monitoring: Prometheus metrics, logging
• Auto-scaling: Scale workers based on load

Basic Workflow

# 1. Archive the model
torch-model-archiver \
    --model-name my_model \
    --version 1.0 \
    --serialized-file model.pt \
    --handler image_classifier \
    --export-path model_store

# 2. Start TorchServe
torchserve --start \
    --model-store model_store \
    --models my_model=my_model.mar

# 3. Send requests
curl http://localhost:8080/predictions/my_model -T input.jpg

---

ExecuTorch

ExecuTorch is PyTorch's solution for on-device inference on mobile phones, wearables, and embedded systems. It produces small, efficient models that run without a full PyTorch runtime.

Key Features

• Small footprint: Runtime is ~100s of KB (vs PyTorch's ~100s of MB)
• Delegate system: Hardware-specific optimizations (Core ML, XNNPACK,

Qualcomm QNN, etc.)

• No Python: Runs on C++ runtime

Workflow

import torch
from executorch.exir import to_edge_transform_and_lower

model = MyModel().eval()
example_input = (torch.randn(1, 3, 224, 224),)

# Export
exported = torch.export.export(model, example_input)

# Lower to edge
edge_program = to_edge_transform_and_lower(exported)

# Save
edge_program.save("model.pte")

Target Platforms

---

Quantization for Deployment

Quantization reduces model size and increases inference speed by using lower precision arithmetic (e.g., INT8 instead of FP32).

PT2E Quantization Flow

The modern quantization flow is built on torch.export:

import torch
from torch.ao.quantization.quantize_pt2e import (
    prepare_pt2e,
    convert_pt2e,
)
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)

model = MyModel().eval()
example_input = (torch.randn(1, 3, 224, 224),)

# Step 1: Export
exported = torch.export.export(model, example_input)

# Step 2: Prepare for quantization (inserts observers)
quantizer = XNNPACKQuantizer().set_global(
    get_symmetric_quantization_config()
)
prepared = prepare_pt2e(exported, quantizer)

# Step 3: Calibrate with representative data
with torch.no_grad():
    for data in calibration_dataloader:
        prepared(data)

# Step 4: Convert to quantized model
quantized = convert_pt2e(prepared)

# Step 5: Deploy (e.g., with AOTInductor or ExecuTorch)

Quantization Types

Dynamic vs Static Quantization

• Static quantization: Calibrate ranges with representative data. Best

accuracy. Requires calibration dataset.

• Dynamic quantization: Compute ranges at runtime. Simpler setup. Slightly

worse performance on some workloads.

---

Practical Workflow

The End-to-End Journey

1. RESEARCH & TRAINING
   ├── Train in eager mode (nn.Module, autograd)
   ├── Validate accuracy
   └── Save checkpoint

2. EXPORT
   ├── torch.export.export(model, example_inputs)
   ├── Add dynamic shapes for variable inputs
   ├── Fix export issues (torch.cond for control flow, etc.)
   ├── Use draft_export to debug failures
   └── Save: torch.export.save(exported, "model.pt2")

3. OPTIMIZE
   ├── Quantize (PT2E flow) for smaller/faster model
   ├── Profile to identify bottlenecks
   └── Choose deployment target

4. DEPLOY
   ├── Server GPU → AOTInductor (.so) or NativeRT
   ├── Server CPU → ONNX Runtime or AOTInductor
   ├── Model serving → TorchServe
   ├── Mobile/Edge → ExecuTorch (.pte)
   └── Cross-framework → ONNX

Common Patterns

Pattern 1: Quick Server Deployment

model = load_trained_model()
exported = torch.export.export(model, example_input)
torch.export.save(exported, "model.pt2")
# Load in production with torch.export.load("model.pt2")

Pattern 2: Maximum Performance Server

model = load_trained_model()
exported = torch.export.export(model, example_input)
torch._inductor.aoti_compile_and_package(exported, "model_pkg.pt2")
# Deploy .so or package without Python

Pattern 3: Mobile Deployment

model = load_trained_model()
exported = torch.export.export(model, example_input)
# Quantize → Lower to edge → Save .pte

Pattern 4: Cross-Platform

model = load_trained_model()
onnx_program = torch.onnx.export(model, example_input, dynamo=True)
onnx_program.save("model.onnx")
# Deploy with ONNX Runtime on any platform

Debugging Export Failures

• Start with draft_export: Get a partial graph and diagnostic report.
• Check control flow: Replace data-dependent branches with torch.cond.
• Check dynamic shapes: Use Dim for variable dimensions.
• Simplify: Export a smaller part of the model first to isolate issues.
• Check operators: Some custom ops may need registration for export.

Performance Comparison (Rough Guidelines)

---

Files in This Module

---

📓 Open Notebook — Interactive version of this module

Module 12: Model Architectures — From Paper to PyTorch

Building real neural network architectures is the bridge between understanding PyTorch basics and doing real deep learning research or engineering. This module walks through several landmark architectures, explaining the why behind each design choice and showing you how to translate ideas from papers into working PyTorch code.

---

How to Read an Architecture Paper and Translate It to Code

Most deep learning papers follow a predictable structure. Learning to read them systematically will save you enormous amounts of time.

Step 1: Understand the Problem Statement

Before looking at the architecture, understand what problem the paper solves. For ResNet, it's the "degradation problem" — deeper networks were performing worse than shallower ones, even on training data. For Transformers, it was the sequential bottleneck of RNNs that prevented parallelization.

Step 2: Identify the Core Innovation

Every architecture paper has one or two key ideas. Everything else is engineering around those ideas:

Step 3: Map the Architecture Diagram to nn.Modules

Papers usually have an architecture diagram. Each box becomes either:

• An nn.Module subclass (if it's a reusable block)
• A line of code inside forward() (if it's a simple operation)

Step 4: Match Dimensions

Papers describe tensor shapes. Track them through the network. A common approach: write comments with shapes at each step in your forward() method during development, then remove them once the code is tested.

Step 5: Implement, Test with Random Data, Then Train

Always verify your architecture with random tensors before training:

model = MyModel()
x = torch.randn(2, 3, 224, 224)  # batch=2, channels=3, 224x224
out = model(x)
print(out.shape)  # should be (2, num_classes)

---

ResNet (Residual Networks)

Paper: "Deep Residual Learning for Image Recognition" (He et al., 2015)

The Degradation Problem

Before ResNet, researchers observed a paradox: adding more layers to a neural network made it worse, even on the training set. This wasn't overfitting — it was a fundamental optimization problem. Deeper networks were harder to optimize because gradients had to flow through many layers.

The Key Insight: Skip Connections

Instead of learning H(x) directly, learn the residual F(x) = H(x) - x, then compute H(x) = F(x) + x. If the optimal transformation is close to identity, it's easier to learn a small residual than to learn identity from scratch.

Input x ──────────────────────┐
   │                          │
   ▼                          │
┌──────────┐                  │
│  Conv-BN  │                  │
│   ReLU    │                  │
│  Conv-BN  │                  │
└──────────┘                  │
   │                          │
   ▼                          │
  F(x)  ─────── + ◄──────────┘
   │
   ▼
  ReLU
   │
   ▼
 Output = ReLU(F(x) + x)

BasicBlock vs Bottleneck

ResNet uses two block types:

BasicBlock (for ResNet-18 and ResNet-34):

• Two 3x3 convolutions
• Each followed by BatchNorm
• A skip connection that adds the input to the output
• If dimensions don't match, a 1x1 conv "projection" shortcut is used

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3,
                               stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3,
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample

    def forward(self, x):
        identity = x
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        return F.relu(out)

Bottleneck (for ResNet-50, 101, 152):

• Three convolutions: 1x1 (reduce), 3x3 (process), 1x1 (expand)
• The 1x1 convolutions reduce and restore channel dimensions
• This "bottleneck" design is more parameter-efficient for deep networks

The expansion factor controls how much the Bottleneck expands channels:

• BasicBlock: expansion = 1 (output channels == internal channels)
• Bottleneck: expansion = 4 (output channels == 4 * internal channels)

Pre-Activation ResNet

The original ResNet applies BN and ReLU after each convolution. A later paper ("Identity Mappings in Deep Residual Networks", He et al., 2016) showed that applying BN and ReLU before the convolution (pre-activation) improves both optimization and generalization:

# Post-activation (original): Conv -> BN -> ReLU
# Pre-activation (improved):  BN -> ReLU -> Conv

The intuition: in the pre-activation design, the skip connection is a true identity mapping (no BN or ReLU on the shortcut path), allowing gradients to flow unimpeded.

ResNet Family Configurations

See resnet.py for the complete implementation.

---

Transformer

Paper: "Attention Is All You Need" (Vaswani et al., 2017)

Motivation

RNNs process sequences one token at a time — you can't compute step t until step t-1 is done. This sequential bottleneck limits parallelism and makes it hard to learn long-range dependencies. The Transformer replaces recurrence entirely with attention mechanisms.

Scaled Dot-Product Attention

The core operation. Given queries Q, keys K, and values V:

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

• Q @ K^T computes similarity scores between every query and every key
• Division by sqrt(d_k) prevents dot products from becoming too large (which

would push softmax into saturation, giving near-zero gradients)

• Softmax normalizes scores to a probability distribution
• Multiplication with V produces a weighted combination of values

Multi-Head Attention

Instead of one big attention computation, split Q, K, V into h heads:

# Instead of d_model-dimensional attention:
# Split into h heads, each of dimension d_k = d_model // h
# Compute attention independently per head
# Concatenate and project back to d_model

Why multiple heads? Each head can attend to different aspects of the input — one head might focus on syntactic relationships, another on semantic ones.

Self-Attention vs Cross-Attention

• Self-attention: Q, K, V all come from the same sequence. Each token

attends to all other tokens in the same sequence.

• Cross-attention: Q comes from one sequence (decoder), K and V come from

another (encoder output). This is how the decoder "reads" the encoder.

Positional Encoding

Attention is permutation-invariant — it has no notion of order. Positional encodings inject position information. The original paper uses sinusoidal encodings:

PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

These have a nice property: the encoding of position pos+k can be expressed as a linear function of the encoding of position pos, making it possible for the model to learn relative positions.

Encoder-Decoder Architecture

Encoder: Stack of N identical layers, each containing:

• Multi-head self-attention (with residual + LayerNorm)
• Position-wise feed-forward network (with residual + LayerNorm)

Decoder: Stack of N identical layers, each containing:

• Masked multi-head self-attention (causal mask prevents attending to future)
• Multi-head cross-attention (attends to encoder output)
• Position-wise feed-forward network

Each sub-layer uses the pattern: LayerNorm(x + Sublayer(x)) (post-norm) or x + Sublayer(LayerNorm(x)) (pre-norm, more common now).

See transformer.py for the complete implementation.

---

GPT (Decoder-Only Transformer)

Paper: "Language Models are Unsupervised Multitask Learners" (Radford et al., 2019)

Design Philosophy

GPT simplifies the Transformer by keeping only the decoder (with causal attention), removing the encoder and cross-attention entirely. The insight: a powerful enough language model trained to predict the next token can learn to perform many tasks without explicit task-specific architecture.

Causal (Autoregressive) Attention

In a decoder-only model, each token can only attend to itself and previous tokens. This is enforced by a triangular mask:

# For sequence length 4:
mask = [[1, 0, 0, 0],   # token 0 sees only itself
        [1, 1, 0, 0],   # token 1 sees tokens 0-1
        [1, 1, 1, 0],   # token 2 sees tokens 0-2
        [1, 1, 1, 1]]   # token 3 sees tokens 0-3

Positions with 0 are set to -infinity before softmax, effectively zeroing those attention weights.

Weight Tying

GPT ties the token embedding matrix with the output projection (language model head). If embedding maps token IDs to vectors, the LM head maps vectors back to logits over the vocabulary — these are inverse operations, so sharing weights makes sense and reduces parameters significantly:

self.token_embedding = nn.Embedding(vocab_size, d_model)
self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
self.lm_head.weight = self.token_embedding.weight  # weight tying

Pre-Norm vs Post-Norm

Original Transformer: x + Sublayer(LayerNorm(x)) — "post-norm" GPT-2 and later: x + Sublayer(LayerNorm(x)) — "pre-norm"

Wait, the formulas look the same? The difference is subtle:

• Post-norm: LayerNorm(x + Sublayer(x)) — norm is outside the residual
• Pre-norm: x + Sublayer(LayerNorm(x)) — norm is inside the residual

Pre-norm is more stable for training deep networks because the residual connection carries un-normalized values, preserving gradient magnitude.

Generation Strategies

Given a trained model, how do you generate text?

Greedy: Always pick the most probable next token. Fast but repetitive.

Temperature: Divide logits by temperature T before softmax:

• T < 1: sharper distribution, more confident (less random)
• T = 1: original distribution
• T > 1: flatter distribution, more random

Top-k sampling: Keep only the top k most probable tokens, zero out the rest, renormalize, then sample. Prevents sampling very unlikely tokens.

Top-p (nucleus) sampling: Keep the smallest set of tokens whose cumulative probability exceeds p. Adapts the number of candidates dynamically — when the model is confident, few tokens are kept; when uncertain, more are kept.

def top_p_sample(logits, p=0.9):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
    # Remove tokens with cumulative probability above the threshold
    mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= p
    sorted_logits[mask] = float('-inf')
    # Sample from the filtered distribution
    probs = F.softmax(sorted_logits, dim=-1)
    idx = torch.multinomial(probs, 1)
    return sorted_indices.gather(-1, idx)

KV Cache

During autoregressive generation, each new token only needs to attend to all previous tokens. Without caching, you'd recompute K and V for the entire sequence at every step. The KV cache stores previously computed K and V tensors and only computes the new token's K and V, then concatenates:

# Step 1: compute K, V for all tokens
# Step 2: only compute K_new, V_new for the new token
#          K = cat(K_cached, K_new)
#          V = cat(V_cached, V_new)

This reduces generation from O(n^2) to O(n) per token.

See gpt.py for the complete implementation.

---

Vision Transformer (ViT)

Paper: "An Image Is Worth 16x16 Words" (Dosovitskiy et al., 2020)

Core Idea

Treat an image as a sequence of patches and process them with a standard Transformer encoder. No convolutions needed.

Patch Embedding

Split the image into non-overlapping patches (e.g., 16x16 pixels), flatten each patch into a vector, then project to the model dimension:

# Image: (B, 3, 224, 224)
# Patches: 224/16 = 14 patches per side, 14*14 = 196 patches
# Each patch: 16*16*3 = 768 pixels

# Efficient implementation using Conv2d:
self.patch_embed = nn.Conv2d(3, d_model, kernel_size=16, stride=16)
# Output: (B, d_model, 14, 14) -> reshape to (B, 196, d_model)

CLS Token

A special learnable token prepended to the sequence. After processing through the Transformer, the CLS token's representation is used for classification:

self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
# Prepend to patch sequence: (B, 196, d_model) -> (B, 197, d_model)

Why not just average-pool all patch representations? The CLS token provides a single, fixed-position summary token that the model can learn to aggregate global information into.

Position Embedding

Since patches have spatial relationships, learnable position embeddings are added (ViT uses learned embeddings, not sinusoidal):

self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, d_model))
# +1 for the CLS token

Classification Head

After the Transformer encoder, take the CLS token's output and pass it through a simple MLP head:

cls_output = transformer_output[:, 0]  # CLS token at position 0
logits = self.mlp_head(cls_output)

See vit.py for the complete implementation.

---

VAE (Variational Autoencoder)

Paper: "Auto-Encoding Variational Bayes" (Kingma & Welling, 2013)

Autoencoders vs Variational Autoencoders

A regular autoencoder learns to compress data to a latent code and reconstruct it. The latent space, however, has no structure — nearby points in latent space may decode to very different outputs, making it useless for generation.

A VAE forces the latent space to be structured (approximately Gaussian) by:

• Encoding to a distribution (mean and variance) instead of a point
• Sampling from that distribution (the reparameterization trick)
• Adding a KL divergence loss that pushes the distribution toward N(0, I)

The Reparameterization Trick

We can't backpropagate through random sampling. The trick: instead of sampling z ~ N(mu, sigma^2), compute z = mu + sigma * epsilon where epsilon ~ N(0, 1). Now the randomness is in epsilon (which doesn't need gradients), and z is a deterministic function of mu and sigma.

def reparameterize(self, mu, log_var):
    std = torch.exp(0.5 * log_var)
    eps = torch.randn_like(std)
    return mu + eps * std

The ELBO Loss

The VAE loss has two terms:

Loss = Reconstruction Loss + KL Divergence
     = E[||x - x_hat||^2]  + KL(q(z|x) || p(z))

• Reconstruction loss: How well the decoder reconstructs the input.

Binary cross-entropy for binary data, MSE for continuous data.

• KL divergence: How far the encoder's distribution is from the prior

N(0, I). For Gaussians, this has a closed-form solution:

kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

The KL term acts as a regularizer, preventing the encoder from collapsing the latent space to a single point (which would defeat the purpose of having a generative model).

See vae.py for the complete implementation.

---

U-Net

Paper: "U-Net: Convolutional Networks for Biomedical Image Segmentation" (Ronneberger et al., 2015)

Architecture Overview

U-Net has a symmetric encoder-decoder structure with skip connections:

Encoder                    Decoder
(downsampling)             (upsampling)
                     
[Input] ──────────────────────► [Output]
   │                               ▲
   ▼                               │
[Down1] ──── skip connection ──► [Up1]
   │                               ▲
   ▼                               │
[Down2] ──── skip connection ──► [Up2]
   │                               ▲
   ▼                               │
[Down3] ──── skip connection ──► [Up3]
   │                               ▲
   ▼                               │
[Down4] ──── skip connection ──► [Up4]
   │                               ▲
   ▼                               │
        [Bottleneck] ──────────────┘

Why Skip Connections?

The encoder captures "what" is in the image (semantic features) but loses spatial precision. The decoder recovers spatial resolution but lacks semantic context. Skip connections combine both: the decoder gets high-resolution features from the encoder concatenated with upsampled semantic features.

Encoder Path

Each encoder block: two 3x3 convolutions + ReLU, then 2x2 max pooling. Channels double at each level: 64 -> 128 -> 256 -> 512 -> 1024.

Decoder Path

Each decoder block: 2x2 transposed convolution (upsample), concatenate with the corresponding encoder features, then two 3x3 convolutions + ReLU.

Key Difference from ResNet Skip Connections

• ResNet: skip connections add the input (element-wise addition)
• U-Net: skip connections concatenate encoder features with decoder features

This is because U-Net needs to preserve both the high-resolution spatial information from the encoder and the semantic information from the decoder, while ResNet just needs to facilitate gradient flow.

---

Common Architectural Patterns

Residual Connections

Found in almost every modern architecture. The core pattern:

output = x + f(x)  # residual connection

Benefits: easier optimization, better gradient flow, ability to train very deep networks.

Pre-Norm vs Post-Norm

Post-norm (original Transformer):

x = layer_norm(x + sublayer(x))

Pre-norm (GPT-2, modern practice):

x = x + sublayer(layer_norm(x))

Pre-norm is more stable for training. Post-norm can achieve slightly better performance but requires careful learning rate warmup.

Weight Tying

Sharing parameters between the input embedding and the output projection:

self.embed = nn.Embedding(vocab_size, d_model)
self.output_proj = nn.Linear(d_model, vocab_size, bias=False)
self.output_proj.weight = self.embed.weight  # shared!

Reduces parameters by vocab_size * d_model and acts as regularization. Used in GPT, BERT, T5, and most modern language models.

Layer Scaling

Introduced in CaiT (Going Deeper with Image Transformers). Scale the output of each residual block by a learnable scalar, initialized to a small value (e.g., 0.1):

self.gamma = nn.Parameter(torch.ones(d_model) * 0.1)
# In forward:
x = x + self.gamma * sublayer(x)

This helps stabilize training of very deep Transformers by starting with near-identity blocks.

GELU Activation

Most modern Transformers use GELU (Gaussian Error Linear Unit) instead of ReLU:

F.gelu(x)  # smooth approximation: x * Phi(x)

GELU is smoother than ReLU and has been empirically shown to work better for Transformers. It's the default in GPT, BERT, ViT, and most modern architectures.

Dropout Patterns in Transformers

Dropout is typically applied in three places:

• After attention weights (before multiplying with V)
• After the feed-forward network's output projection
• After adding positional embeddings (in some architectures)

attn_weights = F.softmax(scores, dim=-1)
attn_weights = F.dropout(attn_weights, p=0.1, training=self.training)

Weight Initialization

Different architectures use different initialization strategies:

# Xavier/Glorot (good for tanh/sigmoid): used in Transformer
nn.init.xavier_uniform_(self.weight)

# Kaiming/He (good for ReLU): used in ResNet
nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='relu')

# Normal with small std: used in GPT
nn.init.normal_(self.weight, mean=0.0, std=0.02)

The right initialization prevents vanishing/exploding gradients at the start of training and can significantly affect convergence speed.

---

Summary

Files in This Module

• resnet.py — Complete ResNet with BasicBlock, Bottleneck, configs for 18/34/50/101
• transformer.py — Full encoder-decoder Transformer from scratch
• gpt.py — GPT with generation: greedy, temperature, top-k, top-p sampling
• vit.py — Vision Transformer for image classification
• vae.py — Variational Autoencoder with reparameterization trick and ELBO loss

---

📓 Open Notebook — Interactive version of this module

Module 13: Advanced PyTorch Features

This module covers advanced PyTorch capabilities that go beyond standard model building and training. These are the tools that separate "I can train a model" from "I can build production-quality ML systems."

---

Functorch (torch.func): Functional Transformations

torch.func (formerly the standalone functorch library) provides composable function transforms inspired by JAX. The core idea: transform a plain Python function into a new function with different behavior.

vmap — Vectorized Map

vmap automatically vectorizes a function over a batch dimension. Instead of writing explicit batch loops or reshaping tensors, you write the function for a single example and vmap handles batching:

import torch
from torch.func import vmap

def compute_norm(x):
    """Compute L2 norm of a single vector."""
    return torch.sqrt(torch.sum(x ** 2))

# Without vmap: need to handle batch dimension explicitly
batch = torch.randn(32, 10)
norms_loop = torch.stack([compute_norm(batch[i]) for i in range(32)])

# With vmap: automatic batching
norms_vmap = vmap(compute_norm)(batch)

Why use vmap instead of just writing batched code? Three reasons:

• Clarity: Write single-example logic, get batched execution
• Correctness: No batch dimension bugs
• Composition: Combine with grad, jacrev, etc.

grad — Functional Gradient

torch.func.grad computes gradients functionally, without modifying tensors in-place or using .backward():

from torch.func import grad

def f(x):
    return torch.sin(x).sum()

# grad returns a function that computes the gradient
grad_f = grad(f)
x = torch.tensor([1.0, 2.0, 3.0])
print(grad_f(x))  # cos(x) = [0.5403, -0.4161, -0.9900]

grad is particularly useful when composed with other transforms.

jacrev / jacfwd — Jacobian Computation

The Jacobian matrix contains all partial derivatives of a vector-valued function:

from torch.func import jacrev, jacfwd

def f(x):
    return torch.stack([x[0]**2 + x[1], x[0] * x[1]**2])

x = torch.tensor([1.0, 2.0])
J_rev = jacrev(f)(x)   # Reverse-mode: efficient when output dim < input dim
J_fwd = jacfwd(f)(x)   # Forward-mode: efficient when input dim < output dim

Rule of thumb:

• Use jacrev when output dimension is smaller than input dimension
• Use jacfwd when input dimension is smaller than output dimension

hessian — Second Derivatives

The Hessian matrix is the Jacobian of the gradient. torch.func.hessian is syntactic sugar for jacrev(jacrev(f)) or jacfwd(jacrev(f)):

from torch.func import hessian

def f(x):
    return (x ** 3).sum()

H = hessian(f)(torch.tensor([1.0, 2.0, 3.0]))
# H[i,j] = d^2f / dx_i dx_j

Composing Transforms

The real power is in composition:

from torch.func import vmap, grad, jacrev

# Per-sample gradients: grad of loss for each sample in a batch
def loss_fn(params, x, y):
    pred = model_fn(params, x)
    return ((pred - y) ** 2).sum()

# Batched Jacobian: Jacobian for each sample in a batch
batched_jacobian = vmap(jacrev(model_fn), in_dims=(None, 0))

See functorch_transforms.py for complete examples.

---

Per-Sample Gradients

The classic vmap + grad use case. Standard training computes the average gradient across a batch. But sometimes you need the gradient for each sample individually — for example, in differential privacy (DP-SGD), where you need to clip per-sample gradients before averaging.

Without vmap, you'd need to loop over samples or use inefficient tricks. With vmap:

from torch.func import vmap, grad
from torch import nn

model = nn.Linear(10, 1)
params = dict(model.named_parameters())

def compute_loss(params, x, y):
    # Stateless function: takes params explicitly
    pred = torch.func.functional_call(model, params, (x,))
    return ((pred - y) ** 2).squeeze()

# grad w.r.t. params for a single sample
grad_fn = grad(compute_loss)

# vmap over the batch dimension of x and y
per_sample_grads = vmap(grad_fn, in_dims=(None, 0, 0))(params, X_batch, Y_batch)

See per_sample_gradients.py for a complete walkthrough.

---

Sparse Tensors

Sparse tensors store only nonzero elements, saving memory and computation for data that is mostly zeros (e.g., adjacency matrices, bag-of-words features).

COO (Coordinate) Format

Stores row and column indices alongside values. Good for construction and conversion, less efficient for arithmetic:

indices = torch.tensor([[0, 1, 2], [1, 0, 2]])  # (2, nnz)
values = torch.tensor([3.0, 4.0, 5.0])
sparse_coo = torch.sparse_coo_tensor(indices, values, size=(3, 3))

CSR (Compressed Sparse Row) Format

Stores row pointers, column indices, and values. Efficient for row-slicing and matrix-vector products:

crow_indices = torch.tensor([0, 1, 2, 3])  # row pointers
col_indices = torch.tensor([1, 0, 2])       # column indices
values = torch.tensor([3.0, 4.0, 5.0])
sparse_csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, size=(3, 3))

BSR (Block Sparse Row) Format

Like CSR but stores dense blocks instead of individual elements. Useful when sparsity has block structure (common in structured pruning):

crow_indices = torch.tensor([0, 1, 2])
col_indices = torch.tensor([0, 1])
values = torch.randn(2, 2, 2)  # two 2x2 blocks
sparse_bsr = torch.sparse_bsr_tensor(crow_indices, col_indices, values, size=(4, 4))

When to use each:

• COO: Building sparse tensors, format conversion, unstructured updates
• CSR: Sparse matrix-vector products, row-based access patterns
• BSR: Block-structured sparsity, GPU-friendly operations

Sparse Operations

# Matrix multiply (sparse @ dense)
result = torch.sparse.mm(sparse_csr, dense_matrix)

# Element-wise operations
sparse_sum = sparse_coo + sparse_coo
sparse_scaled = sparse_coo * 2.0

# Convert between formats
dense = sparse_coo.to_dense()
csr = sparse_coo.to_sparse_csr()

---

Complex Numbers

PyTorch natively supports complex tensors, essential for signal processing, quantum computing simulations, and Fourier analysis.

# Creating complex tensors
z = torch.complex(torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]))
z = torch.tensor([1+3j, 2+4j])  # Python complex literals

# Operations
z.real      # real part
z.imag      # imaginary part
z.abs()     # magnitude
z.angle()   # phase angle
z.conj()    # complex conjugate

# FFT
signal = torch.randn(1000)
spectrum = torch.fft.fft(signal)
freqs = torch.fft.fftfreq(1000)

# Inverse FFT
reconstructed = torch.fft.ifft(spectrum).real

See sparse_and_complex.py for complete examples.

---

Quantization Overview

Quantization reduces model size and increases inference speed by using lower precision (e.g., INT8 instead of FP32).

Why Quantize?

• Memory: INT8 uses 4x less memory than FP32
• Speed: Integer arithmetic is faster, especially on mobile/edge devices
• Accuracy: Modern quantization preserves most of the model's accuracy

Quantization Approaches

Post-Training Quantization (PTQ): Quantize a pre-trained model without retraining. Fast but may lose more accuracy.

Quantization-Aware Training (QAT): Simulate quantization during training so the model learns to be robust to quantization noise. Better accuracy but requires retraining.

PT2E Quantization Flow: The modern PyTorch 2 Export quantization approach. Uses torch.export to capture the model as a graph, applies quantization annotations, and lowers to optimized backends:

# Conceptual PT2E flow (simplified)
import torch
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e

model = MyModel()
exported = torch.export.export(model, example_inputs)
prepared = prepare_pt2e(exported, quantizer)
# Calibrate with representative data
for batch in calibration_data:
    prepared(batch)
quantized = convert_pt2e(prepared)

torchao: A library for architecture optimization, including quantization, sparsity, and low-precision training. Offers simple APIs:

# Conceptual torchao usage
import torchao
torchao.quantize_(model, torchao.quantization.int8_weight_only())

---

Custom Operators (torch.library)

When PyTorch's built-in ops don't cover your needs, you can define custom operators with proper integration into autograd, torch.compile, and torch.export.

import torch
from torch.library import Library, impl

# Create a library namespace for your custom ops
my_lib = Library("myops", "DEF")

# Define the op signature
my_lib.define("my_relu(Tensor x) -> Tensor")

# Register a CPU implementation
@impl(my_lib, "my_relu", "CPU")
def my_relu_cpu(x):
    return x.clamp(min=0)

# Register a Meta (shape-only) implementation for torch.compile
@impl(my_lib, "my_relu", "Meta")
def my_relu_meta(x):
    return torch.empty_like(x)

# Use the custom op
x = torch.randn(5)
result = torch.ops.myops.my_relu(x)

Custom Autograd for Custom Ops

# Register autograd formula
def my_relu_backward(ctx, grad_output):
    x, = ctx.saved_tensors
    return grad_output * (x > 0).float()

torch.library.impl_abstract("myops::my_relu", my_relu_meta)

See custom_operators.py for a complete example.

---

C++ Extensions

For performance-critical code, you can write custom C++ (and CUDA) extensions:

from torch.utils.cpp_extension import load

# JIT compilation: compiles C++ code on first use
my_extension = load(
    name="my_extension",
    sources=["my_extension.cpp"],
    verbose=True,
)

The C++ side uses the PyTorch C++ API (LibTorch):

#include <torch/extension.h>

torch::Tensor my_add(torch::Tensor a, torch::Tensor b) {
    return a + b;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("my_add", &my_add, "Custom add");
}

---

Profiling Deep Dive

PyTorch's profiler helps identify performance bottlenecks.

Basic Profiling

import torch
from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU],
    record_shapes=True,
    profile_memory=True,
) as prof:
    with record_function("model_inference"):
        output = model(input_tensor)

# Print a summary table sorted by CPU time
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

Chrome Trace

Export a trace viewable in Chrome's chrome://tracing or Perfetto UI:

prof.export_chrome_trace("trace.json")

record_function

Annotate specific code regions for fine-grained profiling:

with record_function("my_attention"):
    attn_output = attention(q, k, v)

with record_function("my_ffn"):
    ffn_output = feed_forward(attn_output)

TensorBoard Integration

with profile(
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./log_dir"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(dataloader):
        output = model(batch)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        prof.step()

See profiling.py for runnable examples.

---

Memory Profiling

Basic Memory Tracking

# Note: these require CUDA, shown for reference
torch.cuda.memory_allocated()      # current memory used by tensors
torch.cuda.max_memory_allocated()  # peak memory since last reset
torch.cuda.memory_reserved()       # total memory reserved by allocator
torch.cuda.reset_peak_memory_stats()

Finding Memory Leaks

Common causes:

• Storing tensors in a list that grows: solution is to .detach() or

store .item() for scalar values

• Not clearing gradients: call optimizer.zero_grad() each step
• Keeping computation graph alive: use .detach() or with torch.no_grad()

# BAD: keeps entire computation graph in memory
losses = []
for batch in dataloader:
    loss = model(batch).sum()
    losses.append(loss)  # holds onto graph!

# GOOD: detach the scalar value
losses = []
for batch in dataloader:
    loss = model(batch).sum()
    losses.append(loss.item())  # just a float, no graph

---

Debugging Techniques

Anomaly Detection

Detects the operation that produced a NaN or Inf gradient:

with torch.autograd.detect_anomaly():
    output = model(input)
    loss = criterion(output, target)
    loss.backward()  # will print a traceback if NaN/Inf is detected

Warning: this is SLOW. Only use for debugging.

Gradient Checking

Verify your autograd implementation by comparing with finite differences:

from torch.autograd import gradcheck

func = MyCustomFunction.apply
input = torch.randn(3, 4, dtype=torch.float64, requires_grad=True)
assert gradcheck(func, input, eps=1e-6, atol=1e-4)

Common Error Messages and Fixes

"one of the variables needed for gradient computation has been modified by an inplace operation"

• Cause: in-place operation (like x += 1) on a tensor that requires grad
• Fix: use x = x + 1 (out-of-place) instead

"Trying to backward through the graph a second time"

• Cause: calling .backward() twice without retain_graph=True
• Fix: either use retain_graph=True or restructure to avoid double backward

"Expected all tensors to be on the same device"

• Cause: mixing CPU and GPU tensors in one operation
• Fix: ensure all tensors are on the same device with .to(device)

"RuntimeError: mat1 and mat2 shapes cannot be multiplied"

• Cause: shape mismatch in linear layers
• Fix: print shapes before the operation to identify the mismatch

See debugging_tips.py for practical debugging examples.

---

torch.fx: Symbolic Tracing and Graph Transformation

torch.fx symbolically traces a model to produce a graph IR (intermediate representation) that you can analyze and transform.

Basic Tracing

import torch.fx

class MyModel(nn.Module):
    def forward(self, x):
        x = torch.relu(x)
        x = x + 1
        return x

model = MyModel()
traced = torch.fx.symbolic_trace(model)
print(traced.graph)  # shows the operations as a graph

Writing Custom Passes

def replace_relu_with_gelu(module):
    """Replace all ReLU calls with GELU."""
    traced = torch.fx.symbolic_trace(module)
    for node in traced.graph.nodes:
        if node.op == "call_function" and node.target == torch.relu:
            node.target = torch.nn.functional.gelu
    traced.graph.lint()  # validate the graph
    traced.recompile()
    return traced

Use Cases

• Quantization: Analyze the graph to determine where to insert quant/dequant nodes
• Fusion: Merge compatible operations (e.g., Conv + BN)
• Shape inference: Propagate shapes through the graph without running data
• Visualization: Understand model structure programmatically

---

Meta Device: Shape Inference Without Memory

The meta device lets you analyze models without allocating real memory. Tensors on the meta device have shapes and dtypes but no data:

# Create a model on the meta device (no memory allocated)
with torch.device("meta"):
    model = nn.Linear(1000, 1000)
    # model.weight.shape == (1000, 1000) but uses 0 bytes

# Analyze input/output shapes
x = torch.empty(32, 1000, device="meta")
out = model(x)
print(out.shape)  # torch.Size([32, 1000])

Use cases:

• Model analysis: Count parameters and compute shapes for huge models

that don't fit in memory

• Architecture prototyping: Verify shapes without waiting for memory

allocation

• Deferred initialization: Create model structure on meta, then

materialize weights on the target device

# Count parameters of a huge model without any memory
with torch.device("meta"):
    huge_model = nn.Sequential(
        nn.Linear(10000, 10000),
        nn.ReLU(),
        nn.Linear(10000, 10000),
    )

total_params = sum(p.numel() for p in huge_model.parameters())
memory_gb = total_params * 4 / 1e9  # FP32 = 4 bytes
print(f"Parameters: {total_params:,}, Memory: {memory_gb:.2f} GB")

---

Summary

Files in This Module

• functorch_transforms.py — vmap, grad, jacrev, hessian demonstrations
• per_sample_gradients.py — Per-sample gradient computation with vmap+grad
• custom_operators.py — Defining custom ops with torch.library
• profiling.py — Profiler usage, timing, and analysis
• sparse_and_complex.py — Sparse tensors, complex numbers, and FFT
• debugging_tips.py — Anomaly detection, gradient flow checking, and common fixes

---

📓 Open Notebook — Interactive version of this module

Module 14: Testing and Reproducibility

Writing tests and ensuring reproducibility are essential skills that separate hobby projects from production-quality deep learning code. This module covers PyTorch's testing framework, reproducibility techniques, and benchmarking.

---

PyTorch's Testing Framework

PyTorch has its own testing infrastructure built on top of Python's unittest. The key class is TestCase from torch.testing._internal.common_utils, which provides tensor-aware assertions and convenient utilities.

Basic Test Structure

from torch.testing._internal.common_utils import run_tests, TestCase
import torch

class TestMyFeature(TestCase):
    def test_addition(self):
        a = torch.tensor([1.0, 2.0, 3.0])
        b = torch.tensor([4.0, 5.0, 6.0])
        result = a + b
        expected = torch.tensor([5.0, 7.0, 9.0])
        self.assertEqual(result, expected)

    def test_shape(self):
        x = torch.randn(3, 4, 5)
        self.assertEqual(x.shape, (3, 4, 5))

    def test_dtype(self):
        x = torch.zeros(5, dtype=torch.float32)
        self.assertEqual(x.dtype, torch.float32)

if __name__ == "__main__":
    run_tests()

assertEqual for Tensors

PyTorch's assertEqual is smarter than the standard library version. For tensors, it checks:

• Shape equality
• Dtype equality
• Value equality (with configurable tolerance for floating point)

# Exact equality (for integer tensors)
self.assertEqual(torch.tensor([1, 2, 3]), torch.tensor([1, 2, 3]))

# Approximate equality (for float tensors) — uses default tolerances
self.assertEqual(
    torch.tensor([1.0, 2.0]),
    torch.tensor([1.0 + 1e-7, 2.0 - 1e-7]),
)

# Custom tolerances
self.assertEqual(a, b, atol=1e-4, rtol=1e-4)

Useful Assertions

# Check that a function raises a specific exception
with self.assertRaises(RuntimeError):
    torch.tensor([1, 2]) + torch.tensor([1, 2, 3])

# Check error message content
with self.assertRaisesRegex(RuntimeError, "size mismatch"):
    bad_operation()

# Boolean checks
self.assertTrue(torch.all(x > 0))
self.assertFalse(torch.any(torch.isnan(x)))

Parametrized Tests

Use the @parametrize decorator to run a test with multiple inputs:

from torch.testing._internal.common_utils import parametrize

class TestOps(TestCase):
    @parametrize("dtype", [torch.float32, torch.float64])
    @parametrize("size", [(2, 3), (4, 5)])
    def test_zeros(self, dtype, size):
        x = torch.zeros(size, dtype=dtype)
        self.assertEqual(x.sum().item(), 0.0)
        self.assertEqual(x.dtype, dtype)
        self.assertEqual(x.shape, size)

This generates 4 test cases (2 dtypes x 2 sizes), each with a descriptive name.

Device-Generic Tests

For testing across CPU and (optionally) GPU:

from torch.testing._internal.common_device_type import (
    instantiate_device_type_tests,
    dtypes,
)

class TestMyOp(TestCase):
    @dtypes(torch.float32, torch.float64)
    def test_my_op(self, device, dtype):
        x = torch.randn(10, device=device, dtype=dtype)
        result = my_op(x)
        self.assertEqual(result.device.type, device)
        self.assertEqual(result.dtype, dtype)

instantiate_device_type_tests(TestMyOp, globals())

This creates separate test classes for CPU and CUDA (if available), testing each dtype on each device.

---

OpInfo Framework

PyTorch uses the OpInfo framework for systematic operator testing. Each operator has an OpInfo entry that describes:

• The operator function
• Valid input dtypes
• Sample inputs (for testing)
• Reference implementations (for correctness checking)
• Gradient test configurations

While you probably won't need to write OpInfo entries unless contributing to PyTorch core, understanding the concept helps you write better tests:

# The idea: describe an op's properties declaratively, then auto-generate tests
# OpInfo("torch.add",
#     dtypes=floating_types_and(torch.half),
#     sample_inputs_func=sample_inputs_add,
#     supports_out=True,
# )

---

Reproducibility

Non-deterministic behavior makes debugging nearly impossible. Here's how to control randomness in PyTorch.

Setting All Seeds

import random
import numpy as np
import torch

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    # Also seeds all CUDA devices:
    torch.cuda.manual_seed_all(seed)

Deterministic Mode

Even with fixed seeds, some operations have non-deterministic GPU implementations for performance. To enforce full determinism:

torch.use_deterministic_algorithms(True)
# Also set the CUBLAS workspace config for full CUDA determinism:
import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

Warning: deterministic mode may be slower and some operations will raise errors if no deterministic implementation exists.

torch.backends Settings

# CuDNN: control convolution algorithm selection
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# benchmark=True auto-selects the fastest algorithm, but this selection is
# non-deterministic. Set to False for reproducibility.

DataLoader Reproducibility

def seed_worker(worker_id):
    """Ensure each DataLoader worker has a different but reproducible seed."""
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

g = torch.Generator()
g.manual_seed(42)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    worker_init_fn=seed_worker,
    generator=g,
)

See reproducibility.py for a complete setup.

---

Benchmarking

Reliable benchmarking requires attention to detail: warmup, synchronization, statistical rigor.

torch.utils.benchmark.Timer

The recommended way to benchmark PyTorch code:

from torch.utils.benchmark import Timer

t = Timer(
    stmt="torch.mm(A, B)",
    globals={"A": torch.randn(256, 256), "B": torch.randn(256, 256), "torch": torch},
    label="Matrix Multiplication",
    sub_label="256x256",
)

# blocked_autorange: automatically determines the number of runs
result = t.blocked_autorange(min_run_time=1.0)
print(result)  # prints median, IQR, and other statistics

Why use Timer instead of time.time()?

• Warmup: Automatically warms up JIT compilation, CUDA initialization
• Statistics: Reports median and IQR, not just mean
• Synchronization: Handles CUDA synchronization correctly
• Isolation: Minimizes interference from other processes

Comparing Implementations

from torch.utils.benchmark import Compare

results = []
for size in [64, 128, 256, 512]:
    A = torch.randn(size, size)
    B = torch.randn(size, size)

    for label, stmt in [("mm", "torch.mm(A, B)"), ("@", "A @ B")]:
        t = Timer(
            stmt=stmt,
            globals={"A": A, "B": B, "torch": torch},
            label="matmul",
            sub_label=label,
            description=f"{size}x{size}",
        )
        results.append(t.blocked_autorange(min_run_time=0.5))

compare = Compare(results)
compare.print()

Common Benchmarking Mistakes

• No warmup: First runs include JIT compilation and caching overhead
• No CUDA synchronization: CUDA ops are asynchronous — time.time() measures

only the kernel launch, not execution

• Too few runs: A single measurement is noisy
• Benchmarking in training mode: BatchNorm and Dropout add overhead

See benchmarking.py for complete examples.

---

Common Testing Patterns

Testing Numerical Correctness

def test_softmax_correctness(self):
    x = torch.randn(5, 10)
    result = F.softmax(x, dim=-1)

    # Property: sums to 1 along the softmax dimension
    self.assertTrue(torch.allclose(result.sum(dim=-1), torch.ones(5), atol=1e-6))

    # Property: all values are in [0, 1]
    self.assertTrue((result >= 0).all())
    self.assertTrue((result <= 1).all())

Testing Gradient Correctness

from torch.autograd import gradcheck

def test_custom_op_gradient(self):
    func = my_custom_op
    # Use float64 for numerical gradient checking (better precision)
    input = torch.randn(3, 4, dtype=torch.float64, requires_grad=True)
    self.assertTrue(gradcheck(func, input, eps=1e-6, atol=1e-4))

Testing with Approximate Equality

# torch.testing.assert_close: the modern way to check approximate equality
torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)

# For very loose checks (e.g., stochastic operations)
self.assertTrue(torch.allclose(result, expected, atol=0.1, rtol=0.1))

Testing Model Invariants

def test_model_deterministic_eval(self):
    """In eval mode, same input should produce same output."""
    model = MyModel()
    model.eval()
    x = torch.randn(4, 10)
    with torch.no_grad():
        out1 = model(x)
        out2 = model(x)
    self.assertEqual(out1, out2)

def test_model_output_shape(self):
    """Output shape should be (batch_size, num_classes)."""
    model = MyModel(num_classes=10)
    for batch_size in [1, 4, 16]:
        x = torch.randn(batch_size, 3, 32, 32)
        out = model(x)
        self.assertEqual(out.shape, (batch_size, 10))

---

Summary

Files in This Module

• test_example.py — Complete test file using PyTorch's TestCase
• reproducibility.py — Full reproducibility setup and verification
• benchmarking.py — Benchmarking with torch.utils.benchmark

---

📓 Open Notebook — Interactive version of this module

Module 15: Practical PyTorch Utilities

The Hidden Toolkit Most Tutorials Never Teach

PyTorch ships with a rich set of utility modules that most tutorials skip entirely. These are the tools that separate a beginner from a practitioner: weight parametrization, pruning, normalization techniques, sequence packing, nested tensors, and model fusion. This module covers them all.

---

Table of Contents

• torch.nn.utils.parametrize — Weight Constraints
• torch.nn.utils.prune — Model Pruning
• Spectral Norm & Weight Norm
• torch.nn.utils.rnn — Sequence Packing
• Conv-BN Fusion
• torch.nested — Nested (Jagged) Tensors
• torch.nn.utils.clip_grad — Gradient Clipping Internals
• parameters_to_vector & skip_init

---

1. Weight Parametrization

What it is: A framework to apply constraints or transformations to module parameters. Instead of manually enforcing constraints in forward(), you register a parametrization that automatically transforms the raw weight into a constrained version.

Why it matters: Enforcing constraints like orthogonality, symmetry, or positivity on weights is common in research and production. Without parametrization, you'd need hacky workarounds.

How It Works

import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P

class Symmetric(nn.Module):
    """Parametrization that makes a matrix symmetric."""
    def forward(self, X):
        return X.triu() + X.triu(1).transpose(-1, -2)

linear = nn.Linear(5, 5)
P.register_parametrization(linear, "weight", Symmetric())

# Now linear.weight is ALWAYS symmetric
print(linear.weight)  # Symmetric!
print(torch.allclose(linear.weight, linear.weight.T))  # True

# The raw unconstrained parameter is stored as:
print(linear.parametrizations.weight.original)

Key Concept: Original vs Parametrized

When you register a parametrization:

• The original unconstrained tensor is stored at module.parametrizations.<name>.original
• Accessing module.<name> runs the parametrization on the original and returns the result
• The optimizer updates the original (unconstrained) parameter
• The parametrization is applied on every access (or cached)

Built-in Parametrizations

from torch.nn.utils import parametrizations

# Orthogonal weight matrix (useful for RNNs, preventing vanishing/exploding gradients)
linear = nn.Linear(5, 5)
parametrizations.orthogonal(linear, "weight")
# linear.weight is now always orthogonal: W^T W = I

# Spectral normalization (stabilize GANs and training)
conv = nn.Conv2d(3, 64, 3)
parametrizations.spectral_norm(conv, "weight")
# Constrains the spectral norm (largest singular value) of the weight to 1

# Weight normalization (decouple magnitude from direction)
linear = nn.Linear(10, 5)
parametrizations.weight_norm(linear, "weight")
# Reparametrizes: w = g * (v / ||v||)

Custom Parametrization Example — Positive Weights

class Positive(nn.Module):
    """Ensures weights are always positive via softplus."""
    def forward(self, X):
        return torch.nn.functional.softplus(X)

linear = nn.Linear(3, 3)
P.register_parametrization(linear, "weight", Positive())
print(linear.weight)  # All positive!
print((linear.weight > 0).all())  # True

Caching for Efficiency

If you use a parametrized weight multiple times in forward() (e.g., RNNs sharing the recurrent kernel), use caching to avoid recomputation:

with P.cached():
    output = model(input)  # Parametrizations computed once, cached

Removing Parametrizations

# Remove and keep the parametrized (constrained) value
P.remove_parametrizations(linear, "weight", leave_parametrized=True)

# Remove and go back to unconstrained original
P.remove_parametrizations(linear, "weight", leave_parametrized=False)

---

2. Model Pruning

What it is: Removing (zeroing out) weights from a neural network to make it smaller and faster.

Why it matters: Pruned models can be 2-10x smaller with minimal accuracy loss. Critical for edge deployment.

Pruning Strategies

How Pruning Works in PyTorch

• The original weight is moved to weight_orig
• A binary mask weight_mask is created
• A forward hook computes weight = weight_orig * weight_mask before each forward pass

import torch.nn.utils.prune as prune

linear = nn.Linear(10, 5)

# Prune 30% of weights (smallest magnitude)
prune.l1_unstructured(linear, name="weight", amount=0.3)

print(linear.weight)        # Pruned weight (has zeros)
print(linear.weight_mask)   # Binary mask
print(linear.weight_orig)   # Original weight

# Count sparsity
zeros = (linear.weight == 0).sum().item()
total = linear.weight.numel()
print(f"Sparsity: {zeros}/{total} = {zeros/total:.1%}")

Global Pruning (Prune Across All Layers)

model = nn.Sequential(
    nn.Linear(100, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 10),
)

# Collect all prunable parameters
parameters_to_prune = [
    (model[0], "weight"),
    (model[2], "weight"),
    (model[4], "weight"),
]

# Globally prune 40% of weights (by L1 magnitude across ALL layers)
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.4,
)

Making Pruning Permanent

# Remove the pruning reparametrization (bake the mask into the weight)
prune.remove(linear, "weight")
# Now linear.weight IS the pruned weight directly (no more mask/orig)

---

3. Spectral Norm & Weight Norm

Spectral Normalization

Controls the Lipschitz constant of a layer by normalizing weights by their spectral norm (largest singular value). Essential for stable GAN training.

Math: $\bar{W} = W / \sigma(W)$ where $\sigma(W)$ is the largest singular value.

from torch.nn.utils import spectral_norm

# Apply spectral norm
conv = nn.Conv2d(3, 64, 3, padding=1)
conv = spectral_norm(conv, name="weight")

# The spectral norm is estimated via power iteration (efficient)
# No full SVD needed — just one vector update per forward pass

Weight Normalization

Decouples weight magnitude from direction: $w = g \cdot \frac{v}{\|v\|}$

The optimizer can separately learn the magnitude $g$ and direction $v$, which often leads to faster convergence.

from torch.nn.utils import weight_norm

linear = nn.Linear(10, 5)
linear = weight_norm(linear, name="weight")

# Now has: linear.weight_g (magnitude) and linear.weight_v (direction)
print(linear.weight_g.shape)  # (5, 1)
print(linear.weight_v.shape)  # (5, 10)

---

4. Sequence Packing for RNNs

Problem: Sequences in a batch have different lengths. Padding wastes computation — the RNN processes pad tokens unnecessarily.

Solution: Pack sequences so the RNN only processes real tokens.

from torch.nn.utils.rnn import (
    pack_padded_sequence,
    pad_packed_sequence,
    pad_sequence,
    pack_sequence,
)

# Variable-length sequences
seqs = [torch.randn(5, 10),   # length 5
        torch.randn(3, 10),   # length 3
        torch.randn(8, 10)]   # length 8

# Step 1: Pad to same length
padded = pad_sequence(seqs, batch_first=True)  # (3, 8, 10)
lengths = torch.tensor([5, 3, 8])

# Step 2: Pack (sorts by length internally)
packed = pack_padded_sequence(padded, lengths, batch_first=True, enforce_sorted=False)

# Step 3: Feed to RNN
rnn = nn.LSTM(10, 20, batch_first=True)
output_packed, (h_n, c_n) = rnn(packed)

# Step 4: Unpack
output_padded, output_lengths = pad_packed_sequence(output_packed, batch_first=True)
print(f"Output: {output_padded.shape}")  # (3, 8, 20)

Why Packing Matters

Without packing: RNN processes all 8 timesteps for all 3 sequences (24 steps). With packing: RNN processes 5+3+8=16 steps total. 33% less computation.

For long sequences with high variance in length, the savings are much larger.

---

5. Conv-BN Fusion

What it is: Merging a Conv2d + BatchNorm2d into a single Conv2d for faster inference.

Why it matters: During inference, BatchNorm is a fixed affine transform. Fusing it into the convolution eliminates one entire layer with zero accuracy loss.

from torch.nn.utils.fusion import fuse_conv_bn_eval

conv = nn.Conv2d(3, 64, 3, padding=1)
bn = nn.BatchNorm2d(64)

# Train as normal...
# Then for inference:
conv.eval()
bn.eval()

fused_conv = fuse_conv_bn_eval(conv, bn)
# fused_conv is a single Conv2d that produces identical output
# but is faster (one layer instead of two)

x = torch.randn(1, 3, 32, 32)
print(torch.allclose(fused_conv(x), bn(conv(x)), atol=1e-5))  # True

This is one of the most common inference optimizations and is done automatically by torch.compile and ONNX optimizers.

---

6. Nested Tensors

What it is: A tensor that can hold sequences of different lengths without padding. Also called "jagged tensors" or "ragged tensors."

Why it matters: Eliminates wasted computation on padding tokens in NLP/attention. Flash Attention can directly consume nested tensors for variable-length sequences.

import torch
from torch.nested import nested_tensor, as_nested_tensor

# Create a nested tensor from variable-length sequences
nt = nested_tensor([
    torch.randn(3, 8),   # sequence of length 3, dim 8
    torch.randn(5, 8),   # sequence of length 5, dim 8
    torch.randn(2, 8),   # sequence of length 2, dim 8
])

print(f"Type: {type(nt)}")
print(f"Nested size: {nt.size()}")
# The first dim is the batch, subsequent dims may vary

# Convert to padded tensor when needed
padded = torch.nested.to_padded_tensor(nt, padding=0.0)
print(f"Padded shape: {padded.shape}")  # (3, 5, 8) — padded to max length

# Convert back
nt2 = as_nested_tensor(padded)

Nested Tensors with SDPA

The real power is using nested tensors with F.scaled_dot_product_attention — Flash Attention handles the variable lengths natively, avoiding wasted computation on padding:

# Instead of padding and masking:
# attn = F.scaled_dot_product_attention(Q_padded, K_padded, V_padded, attn_mask=mask)

# With nested tensors, no padding needed:
# Q_nested, K_nested, V_nested are NestedTensors
# attn = F.scaled_dot_product_attention(Q_nested, K_nested, V_nested)
# Flash Attention processes only real tokens!

---

7. Gradient Clipping Internals

We covered gradient clipping in Module 05, but here's the deeper story:

from torch.nn.utils import clip_grad_norm_, clip_grad_value_, get_total_norm

model = nn.Linear(10, 5)
loss = model(torch.randn(3, 10)).sum()
loss.backward()

# Get the total gradient norm BEFORE clipping
total_norm = get_total_norm(model.parameters(), norm_type=2.0)
print(f"Total gradient norm: {total_norm:.4f}")

# Clip by norm (scales all gradients proportionally if norm exceeds max)
clipped_norm = clip_grad_norm_(model.parameters(), max_norm=1.0, norm_type=2.0)
print(f"Returned (pre-clip) norm: {clipped_norm:.4f}")

# Clip by value (clamps each gradient element independently)
clip_grad_value_(model.parameters(), clip_value=0.5)

Key insight: clip_grad_norm_ returns the total norm before clipping — useful for monitoring gradient health during training.

---

8. Parameter Utilities

parameters_to_vector / vector_to_parameters

Flatten all model parameters into a single vector (useful for L-BFGS, evolutionary methods, or model comparison):

from torch.nn.utils import parameters_to_vector, vector_to_parameters

model = nn.Sequential(nn.Linear(10, 5), nn.Linear(5, 2))

# Flatten all parameters
vec = parameters_to_vector(model.parameters())
print(f"All params as vector: {vec.shape}")  # (77,)

# Modify and write back
vec *= 0.5
vector_to_parameters(vec, model.parameters())

skip_init — Create Modules Without Initializing Weights

For very large models, the default weight initialization can be slow and wasteful (especially if you're loading a checkpoint immediately):

from torch.nn.utils import skip_init

# Normal: allocates + initializes weights (slow for large models)
linear = nn.Linear(10000, 10000)

# Skip init: allocates uninitialized memory (fast)
linear = skip_init(nn.Linear, 10000, 10000)
# Then load your checkpoint:
# linear.load_state_dict(torch.load("checkpoint.pt"))

---

Summary

---

Further Reading

• torch.nn.utils.parametrize docs: pytorch.org/docs/stable/generated/torch.nn.utils.parametrize.register_parametrization
• Pruning tutorial: pytorch.org/tutorials/intermediate/pruning_tutorial
• NestedTensor: pytorch.org/docs/stable/nested.html

---

📓 Open Notebook — Interactive version

Module 16: Activation Checkpointing — Trading Memory for Compute

Day 2 of the incremental learning series

---

The Problem: Activations Eat Your GPU Memory

When training a neural network, PyTorch stores all intermediate activations (outputs of each layer) during the forward pass. These are needed for the backward pass to compute gradients. For large models, activations consume far more memory than the model parameters themselves.

Example memory breakdown for a 7B parameter Transformer:

Model parameters:     14 GB  (7B × 2 bytes in BF16)
Optimizer state:      28 GB  (Adam: 2× params in FP32)
Activations:         ~60 GB  (scales with batch_size × seq_len × layers)
Total:              ~102 GB  — doesn't fit on an 80GB GPU!

Activation checkpointing solves this by not saving activations during forward. During backward, it recomputes them on the fly. This trades ~33% more compute for ~60% less memory.

---

Table of Contents

• How Activation Checkpointing Works
• Basic Usage: torch.utils.checkpoint
• checkpoint_sequential for Sequential Models
• Selective Activation Checkpointing (SAC)
• CheckpointPolicy Options
• Integration with torch.compile
• Practical Guidelines

---

1. How Activation Checkpointing Works

Normal Training (No Checkpointing)

Forward pass:  Input → [Layer 1] → a₁ → [Layer 2] → a₂ → [Layer 3] → a₃ → Loss
                         save a₁      save a₂       save a₃

Backward pass: Uses saved a₁, a₂, a₃ to compute gradients
Memory: O(N) — stores all N layer activations

With Activation Checkpointing

Forward pass:  Input → [Layer 1] → [Layer 2] → [Layer 3] → Loss
                        (discard)   (discard)    (discard)

Backward pass: 
  Need a₃ → Recompute: Input → Layer 1 → Layer 2 → Layer 3 → a₃ ✓
  Need a₂ → Recompute: Input → Layer 1 → Layer 2 → a₂ ✓
  Need a₁ → Recompute: Input → Layer 1 → a₁ ✓

Memory: O(1) per checkpointed segment
Compute: ~1.33× (one extra forward pass)

In practice, you checkpoint segments of the model (e.g., each Transformer layer), not the entire model. This gives a good memory/compute tradeoff.

---

2. Basic Usage

import torch
from torch.utils.checkpoint import checkpoint

class TransformerLayer(torch.nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.attn = torch.nn.MultiheadAttention(d_model, 8, batch_first=True)
        self.norm1 = torch.nn.LayerNorm(d_model)
        self.ffn = torch.nn.Sequential(
            torch.nn.Linear(d_model, 4 * d_model),
            torch.nn.GELU(),
            torch.nn.Linear(4 * d_model, d_model),
        )
        self.norm2 = torch.nn.LayerNorm(d_model)
    
    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.ffn(self.norm2(x))
        return x


class TransformerModel(torch.nn.Module):
    def __init__(self, d_model=512, n_layers=12):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            TransformerLayer(d_model) for _ in range(n_layers)
        ])
    
    def forward(self, x, use_checkpoint=False):
        for layer in self.layers:
            if use_checkpoint:
                # Checkpoint each layer — its activations are NOT saved
                x = checkpoint(layer, x, use_reentrant=False)
            else:
                x = layer(x)
        return x

Key Parameter: use_reentrant=False

Always use use_reentrant=False (the modern, recommended path):

• use_reentrant=True (legacy): Uses a different autograd mechanism, has subtle bugs with certain ops
• use_reentrant=False (recommended): Works correctly with all ops, torch.compile, and distributed training

---

3. checkpoint_sequential

For nn.Sequential models, there's a convenience wrapper:

from torch.utils.checkpoint import checkpoint_sequential

model = torch.nn.Sequential(
    TransformerLayer(512),
    TransformerLayer(512),
    TransformerLayer(512),
    TransformerLayer(512),
)

x = torch.randn(8, 32, 512, requires_grad=True)

# Divide into 2 segments — each segment is checkpointed
output = checkpoint_sequential(model, segments=2, input=x, use_reentrant=False)

segments controls how many checkpointed groups to split the sequential into. More segments = more memory savings but more recomputation.

---

4. Selective Activation Checkpointing (SAC)

The problem with basic checkpointing: It recomputes EVERYTHING, including expensive operations like matrix multiplies and attention. Ideally, we'd save expensive outputs and only recompute cheap ones (like activations, norms).

SAC lets you choose per-operation whether to save or recompute:

from torch.utils.checkpoint import (
    checkpoint,
    CheckpointPolicy,
    create_selective_checkpoint_contexts,
)

# Policy function: decides per-op whether to save or recompute
def policy_fn(ctx, op, *args, **kwargs):
    # Save expensive ops (matmul, attention)
    if op in (torch.ops.aten.mm.default, 
              torch.ops.aten.bmm.default,
              torch.ops.aten._scaled_dot_product_flash_attention.default):
        return CheckpointPolicy.MUST_SAVE
    # Recompute cheap ops (relu, add, norm, etc.)
    return CheckpointPolicy.PREFER_RECOMPUTE

# Use with checkpoint's context_fn parameter
context_fn = create_selective_checkpoint_contexts(policy_fn)

x = checkpoint(
    layer, x,
    use_reentrant=False,
    context_fn=context_fn,
)

Shortcut: Pass a List of Ops to Save

# Instead of a policy function, just list the ops you want to save
ops_to_save = [
    torch.ops.aten.mm.default,
    torch.ops.aten.bmm.default,
]

context_fn = create_selective_checkpoint_contexts(ops_to_save)

x = checkpoint(layer, x, use_reentrant=False, context_fn=context_fn)

---

5. CheckpointPolicy Options

The PREFER_ variants allow torch.compile to override the decision based on its global optimization analysis. The MUST_ variants are strict.

---

6. Integration with torch.compile

Activation checkpointing works with torch.compile. When compiled, the compiler can make even smarter decisions about what to save vs recompute:

model = TransformerModel(d_model=512, n_layers=12)

# Compile with checkpointing
compiled_model = torch.compile(model)

x = torch.randn(8, 32, 512)
output = compiled_model(x, use_checkpoint=True)

When using PREFER_* policies with compiled code, the compiler may override your suggestions if it determines a different strategy is more efficient.

---

7. Practical Guidelines

When to Use Activation Checkpointing

Rules of Thumb

• Checkpoint at the Transformer layer granularity — each TransformerEncoderLayer or decoder layer is one checkpoint boundary
• Always use use_reentrant=False — the legacy reentrant mode has known issues
• With SAC, save matmuls and attention — these are 10-100x more expensive than activations/norms
• Memory savings: ~50-70% activation memory reduction with basic checkpointing
• Compute overhead: ~30% more training time (one extra forward pass per checkpointed segment)
• Stacks with everything: Works with AMP, DDP, FSDP2, torch.compile

Memory Estimation

Without checkpointing:
  Activation memory ≈ batch_size × seq_len × hidden_dim × num_layers × 2 bytes (BF16)

With checkpointing (per layer):
  Activation memory ≈ batch_size × seq_len × hidden_dim × 2 × 2 bytes
  (Only stores input/output of each checkpointed segment)

Savings = (num_layers - 2) / num_layers ≈ 90%+ for 24+ layer models

---

Further Reading

• PyTorch Activation Checkpointing Tutorial
• Selective Activation Checkpointing for torch.compile
• Source code: torch/utils/checkpoint.py

---

📓 Open Notebook — Interactive version

Module 17: torch.compile Decorators & Control APIs

Day 3 of the incremental learning series

---

Beyond torch.compile() — Fine-Grained Compilation Control

Module 08 taught you the basics of torch.compile. This module dives into the decorator and control APIs that give you precise control over what gets compiled, how shapes are handled, and how to debug compilation issues.

---

Table of Contents

• Compiler Stances — Global Compilation Behavior
• torch.compiler.disable — Skip Compilation
• allow_in_graph / disallow_in_graph — Control Tracing
• substitute_in_graph — Replace Functions for Tracing
• mark_dynamic / mark_static — Shape Control
• graph_break / error_on_graph_break — Break Control
• assume_constant_result — Constant Folding
• comptime — Compile-Time Debugging
• torch._dynamo.explain — Understanding Compilation
• TORCH_LOGS — Logging and Debugging
• What's New: Recent Upstream Changes

---

1. Compiler Stances

Stances control the global compilation behavior — useful for debugging and gradual adoption:

import torch

torch.compiler.set_stance("default")             # Normal compilation
torch.compiler.set_stance("force_eager")          # Skip all compilation
torch.compiler.set_stance("eager_on_recompile")   # Compile once, eager on recompile
torch.compiler.set_stance("fail_on_recompile")    # Error on recompilation
torch.compiler.set_stance("eager_then_compile")   # Eager first, compile on second call

# Context manager form
with torch.compiler.set_stance("force_eager"):
    output = compiled_model(x)  # Runs eagerly despite being compiled

---

2. Disable — Skip Compilation for Specific Functions

@torch.compiler.disable
def preprocessing(x):
    """Uses Python features that don't compile well."""
    results = []
    for i in range(x.shape[0]):
        if x[i].item() > 0:
            results.append(x[i] * 2)
        else:
            results.append(x[i])
    return torch.stack(results)

@torch.compile
def model_fn(x):
    x = preprocessing(x)     # Runs eagerly (disabled)
    return x.relu().sum()     # This part is compiled

Non-recursive disable (only disables the decorated function, not functions it calls):

@torch.compiler.disable(recursive=False)
def outer(x):
    return inner(x)  # inner() will still be compiled

---

3. allow_in_graph / disallow_in_graph

@torch.compiler.allow_in_graph   # Opaque node in graph
@torch._dynamo.disallow_in_graph  # Forces graph break
@torch._dynamo.forbid_in_graph    # Raises error during tracing

---

4. substitute_in_graph — Replace Functions for Tracing

def original_fn(x):
    result = x.tolist()  # Not traceable!
    return torch.tensor(result)

def traceable_fn(x):
    return x.clone()

torch.compiler.substitute_in_graph(original_fn, traceable_fn)

---

5. mark_dynamic / mark_static — Shape Control

x = torch.randn(batch_size, 512)
torch._dynamo.mark_dynamic(x, 0)  # Dim 0 is dynamic — no recompile for different batch sizes
torch._dynamo.mark_static(x, 1)   # Dim 1 is always 512
torch._dynamo.mark_static_address(weight)  # Data pointer won't change

---

6. Graph Break Control

torch._dynamo.graph_break()       # Force a graph break
torch.compile(fullgraph=True)     # Error on any break

@torch._dynamo.error_on_graph_break
def must_be_one_graph(x):
    return x + 1

---

7. assume_constant_result

@torch._dynamo.assume_constant_result
def get_config():
    return load_config()["lr"]  # Folded to constant at compile time

---

8. comptime — Compile-Time Debugging

from torch._dynamo.comptime import comptime

@torch.compile
def fn(x):
    comptime.breakpoint()  # pdb during COMPILATION
    return x + 1
# In pdb: ctx.print_locals(), ctx.print_graph(), ctx.print_bt()

---

9. torch._dynamo.explain

explanation = torch._dynamo.explain(my_fn)(torch.randn(10))
print(explanation)  # Shows graphs, breaks, guards

---

10. TORCH_LOGS — Logging and Debugging

TORCH_LOGS="graph_breaks" python train.py       # Graph break reasons
TORCH_LOGS="guards,recompiles" python train.py  # Guard failures
TORCH_LOGS="graph_code" python train.py         # Captured FX graph
TORCH_LOGS="output_code" python train.py        # Generated Triton/C++
TORCH_LOGS="+dynamo" python train.py            # Full debug
TORCH_TRACE=/tmp/trace python train.py          # Structured tracing

---

11. What's New: Recent Upstream Changes (June 4-8, 2026)

194 commits landed on PyTorch main in the last 4 days:

QuACK GEMM Kernels Vendored

PyTorch now vendors the QuACK library from Dao-AILab — high-performance CuTeDSL GEMM epilogue adapters with fused RMSNorm. Located at torch/_vendor/quack/.

CUPTI Monitor — Continuous GPU Profiling

New torch.profiler._cupti_monitor for continuous CUPTI activity monitoring across the entire program (not just a profiling window).

Optimized _foreach_mm — Grouped GEMMs

New Python override dispatching to nvmath cublasLt grouped GEMM (bf16) or CUTLASS. At torch/_native/ops/foreach_mm/.

AArch64 torch.compile

Armv9-A target support — compiled models now work on ARM servers and edge devices.

DTensor Autogen Ops

Auto-generated sharding strategies expanding DTensor op coverage. At torch/distributed/tensor/_ops/autogen.py.

NCCL Symmetric Memory Registration

External NCCL comm registration for symmetric memory at torch/distributed/_symmetric_memory/_nccl.py.

Inductor Heuristics Module

Refactored Triton template heuristics into torch/_inductor/heuristics/.

---

Quick Reference

---

No dedicated notebook — covered in Module 08 notebook

Module 18: torch.package — Self-Contained Model Packaging

Day 4 of the incremental learning series

---

The Problem: "It Works on My Machine"

You train a model. You save model.pt. You send it to a colleague. It fails because:

• They don't have the same version of your custom modules
• An import path changed between your environments
• A dependency you forgot about isn't installed on their machine

torch.package solves this by bundling the model and all its Python dependencies into a single .pt archive.

---

Table of Contents

• What is torch.package?
• PackageExporter — Creating Packages
• PackageImporter — Loading Packages
• Module Actions: intern, extern, mock, deny
• Packaging Models with Weights
• Inspecting Package Contents
• Re-Packaging and Dependency Analysis
• torch.package vs torch.save vs torch.export
• Practical Workflow
• What's New Upstream (June 8-9, 2026)

---

1. What is torch.package?

torch.package creates a hermetic zip archive containing:

• Your model's Python source code (the actual .py files)
• The model's pickled state (weights, buffers, config)
• A manifest of external dependencies

When someone loads the package, it uses its own import system — code is loaded from inside the archive, not from the local Python installation. This means:

• The exact code you packaged runs, regardless of what's installed locally
• Only explicitly listed external dependencies are loaded from the system
• No "accidental" dependencies can sneak in

┌──────────────────────────────────┐
│        my_model.pt (zip)         │
├──────────────────────────────────┤
│  .data/                          │
│    model.pkl        (pickled)    │
│    weights.pt       (tensors)    │
│  my_module/                      │
│    model.py         (source)     │
│    layers.py        (source)     │
│    config.py        (source)     │
│  extern_modules     (manifest)   │
│    torch                         │
│    numpy                         │
└──────────────────────────────────┘

---

2. PackageExporter — Creating Packages

from torch.package import PackageExporter

# Create a package
with PackageExporter("my_model.pt") as exporter:
    # INTERN: Include this module's source inside the package
    exporter.intern("my_module.**")

    # EXTERN: This module is expected to exist on the target machine
    exporter.extern("torch.**")
    exporter.extern("numpy.**")

    # MOCK: Replace this module with a stub (for unused optional deps)
    exporter.mock("matplotlib.**")

    # Save the model object
    exporter.save_pickle("model", "model.pkl", my_model)

The Four Module Actions

Patterns use glob syntax: "my_module.**" matches my_module and all submodules.

---

3. PackageImporter — Loading Packages

from torch.package import PackageImporter

# Load a package
importer = PackageImporter("my_model.pt")

# Load the pickled model
model = importer.load_pickle("model", "model.pkl")

# The model runs using code FROM the package, not your local installation
output = model(torch.randn(1, 3, 224, 224))

# Import a module from the package (hermetic import)
my_config = importer.import_module("my_module.config")

Key property: The loaded code runs in an isolated namespace. If my_module/model.py inside the package says import my_module.layers, it loads layers.py from the package, not from your filesystem.

---

4. Module Actions in Detail

intern — Bundle Source Code

# Include specific modules
exporter.intern("my_project.models.**")
exporter.intern("my_project.utils.**")

# Include everything in your project
exporter.intern("my_project.**")

What happens: The .py source files are copied into the zip archive. When loaded, Python reads them from the archive.

extern — External Dependencies

# Standard externals
exporter.extern("torch.**")
exporter.extern("torchvision.**")
exporter.extern("numpy.**")

# Stdlib is automatically handled, but you can be explicit:
exporter.extern("os")
exporter.extern("json")

What happens: A list of external modules is saved in extern_modules. When loading, these are imported from the system Python.

mock — Stub Out Dependencies

# Mock out modules that aren't needed at inference time
exporter.mock("wandb.**")          # Logging library
exporter.mock("matplotlib.**")     # Plotting
exporter.mock("tensorboard.**")    # TensorBoard

What happens: A lightweight _mock module replaces the real one. Any attribute access on a mocked module returns MockedObject.

deny — Prevent Inclusion

# Error if these are encountered
exporter.deny("secret_module.**")
exporter.deny("credentials.**")

---

5. Packaging Models with Weights

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

# Train your model...
model = MyModel(784, 256, 10)

# Package it with all dependencies
with PackageExporter("my_model_package.pt") as exporter:
    exporter.intern("__main__")    # Include the current module
    exporter.extern("torch.**")
    exporter.extern("numpy.**")

    # Save model
    exporter.save_pickle("model", "model.pkl", model)

    # You can also save arbitrary data
    exporter.save_pickle("config", "config.pkl", {
        "input_dim": 784,
        "hidden_dim": 256,
        "output_dim": 10,
        "version": "1.0",
    })

    # Save raw text/binary files
    exporter.save_text("info", "README.txt", "My model v1.0")

# Load on another machine
importer = PackageImporter("my_model_package.pt")
model = importer.load_pickle("model", "model.pkl")
config = importer.load_pickle("config", "config.pkl")
readme = importer.load_text("info", "README.txt")

---

6. Inspecting Package Contents

importer = PackageImporter("my_model_package.pt")

# View the file structure
file_structure = importer.file_structure()
print(file_structure)
# Prints a tree of all files in the archive

# List all extern modules
print(file_structure.has_file("extern_modules"))

You can also inspect with standard zip tools since .pt files are zip archives:

unzip -l my_model_package.pt
python -m zipfile -l my_model_package.pt

---

7. Re-Packaging

You can load a package and re-export it (e.g., to add/remove dependencies):

importer = PackageImporter("model_v1.pt")

with PackageExporter("model_v2.pt", importer=(importer,)) as exporter:
    exporter.intern("my_module.**")
    exporter.extern("torch.**")

    # Load the old model and save it in the new package
    model = importer.load_pickle("model", "model.pkl")
    exporter.save_pickle("model", "model.pkl", model)

---

8. torch.package vs torch.save vs torch.export

When to use each:

• torch.save: Quick checkpoints during training
• torch.package: Ship Python models with all dependencies, research sharing
• torch.export: Production deployment, C++ inference, mobile

---

9. Practical Workflow

Research → Deployment Pipeline

# 1. Researcher trains model (research/train.py)
model = train_my_model()

# 2. Researcher packages model with all custom code
with PackageExporter("model_v1.pt") as pe:
    pe.intern("my_research_code.**")
    pe.extern("torch.**")
    pe.extern("torchvision.**")
    pe.mock("wandb.**")   # Don't need logging in production
    pe.mock("matplotlib.**")
    pe.save_pickle("model", "model.pkl", model)

# 3. Engineer loads on a different machine (no my_research_code installed!)
importer = PackageImporter("model_v1.pt")
model = importer.load_pickle("model", "model.pkl")
output = model(input_data)  # Just works!

Tips

• Always extern torch — it must match the installed version
• Mock unused dependencies — logging, visualization, experiment tracking
• Test the package — load it in a clean environment to verify
• Version your packages — include version info as saved text/pickle
• Inspect before shipping — use file_structure() to verify contents

---

10. Upstream Updates (June 8-9, 2026)

Recent PyTorch main commits (since last update):

• Inductor CUTLASS GELU fusion — Folding decomposed GELU back into native CUTLASS functor for better performance (#185824)
• Inductor TP pattern fusion — Fusing slice-cat tensor parallel collective patterns (#184911)
• Dynamo Python 3.15 support — Build dynamo with Python 3.15, including updated IMPORT_NAME generation (#186402)
• Dynamo operator support — Added divmod, remainder, true_divide, floor_divide operators (#185652-#185655)
• Deterministic topk — torch.topk now respects torch.use_deterministic_algorithms() (#186653)
• XPU oneDNN LSTM — Intel GPU LSTM inference via oneDNN primitives (#185531)
• Stable ABI generator — New torch/csrc/stable/generator.h for stable C API

---

Further Reading

• Source: torch/package/package_exporter.py, torch/package/package_importer.py
• PyTorch docs: torch.package
• Tutorial: Loading a TorchScript Model in C++ (comparison)

---

No dedicated notebook — see Practical Workflow above

Module 19: __torch_function__ & __torch_dispatch__ — Tensor Subclassing

Day 5 of the incremental learning series

---

Why This Matters

Every time you call torch.add(x, y) or x + y, PyTorch checks: does this tensor have a custom dispatch protocol? Two protocols exist:

• __torch_function__ — Python-level override. Intercepts any PyTorch function call. Like NumPy's __array_function__.
• __torch_dispatch__ — Lower-level override. Intercepts at the ATen operator level (after decompositions). More powerful, used by DTensor, FakeTensor, and torch.compile internals.

These are the extension points that power:

• DTensor (distributed tensor) — sharding logic via __torch_dispatch__
• FakeTensor (torch.compile) — shape-only tensors via __torch_dispatch__
• Logging/profiling — intercept all operations without modifying model code
• Custom tensor types — sparse, quantized, masked tensors
• Unit conversion — tensors that carry physical units

---

Table of Contents

• __torch_function__ — Python-Level Override
• TorchFunctionMode — Override Without Subclassing
• __torch_dispatch__ — ATen-Level Override
• TorchDispatchMode — Mode-Based Dispatch
• Practical Examples
• When to Use Which Protocol
• Upstream Updates (June 9-10, 2026)

---

1. __torch_function__ — Python-Level Override

When you define __torch_function__ on a class, PyTorch calls it instead of the normal implementation for any torch function that receives your object as an argument.

import torch

class ScaledTensor:
    """A tensor wrapper that tracks a scaling factor."""

    def __init__(self, data, scale=1.0):
        self.data = data
        self.scale = scale

    def __repr__(self):
        return f"ScaledTensor(data={self.data}, scale={self.scale})"

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        """Called for any torch.* function involving this type."""
        if kwargs is None:
            kwargs = {}

        # Extract ScaledTensors from args, replace with raw data
        new_args = []
        scale = 1.0
        for a in args:
            if isinstance(a, ScaledTensor):
                new_args.append(a.data)
                scale = a.scale
            else:
                new_args.append(a)

        # Call the original function on raw tensors
        result = func(*new_args, **kwargs)

        # Wrap the result back
        if isinstance(result, torch.Tensor):
            return ScaledTensor(result, scale)
        return result

# Usage
x = ScaledTensor(torch.tensor([1.0, 2.0, 3.0]), scale=0.5)
y = ScaledTensor(torch.tensor([4.0, 5.0, 6.0]), scale=0.5)

z = torch.add(x, y)       # Calls ScaledTensor.__torch_function__!
print(z)                    # ScaledTensor(data=tensor([5., 7., 9.]), scale=0.5)

w = torch.mul(x, 2)        # Also intercepted
print(w)                    # ScaledTensor(data=tensor([2., 4., 6.]), scale=0.5)

How the Protocol Works

• PyTorch checks if any argument has __torch_function__
• If yes, it calls __torch_function__(func, types, args, kwargs) where:
• func — the original function (e.g., torch.add)
• types — tuple of types that implement __torch_function__
• args — positional arguments
• kwargs — keyword arguments
• Your implementation decides what to do and returns the result

---

2. TorchFunctionMode — Override Without Subclassing

Modes let you override all torch operations within a context manager — no tensor subclass needed:

from torch.overrides import TorchFunctionMode

class LoggingMode(TorchFunctionMode):
    """Logs every torch operation."""

    def __torch_function__(self, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        print(f"  Called: {func.__module__}.{func.__name__}")
        return func(*args, **kwargs)

# All torch ops inside the context are logged
with LoggingMode():
    x = torch.randn(3, 4)      # Logged
    y = x + 1                   # Logged
    z = torch.relu(y)           # Logged
    w = z.mean()                # Logged

Use Cases for Modes

• Logging/debugging — see every operation a model performs
• Profiling — count operations, measure shapes
• Mocking — override factory functions (torch.randn, torch.zeros)
• Validation — check all inputs are on the correct device

---

3. __torch_dispatch__ — ATen-Level Override

__torch_dispatch__ intercepts at a lower level — after Python function dispatch, at the ATen operator level. This is where the real computation happens.

import torch
from torch.utils._python_dispatch import return_and_correct_aliasing

class LoggingTensor(torch.Tensor):
    """A tensor subclass that logs all ATen operations."""

    @staticmethod
    def __new__(cls, data):
        return torch.Tensor._make_subclass(cls, data)

    @classmethod
    def __torch_dispatch__(cls, func, types, args, kwargs=None):
        """Called for every ATen operator."""
        if kwargs is None:
            kwargs = {}

        # Unwrap LoggingTensors to plain tensors
        def unwrap(t):
            return t.elem if isinstance(t, LoggingTensor) else t

        print(f"  dispatch: {func.__name__}")

        # Call the actual ATen op
        result = func(*args, **kwargs)
        return result

x = LoggingTensor(torch.randn(3, 4))
y = x + 1       # Dispatches through __torch_dispatch__
z = y.relu()     # Also dispatched

Key Differences from __torch_function__

---

4. TorchDispatchMode — Mode-Based Dispatch

Like TorchFunctionMode, but at the ATen operator level:

from torch.utils._python_dispatch import TorchDispatchMode

class CountOps(TorchDispatchMode):
    """Count all ATen operations in a scope."""

    def __init__(self):
        super().__init__()
        self.ops = {}

    def __torch_dispatch__(self, func, types, args, kwargs=None):
        name = str(func.name())
        self.ops[name] = self.ops.get(name, 0) + 1
        if kwargs is None:
            kwargs = {}
        return func(*args, **kwargs)

# Count ops in a forward pass
counter = CountOps()
model = torch.nn.Sequential(
    torch.nn.Linear(10, 20),
    torch.nn.ReLU(),
    torch.nn.Linear(20, 5),
)

with counter:
    output = model(torch.randn(4, 10))

print("Operations performed:")
for op, count in sorted(counter.ops.items()):
    print(f"  {op}: {count}x")

---

5. Practical Examples

Example 1: Device-Checking Mode

class DeviceCheckMode(TorchFunctionMode):
    """Error if any tensor is on the wrong device."""

    def __init__(self, expected_device):
        self.expected_device = torch.device(expected_device)

    def __torch_function__(self, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        for a in args:
            if isinstance(a, torch.Tensor) and a.device != self.expected_device:
                raise RuntimeError(
                    f"{func.__name__}: tensor on {a.device}, "
                    f"expected {self.expected_device}"
                )
        return func(*args, **kwargs)

# This catches CPU/GPU mismatches early
with DeviceCheckMode("cpu"):
    x = torch.randn(3, 4)  # OK
    y = x + 1               # OK

Example 2: Shape Logging Mode

class ShapeTracer(TorchDispatchMode):
    """Track input/output shapes of all ops."""

    def __init__(self):
        super().__init__()
        self.traces = []

    def __torch_dispatch__(self, func, types, args, kwargs=None):
        if kwargs is None:
            kwargs = {}
        result = func(*args, **kwargs)

        in_shapes = [a.shape for a in args if isinstance(a, torch.Tensor)]
        out_shape = result.shape if isinstance(result, torch.Tensor) else "N/A"
        self.traces.append((func.name(), in_shapes, out_shape))
        return result

Example 3: Tensor with Units (Physics)

class UnitTensor:
    """Tensor that tracks physical units (e.g., meters, seconds)."""

    def __init__(self, data, unit=""):
        self.data = data
        self.unit = unit

    def __repr__(self):
        return f"{self.data} [{self.unit}]"

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        tensors = [a for a in args if isinstance(a, UnitTensor)]
        raw_args = [a.data if isinstance(a, UnitTensor) else a for a in args]
        result = func(*raw_args, **kwargs)

        if func == torch.mul and len(tensors) == 2:
            unit = f"{tensors[0].unit}*{tensors[1].unit}"
        elif func == torch.div and len(tensors) == 2:
            unit = f"{tensors[0].unit}/{tensors[1].unit}"
        else:
            unit = tensors[0].unit if tensors else ""

        if isinstance(result, torch.Tensor):
            return UnitTensor(result, unit)
        return result

distance = UnitTensor(torch.tensor(100.0), "m")
time = UnitTensor(torch.tensor(9.58), "s")
speed = torch.div(distance, time)
print(f"Speed: {speed}")  # 10.44 [m/s]

---

6. When to Use Which Protocol

The Dispatch Stack

User code: torch.nn.functional.relu(x)
     |
     v
__torch_function__     <- Python-level, sees relu
     |
     v
Decompositions         <- relu -> clamp(x, min=0)
     |
     v
__torch_dispatch__     <- ATen-level, sees aten.clamp.default
     |
     v
C++ dispatcher         <- Routes to CPU/CUDA/etc. kernel

---

7. Upstream Updates (June 9-10, 2026)

Recent PyTorch main commits:

• FSDP2 separate reduce-scatter group — Opt-in all-gather/reduce-scatter overlap via set_separate_reduce_scatter_group (#186335)
• Activation offloading pinned memory pool — Dedicated pinned memory pool for activation offloading ops (#186162)
• Activation offloading stride preservation — Preserves original tensor strides across offload/reload (#186396)
• BERT SDPA pattern on CUDA — Enables BERT attention pattern for SDPA on CUDA (#184417)
• DTensor group_norm fix — Fixes crash when weight=None in group_norm under DTensor (#184819)
• Pipeline parallel backward fix — Fixes None gradient handling in pipeline backward send/recv (#182182)
• Torch.cuda.stream round-trip — Dynamo now correctly handles torch.cuda.stream context managers across graph breaks (#184487)
• TORCH_TRACE fork-safety — Structured tracing logs now preserved across forks (#184772)
• Open registration profiler — Activity profiler support for custom backend devices via open registration (new test files)

---

Quick Reference

# __torch_function__ -- Python-level override (any class)
class MyType:
    @classmethod
    def __torch_function__(cls, func, types, args, kwargs=None):
        ...

# TorchFunctionMode -- scope-based override (no subclass needed)
class MyMode(TorchFunctionMode):
    def __torch_function__(self, func, types, args=(), kwargs=None):
        ...

# __torch_dispatch__ -- ATen-level override (tensor subclass)
class MyTensor(torch.Tensor):
    @classmethod
    def __torch_dispatch__(cls, func, types, args, kwargs=None):
        ...

# TorchDispatchMode -- scope-based ATen override
class MyDispatchMode(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args, kwargs=None):
        ...

---

Further Reading

• Source: torch/overrides.py (torch_function), torch/utils/_python_dispatch.py (torch_dispatch)
• Extending PyTorch docs
• __torch_function__ protocol

---

No dedicated notebook — see examples in torch_function_examples.py

Module 20: torch.backends — Performance Tuning

Prerequisites: Module 07 (Training Pipelines)
Time: ~2 hours
Level: Intermediate → Advanced

---

Overview

PyTorch's torch.backends module is a configuration layer that controls hardware-specific optimizations. These settings determine which algorithms run under the hood for convolutions, matrix multiplications, attention, and parallelism—often yielding 2–10x speedups with a single line of code.

Most tutorials never mention these knobs. This module changes that.

Files in This Module

---

1. What Are torch.backends?

torch.backends exposes runtime configuration for the hardware libraries PyTorch uses:

torch.backends
├── cudnn          # NVIDIA cuDNN (convolutions, RNNs)
├── cuda           # NVIDIA CUDA (matmul, SDPA)
├── mkldnn         # Intel oneDNN (CPU conv, linear)
├── mkl            # Intel MKL (BLAS/LAPACK)
├── openmp         # OpenMP (CPU threading)
├── opt_einsum     # Optimized einsum contraction
└── mps            # Apple Metal (M1/M2/M3)

Each backend exposes flags you can toggle at runtime. No recompilation needed.

Key principle: backends control the how, not the what. The mathematical operation stays the same; the algorithm, precision, or parallelism strategy changes.

---

2. torch.backends.cudnn

cuDNN is NVIDIA's deep neural network library. It provides optimized implementations for convolutions, pooling, normalization, and RNNs.

2.1 cudnn.enabled

torch.backends.cudnn.enabled  # default: True

When True, PyTorch uses cuDNN for supported operations. Disabling it falls back to slower native implementations. You almost never want to disable this.

2.2 cudnn.benchmark

torch.backends.cudnn.benchmark = True  # default: False

What it does: Before the first convolution at each input size, cuDNN runs multiple algorithm variants and selects the fastest. Results are cached for the session.

When it helps:

• Fixed input sizes (standard training with constant batch size and image dimensions)
• Repeated convolutions with the same shapes
• Training CNNs (ResNet, EfficientNet, etc.)

When it hurts:

• Variable input sizes (NLP with different sequence lengths, object detection with varying image sizes)
• Short-lived scripts (benchmarking overhead > savings)
• First iteration is slower (paying the auto-tuning cost)

# Typical training setup for CNNs with fixed input
torch.backends.cudnn.benchmark = True

# Disable for variable-size inputs
torch.backends.cudnn.benchmark = False

2.3 cudnn.deterministic

torch.backends.cudnn.deterministic = True  # default: False

Forces cuDNN to use deterministic algorithms. Non-deterministic algorithms are often faster because they can exploit parallelism without worrying about reduction order.

Tradeoffs:

For full determinism, also call torch.use_deterministic_algorithms(True).

2.4 cudnn.allow_tf32

torch.backends.cudnn.allow_tf32 = True  # default: True (PyTorch 2.x)

Controls whether cuDNN can use TF32 precision for convolutions on Ampere+ GPUs. See Section 9 for details on TF32.

---

3. torch.backends.cuda

Configuration for CUDA operations beyond cuDNN (matrix multiplications, attention).

3.1 matmul.allow_tf32

torch.backends.cuda.matmul.allow_tf32 = True  # default: True (PyTorch 2.x)

Allows cuBLAS to use TF32 for float32 matrix multiplications. On Ampere (A100) and later, this can double matmul throughput with minimal precision loss.

3.2 matmul.allow_fp16_reduced_precision_reduction

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True

Controls whether fp16 GEMMs can use reduced precision for internal accumulation. Faster but potentially less accurate for very large matrices.

3.3 matmul.allow_bf16_reduced_precision_reduction

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True

Same as above but for bfloat16 operations. Relevant for training on Ampere+ GPUs where bf16 is preferred over fp16.

3.4 Flash SDP (Scaled Dot-Product Attention)

# Check if Flash Attention is enabled
torch.backends.cuda.flash_sdp_enabled()  # True/False

# Enable/disable Flash Attention
torch.backends.cuda.enable_flash_sdp(True)

# Other SDPA backends
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(True)

Flash Attention is the fastest SDPA implementation for most cases. You might disable it for debugging or to force a specific backend.

3.5 preferred_blas_library

torch.backends.cuda.preferred_blas_library()           # get current
torch.backends.cuda.preferred_blas_library("cublas")   # set preference
# Options: "cublas", "cublaslt", "hipblaslt" (AMD)

Select which BLAS library handles matrix multiplications. cuBLASLt supports more fused operations and epilogues.

---

4. torch.backends.mkldnn

Intel oneDNN (formerly MKL-DNN) provides optimized CPU implementations for convolutions, linear layers, and batch normalization.

torch.backends.mkldnn.enabled  # default: True on x86

# Check if available
torch.backends.mkldnn.is_available()

When enabled, PyTorch automatically routes supported operations through oneDNN for faster CPU execution. This is especially impactful for inference on Intel CPUs.

---

5. torch.backends.mkl

Intel Math Kernel Library (MKL) provides optimized BLAS/LAPACK routines.

# Enable verbose mode to see which MKL routines are called
torch.backends.mkl.verbose(torch.backends.mkl.VERBOSE_ON)
# VERBOSE_OFF, VERBOSE_ON

Use case: Profiling CPU performance to confirm MKL is being used for linear algebra operations.

---

6. torch.backends.openmp

Controls OpenMP threading for CPU parallelism.

import torch

# Get/set number of threads
torch.get_num_threads()        # current thread count
torch.set_num_threads(4)       # set to 4 threads

# Also controlled via environment variable (before import):
# OMP_NUM_THREADS=4
# MKL_NUM_THREADS=4

Guidelines:

• For training: set to number of physical cores (not hyperthreads)
• For inference with batching: reduce threads to allow concurrent requests
• For DataLoader workers: reduce to avoid oversubscription

import os
os.environ["OMP_NUM_THREADS"] = "4"  # must be set BEFORE importing torch

---

7. torch.backends.opt_einsum

Optimized path planning for torch.einsum operations.

torch.backends.opt_einsum.enabled       # default: True if opt_einsum installed
torch.backends.opt_einsum.strategy      # default: "auto"
# Strategies: "auto", "greedy", "optimal", "branch-all", "branch-2", "dp"

What it does: For complex einsum expressions with 3+ tensors, finding the optimal contraction order is NP-hard. opt_einsum uses heuristics to find near-optimal orderings that can be orders of magnitude faster.

# Example: without optimization, this could be O(N^5)
# With optimal contraction order, it's O(N^3)
result = torch.einsum("ij,jk,kl->il", A, B, C)

Strategies:

---

8. torch.backends.mps

Apple Metal Performance Shaders backend for M1/M2/M3/M4 chips.

torch.backends.mps.is_available()  # True on Apple Silicon with macOS 12.3+
torch.backends.mps.is_built()      # True if PyTorch was compiled with MPS

MPS provides GPU acceleration on Apple hardware. While not as fast as CUDA for large models, it enables GPU training on Apple laptops and desktops.

if torch.backends.mps.is_available():
    device = torch.device("mps")
    tensor = torch.randn(1000, 1000, device=device)

---

9. TF32 Precision

What is TF32?

TF32 (TensorFloat-32) is a 19-bit floating point format introduced with NVIDIA Ampere (A100):

Format comparison:
FP32:  1 sign + 8 exponent + 23 mantissa = 32 bits
TF32:  1 sign + 8 exponent + 10 mantissa = 19 bits
FP16:  1 sign + 5 exponent + 10 mantissa = 16 bits
BF16:  1 sign + 8 exponent +  7 mantissa = 16 bits

TF32 has the range of FP32 (8-bit exponent) with the precision of FP16 (10-bit mantissa). It's used internally by tensor cores—inputs are read as FP32, rounded to TF32 for computation, and results are accumulated in FP32.

How TF32 Affects Operations

Enable/Disable TF32

# Enable TF32 everywhere (recommended for training)
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True

# Disable TF32 for full FP32 precision (validation, debugging)
torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False

When to Disable TF32

• Numerical validation (comparing against reference implementations)
• Scientific computing requiring full FP32 precision
• Debugging convergence issues
• Unit tests checking exact numerical equality

---

10. torch.set_float32_matmul_precision()

A high-level API that controls matmul precision globally:

torch.set_float32_matmul_precision("highest")  # No TF32, full FP32
torch.set_float32_matmul_precision("high")     # TF32 on Ampere+
torch.set_float32_matmul_precision("medium")   # TF32 + reduced precision reductions

# Recommended: set once at the top of your training script
torch.set_float32_matmul_precision("high")

This is what torch.compile hints at when it logs: "TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled."

---

11. torch.backends.flags() Context Manager

Temporarily override backend settings within a scope:

with torch.backends.flags(
    cudnn_benchmark=True,
    cudnn_deterministic=False,
    cudnn_enabled=True,
    allow_tf32=True
):
    # These settings active only inside this block
    output = model(input)

# Original settings restored here

Use cases:

• Temporarily disabling TF32 for a validation step
• Enabling benchmark mode for a specific layer
• Writing tests that need specific backend states

---

12. Performance Checklist

Training Settings (Maximum Speed)

# Set at top of training script
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")

Inference Settings (Maximum Throughput)

# Set before inference
torch.backends.cudnn.benchmark = True       # if input sizes are fixed
torch.backends.cudnn.deterministic = False
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision("high")
torch.set_num_threads(num_physical_cores)

Debugging/Reproducibility Settings

# Set for exact reproducibility
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.use_deterministic_algorithms(True)
torch.manual_seed(42)

Quick Reference Table

---

13. Upstream Updates (June 10–11, 2026)

Recent PyTorch commits that affect backend behavior and performance:

cuDNN SDPA d=256 Support (#185553)

cuDNN's SDPA backend now supports head dimension d=256, enabling Flash Attention-like performance for larger attention heads without falling back to the math kernel. This benefits models using large head dimensions (e.g., certain MoE architectures).

ARM scatter/gather Optimization (#156161)

Scatter and gather operations now use optimized ARM NEON intrinsics on aarch64, improving CPU performance on ARM servers (AWS Graviton, Apple M-series) by 2–4x for these ops.

Dynamo Polyfills for itertools (#186240)

torch.compile now handles itertools.chain, itertools.islice, and other itertools functions as graph-safe polyfills, reducing graph breaks in models that use standard library iteration patterns.

DTensor Single-Dim Strategies Migration (#186667)

Internal migration of DTensor sharding strategies to a single-dimension representation, improving compilation speed and reducing memory for distributed models using DeviceMesh.

AOTI torch.cond/while_loop Support (#184736)

AOTInductor (the ahead-of-time compiler) now supports torch.cond and torch.while_loop control flow operations, enabling export of models with conditional branches and loops.

---

Summary

torch.backends — what to remember:

1. cudnn.benchmark = True     → auto-tune conv algorithms (fixed sizes)
2. allow_tf32 = True          → 2–3x matmul/conv speedup on Ampere+
3. set_float32_matmul_precision("high") → same as TF32, cleaner API
4. set_num_threads(N)         → match physical cores for CPU
5. opt_einsum.strategy        → optimize multi-tensor contractions
6. deterministic = True       → reproducibility at the cost of speed
7. backends.flags()           → scope backend settings temporarily

---

Further Reading

• PyTorch backends documentation
• NVIDIA TF32 documentation
• cuDNN developer guide
• Intel oneDNN

---

Notebook: 20_backends_tuning.ipynb

Module 21: CUDA Graphs — Eliminating CPU Launch Overhead

Day 7 of the incremental learning series

---

Table of Contents

• What Are CUDA Graphs?
• Why CUDA Graphs Matter
• Basic API: torch.cuda.CUDAGraph
• The Static Inputs Requirement
• Warmup — Why It's Mandatory
• CUDA Graph Pools
• torch.compile with CUDA Graphs
• Limitations — What Breaks
• CUDA Graphs with AMP
• torch.cuda.make_graphed_callables
• Practical Patterns
• When to Use / Not Use
• Upstream Updates (June 11–12, 2026)
• Further Reading

---

1. What Are CUDA Graphs?

When you run PyTorch code on a GPU, every operation — matrix multiply, activation, copy — is a kernel launch. Each launch requires the CPU to:

• Prepare kernel arguments
• Submit work to the CUDA driver
• The driver queues it on the GPU stream
• Return control to the CPU

For large kernels (e.g., a big GEMM), the GPU execution time dwarfs this overhead. But for small/medium operations, the CPU launch latency (5–15 microseconds per kernel) can dominate total runtime.

CUDA Graphs solve this by recording a sequence of GPU operations into a graph during a capture phase, then replaying the entire graph with a single CPU-side launch.

Normal execution:            CUDA Graph replay:

CPU: launch K1               CPU: launch graph ─────────────────┐
     wait...                                                     │
     launch K2               GPU: K1 → K2 → K3 → K4 → K5       │
     wait...                       (all pre-recorded, no waits)  │
     launch K3                                                   │
     wait...                 Result: 1 CPU launch instead of 5   │
     launch K4               ────────────────────────────────────┘
     wait...
     launch K5

GPU: ▓░░▓░░▓░░▓░░▓          GPU: ▓▓▓▓▓
     (gaps = idle)                 (no gaps)

The graph captures:

• Which kernels to run and in what order
• Memory addresses of all inputs and outputs
• Kernel launch parameters

It does not capture tensor values — only the operations and where they read/write.

---

2. Why CUDA Graphs Matter

The CPU Bottleneck

Modern GPUs execute small kernels in microseconds. If each kernel takes 3 µs on the GPU but 10 µs for the CPU to launch, you spend 77% of your time on CPU overhead.

Without CUDA Graphs (100 kernels):
  CPU overhead:  100 × 10 µs = 1,000 µs
  GPU compute:   100 ×  3 µs =   300 µs
  Total:                        1,300 µs
  GPU utilization:                  23%

With CUDA Graphs (100 kernels):
  CPU overhead:    1 × 10 µs =    10 µs
  GPU compute:   100 ×  3 µs =   300 µs
  Total:                          310 µs
  GPU utilization:                  97%
  Speedup:                        4.2x

When Speedups Are Largest

The pattern: the smaller the model and the more kernels per unit of compute, the bigger the win.

---

3. Basic API: torch.cuda.CUDAGraph

Minimal Example

import torch

model = torch.nn.Linear(512, 512).cuda()
static_input = torch.randn(64, 512, device="cuda")

# Step 1: Warmup (covered in Section 5)
with torch.no_grad():
    for _ in range(3):
        _ = model(static_input)

# Step 2: Capture
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = model(static_input)

# Step 3: Replay
static_input.copy_(torch.randn(64, 512, device="cuda"))
g.replay()
print(static_output)  # Contains result for the new input

What Happens During Capture

When you enter torch.cuda.graph(g):

• PyTorch switches to a capturing stream
• Every CUDA operation is recorded (not executed normally)
• Memory allocations inside the block come from a special graph pool
• On exit, the graph is finalized and ready for replay

What Happens During Replay

g.replay() submits the entire recorded sequence to the GPU in one shot. The GPU executes all kernels using the same memory addresses that were captured.

This is why inputs must be static — the graph hardcodes memory pointers.

---

4. The Static Inputs Requirement

This is the most important concept to understand. CUDA Graphs capture memory addresses, not tensor values. When you replay, the GPU reads from and writes to the exact same addresses.

Correct Pattern

# Pre-allocate (addresses are fixed)
static_input = torch.zeros(batch_size, features, device="cuda")
static_output = None

# Capture
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = model(static_input)

# For each new input: copy data IN, replay, read data OUT
for batch in dataloader:
    static_input.copy_(batch.cuda())  # Copy INTO the pre-allocated tensor
    g.replay()                        # GPU reads from same address
    result = static_output.clone()    # Copy OUT (or use in-place)

Wrong Pattern

# DON'T DO THIS — creates new tensor each iteration
for batch in dataloader:
    new_input = batch.cuda()       # New memory address each time!
    g.replay()                     # Graph still reads from OLD address
    # Result: stale data, wrong answers

Why This Design?

Recording memory addresses (not values) is what makes replay so fast. If the graph had to remap pointers each time, it would lose most of its advantage. The tradeoff: you manage input/output buffers manually.

---

5. Warmup — Why It's Mandatory

Before capturing a CUDA Graph, you must run the model at least once (usually 3 times for safety). Warmup triggers:

Warmup Pattern

model = MyModel().cuda().eval()
static_input = torch.randn(B, C, H, W, device="cuda")

# Warmup: run several forward passes
with torch.no_grad():
    for _ in range(3):
        _ = model(static_input)

# Now safe to capture
torch.cuda.synchronize()
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = model(static_input)

What Happens Without Warmup?

If you skip warmup, the capture records the lazy initialization itself — memory allocations, algorithm searches, handle creation. This means:

• The graph becomes bloated with one-time setup
• Replaying re-executes setup on every call (pointless overhead)
• Some lazy ops allocate memory dynamically, which is forbidden inside capture and will raise a RuntimeError

---

6. CUDA Graph Pools

When multiple graphs share the same model or intermediate buffers, you can share their memory pools to avoid redundant allocations.

Default: Separate Pools

g1 = torch.cuda.CUDAGraph()
g2 = torch.cuda.CUDAGraph()

with torch.cuda.graph(g1):
    out1 = model(input1)

with torch.cuda.graph(g2):
    out2 = model(input2)
# g1 and g2 each have their own memory pool

Shared Pool

g1 = torch.cuda.CUDAGraph()
g2 = torch.cuda.CUDAGraph()

# Capture both with the same pool
with torch.cuda.graph(g1):
    out1 = model(input1)

# Share g1's pool — no extra memory allocated for g2
with torch.cuda.graph(g2, pool=g1.pool()):
    out2 = model(input2)

Pool sharing works when graphs don't execute concurrently. The shared pool means they reuse the same memory, so overlapping execution would corrupt data.

When to Share Pools

• Multiple batch-size variants of the same model
• Encoder and decoder graphs that run sequentially
• A/B model comparisons that never overlap

---

7. torch.compile with CUDA Graphs

The easiest way to use CUDA Graphs is through torch.compile:

model = MyModel().cuda()

# reduce-overhead mode automatically uses CUDA Graphs
compiled = torch.compile(model, mode="reduce-overhead")

# Just use it normally — graphs are managed for you
output = compiled(input_tensor)

How It Works: cudagraph_trees

Under the hood, reduce-overhead mode uses Inductor's cudagraph_trees system:

torch.compile(mode="reduce-overhead")
  └─ Dynamo traces the Python code
      └─ AOTAutograd generates forward/backward
          └─ Inductor generates optimized CUDA code
              └─ cudagraph_trees wraps each compiled region in a graph

cudagraph_trees manages:

• Automatic warmup — runs the compiled code once before capturing
• Graph caching — one graph per unique input shape
• Memory management — pools are handled internally
• Fallback — if a region can't be graphed, it runs eagerly

Advantages Over Manual Graphs

Checking What Got Graphed

import torch._dynamo as dynamo

compiled = torch.compile(model, mode="reduce-overhead")
explanation = dynamo.explain(compiled)(sample_input)
print(explanation)

---

8. Limitations — What Breaks

CUDA Graphs capture a fixed sequence of GPU operations. Anything that deviates from this at replay time will fail or produce wrong results.

Hard Failures (RuntimeError)

These operations cannot be captured and will raise errors:

Silent Failures (Wrong Results)

These won't crash but will produce incorrect output:

Operations That Prevent Capture

# These will fail during capture:

# 1. CPU sync
with torch.cuda.graph(g):
    out = model(x)
    print(out.sum().item())  # ERROR: .item() syncs to CPU

# 2. Dynamic allocation
with torch.cuda.graph(g):
    out = model(x)
    mask = out > 0          # OK so far
    filtered = out[mask]    # ERROR: output size depends on data

# 3. CPU tensor creation
with torch.cuda.graph(g):
    scale = torch.tensor(2.0)  # ERROR: creates CPU tensor
    out = model(x) * scale

# Fix: pre-allocate the scale on GPU
scale = torch.tensor(2.0, device="cuda")
with torch.cuda.graph(g):
    out = model(x) * scale     # OK: scale already on GPU

NCCL and Distributed

Most NCCL collective operations (all-reduce, all-gather, etc.) are not compatible with CUDA Graphs. This is why CUDA Graphs are primarily used for inference and single-GPU training.

Exception: PyTorch's torch.distributed has experimental support for graphing some collectives on newer NCCL versions.

---

9. CUDA Graphs with AMP

Automatic Mixed Precision works inside CUDA Graph capture, but you must set up the autocast context inside the capture block:

model = MyModel().cuda()
static_input = torch.randn(64, 512, device="cuda")

# Warmup with AMP
with torch.no_grad(), torch.cuda.amp.autocast():
    for _ in range(3):
        _ = model(static_input)

# Capture with AMP
g = torch.cuda.CUDAGraph()
with torch.cuda.amp.autocast():
    with torch.cuda.graph(g):
        static_output = model(static_input)

# Replay (autocast context not needed — types are baked into the graph)
static_input.copy_(new_data)
g.replay()

Key Points

• The autocast context must be inside capture (or wrapping it) so the graph records the mixed-precision kernel variants
• After capture, replay uses whatever dtypes were recorded — no autocast needed at replay time
• GradScaler is harder — avoid it with CUDA Graphs. If you need training with AMP + graphs, prefer torch.compile(mode="reduce-overhead")

---

10. torch.cuda.make_graphed_callables

make_graphed_callables is a convenience wrapper that handles warmup, capture, and static-input management for nn.Module or simple callables:

model = MyModel().cuda()
sample_input = torch.randn(64, 512, device="cuda")

# Wrap the model — handles warmup and capture automatically
graphed_model = torch.cuda.make_graphed_callables(
    model,
    sample_args=(sample_input,),
    num_warmup_iters=3,
)

# Use like a normal callable
output = graphed_model(sample_input)

Multiple Callables

You can graph multiple modules together, sharing a pool:

encoder = Encoder().cuda()
decoder = Decoder().cuda()

graphed_encoder, graphed_decoder = torch.cuda.make_graphed_callables(
    (encoder, decoder),
    sample_args=(
        (encoder_input,),
        (decoder_input,),
    ),
)

When to Use

• Quick experiments where you want CUDA Graph speedups without manual buffer management
• Models with simple input/output signatures
• Inference-only paths

Caveats

• Only supports fixed shapes (like manual graphs)
• The returned callable replaces the original forward pass
• Multiple return values need careful handling

---

11. Practical Patterns

Pattern 1: Inference Server

The most common CUDA Graph use case — a model serving predictions with fixed batch size:

class GraphedInferenceServer:
    def __init__(self, model, batch_size, input_dim):
        self.model = model.cuda().eval()
        self.static_input = torch.zeros(
            batch_size, input_dim, device="cuda"
        )
        self.static_output = None
        self.graph = torch.cuda.CUDAGraph()

        # Warmup
        with torch.no_grad():
            for _ in range(3):
                _ = self.model(self.static_input)

        # Capture
        with torch.no_grad():
            with torch.cuda.graph(self.graph):
                self.static_output = self.model(self.static_input)

    def predict(self, input_tensor):
        self.static_input.copy_(input_tensor)
        self.graph.replay()
        return self.static_output.clone()

Pattern 2: Multiple Batch Sizes

For serving with variable batch sizes, capture one graph per batch size:

class MultiBatchServer:
    def __init__(self, model, batch_sizes, input_dim):
        self.model = model.cuda().eval()
        self.graphs = {}

        for bs in batch_sizes:
            static_in = torch.zeros(bs, input_dim, device="cuda")

            # Warmup
            with torch.no_grad():
                for _ in range(3):
                    _ = self.model(static_in)

            g = torch.cuda.CUDAGraph()
            with torch.no_grad():
                with torch.cuda.graph(g):
                    static_out = self.model(static_in)

            self.graphs[bs] = (g, static_in, static_out)

    def predict(self, input_tensor):
        bs = input_tensor.shape[0]
        g, static_in, static_out = self.graphs[bs]
        static_in.copy_(input_tensor)
        g.replay()
        return static_out.clone()

Pattern 3: Partial Graph Capture for Training

Full training steps are hard to graph (optimizer state updates, gradient scaling). Instead, graph just the forward pass:

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())
static_input = torch.randn(32, 512, device="cuda")
static_target = torch.randn(32, 10, device="cuda")

# Warmup
for _ in range(3):
    out = model(static_input)
    loss = torch.nn.functional.mse_loss(out, static_target)
    loss.backward()
    optimizer.zero_grad()

# Capture forward + backward (not optimizer step)
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    out = model(static_input)
    loss = torch.nn.functional.mse_loss(out, static_target)
    loss.backward()

# Training loop
for batch_input, batch_target in dataloader:
    static_input.copy_(batch_input)
    static_target.copy_(batch_target)
    g.replay()                   # Forward + backward replayed
    optimizer.step()             # Optimizer runs eagerly (safe)
    optimizer.zero_grad()

---

12. When to Use / Not Use

Decision Tree

Is your model running on NVIDIA GPU?
├─ No  → CUDA Graphs not applicable
└─ Yes
    │
    Are input shapes fixed (or a small set of fixed shapes)?
    ├─ No  → Use torch.compile (handles dynamic shapes)
    └─ Yes
        │
        Is CPU launch overhead significant?
        (small/medium model, many small kernels, high throughput needed)
        ├─ No  → Graphs won't help much, use torch.compile for fusion
        └─ Yes
            │
            Is it inference only?
            ├─ Yes → CUDA Graphs are ideal
            │        Use torch.compile(mode="reduce-overhead")
            │        or manual CUDAGraph API
            └─ No (training)
                │
                Can you graph forward+backward only?
                ├─ Yes → Partial graph capture (optimizer runs eagerly)
                └─ No  → torch.compile(mode="reduce-overhead") is safer

Quick Reference

---

13. Upstream Updates (June 11–12, 2026)

Recent PyTorch development activity relevant to CUDA Graphs and the broader ecosystem:

Version Bump to 2.14.0a0 (#187070)

The main branch has been bumped to version 2.14.0a0+, signaling the start of the next development cycle. This is the version that includes the latest CUDA Graph improvements and new features described below.

FlexGEMM Higher-Order Op (torch/_higher_order_ops/flex_gemm.py)

A new higher-order operation for flexible GEMM execution has been added. FlexGEMM allows user-defined epilogues on GEMM results, similar to how FlexAttention allows custom attention score modifications. This interacts with CUDA Graphs through the reduce-overhead compilation path.

c10d Window Interfaces for One-Sided Communication (#186299)

New window-based interfaces in PyTorch's distributed backend (c10d) enable one-sided MPI-style communication patterns (put, get, accumulate). While most NCCL collectives remain incompatible with CUDA Graphs, these new primitives expand the distributed toolkit.

cuSOLVER Workspace Optimization (#181998)

Workspace allocation for cuSOLVER operations has been optimized, reducing memory overhead for linear algebra operations. This is relevant to CUDA Graphs because workspace allocations during graph capture can cause issues — the optimization makes these allocations more predictable and graph-friendly.

Dynamo itertools.permutations Polyfill (#186937)

torch._dynamo now supports itertools.permutations during tracing, allowing more Python code to be captured without graph breaks. Fewer graph breaks mean larger graphable regions, which directly benefits CUDA Graph capture through torch.compile(mode="reduce-overhead").

uint16/uint32/uint64 Test Coverage Extended (#183473)

Test infrastructure now covers unsigned integer types more broadly. While not directly related to CUDA Graphs, this improves the reliability of quantized models that may be served with CUDA Graph-accelerated inference.

FlexAttention INDEX_DTYPE for Pointer Arithmetic (#185264)

FlexAttention now uses a dedicated INDEX_DTYPE for pointer arithmetic in block-sparse patterns, improving numerical stability on different GPU architectures. FlexAttention kernels are common targets for CUDA Graph capture in Transformer inference.

---

14. Further Reading

• CUDA Graphs documentation (PyTorch)
• NVIDIA CUDA Graphs guide
• torch.compile reduce-overhead mode
• make_graphed_callables API
• Accelerating PyTorch with CUDA Graphs (NVIDIA blog)

---

Notebook: 21_cuda_graphs.ipynb