style(tinytorch): apply pre-push hook markdown formatting

Author: profvjreddi

books/tinytorch/02_activations.qmd

@@ -31,7 +31,7 @@ Substituting the first equation into the second yields:
 
 $$\mathbf{y} = (\mathbf{x} \mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2 = \mathbf{x}(\mathbf{W}_1 \mathbf{W}_2) + (\mathbf{b}_1 \mathbf{W}_2 + \mathbf{b}_2)$$
 
-Because matrix multiplication is associative and distributive, the product $\mathbf{W}_1 \mathbf{W}_2$ can be precomputed as a single matrix $\mathbf{W}_{\text{comb}} \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$, and the combined bias vector $\mathbf{b}_1 \mathbf{W}_2 + \mathbf{b}_2$ as $\mathbf{b}_{\text{comb}} \in \mathbb{R}^{d_{\text{out}}}$. 
+Because matrix multiplication is associative and distributive, the product $\mathbf{W}_1 \mathbf{W}_2$ can be precomputed as a single matrix $\mathbf{W}_{\text{comb}} \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$, and the combined bias vector $\mathbf{b}_1 \mathbf{W}_2 + \mathbf{b}_2$ as $\mathbf{b}_{\text{comb}} \in \mathbb{R}^{d_{\text{out}}}$.
 
 ```
 ┌─────────────────────────────────────────────────────────────┐
@@ -97,10 +97,11 @@ $$\sigma(x) = \frac{1}{1 + e^{-x}}$$
 In standard IEEE 754 single-precision arithmetic (`float32`), numbers are encoded with 1 sign bit, 8 exponent bits (with bias 127), and 23 mantissa bits.[^ieee754-spec] This yields a dynamic range with maximum finite float $\approx (2 - 2^{-23}) \times 2^{127} \approx 3.4028 \times 10^{38}$.
 
 Evaluating $e^{-x}$ on hardware triggers numerical hazards at extreme boundaries:
+
 - If $x < -88.72$, $-x > 88.72$, causing $e^{-x} > 3.4028 \times 10^{38}$, which overflows to `+inf`. Evaluating $1 / (1 + \text{inf})$ can cause floating-point exception traps.
 - If $x > 88.72$, $e^{-x}$ drops below the minimum positive normal float ($2^{-126} \approx 1.1755 \times 10^{-38}$), underflowing to subnormal numbers or `0.0`.
 
-[^ieee754-spec]: IEEE 754 `float32` allocates 8 bits to the exponent ($E \in [0, 255]$). An exponent of $E=255$ is reserved for $\pm \infty$ (when mantissa is 0) and `NaN` (when mantissa is non-zero). An exponent of $E=0$ represents subnormal numbers or zero. When subnormal numbers occur, CPU pipelines may switch to microcode exception handling, slowing execution by up to 100× unless Flush-To-Zero (FTZ) and Denormals-Are-Zero (DAZ) hardware control flags are enabled.
+[^ieee754-spec]: IEEE 754 `float32` allocates 8 bits to the exponent ($E \in [0, 255]$). An exponent of $E=255$ is reserved for $\pm \infty$ (when mantissa is 0) and `NaN` (when mantissa is nonzero). An exponent of $E=0$ represents subnormal numbers or zero. When subnormal numbers occur, CPU pipelines may switch to microcode exception handling, slowing execution by up to 100× unless Flush-To-Zero (FTZ) and Denormals-Are-Zero (DAZ) hardware control flags are enabled.
 
 To guarantee complete numerical stability across the entire real number line without register overflow, TinyTorch uses a branch-stable formulation:
 
@@ -119,14 +120,14 @@ class Sigmoid:
         with np.errstate(over='ignore'):
             positive_mask = (x_data >= 0)
             result = np.zeros_like(x_data, dtype=np.float32)
-            
+
             # x >= 0 branch: no positive exponential overflow
             result[positive_mask] = 1.0 / (1.0 + np.exp(-x_data[positive_mask]))
-            
+
             # x < 0 branch: evaluate exp(x)/(1 + exp(x))
             exp_neg = np.exp(x_data[~positive_mask])
             result[~positive_mask] = exp_neg / (1.0 + exp_neg)
-            
+
         return Tensor(result)
 
     def __call__(self, x: Tensor) -> Tensor:
@@ -280,6 +281,7 @@ Because activation layers contain no internal state or learnable weights, they c
 While matrix multiplications ($\mathbf{x} \mathbf{W}$) have high arithmetic intensity ($O(N^3)$ operations on $O(N^2)$ data), activation functions are **strictly memory-bandwidth bound**.
 
 For a tensor of $N$ elements executing ReLU in 32-bit float:
+
 - Read 4 bytes per element from DRAM into CPU/GPU registers.
 - Perform 1 comparison operation in the ALU.
 - Write 4 bytes per element back to the DRAM memory buffer.

books/tinytorch/03_layers.qmd

@@ -96,6 +96,7 @@ class Layer:
 ```
 
 The `Linear` layer (also termed a *Dense* or *Fully Connected* layer) encapsulates two trainable parameter tensors:
+
 1. **Weight Matrix ($\mathbf{W}$)** of shape $(d_{\text{in}}, d_{\text{out}})$, initialized with LeCun scaled normal noise.
 2. **Bias Vector ($\mathbf{b}$)** of shape $(d_{\text{out}},)$, initialized to zeros.
 
@@ -105,7 +106,7 @@ import numpy as np
 class Linear(Layer):
     """
     Fully connected linear transformation: y = x @ W + b.
-    
+
     Encapsulates trainable weight and bias parameter Tensors.
     """
 
@@ -134,7 +135,7 @@ class Linear(Layer):
     def forward(self, x: Tensor) -> Tensor:
         """
         Compute affine projection: y = x @ W + b.
-        
+
         Args:
             x: Input Tensor of shape (batch_size, in_features)
         Returns:
@@ -176,6 +177,7 @@ Standard vs. Inverted Dropout Execution. During training, a random binary mask z
 In classical dropout (Hinton et al., 2012), activations were multiplied by a Bernoulli mask during training, which reduced the expected activation sum by a factor of $(1 - p)$. To compensate, test-time inference required multiplying all model weights or activations by $(1 - p)$.
 
 In production ML systems, modifying weights or executing scaling operations during inference creates severe systems overhead:
+
 1. **Quantization Invalidation**: Multiplying pre-quantized INT8 weights by arbitrary floating-point factors destroys quantization scales.
 2. **Inference Latency**: Launching an element-wise scaling kernel at every layer wastes memory bandwidth on latency-critical edge devices.
 
@@ -189,7 +191,7 @@ During inference, evaluation mode becomes a **pure cost-free identity pass** ($y
 class Dropout(Layer):
     """
     Inverted Dropout regularization layer.
-    
+
     Randomly zeroes elements with probability p during training.
     Scales surviving elements by 1 / (1 - p) to preserve expectation.
     """
@@ -318,10 +320,10 @@ $$\text{Memory} = \text{Total Parameters} \times \text{sizeof}(\text{dtype})$$
 ## Check Your Understanding: Calculating Network Weight Footprint
 
 Consider the 3-layer MLP constructed above:
+
 - Layer 1: `Linear(784, 128)`
 - Layer 2: `Linear(128, 64)`
 - Layer 3: `Linear(64, 10)`
-
 1. Calculate the total number of learnable parameters in the network.
 2. How many bytes of RAM are required to store these parameters in 32-bit float (`float32`, 4 bytes per element)?
 

books/tinytorch/04_losses.qmd

@@ -43,6 +43,7 @@ $$\mathcal{L} = \text{Loss}(\hat{\mathbf{y}}, \mathbf{y})$$
 ```
 
 The computed loss scalar serves as the definitive numerical score of prediction quality:
+
 - If predictions match targets perfectly, $\mathcal{L} \to 0$.
 - If predictions diverge from targets, $\mathcal{L}$ grows, providing the magnitude of error that will subsequently guide optimization.
 
@@ -66,7 +67,7 @@ class MSELoss:
     def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
         """
         Compute mean squared error between predictions and targets.
-        
+
         Args:
             predictions: Model output Tensor of shape (B, ...)
             targets: Ground-truth Tensor of identical shape (B, ...)
@@ -82,6 +83,7 @@ class MSELoss:
 ```
 
 **Properties of MSE**:
+
 1. **Non-negativity**: $\mathcal{L}_{\text{MSE}} \ge 0$, with equality if and only if $\hat{\mathbf{y}} = \mathbf{y}$.
 2. **Outlier Sensitivity**: Because errors are squared, a prediction error of $10.0$ generates a penalty 100× greater than an error of $1.0$, heavily discouraging large deviations.
 
@@ -157,6 +159,7 @@ Substituting back into the log-softmax formula gives the exact **Log-Sum-Exp Inv
 $$\log(\text{Softmax}(\mathbf{z})_i) = (z_i - M) - \log\left(\sum_{j=1}^C e^{z_j - M}\right)$$
 
 This formula possesses two vital numerical stability guarantees:
+
 1. **No Overflow**: Because $z_j - M \le 0$ for all $j$, the exponential terms $e^{z_j - M} \in (0, 1]$ **never exceed 1.0**.
 2. **No Underflow to Zero**: Because the maximum element satisfies $z_{\text{argmax}} - M = 0$, the sum always contains $e^0 = 1.0$, guaranteeing $\sum_{j=1}^C e^{z_j - M} \ge 1.0$. The logarithm argument is strictly bounded away from zero ($\log(\ge 1.0) \ge 0.0$), making $\log(0) \to -\infty$ mathematically impossible.
 :::
@@ -165,7 +168,7 @@ This formula possesses two vital numerical stability guarantees:
 def log_softmax(x: Tensor, dim: int = -1) -> Tensor:
     """
     Compute log-softmax using the Log-Sum-Exp numerical stability trick.
-    
+
     Args:
         x: Input Tensor of unnormalized logits
         dim: Dimension along which to compute log-softmax (default: -1)
@@ -201,7 +204,6 @@ Consider a large language model with vocabulary size $C = 128,000$ processing a
    - Softmax reads logits $(4096 \times 128,000)$ from DRAM, computes probabilities, and writes probabilities $(4096 \times 128,000 \times 4\text{ B} \approx 2.1\text{ GB})$ back to DRAM.
    - Cross-entropy reads $2.1\text{ GB}$ of probabilities from DRAM, computes logs, indexes target tokens, and writes loss.
    - **Total DRAM Traffic**: Over $4.2\text{ GB}$ of high-latency memory bus traffic.
-
 2. **Fused Log-Sum-Exp Loss Kernel**:
    - The GPU kernel loads logit tiles directly into on-chip SRAM registers.
    - Threads compute the maximum logit $M$ and sum-of-exponentials $\sum e^{z_j - M}$ using fast in-register warp shuffle reductions.[^warp-shuffle]
@@ -220,14 +222,14 @@ In TinyTorch (as in PyTorch), `CrossEntropyLoss` takes **raw unnormalized logits
 class CrossEntropyLoss:
     """
     Cross-Entropy loss combining log_softmax and Negative Log-Likelihood (NLL).
-    
+
     Expects unnormalized logits directly from the model.
     """
 
     def forward(self, logits: Tensor, targets: Tensor) -> Tensor:
         """
         Compute categorical cross-entropy loss.
-        
+
         Args:
             logits: Unnormalized logits Tensor of shape (B, C)
             targets: Class indices (B,) or one-hot vectors (B, C)

books/tinytorch/07_optimizers.qmd

@@ -47,6 +47,7 @@ where $\eta > 0$ is the learning rate and $\mathbf{u}(\cdot)$ is the transformed
 In high-dimensional deep learning loss surfaces, the local curvature is characterized by the Hessian matrix $\mathbf{H} = \nabla^2 \mathcal{L}(\mathbf{\theta})$. The condition number $\kappa = \frac{\lambda_{\max}}{\lambda_{\min}}$ measures the ratio of maximum to minimum curvature eigenvalues.
 
 When $\kappa \gg 1$ (an ill-conditioned ravine):
+
 - The gradient along the steep perpendicular direction ($\mathbf{g}_\perp$) is massive, causing vanilla SGD to oscillate violently between ravine walls.
 - The gradient along the flat parallel valley floor ($\mathbf{g}_\parallel$) is tiny, causing progress toward the minimum to stall.
 
@@ -69,6 +70,7 @@ $$\mathbf{v}_t = \beta \mathbf{v}_{t-1} + (1 - \beta) \mathbf{g}_t, \quad \mathb
 where $\beta \in [0, 1)$ is the momentum decay factor (typically $\beta = 0.9$).
 
 **Two Critical Physical Properties**:
+
 1. **High-Frequency Damping**: Successive gradient vectors along oscillating directions alternate signs ($\mathbf{g}_{\perp, t} \approx -\mathbf{g}_{\perp, t-1}$), canceling out inside the accumulator ($\sum \mathbf{g}_\perp \to \mathbf{0}$).
 2. **Saddle Point Escape**: At strict saddle points or flat plateaus where local gradient $\nabla \mathcal{L} \approx \mathbf{0}$, standard SGD becomes trapped. Momentum retains stored kinetic velocity ($\mathbf{v}_t = \beta \mathbf{v}_{t-1} \ne \mathbf{0}$), rolling the parameter trajectory through flat regions without stopping.
 
@@ -207,24 +209,24 @@ In large-scale training systems, optimizer state buffers are often the single la
 
 For a model with $N$ parameters, let us calculate the exact memory required to store parameters, gradients, and optimizer states:
 
-| Precision Regime | Component | Precision & Type | Bytes / Param | 7B Model Footprint |
-| :--- | :--- | :--- | :---: | :---: |
-| **Pure FP32 Training** | Model Weights ($\mathbf{\theta}$) | `float32` | $4\text{ bytes}$ | $28.0\text{ GB}$ |
-| | Gradients ($\mathbf{g}$) | `float32` | $4\text{ bytes}$ | $28.0\text{ GB}$ |
-| | AdamW 1st Moment ($\mathbf{m}$) | `float32` | $4\text{ bytes}$ | $28.0\text{ GB}$ |
-| | AdamW 2nd Moment ($\mathbf{v}$) | `float32` | $4\text{ bytes}$ | $28.0\text{ GB}$ |
-| | **Total Static Footprint** | | **$16\text{ bytes}$** | **$112.0\text{ GB}$** |
-| **Mixed Precision (AMP)** | Model Weights ($\mathbf{\theta}$) | `float16` / `bfloat16` | $2\text{ bytes}$ | $14.0\text{ GB}$ |
-| | Gradients ($\mathbf{g}$) | `float16` / `bfloat16` | $2\text{ bytes}$ | $14.0\text{ GB}$ |
-| | Master Weights ($\mathbf{\theta}_{\text{master}}$) | `float32` | $4\text{ bytes}$ | $28.0\text{ GB}$ |
-| | AdamW 1st Moment ($\mathbf{m}$) | `float32` | $4\text{ bytes}$ | $28.0\text{ GB}$ |
-| | AdamW 2nd Moment ($\mathbf{v}$) | `float32` | $4\text{ bytes}$ | $28.0\text{ GB}$ |
-| | **Total Static Footprint** | | **$16\text{ bytes}$** | **$112.0\text{ GB}$** |
+| Precision Regime          | Component                                          | Precision & Type       |     Bytes / Param     |   7B Model Footprint  |
+|:--------------------------|:---------------------------------------------------|:-----------------------|:---------------------:|:---------------------:|
+| **Pure FP32 Training**    | Model Weights ($\mathbf{\theta}$)                  | `float32`              |    $4\text{ bytes}$   |    $28.0\text{ GB}$   |
+|                           | Gradients ($\mathbf{g}$)                           | `float32`              |    $4\text{ bytes}$   |    $28.0\text{ GB}$   |
+|                           | AdamW 1st Moment ($\mathbf{m}$)                    | `float32`              |    $4\text{ bytes}$   |    $28.0\text{ GB}$   |
+|                           | AdamW 2nd Moment ($\mathbf{v}$)                    | `float32`              |    $4\text{ bytes}$   |    $28.0\text{ GB}$   |
+|                           | **Total Static Footprint**                         |                        | **$16\text{ bytes}$** | **$112.0\text{ GB}$** |
+| **Mixed Precision (AMP)** | Model Weights ($\mathbf{\theta}$)                  | `float16` / `bfloat16` |    $2\text{ bytes}$   |    $14.0\text{ GB}$   |
+|                           | Gradients ($\mathbf{g}$)                           | `float16` / `bfloat16` |    $2\text{ bytes}$   |    $14.0\text{ GB}$   |
+|                           | Master Weights ($\mathbf{\theta}_{\text{master}}$) | `float32`              |    $4\text{ bytes}$   |    $28.0\text{ GB}$   |
+|                           | AdamW 1st Moment ($\mathbf{m}$)                    | `float32`              |    $4\text{ bytes}$   |    $28.0\text{ GB}$   |
+|                           | AdamW 2nd Moment ($\mathbf{v}$)                    | `float32`              |    $4\text{ bytes}$   |    $28.0\text{ GB}$   |
+|                           | **Total Static Footprint**                         |                        | **$16\text{ bytes}$** | **$112.0\text{ GB}$** |
 
 ::: {.callout-principle}
 ## The 16-Byte Invariant of AdamW Training
 
-A common systems misconception is that switching from FP32 to 16-bit Mixed Precision (FP16/BF16) cuts total training memory in half. 
+A common systems misconception is that switching from FP32 to 16-bit Mixed Precision (FP16/BF16) cuts total training memory in half.
 
 In reality, because multiplying 16-bit gradients by small learning rates ($\eta = 10^{-4}$) causes underflow below the dynamic range of FP16 ($2^{-14} \approx 6 \times 10^{-5}$), adaptive optimizers **must maintain master weights and moment accumulators in full 32-bit float**. Consequently, both FP32 and AMP AdamW training require exactly **$16\text{ bytes per parameter}$** of static VRAM—excluding all activation memory.[^zero-mem][^eight-bit-adam]
 :::

books/tinytorch/08_training.qmd

@@ -105,7 +105,7 @@ def clip_grad_norm(params: List[Tensor], max_norm: float = 1.0) -> float:
     for p in params:
         if p.grad is not None:
             total_norm_sq += float(np.sum(p.grad.data ** 2))
-    
+
     total_norm = np.sqrt(total_norm_sq)
     clip_coef = max_norm / max(total_norm, 1e-6)
 
@@ -228,7 +228,7 @@ def load_checkpoint(model: Layer, optimizer: Optimizer, filepath: str) -> int:
     """Restore model parameters and optimizer state from disk."""
     with open(filepath, 'rb') as f:
         state = pickle.load(f)
-    
+
     for p, p_data in zip(model.parameters(), state['model_params']):
         p.data = p_data.copy()
 
@@ -245,6 +245,7 @@ def load_checkpoint(model: Layer, optimizer: Optimizer, filepath: str) -> int:
 ## Checkpoint Provenance and Determinism
 
 Saving only model weights ($\mathbf{\theta}$) is insufficient to resume training faithfully. If optimizer momentum buffers ($\mathbf{m}, \mathbf{v}$) are discarded, the resumed optimizer starts with zero momentum, triggering massive gradient shocks and loss spikes. A production checkpoint must serialize:
+
 1. **Model parameters** (`state_dict`).
 2. **Optimizer state buffers** (momentum, variance, step counters).
 3. **Learning rate scheduler state** (current decay factor and step index).
@@ -263,6 +264,7 @@ In modern AI clusters, the single-GPU `Trainer` scales horizontally across hundr
 
 1. **`torch.nn.DataParallel` (Single-Process Multi-Threaded)**:
    Runs in a single Python process across multiple threads. It is heavily bottlenecked by Python's Global Interpreter Lock (GIL). On every batch, GPU 0 must broadcast the model weights to all other GPUs, gather all predictions to GPU 0 to compute the loss, and scatter gradients. This causes severe **GPU 0 VRAM memory imbalance** and stalls compute.
+
 2. **`torch.nn.parallel.DistributedDataParallel` (Multi-Process Multi-GPU)**:
    Spawns an independent OS process per GPU (`torchrun`), completely eliminating GIL contention. Each process maintains its own independent model replica and optimizer in local GPU VRAM.
 
@@ -283,6 +285,7 @@ In modern AI clusters, the single-GPU `Trainer` scales horizontally across hundr
 ## PyTorch Systems Bridge: DDP Gradient Bucketing in C++
 
 In PyTorch, DDP communication is handled in native C++ by `torch/csrc/distributed/c10d/reducer.cpp`:
+
 - **Gradient Bucketing**: PyTorch allocates flat, contiguous memory buffers called **buckets** (default: `bucket_cap_mb=25`).
 - As backpropagation evaluates parameter gradients in reverse DAG order, the autograd engine writes gradients directly into the active bucket.
 - As soon as all parameters assigned to a bucket finish computing gradients, the `Reducer` dispatches an asynchronous, non-blocking **NCCL Ring-AllReduce** kernel over NVLink or InfiniBand on a dedicated CUDA communication stream.

books/tinytorch/12_attention.qmd

@@ -262,7 +262,6 @@ class MultiHeadAttention(Layer):
 
 ---
 
-
 ::: {.callout-note}
 ## PyTorch Systems Bridge: `torch.nn.functional.scaled_dot_product_attention` (SDPA)