Memory Optimization

Techniques and strategies for managing GPU memory efficiently during neural network training, particularly for large language models where memory constraints often represent the primary bottleneck in scaling training to larger models and batch sizes.

Memory as Hard Constraint

Memory represents a hard constraint in LLM training - if a single training step doesn't fit in GPU memory, training simply cannot proceed. This makes memory optimization the most critical aspect of distributed training, as established in the ultra-scale-playbook.

Unlike compute or communication which can be optimized for efficiency, memory has absolute limits that must be respected for training to function at all.

Four Memory Components

GPU memory during training consists of four primary components:

1. Model Weights

• Neural network parameters
• Stored in various precisions (FP32, BF16, FP8)
• Size determined by model architecture

2. Gradients

• Computed during backward pass
• Typically same size as model weights
• Temporarily stored before optimization step

3. Optimizer States

• Often the largest component
• Adam optimizer stores momentum and variance (2x model size)
• Can dominate total memory usage

4. Activations

• Intermediate values from forward pass
• Needed for gradient computation during backward pass
• Can be traded for computation through activation-recomputation

Memory Profiling and Analysis

Empirical Measurement Approach

The ultra-scale-playbook emphasizes empirical memory profiling over theoretical calculations:

• PyTorch Memory Profiler: Step-by-step memory allocation tracking
• Dynamic Patterns: Memory usage varies significantly during training steps
• First Step Anomaly: Initial step shows different patterns due to caching allocator preparation

Training Step Memory Anatomy

1. Forward Pass: Activations build up progressively
2. Backward Pass: Gradients accumulate while activations are cleared
3. Optimization: All gradients needed, optimizer states updated

Memory Optimization Techniques

Precision Management

• Mixed Precision Training: Use BF16/FP16 instead of FP32
• FP8 Training: Cutting-edge precision for maximum memory savings
• Gradient Scaling: Maintain numerical stability with lower precision

Activation Management

• activation-recomputation: Trade computation for memory (50-90% memory reduction)
• Gradient Checkpointing: Strategic activation storage points
• Progressive Clearing: Clear activations as soon as gradients computed

Optimizer Optimization

• zero-optimizer: Partition optimizer states across devices
• AdamW vs Adam: More memory-efficient optimizer variants
• State Precision: Lower precision for optimizer states

Batch Size Management

• gradient-accumulation: Simulate larger batches without memory increase
• Micro-batching: Process smaller chunks within larger logical batches
• Dynamic Batching: Adjust batch size based on sequence length

Memory Prediction and Tools

Theoretical Calculation

Memory usage can be estimated from:

• Tensor shapes (batch size, sequence length, hidden dimensions)
• Precision formats (4 bytes for FP32, 2 for BF16, 1 for FP8)
• Model architecture parameters

Empirical Tools

• Memory Prediction Tools: Hugging Face's memory estimation widgets
• Profiling Dashboards: Real-time memory usage visualization
• Benchmarking Suites: Systematic memory usage measurement

Memory Fragmentation and Allocation

PyTorch Caching Allocator

• Pre-allocates memory blocks to speed up subsequent allocations
• Causes first step anomaly in memory patterns
• Can lead to fragmentation reducing usable memory

CUDA Kernel Overhead

• Kernels typically require 1-2 GB of GPU memory
• Constant overhead independent of model size
• Must be factored into memory budget

Trade-offs and Strategies

Computation-Memory Trade-offs

• Recomputation: Use more computation to reduce memory storage
• Batching: Larger batches improve efficiency but increase memory
• Precision: Lower precision saves memory but may affect convergence

Memory-Communication Balance

• Smaller models per device reduce memory but increase communication
• Optimal balance depends on interconnect bandwidth
• Different strategies for intra-node vs inter-node communication

Scaling Implications

Memory optimization becomes increasingly critical at scale:

• Single GPU: Focus on activation recomputation and mixed precision
• Multi-GPU: Add optimizer state sharding and gradient compression
• Ultra-Scale: Combine all techniques with sophisticated parallelism strategies

Understanding memory patterns through empirical profiling enables informed decisions about which optimization techniques to apply for specific training configurations.

See also

• training-step-anatomy
• activation-recomputation
• gradient-accumulation
• zero-optimizer
• Mixed Precision Training
• ultra-scale-playbook
• gpu-cluster-training

---