Memory components in LLM

LLMs are memory hungry. GPUs are limited by memory more compared to compute.

Forward: [Parameters] -> [Activations]

Backward: [Activations] -> [Gradients]

[Gradients] + [Optimizer States] -> [Updated Parameters]

Memory usage across sequence length in different LLM:

(“dotted line” is the memory available in 1 GPU; Note 8B parameter models do not fit in one GPU.)

Note:

• Parameters, gradients and optimizer states remain constant.
• Only, activation is dependent on sequence length.
• In inference, only parameter and activation exist.

(1) Parameters (weights)

• Model's learnable variables. Example: weight matrix for Q, K, V; FFN. Shape: mostly (hidden_dim × hidden_dim)
• Role: Used in both forward and backward pass.
• Memory: Static during one iteration.
• Communication:
  • Sharded in Tensor Parallelism
  • Replicated in Data Parallelism (synced only via gradients)

(2) Activations

• Intermediate outputs produced during forward pass. Example: attention output, FFN output. Needed for: Backward pass (to compute gradients)
• Shape: (batch_size × sequence_length × hidden_dim)
• Communication:
  • Passed between GPUs in Pipeline Parallelism.
  • Memory: Large and transient (cause of activation checkpointing trade-offs)

(3) Gradients

• Derivatives of loss w.r.t. parameters.
• Used to update parameters.
• Computed during: Backward pass.
• Communication:
  • In Data Parallelism, all GPUs compute different gradients and then do an AllReduce so everyone has the average.
  • Size = parameters size.

(4) Optimizer States

• Extra tensors used by optimizers (like Adam or AdamW) Example: momentum, variance. Usually, 2× or 3× the parameter size
• Communication:
  • Synced or partitioned in ZeRO optimizations
  • Memory-heavy, not needed for inference.