So I built mini-vllm: a small inference loop around Llama-2-7B on a single GPU. Scheduler, block allocator, paged KV, generation loop. Not a production server, just enough to see how the ideas connect. Code is on GitHub if you want the implementation.

What an inference engine actually is

From the outside it looks like one model call. Inside it is three boring layers plus weights:

• Per-request state on the CPU: prompt, tokens generated so far, when to stop, which memory blocks you own.
• Memory management: a fixed pool of KV blocks on the GPU, handed out like parking spots.
• Scheduling: who runs this step, who waits, who just finished and must give blocks back.

The transformer math is the easy part to read about. Most of my confusion was in those three layers talking to each other.

PagedAttention, in one picture

Naive serving often reserves a long KV buffer per request. Short prompts waste most of it, so you fit fewer concurrent users.

PagedAttention is the opposite idea: allocate one big KV tensor up front, then give each request a short list of block ids. Token 47 does not mean "grow a tensor." It means "block 2, slot 15 in the shared buffer."

The allocator is just bookkeeping on those ids. The GPU tensor never changes shape during inference. That was the click for me: all the "paging" drama is CPU-side lists and integers, not new CUDA allocations every token.

Things I had backwards at first

I thought the block table stored vectors. It stores handles into the pool. The vectors live in one place; the table is a map.

I mixed up fleet metrics with request state. Before writing code I listed SLOs and counts. Step six of generation only needs identity, full output history, stop rules, and the prompt.

I reordered the scheduler wrong. If you grow block tables before you drop finished sequences, you can allocate for a request that already ended and free it in the same pass. Small bug, very vLLM-shaped lesson: lifecycle order matters as much as the algorithm on paper.

I treated prefill and decode as different models. Same weights. Different input shapes. Decode is one new token with a position index in the full sequence, not "position 0 in this one-row tensor." I lost time on RoPE before I believed that.

I underestimated memory layout vs math. Growing KV with concat-every-step means allocate, copy, free in the hot loop. Writing into a fixed slot is why serving engines feel so much faster in practice. The gap is often memory traffic, not a smarter softmax.

What surprised me

Papers teach the what. Building taught the why behind small choices.

• Refcounts on blocks only make sense once you imagine two requests sharing a prefix later.
• Most bugs were state machines and indexing, not attention.
• File layout followed dependencies: engine owns everything, nothing imports the engine. I did not design that upfront; it fell out of "who creates whom at startup."

If you want to try this yourself

Start from the PagedAttention paper with a sketch of logical tokens vs physical blocks. Get one sequence generating before you turn on batching. Then read Inside vLLM or peek at nano-vllm once your own loop works. Comparisons land better when you have something to compare against.

mini-vllm is my version of that exercise. The repo has benchmarks and the full wiring; this post is only the flavour I wish I had before I opened the editor.

When people say a model "supports 128K context," that sounds like a single setting you can crank up. It's not. Context length is actually limited by three things at once: how the model tracks token position, how expensive attention gets as sequences grow, and how much memory you need to store past tokens during inference. Each one puts a ceiling on how far you can go.

Constraint 1: Positional Encoding Scheme

Transformers have no built-in sense of order. Self-attention is permutation-invariant, meaning it treats "the cat sat" the same as "sat cat the" unless you explicitly tell it which token came first. Positional encodings inject that order. The scheme you pick dictates the ceiling.

RoPE: The Dominant Approach (and Its Limitation)

Most modern models use Rotary Position Embedding (RoPE). For each dimension pair (i), the rotation angle at position m is:

The problem: If you train on positions 0 to 4096, the model has only ever seen rotation angles in that range. Position 8000 produces rotations the model has never encountered. It's like teaching someone a 12-hour clock then asking them to read a 24-hour clock. The mechanism is the same but the values are foreign. This is a distribution shift, and the model's output degrades unpredictably.

Constraint 2: Attention's Quadratic Cost

Self-attention computes pairwise scores between every token pair:

This is the fundamental reason you can't "just make context longer." The cost grows with the square of the length.

Constraint 3: KV Cache at Inference

At inference, you store Key and Value vectors for every past token, every layer. For a 70B model with GQA (8 KV heads):

Why You Can't "Just Extend" the Context

Problem 1: Training Distribution Mismatch (the deepest reason)

The model learns attention patterns from training data. If most training documents are 2-8K tokens:

• It learns "the answer is usually within a few thousand tokens of the question"
• It learns attention heads that specialize in local patterns
• It never practices retrieving a fact from 50K tokens away

Even if you architecturally support 128K, the model hasn't learned when and how to attend across that distance. This is a learned behavior problem, not an architecture problem.

Problem 2: Lost in the Middle (Liu et al., 2023)

Even models with "long context" show a U-shaped retrieval curve:

Softmax attention distributes probability mass. Over long sequences, the middle tokens get starved because they don't benefit from primacy or recency bias.

Problem 3: Effective vs. Claimed Context

A model "supporting" 128K and effectively using 128K are different things:

The spec sheet number is the architectural limit. The effective limit is much lower.

The Fundamental Trilemma

Pick at most two:

• Long + Quality = massive compute (full attention, long training)
• Long + Efficient = sparse/approximate attention (loses quality)
• Quality + Efficient = short context (what most models do best)

2. Strategies for Extending and Managing Context

There are two fundamentally different approaches: make the window bigger (architectural) or make the model smarter within the window it has (behavioral). Both have their place.

Architectural Strategies: Making the Window Bigger

Behavioral Strategies: Making the Model Smarter Within Its Window

Instead of making the window bigger, teach the model to use its existing window more intelligently. These are agent-level strategies.

Strategy 1: Surgical Retrieval. Read functions, not files

The agent already knows (from the stack trace or grep) which function matters. Reading the whole file dumps 1800 lines of noise into context that competes for attention with the 30 lines that matter.

Strategy 2: Context Eviction. Proactively discard stale information

After you've acted on information, you rarely need it verbatim. A summary preserves the decision-relevant facts while freeing context for new information.

Strategy 3: Structured Memory. Notes over verbatim copies

Strategy 4: Plan Before Reading. Avoid wrong leads entirely

200 tokens of planning saves 8K tokens of unnecessary reads.

Strategy 5: Compaction as a Tool

Compaction doesn't have to be an external system decision. You can give the agent a compact() tool and let it learn when to use it through reinforcement learning:

The model learns through trial and error: compact too early and you lose critical details. Compact too late and context fills up with noise. Extract key facts before compacting and you get the best of both worlds. This is trainable with RLVR.

The Core Principle

3. At Which Stage Can You Do This?

LLM development has four stages. Each one offers different knobs for improving context, at very different cost levels. Understanding which stage to operate at is critical.

Stage 1: Pretraining

What you can do:

• Choose the positional encoding scheme (RoPE, ALiBi)
• Choose the attention architecture (full, sparse, sliding window)
• Train on progressively longer sequences (4K → 16K → 64K → 128K)
• Curate long-context training data (full codebases, books, legal documents)
• Build in GQA for KV cache efficiency

Who does this: Only frontier labs (OpenAI, Anthropic, Google, Meta, DeepSeek). This is where the fundamental context capacity is set. Everything downstream is constrained by what happens here.

What it costs: Tens of millions to billions of dollars. Training on 128K sequences costs 1,024x more per sample than 4K (quadratic attention). You need thousands of GPUs for weeks.

What you get: A model that can physically attend to 128K-200K tokens and has actually practiced using long-range attention patterns during training. This is what makes Claude, GPT-4, and Gemini good at long context. They spent enormous compute on long-context pretraining.

Stage 2: Mid-Training (Context Extension)

What you can do:

• Take a model pretrained at 4K-8K and extend to 32K-128K
• Adjust RoPE parameters (YaRN, NTK scaling, change base frequency)
• Continue training on 10-100B tokens of long-context data

Real examples:

What it costs: Still expensive. You need long-context training data (books, repos, long docs) and significant GPU time. Context extension works well up to ~4x the original length. 32K → 128K is fine. 32K → 1M is sketchy.

Stage 3: Post-Training (SFT and RL)

What you can do:

• Fine-tune on long-context tasks (needle-in-a-haystack, multi-doc QA)
• SFT on full multi-turn agent trajectories with tool use
• RLVR to teach the model when to compact, extract, and manage context
• Train with compaction/memory tools as available actions

This is where you teach the model to be smart within its window. The context capacity is already set by pretraining. Post-training teaches the model behavioral strategies:

RLVR is particularly powerful here because optimal context management is task-dependent. Sometimes you should compact aggressively. Sometimes you need every detail. The model learns the strategy through trial and error, not imitation.

What it costs:

Stage 4: Agent Scaffolding

What you can do:

• Automatic context compaction (summarize old turns when approaching limits)
• Prompt caching (don't reprocess the same system prompt every turn)
• Selective file reading (tools load specific functions, not entire files)
• Conversation summarization (compress old turns into key decisions)
• KV cache management across turns

This is what Claude Code, Cursor, and Codex actually do. The model itself doesn't change. The system around the model manages context intelligently.

What it costs: Engineering time only. No GPU costs beyond normal inference.

Summary: The Four Stages

4. Honest Cost Analysis

Everyone talks about long context. Nobody talks about what it actually costs. Here's the full picture.

Cost of Training with Long Context

Attention is O(n²). This means training costs scale quadratically with context length:

Cost of Inference with Long Context

Using Claude/GPT-4 API pricing as reference:

At 1M context, a 30-turn coding session costs $75. Most of those tokens are noise the model barely attends to. A smart agent at 32K would cost $2.40 and probably produce better results.

Cost of Self-Hosted Inference

The Cost-Performance Sweet Spot

5. What Reasoning Over Long Context Actually Means

With the mechanics and costs established, let's address what everyone is actually chasing: making models reason across long context, not just hold it.

Long-Horizon Tasks: The Motivation

Long-horizon tasks require many sequential steps, decisions, or intermediate sub-goals before reaching a final outcome. "Reasoning over" them means the model must plan, track state, and make coherent decisions across that entire chain.

Why It's Hard for LLMs

1. Error compounding: A small mistake in step 3 of 20 derails everything downstream.
2. State tracking: The model must remember what it has done, what remains, and what intermediate results it produced.
3. Planning under uncertainty: Early decisions constrain later options.
4. Credit assignment: When the final answer is wrong, it's hard to identify which step failed.

The Difference Between "Holds" and "Reasons Across"

This is the crux. Many models support 128K tokens. The question is what happens to information at various positions:

What "Reasoning Across" Looks Like in Code

Cross-File Causal Tracing

# File A (in context at position ~2K)
def get_user(id):
    return db.query(User).filter(User.user_id == id).first()  # returns None if not found

# File B (in context at position ~15K)
def login(request):
    user = get_user(request.id)
    token = generate_token(user.email)  # ← crashes: NoneType has no attribute 'email'

The model must connect a None return in one file to an unguarded attribute access in another file. These two code fragments might be 13,000 tokens apart in the context. The model's attention mechanism must literally assign high attention weight from the user.email token to the return ... .first() token across that gap.

This is what "lost in the middle" kills. A weaker model sees File B, generates a fix like if user is None: return error, which is correct but shallow. A model reasoning across full context also notices that get_user should probably raise UserNotFoundError instead of returning None, because 6 other call sites (also in context from earlier grep results) all assume it returns a valid user.

Pattern Recognition Across the Codebase

After reading 8-10 files, the model should recognize:

• "This codebase uses the repository pattern"
• "Errors are custom exceptions caught by middleware, not return codes"
• "Every endpoint has a corresponding Pydantic schema in schemas/"
• "Tests use factory_boy fixtures, not raw object creation"

This is not stored in any single file. It emerges from reasoning across the aggregate context. A model that can't do this generates code that is locally correct but stylistically alien to the codebase.

Multi-Step Plan Coherence

By step 5, the model is writing test code that must be simultaneously consistent with decisions made in steps 1-4. If the model "forgets" that it used account_id in step 1 and writes user_id in the step 5 test, the entire chain breaks.

Why This Is So Hard (The Non-Obvious Part)

The real difficulty isn't memory. It's relevance detection at scale.

At 100K tokens of context, there might be 500 function definitions, 200 class attributes, 50 config values, and 100 test assertions. When generating one line of code, maybe 3-4 of those are relevant. The model must:

1. Not attend to the 796 irrelevant items (noise suppression)
2. Strongly attend to the 4 relevant items (signal detection)
3. Do this for every token it generates

How Models Address Long-Horizon Reasoning

• Chain-of-thought (CoT): Explicit step-by-step reasoning reduces the per-step difficulty.
• Tree/graph search (ToT, GoT): Explore multiple reasoning paths, backtrack when stuck.
• Reinforcement learning (GRPO, PPO): Train models to get reward for correct final answers, forcing them to learn robust multi-step strategies. This is what DeepSeek-R1 and OpenAI's o1/o3 do.
• Process reward models (PRMs): Give reward at each step, not just the final answer, so the model learns which intermediate steps are good.
• Decomposition: Break the long task into sub-tasks, solve each, then compose.

How Agentic Tools Put It All Together

No vector database. No knowledge graph. No external memory. Just a flat conversation where the model is simultaneously the retriever, the planner, and the reasoner.

At message 11, the model sees ALL of messages 1-10. "Long context reasoning" = at this point, attention reaches back across every previous message simultaneously to produce a coherent fix.

The Model IS the Retriever

No embedding approximation. No similarity threshold. The model reasons about code dependencies and retrieves based on understanding, not pattern matching. And the same model choosing what to retrieve will use what it retrieves.

Iterative Context Building

Each turn refines understanding. By turn 5, the model has a complete causal chain in context AND has been building its mental model incrementally. This is fundamentally better than dumping 5 files cold. The order of discovery helps the model organize information.

Automatic Context Compression

This mirrors what attention does naturally (recent = high attention, old = low attention), but makes it explicit and honest instead of pretending the model attends well to 180K tokens.

6. Where the Field Is Heading

The endgame is not "2M context." It's an agent that navigates a million-file codebase using a 50K context window, because it knows exactly where to look and what to remember. The next breakthrough won't be bigger windows. It'll be same quality at 64K with 10x less cost and 5x less latency.

Key Concepts Reference

When you want to run an LLM for inference, you have two options: use a cloud API (OpenAI, Anthropic) or host your own. Self-hosting gives you control over costs, latency, and model choice. In this post, I'll break down exactly what happens when you deploy a model on RunPod using vLLM.

Understanding RunPod

RunPod is a cloud GPU provider. Unlike AWS or GCP where you rent full VMs, RunPod specializes in GPU containers. You pay only for GPU time, often at 50-70% lower cost than major cloud providers.

What RunPod Provides

Templates: Pre-configured Recipes

A RunPod template is a pre-configured deployment recipe. It specifies everything needed to run a specific application:

Understanding vLLM

vLLM is a high-performance inference engine for LLMs. It's not a model—it's the software that loads models and serves them efficiently.

Why vLLM Over Plain PyTorch?

Key vLLM Optimizations

1. PagedAttention

Traditional attention stores KV cache in contiguous memory blocks. vLLM uses "paged" memory (like OS virtual memory), allowing dynamic allocation and preventing memory fragmentation.

2. Continuous Batching

Instead of waiting for a batch to complete, vLLM continuously adds new requests and removes completed ones. This maximizes GPU utilization.

The Deployment Flow

Here's exactly what happens when you deploy:

Step 1: Pod Creation

Step 2: Container Startup

Step 3: Model Download

Step 4: Model Loading

Step 5: API Server Ready

Making Requests

vLLM exposes an OpenAI-compatible API. This means you can use the same code you'd use for OpenAI, just change the base URL:

Request Flow

Benchmark Results

I ran benchmarks at different concurrency levels to understand throughput scaling:

Cost Analysis

Self-hosting only makes sense if it's cheaper than API providers. Let's do the math:

When to Self-Host vs Use APIs

Key Takeaways

What's Next?

In the next post, I'll test function calling capabilities with this setup. We'll compare zero-shot function calling accuracy across different open-source models and see how they stack up against GPT-4.

This article builds that foundation. We cover GPU architecture, the roofline model, memory estimation, arithmetic intensity analysis, and latency estimation , everything you need to develop strong intuition about what makes LLM inference fast or slow.

GPU Architecture Fundamentals

NVIDIA A100 Architecture

The A100 is the workhorse GPU for LLM inference and training. Understanding its architecture is essential for reasoning about performance.

The A100-80GB packs 108 SMs, each with 64 CUDA cores and 4 Tensor cores, connected via 40 MB L2 to 80 GB HBM2e at 2 TB/s. Fast on-chip SRAM (192 KB per SM) sits next to each SM; the size and speed gap between SRAM and HBM is what makes memory bandwidth the bottleneck.

Key Numbers to Memorize

Execution Model: Warps and Thread Blocks

• Kernel launch → a grid of thread blocks is scheduled onto SMs.
• Each block has up to 1024 threads, split into warps.
• Each warp = 32 threads executing in lockstep (SIMT); they all run the same instruction at once.

The warp is the fundamental unit of execution. All 32 threads in a warp execute the same instruction at the same time. If threads diverge (different if branches), both paths must be executed serially , this is called warp divergence and it wastes cycles.

SRAM vs. HBM: The Crucial Ratio

• Total SRAM: ~20 MB (fast, on-chip, ~30 cycle latency)
• Total HBM: 80 GB (slow, off-chip, ~400 cycle latency)
• Ratio: HBM is ~4000x larger but ~13x slower

This mismatch is the fundamental reason why memory bandwidth is the bottleneck for most LLM inference workloads. The model weights live in HBM, and for each token generated during decode, we must read ALL weights from HBM through a limited bandwidth pipe.

Memory Hierarchy Latency

Why This Matters for LLM Inference

During autoregressive decode, each generated token requires reading the full model weights from HBM:

• Llama-7B at FP16: 14 GB of weights
• At 2 TB/s bandwidth: Takes 14 GB / 2 TB/s = 7 ms just to read the weights
• Computation: ~14 GFLOP per token, which at 312 TFLOPS takes only 0.045 ms

The weight-loading time is 150x larger than the compute time. This is why decode is memory-bandwidth-bound, and why all the major inference optimizations (quantization, KV caching, speculative decoding, batching) ultimately aim to reduce the bytes-per-useful-FLOP ratio.

The Roofline Model

The roofline model is the single most important mental model for understanding inference performance. It tells you whether a workload is limited by:

1. Memory bandwidth (loading data from HBM) , left side of the plot
2. Compute throughput (doing arithmetic) , right side / top of the plot

Arithmetic Intensity

Ridge Point

The ridge point is the arithmetic intensity at which the compute and memory roofs intersect:

Attainable Performance

Figure 1: The A100-80GB Roofline Model. Decode (batch=1) sits at AI=1, achieving only 0.6% of peak compute. Prefill (seq=2048) sits near the ridge point, approaching full compute utilization. The red region is memory-bound; the green region is compute-bound.

The Key Insight: Prefill is Compute-Bound, Decode is Memory-Bound

This single insight explains nearly every optimization in modern LLM serving.

Why is decode so memory-inefficient? During decode, we generate ONE token. This means we do a matrix-vector product: the weight matrix has d_model × d_model elements, but the input vector has only d_model elements. Each weight is loaded from HBM but used for only one multiply-add (2 FLOPs / 2 bytes at FP16 = AI of 1).

Why is prefill efficient? During prefill, the input is a matrix of shape (seq_len × d_model). The same weight matrix (loaded once from HBM) is multiplied against seq_len vectors. The weight bytes are amortized across seq_len tokens, giving AI ~ seq_len.

Optimization Implications

Memory Estimation

Understanding where GPU memory goes during inference is critical for capacity planning. Let's break down the memory components for Llama-7B.

VRAM Breakdown: Llama-7B at FP16

Figure 2: Llama-7B memory breakdown (FP16, batch=1, seq=2048). Weights 87.8% (13.04 GB), KV cache 6.7% (1.00 GB), activations 3.4% (0.51 GB), CUDA overhead 2.1% (0.31 GB). Total ~14.86 GB.

Does Llama-70B Fit on A100-80GB?

When Does KV Cache Dominate?

At batch=1, weights dominate. But at production batch sizes, KV cache quickly overtakes weights. This plot shows the crossover point for Llama-7B.

Figure 3: Llama-7B FP16 (seq_len=2048), KV cache vs weights vs batch size. Weights constant at ~13 GB; KV cache grows linearly. KV cache equals weights at batch=14. Max batch=50 on A100-80GB before OOM. KV cache management (PagedAttention, GQA) is critical at scale.

Memory Across Model Sizes and Precisions

Key observations:

1. Weights scale linearly with parameter count and inversely with quantization.
2. INT4 gives a 4x reduction in weight memory vs FP16, making 70B feasible on one GPU.
3. KV cache for 70B is smaller than expected thanks to GQA (8 KV heads vs 64 query heads).
4. At batch=1, weights dominate. At production batch sizes (32+), KV cache dominates.

Arithmetic Intensity Analysis

Let's see how arithmetic intensity changes with different operating conditions, and how that determines whether we're memory-bound or compute-bound.

Decode AI vs Batch Size

Decode at batch=1 has an arithmetic intensity of just 1 , only 0.6% of the ridge point. Even at batch=64, we're still deeply memory-bound. It takes a batch size of ~156 to fully saturate the A100's compute capability during decode.

Figure 4: Decode arithmetic intensity grows linearly with batch size. The blue dashed line marks the A100 ridge point at 156 FLOPs/byte. Below this line, the GPU's compute cores are underutilized , adding more sequences to the batch is essentially "free" in terms of latency, because we're bottlenecked on memory bandwidth anyway.

Prefill AI vs Sequence Length

Prefill crosses the ridge point at relatively short sequence lengths (~300-400 tokens). For typical prompt lengths (1K+ tokens), prefill is solidly compute-bound , the GPU cores become the bottleneck, not memory bandwidth.

Figure 5: Prefill arithmetic intensity vs sequence length (Llama-7B, FP16). Crosses A100 ridge (156 FLOPs/byte) at seq_len ~192; below that, prefill is memory-bound; above, compute-bound. For short prompts, latency is dominated by weight loading; for long prompts, by computation.

Time Estimates

Using the roofline model, we can estimate latencies for prefill (TTFT = time to first token) and decode (per-token generation time). These back-of-the-envelope calculations are exactly the kind of reasoning expected in interviews.

Prefill Latency (TTFT)

Prefill latency scales roughly linearly with prompt length. For Llama-7B on A100, a 2048-token prompt takes about 9 ms , fast enough to be imperceptible.

Figure 6: Prefill latency (TTFT) vs prompt length for Llama-7B, 13B, and 70B on A100 (FP16, batch=1). Horizontal reference lines mark 100ms (imperceptible), 500ms (noticeable), and 1000ms (slow). Larger models have proportionally higher TTFT. For Llama-70B, prompts beyond 4K tokens push TTFT past 1 second on a single A100.

Decode Latency and Throughput

Decode time stays ~7 ms per step while memory-bound (batch < ~164); total tokens/sec scales with batch. Past the ridge, latency rises and throughput plateaus (~22K tokens/sec). See the diagram below.

The key insight: decode time per step stays nearly constant as batch size increases (while memory-bound). This means total throughput scales linearly with batch size , for free! This is the fundamental insight behind continuous batching.

Figure 7: Left: Llama-7B decode latency per step stays flat (~7 ms) while memory-bound, then rises after the inflection at batch=164 (transition to compute-bound). Right: Throughput scales linearly with batch size in the memory-bound regime ("free" batching), then plateaus at ~22K tokens/sec. This is why production serving uses continuous batching.

Key Takeaways

5 Things You Should Know

1. "Autoregressive decode is memory-bandwidth-bound, not compute-bound." Each decode step requires reading all model weights from HBM to generate a single token. The arithmetic intensity is ~1 FLOP/byte (at batch=1, FP16), far below the A100's ridge point of 156. This means the Tensor Cores sit idle ~99% of the time during decode.
2. "Prefill is compute-bound for reasonable sequence lengths." During prefill, the same weights are reused across all tokens in the prompt, giving arithmetic intensity that scales with seq_len. For seq_len > ~300 on A100, prefill becomes compute-bound.
3. "Batching is the primary lever for decode throughput." Since decode is memory-bound, adding more sequences to a batch reuses the same weight data already being streamed from HBM. The per-step latency barely changes, but throughput scales linearly. This is why continuous batching (vLLM, TGI) is transformative.
4. "KV cache memory grows as O(batch_size × seq_len × num_layers × d_head × num_kv_heads)." At production batch sizes (32+), KV cache can dominate total GPU memory. This is why PagedAttention, GQA (grouped-query attention), and KV cache quantization are critical.
5. "Quantization helps decode more than prefill." Since decode is memory-bound, reducing the number of bytes per weight (FP16 → INT4 = 4x fewer bytes) directly translates to ~4x faster decode. For prefill (compute-bound), quantization helps less unless it also reduces the compute cost.

Quick-Reference Formulas

Keep these formulas in your head for back-of-the-envelope calculations:

💻 Fun disclaimer: Used GPT to get all the beautiful visual gradients, but the content is mine!

When training large-scale machine learning models, optimization goes beyond just hyperparameter tuning. This guide covers the essential aspects of efficient ML training: computational constraints, performance profiling, and learning strategies that can save you significant costs and time.

1. FLOPs and Chinchilla Scaling Law

When training large-scale ML models, you typically have FLOPs (Floating Point Operations) constraints. The Chinchilla scaling law provides crucial guidance on how to allocate your compute budget effectively.

Two Critical Cases to Avoid

1. Compute Inefficient Training

2. Data Inefficient Training

2. Profiling Your Code

Profiling your training code is essential for maximizing GPU utilization and getting the best performance for your investment. This is different from hyperparameter tuning, which focuses on model learning rather than computational efficiency.

Key Bottlenecks to Monitor

I/O Bottleneck

Don't assume that just because your GPU can handle a larger batch size, you should use it. PyTorch data loaders work on CPU threads, and if your GPU finishes processing batch 1 but your data loader isn't ready with batch 2, your GPU sits idle.

Memory Bottleneck

Good profiling reveals what's consuming your memory. Common culprits include:

• Per-layer activations
• Gradients storage
• Temporary tensor assignments
• Optimizer states

Memory Optimization Techniques

CPU ↔ GPU Transfer Bottleneck

Moving data between CPU and GPU is often a major bottleneck due to bandwidth limitations. Common scenarios that cause this issue:

• Using .item() to extract scalar values
• Checkpointing weights to CPU
• Frequent data transfers during training

Kernel Overhead

Launching many small kernels can create overhead. The CPU tells the GPU to launch numerous kernels, and the GPU may struggle to keep up with the launch rate.

3. Learning Strategies

Once you've optimized your computational efficiency, focus on improving model performance through effective learning strategies.

Batch Size Selection

Choose the highest batch size your GPU and data loader can handle, but ensure you maintain some stochasticity in your updates. When you change batch size, adjust your learning rate accordingly (usually linearly).

Gradient Accumulation

If your learning is too noisy (loss oscillates up and down), consider gradient accumulation to smooth the updates:

Frequently Asked Questions

When do you stop training? What is the ideal loss?

What if training loss keeps dropping but validation loss increases?

This is classic overfitting. Solutions include:

• Add regularization (dropout, weight decay)
• Collect more training data
• Implement early stopping
• Reduce model complexity

How do I know if my learning rate is too high or low?

Understanding Fine-tuning vs Training

The Memory Challenge

When running large models on available hardware, you typically have two main options:

Evolution of Fine-tuning Approaches

1. Full Fine-tuning

Figure 1: Traditional full fine-tuning approach where all model parameters are updated

In the early days of deep learning, when models had fewer parameters, full fine-tuning was the standard approach. This method updates all model weights during the adaptation process.

2. Partial Fine-tuning

As models grew larger, researchers began experimenting with partial fine-tuning, freezing about 10% of the model weights. While this reduced memory footprint, the performance benefits were limited. Meaningful improvements typically required fine-tuning at least 30% of the model parameters.

The critical question emerged: Could we achieve significant performance improvements while updating only a tiny fraction of parameters? This led to the development of Parameter-Efficient Fine-Tuning (PEFT) methods.

Parameter-Efficient Fine-Tuning (PEFT)

The core idea behind PEFT is strategic: introduce small trainable neural networks called "adapters" at carefully selected locations within the model. These adapters act as learnable interfaces between the model's frozen layers. During fine-tuning, the original pre-trained weights remain unchanged, and only these adapter parameters are updated, making the process highly efficient.

Figure 2: Overview of Parameter-Efficient Fine-Tuning approaches

Understanding Prompt-Based Methods

Hard Prompting

Soft Prompting

Key Prompt-Based Methods

• Prefix Tuning: Adds trainable continuous tokens before specific layers
• Prompt Tuning: Prepends trainable embeddings to the input
• P-Tuning: Introduces trainable prompts at multiple positions

While these prompt-based methods showed promise, they haven't gained as much widespread adoption as more recent approaches like LoRA, which offers better efficiency and easier implementation.

Understanding SVD and Low-Rank Decomposition

Before diving into LoRA, let's understand the key concept behind it: low-rank matrix decomposition. Through a practical example using SVD (Singular Value Decomposition), we'll see how a large matrix can be represented using fewer parameters. This same principle is what makes LoRA efficient - it essentially adds a low-rank update (product of two smaller matrices) to the original weight matrix.

As expected, the matrix has rank 2, confirming that all its information can be represented using just two dimensions, despite being a 10×10 matrix. This is a key insight into why low-rank methods work.

SVD decomposes our matrix into three components, but remarkably, we only need to keep the first two singular values and their corresponding vectors. This is because these capture the essential structure of our matrix, while the remaining values are effectively zero.

This example demonstrates the power of low-rank decomposition: we could perfectly replicate the behavior of the original weights using far fewer parameters. A and B combined have only 40 parameters (10×2 + 2×10 = 40), while the original W matrix had 100 parameters (10×10 = 100). This 60% reduction in parameters is exactly the kind of efficiency that makes LoRA so powerful.

Low-Rank Adaptation (LoRA)

The main idea of Low-Rank Adaptation (LoRA) is to decompose weight updates into low-rank matrices, train these smaller matrices, and then add their product back to the original weights.

Let's break down how LoRA works: Instead of directly updating the large weight matrices of the model, LoRA introduces two smaller matrices (A and B) whose product approximates the weight update. This approach significantly reduces the number of trainable parameters while maintaining model quality.

Figure 3: LoRA's low-rank matrix decomposition and update process

Key Advantages of LoRA

Unlike traditional adapter methods that add extra layers and increase inference latency, LoRA's design offers a unique advantage: the trained matrices can be merged with the original weights at inference time, resulting in zero additional latency.

Serving LoRA Models

One of LoRA's most powerful features is its flexibility during inference. There are two main approaches:

This flexibility allows for interesting deployment scenarios. For example, you could have a base model with different LoRA adaptations for different languages or tasks, and dynamically choose or even combine them at inference time.

Figure 4: Different approaches to serving LoRA models - merged vs separate weights

QLoRA: Quantized LoRA

QLoRA combines the efficiency of LoRA with the memory benefits of quantization. It's a powerful approach that makes fine-tuning possible on consumer GPUs while maintaining model quality.

Thanks for sticking till the end! In Part 2, we'll explore advanced topics including model merging and multitask fine-tuning. Stay tuned!

When working with deep learning models, understanding memory usage is crucial. The total memory consumption can be broken down into two main components:

Example: Calculating Model Memory

Memory Optimization Techniques

Let's explore key techniques for reducing memory usage in deep learning models, starting with an important but often overlooked approach:

1. Gradient Checkpointing

Gradient checkpointing is a powerful technique that trades computation time for memory savings. Instead of storing all activations in memory during the forward pass, we do the following:

Strategy:

1. Store activations at checkpoints only
2. Recompute intermediate activations when needed
3. Free memory after gradients are computed

Trade-offs:

• ✓ Reduced memory footprint
• ✗ Increased training time (recomputation)

Basics and Lookup Table

Understanding INT4 Range

When we say INT4 (4-bit integer) has a range of -8 to 7, we're describing the minimum and maximum values that can be represented using 4 bits in signed integer format. Let's break this down:

2. Understanding Quantization

Figure: Quantization can be applied to four key areas: Weights, Training Time, Inference Time, and Activations

From the figure, we now understand that:

1. Quantization can be applied to weights, Activations.
2. It can also be applied in Inference time and Training time.

Let's see one by one.

1. Quantizing Weights (Static and Stable)

Weights are ideal candidates for quantization because they change less frequently. Once trained, weights remain constant unless the model is fine-tuned, making them perfect for one-time quantization.

2. Quantizing Activations (Dynamic Values)

Unlike weights, activations change with every inference because they depend on the input data. This makes activation quantization more challenging and requires careful consideration of the dynamic range.

Inference Time Quantization

Inference time quantization focuses on serving the model in low precision to accelerate computation. Modern approaches have moved beyond simple quantization to mixed precision strategies, which offer a better balance between performance and accuracy.

In a typical mixed precision setup:

• Model weights are stored in FP16 or FP8 format for memory efficiency
• Activations and gradients use FP32 or FP16 for better numerical stability
• Critical operations may dynamically switch between precisions as needed

Quantization-Aware Training (QAT)

QAT is a training-time technique designed to maintain high accuracy when models are deployed with low-bit quantization (like INT8, INT4). Unlike post-training quantization, QAT allows the model to adapt to quantization effects during the training process itself.

How QAT Works

Figure: The QAT process showing fake quantization during training and real quantization for deployment

The QAT process involves four key steps:

1. Simulate Quantization During Training

During the forward pass, weights and activations are "fake quantized" to simulate deployment conditions. This involves rounding and clipping values based on the target precision (like INT8), helping the model learn to work within quantization constraints.

2. Backpropagate Using High Precision

The backward pass maintains high-precision gradients (typically FP32) to ensure accurate learning. This dual approach allows stable gradient updates while still preparing the model for quantized deployment.

Common Challenges

When implementing QAT, teams typically face several challenges:

• Balancing training time with quantization accuracy
• Choosing appropriate quantization parameters
• Handling layers with different sensitivity to quantization
• Managing the increased complexity of the training pipeline

What's Next?

In our next article, we'll explore advanced memory-efficient techniques like LoRA (Low-Rank Adaptation) and other parameter-efficient fine-tuning methods that are revolutionizing how we train large language models.

🚧 Blog Post Coming Soon

This blog post will dive deep into my complete from-scratch implementation of the ResShift paper for image super-resolution. I'll explain how I built an efficient diffusion-based super-resolution model using a U-Net architecture with Swin Transformer blocks, including the residual shifting mechanism that reduces diffusion steps to just 15 timesteps. The post will cover architecture design, training on DIV2K dataset, and lessons learned from implementing this state-of-the-art approach to image enhancement.

Check out the implementation: GitHub Repository

Stay tuned for detailed explanations, code walkthroughs, and practical insights!