Memory Management › Code Implementation
Language AI HandbookAll chapters Memory Management › Code ImplementationOne section · Sections in this chapter4 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement a memory estimator that computes these values programmatically and validates them against PyTorch's reported memory usage. The goal is a tool you can run before any training job to predict GPU requirements and identify which component is the bottleneck.
Setup and Model Architecture
We start by defining the model configuration and implementing functions to count parameters and estimate memory:
# Model configuration: small GPT-style transformer
config = {
"vocab_size": 50257,
"hidden_dim": 512,
"num_layers": 6,
"num_heads": 8,
"ff_dim": 2048,
"max_seq_len": 512,
}
def count_parameters(cfg):
"""Count total trainable parameters for a GPT-style transformer."""
V = cfg["vocab_size"]
d = cfg["hidden_dim"]
L = cfg["num_layers"]
ff = cfg["ff_dim"]
# Token embeddings + position embeddings
embedding_params = V * d + cfg["max_seq_len"] * d
# Per-layer: 4 attention projection matrices (Q, K, V, O) + 2 FF layers + biases + layer norms
attn_params = 4 * d * d + 4 * d # projections + biases
ff_params = d * ff + ff + ff * d + d # two linear layers with biases
ln_params = 2 * 2 * d # two layer norms per layer (scale + bias each)
per_layer = attn_params + ff_params + ln_params
# Final layer norm + output projection (often tied with embedding, but counted separately here)
head_params = d * V + 2 * d
total = embedding_params + L * per_layer + head_params
return total
total_params = count_parameters(config)Total parameters: 70,640,640 Total parameters: 70.6M
This gives us the parameter count for our toy model. The parameter range in the tens of millions is appropriate for single-GPU experimentation, where we can validate our estimation formulas against actual PyTorch measurements without requiring expensive hardware.
Memory Breakdown Estimator
Now we implement the full memory estimator using the formulas derived above:
def estimate_memory(
cfg,
batch_size,
seq_len,
use_checkpointing=False,
optimizer="adam",
param_dtype_bytes=2, # FP16 parameters
grad_dtype_bytes=4, # FP32 gradients
optim_dtype_bytes=4,
): # FP32 optimizer states
"""
Estimate GPU memory requirements in bytes for transformer training.
Returns a breakdown dict with each memory component.
"""
P = count_parameters(cfg)
d = cfg["hidden_dim"]
L = cfg["num_layers"]
h = cfg["num_heads"]
# Static memory (independent of batch size)
param_mem = P * param_dtype_bytes
grad_mem = P * grad_dtype_bytes
# Optimizer state memory
optim_multipliers = {
"adam": 2, # first and second moment
"lion": 1, # only momentum
"sgd": 1, # momentum buffer
"adafactor": 0.1, # approximate (factored)
}
optim_states_per_param = optim_multipliers.get(optimizer, 2)
optim_mem = P * optim_states_per_param * optim_dtype_bytes
# Activation memory (scales with batch_size * seq_len)
# Approximate formula per layer: BxTx(34d + 5*h*T) bytes (FP16, 2 bytes each)
# Factor of 2 bytes already included in constants assuming FP16 activations
act_per_layer = batch_size * seq_len * (34 * d + 5 * h * seq_len)
if use_checkpointing:
# Only one layer's activations needed at a time (plus boundary states)
act_mem = act_per_layer # approximately one layer
else:
act_mem = L * act_per_layer
# CUDA/framework overhead (rough estimate)
overhead_mem = 512 * 1024 * 1024 # 512 MB baseline
total_mem = param_mem + grad_mem + optim_mem + act_mem + overhead_mem
return {
"parameters_bytes": param_mem,
"gradients_bytes": grad_mem,
"optimizer_states_bytes": optim_mem,
"activations_bytes": act_mem,
"overhead_bytes": overhead_mem,
"total_bytes": total_mem,
"total_gb": total_mem / (1024**3),
}
# Estimate for batch_size=4, seq_len=512, no checkpointing
estimate_no_ckpt = estimate_memory(
config, batch_size=4, seq_len=512, use_checkpointing=False, optimizer="adam"
)
estimate_with_ckpt = estimate_memory(
config, batch_size=4, seq_len=512, use_checkpointing=True, optimizer="adam"
)Component No Checkpointing With Checkpointing --------------------------------------------------------------------------- Parameters (FP16) 134.74 MB 134.74 MB Gradients (FP32) 269.47 MB 269.47 MB Optimizer states (Adam FP32) 538.95 MB 538.95 MB Activations 444.00 MB 74.00 MB Framework overhead 512.00 MB 512.00 MB TOTAL 1.85 GB 1.49 GB Gradient checkpointing saves 83% of activation memory. Total memory reduction: 1.85 GB → 1.49 GB
The table shows how the optimizer state is the dominant consumer for this model configuration, a pattern that holds across all transformer scales. Gradient checkpointing has no effect on the static components (parameters, gradients, optimizer states), but it dramatically reduces the activation term. Activations become comparably large relative to static memory at longer sequence lengths, where the attention term dominates.
Validating Against PyTorch
Let us build an actual small model and verify that our estimates align with what PyTorch reports:
import torch
import torch.nn as nn
class SimpleTransformerLayer(nn.Module):
def __init__(self, d_model, nhead, dim_ff):
super().__init__()
self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
self.ff = nn.Sequential(
nn.Linear(d_model, dim_ff), nn.GELU(), nn.Linear(dim_ff, d_model)
)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x):
attn_out, _ = self.self_attn(x, x, x)
x = self.norm1(x + attn_out)
x = self.norm2(x + self.ff(x))
return x
class SimpleTransformer(nn.Module):
def __init__(self, cfg):
super().__init__()
d = cfg["hidden_dim"]
self.embed = nn.Embedding(cfg["vocab_size"], d)
self.pos_embed = nn.Embedding(cfg["max_seq_len"], d)
self.layers = nn.ModuleList(
[
SimpleTransformerLayer(d, cfg["num_heads"], cfg["ff_dim"])
for _ in range(cfg["num_layers"])
]
)
self.ln_f = nn.LayerNorm(d)
self.head = nn.Linear(d, cfg["vocab_size"], bias=False)
def forward(self, input_ids):
B, T = input_ids.shape
positions = torch.arange(T, device=input_ids.device).unsqueeze(0)
x = self.embed(input_ids) + self.pos_embed(positions)
for layer in self.layers:
x = layer(x)
x = self.ln_f(x)
return self.head(x)
# Count actual parameters
model = SimpleTransformer(config)
actual_params = sum(p.numel() for p in model.parameters() if p.requires_grad)Estimated parameters: 70,640,640 (70.6M) Actual parameters: 70,640,640 (70.6M) Estimation error: 0.0%
Our analytical estimate closely matches the ground truth parameter count from PyTorch. The small discrepancy comes from minor architectural differences between our counting formula and the exact layer implementations, particularly in how nn.MultiheadAttention organizes its projection weight tensors.
Measuring GPU Memory Usage
We can instrument actual GPU memory consumption if a CUDA device is available. Even on CPU, we can measure the parameter storage footprint:
def get_model_memory_footprint(model):
"""Calculate actual memory footprint of model parameters in bytes."""
total_bytes = 0
param_info = {}
for name, param in model.named_parameters():
param_bytes = param.numel() * param.element_size()
total_bytes += param_bytes
# Group by component type
component = name.split(".")[0]
param_info[component] = param_info.get(component, 0) + param_bytes
return total_bytes, param_info
fp32_model = SimpleTransformer(config)
fp16_model = SimpleTransformer(config).half()
fp32_bytes, fp32_components = get_model_memory_footprint(fp32_model)
fp16_bytes, fp16_components = get_model_memory_footprint(fp16_model)Model parameter memory by precision: FP32 (float32): 269.5 MB FP16 (float16): 134.7 MB Ratio: 2.0x FP32 memory breakdown by component: embed : 98.16 MB (36.4%) head : 98.16 MB (36.4%) layers : 72.15 MB (26.8%) pos_embed : 1.00 MB (0.4%) ln_f : 0.00 MB (0.0%) Estimated FP16 parameter memory: 134.7 MB Actual FP16 parameter memory: 134.7 MB
The match between estimated and actual parameter memory confirms our formula is accurate. In a real training run, you would observe three to seven times this parameter footprint once gradients and Adam states are added.
Scaling Analysis
Let us compute memory requirements across different model sizes to build intuition for how memory scales:
# Standard model configurations
model_configs = {
"GPT-2 Small (117M)": {
"vocab_size": 50257,
"hidden_dim": 768,
"num_layers": 12,
"num_heads": 12,
"ff_dim": 3072,
"max_seq_len": 1024,
"params": 117e6,
},
"GPT-2 Large (774M)": {
"vocab_size": 50257,
"hidden_dim": 1280,
"num_layers": 36,
"num_heads": 20,
"ff_dim": 5120,
"max_seq_len": 1024,
"params": 774e6,
},
"LLaMA-7B": {
"vocab_size": 32000,
"hidden_dim": 4096,
"num_layers": 32,
"num_heads": 32,
"ff_dim": 11008,
"max_seq_len": 2048,
"params": 7e9,
},
"LLaMA-13B": {
"vocab_size": 32000,
"hidden_dim": 5120,
"num_layers": 40,
"num_heads": 40,
"ff_dim": 13824,
"max_seq_len": 2048,
"params": 13e9,
},
"LLaMA-70B": {
"vocab_size": 32000,
"hidden_dim": 8192,
"num_layers": 80,
"num_heads": 64,
"ff_dim": 28672,
"max_seq_len": 2048,
"params": 70e9,
},
}
scaling_results = {}
for model_name, cfg in model_configs.items():
P = cfg["params"]
# Use the 14x rule for static memory (FP16 params + FP32 grads + FP32 Adam)
static_gb = 14 * P / 1024**3
# Activation memory for batch_size=4, seq_len=2048 (approximate)
d = cfg["hidden_dim"]
L = cfg["num_layers"]
h = cfg["num_heads"]
T = 2048
B = 4
act_bytes = L * B * T * (34 * d + 5 * h * T)
act_gb = act_bytes / 1024**3
act_ckpt_gb = B * T * (34 * d + 5 * h * T) / 1024**3 # one layer only
scaling_results[model_name] = {
"params_b": P / 1e9,
"static_gb": static_gb,
"act_gb": act_gb,
"act_ckpt_gb": act_ckpt_gb,
"total_gb": static_gb + act_gb,
"total_ckpt_gb": static_gb + act_ckpt_gb,
}Model Params Static Act Act+Ckpt Total Total+Ckpt ----------------------------------------------------------------------------------- GPT-2 Small (117M) 0B 1.5G 13.6G 1.1G 15.2G 2.7G GPT-2 Large (774M) 1B 10.1G 68.2G 1.9G 78.3G 12.0G LLaMA-7B 7B 91.3G 114.0G 3.6G 205.3G 94.8G LLaMA-13B 13B 169.5G 178.1G 4.5G 347.6G 174.0G LLaMA-70B 70B 912.7G 570.0G 7.1G 1482.7G 919.8G Notes: batch_size=4, seq_len=2048, mixed precision + Adam Static = FP16 params + FP32 gradients + FP32 Adam states (14P rule) Act = activation memory without gradient checkpointing Act+Ckpt = activation memory with full gradient checkpointing
This table reveals why distributed training is not optional for models above a few billion parameters. A 7B model requires over 100 GB of static memory alone, far exceeding any single GPU's capacity. The 70B model requires over a terabyte of static memory, requiring many GPUs even before considering activations. Gradient checkpointing helps with activations but does nothing for the static component, which is why memory-efficient optimizers and sharding strategies are both essential tools at large scale.
Visualizations
The stacked comparison shows why practitioners typically combine multiple techniques rather than relying on any single one. Mixed precision alone gets you to the 14P static baseline. Adding gradient checkpointing addresses the activation term for long sequences. Memory-efficient optimizers then chip away at the dominant static component. The full combination can make a 7B model trainable on hardware that would otherwise require tensor parallelism.