Training Batch GPU Memory Estimatormedium

Training Batch GPU Memory Estimator

medium ML Systems & MLOps 15 min

NumPy allowed · plain Python earns 10 bonus points

Implement a real, back-of-envelope GPU memory estimator MLOps engineers actually use to pick a batch size before a training run OOMs.

Task

Implement estimate_batch_memory_gb(num_params, batch_size, seq_len, hidden_dim, bytes_per_param=4.0). Model weight memory (with Adam's 2 extra moment buffers, i.e. 3x weights total) is num_params * bytes_per_param * 3. Activation memory (a common simplified estimate) is batch_size * seq_len * hidden_dim * bytes_per_param * 2 (forward + backward). Return (weight_memory + activation_memory) / (1024**3) (bytes to GB).

Function Signature

estimate_batch_memory_gb(num_params: int, batch_size: int, seq_len: int, hidden_dim: int, bytes_per_param: float = 4.0) -> float

Examples

Example 1: small model small batch
Input: num_params = 1000000
batch_size = 8
seq_len = 128
hidden_dim = 256
bytes_per_param = 4
Output: 0.013128995895385742

Example 2: fp16 halves memory
Input: num_params = 1000000
batch_size = 8
seq_len = 128
hidden_dim = 256
bytes_per_param = 2
Output: 0.006564497947692871

Example 3: larger batch scales activations
Input: num_params = 1000000
batch_size = 64
seq_len = 128
hidden_dim = 256
bytes_per_param = 4
Output: 0.026800870895385742

Constraints

all inputs positive

Python3Saved ✓

def estimate_batch_memory_gb(num_params, batch_size, seq_len, hidden_dim, bytes_per_param=4.0):

# Your implementation here

pass

Case 1: small model small batch

Input: num_params = 1000000
batch_size = 8
seq_len = 128
hidden_dim = 256
bytes_per_param = 4

Expected: 0.013128995895385742

Case 2: fp16 halves memory

Input: num_params = 1000000
batch_size = 8
seq_len = 128
hidden_dim = 256
bytes_per_param = 2

Expected: 0.006564497947692871

Case 3: larger batch scales activations

Input: num_params = 1000000
batch_size = 64
seq_len = 128
hidden_dim = 256
bytes_per_param = 4

Expected: 0.026800870895385742