ML Model Memory & Storage Estimator

A Transformer's parameters come from several components: embedding tables (vocab_size × hidden_dim), attention projections (4 × hidden_dim² per layer for Q/K/V/O), feed-forward layers (2 × hidden_dim × ffn_dim per layer), and layer normalization weights. For GPT-style models, the total is roughly: (vocab_size × d_model) + num_layers × (4 × d_model² + 2 × d_model × ffn_dim + small_norms). Larger hidden dimensions and more layers scale quadratically and linearly respectively, so doubling hidden_dim quadruples attention cost.

Quantization reduces the bit-width of each parameter weight: FP32 uses 4 bytes/param, FP16/BF16 uses 2 bytes/param, INT8 uses 1 byte/param, and INT4 uses 0.5 bytes/param. A 7B parameter model needs ~28 GB at FP32, ~14 GB at FP16, ~7 GB at INT8, and ~3.5 GB at INT4. However, quantization can hurt accuracy (especially INT4 for small models), and activations during inference are often kept in higher precision internally. BF16 is generally preferred over FP16 for training stability on modern hardware.

Training requires storing three things beyond the model weights: (1) Gradients — one gradient tensor per parameter, equal in size to the model weights. (2) Optimizer states — Adam stores two additional tensors per parameter (first and second moment estimates), so 2× model size. Combined, AdamW needs roughly 4× the model's weight size in FP32, or 16 bytes per parameter. Mixed-precision training (AMP) stores a FP32 master copy plus FP16 weights and gradients, totaling ~18 bytes/param. (3) Activation checkpointing can reduce activation memory at the cost of extra compute.

Model VRAM is just the memory needed to hold the weight tensors. Inference VRAM also includes: activation tensors (intermediate outputs at each layer, proportional to batch size × sequence length × hidden_dim), KV cache for autoregressive generation (2 × num_layers × num_heads × head_dim × sequence_length × batch_size × bytes_per_element), and a small framework overhead (~0.5–1 GB). For batch inference with long sequences, the KV cache can easily exceed the model weights, so always add headroom when sizing GPU memory.

A practical rule: ensure the GPU VRAM is at least 1.2× your estimated inference memory (for safety headroom). For inference with FP16, a 7B model (~14 GB) fits on a single RTX 4090 (24 GB) or A10G (24 GB). A 13B model (~26 GB) needs at least an A100 40 GB or two RTX 4090s. A 70B model (~140 GB at FP16) requires 2–4× A100 80 GB cards. Quantized INT4 versions of 70B models (~35 GB) can fit on two A100 40 GB cards. For training, multiply the inference requirement by 4–6× to account for gradients and optimizer states.