rustytorch/docs/book/src/troubleshooting/memory.md

Memory Problems

Diagnose and fix memory issues.

Out of Memory (OOM)

GPU OOM During Training

CUDA error: out of memory
Tried to allocate 2.00 GB

Solutions:

1. Reduce batch size:

   let batch_size = batch_size / 2;
   

2. Enable gradient checkpointing:

   use rtx_autograd::CheckpointManager;
   let checkpoint = CheckpointManager::sqrt_n();
   
   let hidden = checkpoint.checkpoint(|| {
       layer.forward(&x)
   })?;
   

3. Use mixed precision:

   let model = model.to_dtype(DType::F16)?;
   

4. Clear unused memory:

   use rtx_tensor::clear_pool;
   clear_pool();
   

5. Process in chunks:

   let mut outputs = vec![];
   for chunk in data.chunks(safe_batch_size) {
       outputs.push(model.forward(&chunk)?);
       clear_pool();  // Free between chunks
   }
   

GPU OOM During Inference

Solutions:

1. Disable gradient tracking:

   let output = no_grad(|| model.forward(&input))?;
   

2. Use streaming inference:

   for token in model.generate_stream(&input)? {
       // Process one token at a time
   }
   

Memory Leaks

Symptoms

• Memory usage grows over time
• Eventually crashes with OOM

Finding Leaks

use rtx_memory::AllocationTracker;

AllocationTracker::enable();

for epoch in 0..100 {
    train_epoch(&model, &data)?;

    let stats = AllocationTracker::snapshot();
    println!("Epoch {}: {} MB", epoch, stats.current_mb);

    if stats.current_mb > threshold {
        println!("Leak detected!");
        AllocationTracker::dump_live();
    }
}

Common Leak Causes

1. Accumulating tensors in loops:

   // Bad - accumulates history
   let mut history = vec![];
   for batch in data {
       let output = model.forward(&batch)?;
       history.push(output);  // Grows forever!
   }
   
   // Good - process and discard
   for batch in data {
       let output = model.forward(&batch)?;
       process(&output)?;
       // output dropped here
   }
   

2. Reference cycles with Arc:

   // Use Weak references to break cycles
   use std::sync::Weak;
   

3. Holding autograd graph:

   // Clear tape after backward
   clear_tape();
   

Memory Fragmentation

Symptoms

• Free memory but allocation fails
• Slow allocations

Solutions

1. Use tensor pooling:

   use rtx_tensor::TensorPool;
   
   let buffer = TensorPool::get(shape, &device)?;
   

2. Pre-allocate fixed-size buffers:

   let mut output = Tensor::zeros([max_batch, hidden], &device)?;
   
   for batch in batches {
       model.forward_out(&batch, &mut output)?;
   }
   

3. Periodic defragmentation:

   if pool_stats().fragmentation > 0.5 {
       clear_pool();
   }
   

CPU Memory Issues

High RAM Usage

Solutions:

1. Use pinned memory for transfers:

   let pinned = PinnedBuffer::new(size)?;
   

2. Stream data instead of loading all:

   for batch in DataLoader::new(&dataset).batch_size(64) {
       // Load one batch at a time
   }
   

3. Memory-map large files:

   use memmap2::Mmap;
   let mmap = unsafe { Mmap::map(&file)? };
   

Debugging Memory

Memory Snapshot

use rtx_memory::Snapshot;

let before = Snapshot::capture(&device)?;

// Run code
model.forward(&input)?;

let after = Snapshot::capture(&device)?;
let diff = after.diff(&before);

println!("Memory delta: {} MB", diff.total_mb);
for alloc in diff.new_allocations() {
    println!("  New: {} at {}", alloc.size_mb, alloc.location);
}

CUDA Memory Debug

# Enable CUDA memory debugging
export CUDA_LAUNCH_BLOCKING=1
export PYTORCH_NO_CUDA_MEMORY_CACHING=1  # Disable caching for debugging

Valgrind (CPU)

valgrind --leak-check=full ./target/release/my-app

Memory Estimation

Estimate Model Memory

fn estimate_model_memory(model: &Model) -> usize {
    let mut total = 0;
    for param in model.parameters() {
        total += param.numel() * param.dtype().size_bytes();
    }
    total
}

// Account for gradients during training
let training_memory = estimate_model_memory(&model) * 2;

// Account for activations (rough estimate)
let batch_memory = batch_size * hidden_size * num_layers * 4;

Memory Budget

let available = device.free_memory()?;
let model_mem = estimate_model_memory(&model);
let safe_batch_mem = available - model_mem - safety_margin;
let max_batch_size = safe_batch_mem / mem_per_sample;

Next Steps

• GPU Issues - GPU-specific problems
• Performance Memory - Memory optimization
• FAQ - Common questions