1.3 Training vs. Inference Computational Profiles

1.3 Training vs. Inference Computational Profiles

Architectural Overview: In enterprise AI infrastructure, training and inference represent fundamentally different computational, memory, and networking paradigms. Training is an offline, throughput-oriented workload designed to learn parameters across large distributed clusters over days or weeks. Inference is an online, latency-sensitive workload designed to evaluate queries within strict Service Level Agreements (SLAs).

---

1. Deep Comparison: Training vs. Inference

Designing optimal GPU infrastructure requires matching hardware capabilities to the specific computational requirements of training or inference workloads.

---

2. Memory Footprint Breakdown: Training vs. Inference

The total memory footprint allocated inside GPU High Bandwidth Memory (HBM) differs dramatically between training and inference.

+-----------------------------------------------------------------------------+
|                     GPU HBM MEMORY ALLOCATION PROFILES                      |
|                                                                             |
|  [ TRAINING MEMORY ALLOCATION ]         [ INFERENCE MEMORY ALLOCATION ]     |
|  +-----------------------------------+  +--------------------------------+  |
|  | Optimizer States (AdamW: 12-16B/p)|  | Model Weights (FP16/FP8: 1-2B/p)|  |
|  +-----------------------------------+  +--------------------------------+  |
|  | Gradients (FP16/BF16: 2-4B/param) |  | Dynamic KV Cache               |  |
|  +-----------------------------------+  | (Grows with Context & Concur.) |  |
|  | Model Parameters (2-4B/param)     |  +--------------------------------+  |
|  +-----------------------------------+  | Ephemeral Activation Buffers   |  |
|  | Stored Forward Pass Activations   |  +--------------------------------+  |
|  | (Required for Backpropagation)    |                                      |
|  +-----------------------------------+                                      |
|  | Working Scratchpad Memory         |                                      |
|  +-----------------------------------+                                      |
+-----------------------------------------------------------------------------+

Training Memory Components

For an $N$-parameter model trained with mixed precision and the AdamW optimizer:

1. Model Parameters: Stored in BF16/FP16 (2 bytes per parameter) or FP32 (4 bytes per parameter).
2. Gradients: Stored in BF16/FP16 (2 bytes per parameter) during backward accumulation.
3. Optimizer States: In standard AdamW, the system maintains FP32 master weights (4 bytes), first momentum estimates (4 bytes), and second momentum variance estimates (4 bytes), requiring 12 to 16 bytes per parameter.
4. Forward Activations: Intermediate layer outputs generated during the forward pass must be retained in memory until the corresponding layer's backward pass calculates parameter gradients. In long-context models, activations frequently exceed model weight memory unless mitigated by Activation Checkpointing (Recomputation).
5. Temporary Workspace: Buffers for cuDNN convolution plans, GEMM workspaces, and NCCL communication rings.

$$\text{Total Training Memory} \approx 16N \text{ (Static Weights/States)} + \text{Memory}_{\text{Activations}} + \text{Memory}_{\text{Workspace}}$$

Inference Memory Components

Inference eliminates gradients, optimizer states, and stored backward activations entirely:

1. Model Weights: Static baseline footprint. For example, an unquantized 70B parameter model in FP16 requires $70 \times 10^9 \times 2 \text{ bytes} \approx 140 \text{ GB}$ of HBM. In FP8, this drops to $\approx 70 \text{ GB}$.
2. Key-Value (KV) Cache (LLMs): Autoregressive models cache previous token Key and Value projection vectors to avoid redundant attention recalculations. The KV cache size per token across batch size $b$, context length $s$, number of layers $l$, number of KV heads $h_{kv}$, and head dimension $d$ is: $$\text{Memory}_{\text{KV}} = 2 \times b \times s \times l \times h_{kv} \times d \times \text{bytes per element}$$ In multi-user serving environments, the KV cache dynamically expands and can easily consume 50–70% of available GPU memory.
3. Activation Working Memory: Minimal ephemeral memory allocated only for the current token or batch forward evaluation, immediately overwritten by subsequent layers.

---

3. Deconstructing LLM Inference: Prefill vs. Decoding Phase

Large Language Model serving consists of two computationally distinct phases that stress different hardware subsystems within the GPU.

   [ USER PROMPT: N Tokens ]
              │
              ▼
   ┌────────────────────────────────────────────────────────────────────────┐
   │                   PHASE 1: PREFILL (PROMPT PROCESSING)                 │
   │   - Ingests all input tokens concurrently in parallel                  │
   │   - Dense Matrix-Matrix Multiplication (GEMM)                          │
   │   - Arithmetic Intensity: HIGH (Hundreds of FLOPs/byte)                │
   │   - Primary Hardware Bound: COMPUTE-BOUND (Tensor Cores)               │
   │   - Key Performance Metric: Time to First Token (TTFT)                 │
   │   - Output: Initial token + Populated Initial KV Cache                 │
   └───────────────────────────────────┬────────────────────────────────────┘
                                       │
                                       ▼
   ┌────────────────────────────────────────────────────────────────────────┐
   │                  PHASE 2: DECODING (AUTOREGRESSIVE GENERATION)         │
   │   - Generates tokens sequentially, one token at a time                 │
   │   - Matrix-Vector Multiplication (GEMV)                                │
   │   - Arithmetic Intensity: LOW (Single-digit FLOPs/byte)                │
   │   - Primary Hardware Bound: MEMORY-BANDWIDTH-BOUND (HBM Streaming)     │
   │   - Key Performance Metric: Inter-Token Latency (ITL) / TPOT           │
   │   - Mechanism: Reads entire model weight matrix from HBM for EACH token│
   └────────────────────────────────────────────────────────────────────────┘

Detailed Phase Comparison

---

4. The Roofline Model & Operational Intensity

The Roofline Model is a foundational performance model used by systems architects to diagnose whether an AI kernel or application is bounded by GPU compute capacity (Tensor Core FLOPs) or memory subsystem bandwidth (HBM bytes/sec).

   Attainable Performance (TFLOPS)
         ▲
         │                           /──────────────────────── Peak Compute Ceiling
         │                          /                          (e.g., H100 FP16: 989 TFLOPS)
         │                         /  COMPUTE-BOUND REGION
         │                        /   (Performance limited by Tensor Cores)
         │                       /│
         │                      / │
         │                     /  │
         │  MEMORY-BOUND      /   │
         │  REGION           /    │
         │  (Performance    /     │
         │   limited by    /      │
         │   HBM Bandwidth/       │
         │               /        │
         │              /         │
         │             /          │
         │            /           │
         └───────────┴────────────┴───────────────────────────► Operational Intensity
                                  ▲                             (FLOPs / Byte)
                                  │
                           RIDGE POINT (Knee)
                 = Peak Compute / Peak Memory Bandwidth

Key Mathematical Formulations

1. Operational Intensity (Arithmetic Intensity): The ratio of total floating-point operations executed to the total bytes of data transferred to/from device memory (HBM): $$\text{Operational Intensity} = \frac{\text{Total Floating Point Operations (FLOPs)}}{\text{Total DRAM / HBM Bytes Transferred}}$$
2. Attainable Performance Ceiling: $$\text{Attainable Performance} = \min\left(\text{Peak Compute Capacity}, \text{Operational Intensity} \times \text{Memory Bandwidth}\right)$$
3. The Ridge Point (Knee of the Curve): The operational intensity threshold where a GPU transitions from memory-bandwidth-bound to compute-bound: $$\text{Ridge Point} = \frac{\text{Peak Compute FLOP/s}}{\text{Peak Memory Bandwidth (Bytes/s)}}$$

Practical GPU Example: NVIDIA H100 SXM5

• Peak FP16 Tensor Core Compute: $989 \times 10^{12} \text{ FLOP/s}$ (without structured sparsity).
• Peak HBM3 Memory Bandwidth: $3.35 \times 10^{12} \text{ Bytes/s}$ ($3.35 \text{ TB/s}$).
• Ridge Point Calculation: $$\text{Ridge Point} = \frac{989 \times 10^{12}}{3.35 \times 10^{12}} \approx 295.2 \text{ FLOPs/Byte}$$

Operational Takeaway:

• If an operation (such as single-token LLM decoding or DLRM embedding table lookup) achieves an operational intensity of only $15 \text{ FLOPs/Byte}$ (well below $295.2$), it resides squarely in the memory-bound region. Upgrading compute cores will yield zero performance gain; only increasing memory bandwidth or batch size will improve throughput.
• If an operation (such as large-batch training GEMM or prompt prefill) achieves an operational intensity of $500 \text{ FLOPs/Byte}$ (well above $295.2$), it resides in the compute-bound region, fully saturating Tensor Cores.

---

5. Precision Formats & Numerical Representation

Modern accelerated computing leverages reduced-precision arithmetic formats to maximize Tensor Core execution speed and reduce memory consumption while preserving numerical gradient stability.