Memory Profiling & Optimization

Overview

Memory profiling is critical for production ML systems: - Finding bottlenecks: Where does memory go? - Memory leaks: Long-running inference servers - Optimization strategies: Reduce memory usage by 50-80% - Production monitoring: Track memory in real-time

---

Memory Profiling Tools

tracemalloc (Built-in)

import tracemalloc
import torch

tracemalloc.start()

# Code that might leak
for _ in range(1000):
    x = torch.randn(100, 100)
    y = x ** 2
    z = y.sum()

# Get top memory allocations
current, peak = tracemalloc.get_traced_memory()
print(f"Current: {current / 1e6:.1f}MB; Peak: {peak / 1e6:.1f}MB")

# Get detailed snapshot
snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics('lineno')

for stat in top_stats[:3]:
    print(stat)

memory_profiler (Line-by-line)

from memory_profiler import profile
import torch

@profile
def train_epoch():
    """Profile memory line by line."""

    # This line allocated memory
    data = [torch.randn(100, 100) for _ in range(100)]

    # This line processes data
    results = [x.sum() for x in data]

    return results

# Run with: python -m memory_profiler script.py
# Shows memory used after each line

objgraph (Object tracking)

import objgraph
import torch

# Track object growth
objgraph.show_most_common_types(limit=3)

# Create objects
for _ in range(1000):
    x = torch.randn(10, 10)
    y = x ** 2

# Show what grew
objgraph.show_most_common_types(limit=3)

# Find leaks
objgraph.show_refs([x], filename='refs.png')

---

Common Memory Patterns

Pattern 1: Accumulating Lists

# BAD: Accumulating list
results = []
for i in range(10000):
    tensor = torch.randn(1000, 1000)
    results.append(tensor)  # Keeps all tensors in memory!

# Total memory: 10000 * 4MB = 40GB!

# GOOD: Process and discard
for i in range(10000):
    tensor = torch.randn(1000, 1000)
    result = process(tensor)
    save_to_disk(result)
    # Tensor freed after processing

Pattern 2: Unclosed Resources

# BAD: File not closed
f = open('data.txt')
lines = f.readlines()
# File still open, memory held

# GOOD: Context manager
with open('data.txt') as f:
    lines = f.readlines()
# File automatically closed, memory freed

Pattern 3: Circular References

# BAD: Circular reference
class Node:
    def __init__(self):
        self.ref = self  # Circular!

node = Node()
del node  # NOT freed due to cycle

# GOOD: Use weak reference
import weakref

class Node:
    def __init__(self):
        self.ref = weakref.ref(self)

node = Node()
del node  # Freed immediately

---

Gradient Checkpointing (50-80% Memory Reduction)

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class ModelWithCheckpointing(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(1000, 1000)
        self.layer2 = nn.Linear(1000, 1000)
        self.layer3 = nn.Linear(1000, 1000)

    def forward(self, x):
        # Standard forward (stores all activations)
        # x1 = self.layer1(x)
        # x2 = self.layer2(x1)
        # return self.layer3(x2)

        # With checkpointing (recompute instead of storing)
        x1 = checkpoint(self.layer1, x)
        x2 = checkpoint(self.layer2, x1)
        return checkpoint(self.layer3, x2)

model = ModelWithCheckpointing()

# Memory savings: ~50-80%
# Speed cost: ~20-30% slower (recomputation overhead)

---

Quantization (75% Memory Reduction)

import torch
from torch.quantization import quantize_dynamic

# Original model
model = torch.nn.Sequential(
    torch.nn.Linear(1000, 1000),
    torch.nn.ReLU(),
    torch.nn.Linear(1000, 10)
)

print(f"Original size: {sum(p.numel() for p in model.parameters()) * 4 / 1e6:.1f}MB")

# Quantized model (int8)
quantized_model = quantize_dynamic(
    model,
    {torch.nn.Linear},
    dtype=torch.qint8
)

print(f"Quantized size: {sum(p.numel() for p in quantized_model.parameters()) * 1 / 1e6:.1f}MB")
# Memory reduction: 75%

---

Model Pruning (30-70% Memory Reduction)

import torch
import torch.nn.utils.prune as prune

model = torch.nn.Sequential(
    torch.nn.Linear(1000, 1000),
    torch.nn.ReLU(),
    torch.nn.Linear(1000, 10)
)

# Prune 50% of weights in layer 1
prune.l1_unstructured(model[0], name='weight', amount=0.5)

# Make pruning permanent
prune.remove(model[0], 'weight')

# Memory reduction: 30-70% depending on sparsity

---

Memory Monitoring Decorators

from functools import wraps
import tracemalloc

def memory_monitor(func):
    """Decorator to track memory usage of function."""
    @wraps(func)
    def wrapper(*args, **kwargs):
        tracemalloc.start()

        result = func(*args, **kwargs)

        current, peak = tracemalloc.get_traced_memory()
        print(f"{func.__name__}: Peak {peak / 1e6:.1f}MB")

        return result

    return wrapper

@memory_monitor
def train_epoch():
    """Function with memory monitoring."""
    data = [torch.randn(100, 100) for _ in range(100)]
    return sum(x.sum() for x in data)

train_epoch()
# Output: train_epoch: Peak 40.5MB

---

Production Monitoring with psutil

import psutil
import torch
import time

class MemoryMonitor:
    """Monitor system and process memory in production."""

    def __init__(self, alert_threshold_mb=8000):
        self.process = psutil.Process()
        self.alert_threshold = alert_threshold_mb

    def check_memory(self):
        """Check current memory usage."""
        info = self.process.memory_info()
        rss_mb = info.rss / 1e6  # Resident set size

        if rss_mb > self.alert_threshold:
            print(f"WARNING: Memory usage {rss_mb:.0f}MB > {self.alert_threshold}MB")

        return rss_mb

    def periodic_check(self, interval=60):
        """Check memory every N seconds."""
        while True:
            self.check_memory()
            time.sleep(interval)

# Usage in inference server
monitor = MemoryMonitor(alert_threshold_mb=8000)

for request in incoming_requests:
    result = model.infer(request)
    memory = monitor.check_memory()
    print(f"Current memory: {memory:.0f}MB")

---

Optimization Checklist

For Training:

• ✓ Use gradient checkpointing for large models
• ✓ Use mixed precision (fp16) training
• ✓ Enable gradient accumulation (simulate larger batches)
• ✓ Use appropriate batch size (not too large)
• ✓ Profile with tracemalloc to find leaks

For Inference:

• ✓ Use quantization (int4/int8) for model weights
• ✓ Implement batch inference (higher throughput, lower per-token memory)
• ✓ Use KV cache for transformers (avoid recomputation)
• ✓ Implement token streaming (don't buffer entire output)
• ✓ Monitor with psutil in production

General:

• ✓ Use context managers for resources
• ✓ Avoid circular references (use weakref if needed)
• ✓ Implement proper caching (bounded cache size)
• ✓ Profile before optimizing (measure, don't guess)

---

Memory Reduction Impact Table

Technique	Memory Reduction	Speed Impact	Difficulty	Use Case
Gradient Checkpointing	50-80%	+20-30% slower	Medium	Training large models
Quantization	75% (int4)	+2-4x faster	Low	Inference, storage
Model Pruning	30-70%	-5-20% slower	High	Model optimization
Mixed Precision	50%	+2-3x faster	Low	Training, inference
Batch Processing	Neutral	+5-10x	Low	Inference throughput
KV Cache	90% (cache)	+10x	Medium	Transformer inference
Token Streaming	99% (output)	Neutral	Medium	Large output tasks

---

Real-World Example: Llama 2 70B

Baseline (FP32):
  - Model weights: 280GB (70B * 4 bytes)
  - Activations: 40GB (batch size 1, seq len 2048)
  - Gradients: 280GB (training)
  - Total: 600GB (training), 320GB (inference)

With Optimizations:
  - Quantization (int4): 70GB (4x reduction)
  - KV Cache: 5GB (cached, not 40GB)
  - Gradient checkpointing: 30GB (training)
  - Total: 105GB (training), 75GB (inference)

Final with all optimizations: 87.5% memory reduction!

---

Related Topics

• 01 Reference Counting & Garbage Collection - Memory management
• 02 Memory Layout & Cache Efficiency - Performance tuning
• 03 Global Interpreter Lock (Gil) - Concurrency
• 02 Autograd Implementation - Gradient computation