PyTorch Fundamentals

I. The Imperative Mental Model

Transitioning from standard programming to deep learning requires a shift in how we view data and operations. As described in Paszke et al. (2019), PyTorch was designed to provide an imperative, "Python-first" experience.

Unlike static graph frameworks (like early TensorFlow), PyTorch constructs the computational graph on the fly during the forward pass. This allows for standard Python control flow (if-statements, loops) to be part of the model logic, making debugging and research much more intuitive.

II. Tensor Semantics & Memory

The **Tensor** is the fundamental unit of data in PyTorch. While it resembles a NumPy array, a data type, but also the device (CPU/CUDA) and layout.

Storage vs. View: A tensor is physically a 1D array of typically contiguous bytes in RAM/VRAM known as the Storage. The "Tensor" object is merely a metadata view (Shape, Stride, Offset) over this storage.

Strides: The stride is a tuple indicating the number of steps needed in the storage to move by 1 in each dimension. For a tensor \( T \), the address of cell \( (i, j) \) is: $$ \text{addr} = \text{base\_addr} + i \cdot \text{stride}[0] + j \cdot \text{stride}[1] $$ Operations like transpose or permute are typically zero-cost because they only modify the stride tuple, not the underlying memory.

III. The Semantics of Broadcasting

Broadcasting allows operations on tensors of different ranks and shapes by logically expanding dimensions of size 1 to match target dimensions. This follows the standard NumPy broadcasting convention.

IV. Automatic Differentiation

**Autograd** is the engine that enables neural networks to learn. It implements a technique known as Reverse-Mode Automatic Differentiation. As you perform operations on tensors, PyTorch builds a directed acyclic graph (DAG) where nodes represent tensors and edges represent functions.

By calling .backward(), the engine traverses this graph in reverse order, applying the chain rule to compute the gradient of the loss with respect to every leaf tensor (parameter). This eliminates the need for researchers to manually derive complex vector derivatives, a process that was the primary bottleneck of 20th-century AI.

The Jacobian-Vector Product (JVP): PyTorch does not fully compute the Jacobian matrix \( J \) (which would be massive). Instead, it computes the product of the Jacobian with an incoming gradient vector \( v \):

This accumulation is efficient and naturally supports chain rule composition: \( v_{in} \cdot J_{total} = (v_{out} \cdot J_n) \cdot J_{n-1} \dots \).

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input) # Save tensor for backward phase
        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 # Mask gradient for negative inputs
        return grad_input

V. CUDA Kernels & Async Dispatch

CUDA (Compute Unified Device Architecture) execution is Asynchronous. When you call a PyTorch function on GPU tensors, the CPU enqueues the kernel into a "Stream" and immediately returns. This hides the launch latency. Synchronization (blocking) only happens when you print values, copy to CPU (.item()), or assume logical dependencies.

Pinned Memory: For faster Host-Device transfer, use tensor.pin_memory(). This effectively "page-locks" the RAM, allowing the DMA (Direct Memory Access) controller to copy data to the GPU without CPU intervention.

Training on a single GPU is rarely sufficient for modern models. Distributed Data Parallel (DDP) is the industry standard for multi-GPU training.

The Mechanism:

1. Replica: The model is replicated across all GPUs (ranks).
2. Scatter: The batch is split (scattered) so each GPU processes a different slice of data.
3. Forward/Backward: Each GPU computes gradients independently.
4. All-Reduce: Gradients are averaged across all GPUs using a Ring-AllReduce algorithm.
5. Update: The optimizer steps on each GPU with the identical average gradient, keeping weights synchronized.

FSDP (Fully Sharded Data Parallel): For models that don't fit on one GPU (e.g., Llama-70B), FSDP shards parameters, gradients, and optimizer states across GPUs, gathering them only when needed for computation.

VII. The Compiler: JIT & TorchScript

Python's overhead (GIL, interpreter dispatch) can be significant for small operations. TorchScript and torch.compile (PyTorch 2.0) aim to decouple the model from the Python runtime.

PyTorch 2.0 (Dynamo & Inductor): Use model = torch.compile(model). Dynamo captures the graph by overriding Python bytecode execution. Inductor generates tuned Triton kernels for the specific GPU architecture, often yielding 30-200% speedups.

Primary Sources & Further Reading