Model Optimization & Efficiency

I. The Memory Wall & VRAM Bottlenecks

In the current era of LLMs, the primary constraint is not compute power (FLOPS), but Memory Bandwidth. Training and serving models with billions of parameters requires moving massive amounts of data from High Bandwidth Memory (HBM) to the GPU cores.

We define the "Memory Wall" as the point where the speed of data transfer becomes the limiting factor in AI performance. To overcome this, we must employ compression techniques that reduce the model's footprint without severely degrading its reasoning capabilities.

Sharpness-Aware Minimization (SAM) Modern optimization isn't just about finding the lowest loss, but finding the **flattest** minima. Foret et al. (2020) introduced SAM, which simultaneously minimizes loss value and loss sharpness. This improves generalization by ensuring the model stays in a "safe" region even if weights are slightly perturbed (e.g., via quantization).

Hessian Curvature

Second-order derivatives define if the loss landscape is a narrow ravine or a flat valley.

II. Adaptive Optimization: Beyond SGD

Training Large Language Models requires navigating a highly non-convex loss landscape. Vanilla **Stochastic Gradient Descent (SGD)** struggles with saddle points and ravines. Modern training relies on **Adaptive Moment Estimation (Adam)**.

AdamW (Loshchilov & Hutter, 2017) decouples weight decay from the gradient update. It maintains per-parameter learning rates based on the first moment (mean) $m_t$ and second moment (variance) $v_t$ of the 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 $$ $$ \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta \lambda \theta_t $$

Why AdamW? In standard Adam with L2 regularization, the weight decay is applied to the gradients, which interacts poorly with the adaptive learning rate scaling. AdamW applies decay directly to the weights, leading to better generalization.

III. PEFT: Low-Rank Adaptation (LoRA)

Full-parameter fine-tuning of LLMs is prohibitively expensive. LoRA (Hu et al., 2021) solves this by freezing the original weights and training two small, low-rank matrices that represent the delta of the weights.

Mathematically, if the original weight matrix is $W$, we learn its update as $\Delta W = A \times B$, where A and B have a very small rank $r$. This reduces the number of trainable parameters by over 10,000x, enabling fine-tuning on a single GPU.

$$ \mathcal{L} = (1-\alpha)\mathcal{L}_{CE}(y, \hat{y}) + \alpha T^2 \mathcal{L}_{KL}(p_T, p_S) $$

IV. Learning Rate & Scheduler Dynamics

The **Learning Rate ($\eta$)** is the most critical hyperparameter. Schedulers like 'Cosine Annealing' or 'OneCycleLR' fluctuate the rate to escape local minima and ensure convergence to flat, generalizable regions.

Scheduler Workbench

Compare Cosine Decay vs. Linear Warmup. Warmup is essential to stabilize the first few thousand training steps.

Scheduler Type

LR Scale Factor

1.0

V. Quantization: INT8, FP4 & NF4

Affine Quantization maps the floating point range $[min, max]$ to the integer range $[0, 2^b-1]$ using a scale factor $S$ and zero-point $Z$.

$$ Q(x) = \text{clamp}\left(\text{round}\left(\frac{x}{S} + Z\right), 0, 2^B-1\right) $$ $$ x_{dequant} = S(Q(x) - Z) $$

Normal Float 4 (NF4): Introduced in QLoRA (Dettmers et al., 2023), NF4 is an information-theoretically optimal data type for weights that are normally distributed. It defines the quantization levels based on the quantiles of a standard normal distribution $\mathcal{N}(0, 1)$.

Double Quantization QLoRA goes further by quantizing the quantization constants (scales) themselves, saving an additional 0.37 bits per parameter.

VI. Teacher-Student Knowledge Distillation

A compact "Student" model can achieve higher performance by learning from the **soft probabilities** of a massive "Teacher" model rather than the raw hard labels.

$$ \mathcal{L}_{total} = (1-\alpha)\mathcal{L}_{CE}(y, \hat{y}) + \alpha \tau^2 \text{KL}( p(\frac{z_T}{\tau}) || p(\frac{z_S}{\tau}) ) $$

The "Temperature" parameter $\tau$ softens the probability distribution, forcing the student to learn the structual relationships between incorrect classes (e.g., "Dog" is more similar to "Cat" than "Car"), known as Dark Knowledge.

VII. GPU Kernels & Triton

Writing raw CUDA (C++) is hard. OpenAI Triton is a Python-like DSL that compiles to highly efficient block-based GPU kernels. It is the backend for torch.compile (Inductor).

Tiled Access: GPUs have massive global memory (HBM) but small, fast cache (SRAM). Triton optimizes Memory Coalescing—loading data in aligned blocks to minimize memory transactions.

Operator Fusion The biggest speedup comes from fusion. Instead of writing `x = MatMul(a, b)` then `y = ReLU(x)` to HBM (2 reads, 2 writes), Triton generates a single kernel `ReLU(MatMul(...))` (1 read, 1 write), running entirely in SRAM.

VIII. Pruning & The Lottery Ticket Hypothesis

Frankle & Carbin (2018) proved that dense networks contain sparse subnetworks ("winning tickets") that can match the original performance when trained in isolation.

Iterative Magnitude Pruning (IMP) The "winning tickets" are identified by training a network, pruning the smallest weights, and then **rewinding** the remaining weights to their original initializations. This suggests that initialization is as critical as architecture.

Unstructured Pruning

Zeroing out individual weights. Easiest to implement, but requires specialized Sparse Matrix kernels to see wall-clock speedups.

Structured Pruning

Removing entire neurons, heads, or channels (e.g., N:M sparsity). Naturally maps to hardware vector instructions.

IX. Advanced Number Formats

Precision affects both memory footprint and compute throughput.

FP32 (Single): Standard. 1 sign, 8 exponent, 23 mantissa.
TF32 (TensorFloat): 19 bits. 8 exponent (same range as FP32), 10 mantissa. Native on Ampere.
BF16 (BFloat16): 16 bits. 8 exponent (same range as FP32), 7 mantissa. The gold standard for LLM training.
FP8 (Hopper): The new frontier. Split into E4M3 (Weights, higher precision) and E5M2 (Gradients, higher dynamic range).

Primary Sources & Further Reading

Optimization Papers

Hu et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models.
Dettmers et al. (2023). QLoRA: Efficient Finetuning of Quantized LLMs.
Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention.

Tooling & Guides

Hugging Face (2023). PEFT: State-of-the-art Parameter-Efficient Fine-Tuning.
NVIDIA (2024). TensorRT-LLM Documentation.
Unsloth.ai. Fast Llama-3 Fine-tuning Guide.