trustformers_training/

memory_optimization.rs

/// Memory optimization techniques for training large models
///
/// This module provides advanced memory management strategies including:
/// - Gradient checkpointing for reducing memory usage
/// - CPU offloading for large tensors
/// - Dynamic memory management with automatic cleanup
/// - Tensor rematerialization strategies
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex};
use trustformers_core::tensor::Tensor;

/// Configuration for memory optimization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryOptimizationConfig {
    /// Enable gradient checkpointing
    pub gradient_checkpointing: bool,
    /// Enable CPU offloading for large tensors
    pub cpu_offloading: bool,
    /// Enable dynamic memory management
    pub dynamic_memory: bool,
    /// Enable tensor rematerialization
    pub tensor_rematerialization: bool,
    /// Memory threshold for triggering optimizations (in bytes)
    pub memory_threshold: usize,
    /// Maximum memory usage before aggressive cleanup (in bytes)
    pub max_memory_usage: usize,
    /// Checkpoint interval (number of layers)
    pub checkpoint_interval: usize,
    /// CPU offloading threshold (tensor size in bytes)
    pub offload_threshold: usize,
}

impl Default for MemoryOptimizationConfig {
    fn default() -> Self {
        Self {
            gradient_checkpointing: false,
            cpu_offloading: false,
            dynamic_memory: false,
            tensor_rematerialization: false,
            memory_threshold: 1_000_000_000, // 1GB
            max_memory_usage: 8_000_000_000, // 8GB
            checkpoint_interval: 4,
            offload_threshold: 100_000_000, // 100MB
        }
    }
}

/// Checkpoint information for gradient checkpointing
#[derive(Debug, Clone)]
pub struct Checkpoint {
    pub layer_index: usize,
    pub activations: Vec<Tensor>,
    pub timestamp: std::time::Instant,
}

/// Memory optimization manager
#[allow(dead_code)]
pub struct MemoryOptimizer {
    config: MemoryOptimizationConfig,
    checkpoints: VecDeque<Checkpoint>,
    offloaded_tensors: HashMap<String, (Tensor, std::time::Instant)>,
    memory_usage: Arc<Mutex<usize>>,
    recompute_cache: HashMap<String, Vec<Tensor>>,
}

impl MemoryOptimizer {
    pub fn new(config: MemoryOptimizationConfig) -> Self {
        Self {
            config,
            checkpoints: VecDeque::new(),
            offloaded_tensors: HashMap::new(),
            memory_usage: Arc::new(Mutex::new(0)),
            recompute_cache: HashMap::new(),
        }
    }

    /// Create a checkpoint for gradient checkpointing
    pub fn create_checkpoint(
        &mut self,
        layer_index: usize,
        activations: Vec<Tensor>,
    ) -> Result<()> {
        if !self.config.gradient_checkpointing {
            return Ok(());
        }

        let checkpoint = Checkpoint {
            layer_index,
            activations,
            timestamp: std::time::Instant::now(),
        };

        self.checkpoints.push_back(checkpoint);

        // Limit checkpoint buffer size
        while self.checkpoints.len() > self.config.checkpoint_interval * 2 {
            self.checkpoints.pop_front();
        }

        Ok(())
    }

    /// Retrieve activations from checkpoint
    pub fn get_checkpoint_activations(&self, layer_index: usize) -> Option<Vec<Tensor>> {
        if !self.config.gradient_checkpointing {
            return None;
        }

        // Find the most recent checkpoint for this layer
        for checkpoint in self.checkpoints.iter().rev() {
            if checkpoint.layer_index == layer_index {
                return Some(checkpoint.activations.clone());
            }
        }

        None
    }

    /// Offload tensor to CPU to free GPU memory
    pub fn offload_to_cpu(&mut self, name: String, tensor: Tensor) -> Result<()> {
        if !self.config.cpu_offloading {
            return Ok(());
        }

        let tensor_size = self.estimate_tensor_size(&tensor)?;

        if tensor_size >= self.config.offload_threshold {
            // Move tensor to CPU (simplified - in real implementation would use actual CPU/GPU transfer)
            self.offloaded_tensors.insert(name, (tensor, std::time::Instant::now()));
        }

        Ok(())
    }

    /// Retrieve tensor from CPU offloading
    pub fn retrieve_from_cpu(&mut self, name: &str) -> Option<Tensor> {
        if !self.config.cpu_offloading {
            return None;
        }

        self.offloaded_tensors.remove(name).map(|(tensor, _)| tensor)
    }

    /// Estimate tensor memory usage
    fn estimate_tensor_size(&self, tensor: &Tensor) -> Result<usize> {
        // Simplified tensor size estimation
        let shape = tensor.shape();
        let element_size = 4; // Assume f32 elements
        let total_elements: usize = shape.iter().product();
        Ok(total_elements * element_size)
    }

    /// Update memory usage tracking
    pub fn update_memory_usage(&self, delta: isize) {
        let mut usage = self.memory_usage.lock().expect("lock should not be poisoned");
        if delta < 0 {
            *usage = usage.saturating_sub((-delta) as usize);
        } else {
            *usage += delta as usize;
        }
    }

    /// Get current memory usage
    pub fn get_memory_usage(&self) -> usize {
        *self.memory_usage.lock().expect("lock should not be poisoned")
    }

    /// Check if memory cleanup is needed
    pub fn should_cleanup(&self) -> bool {
        let usage = self.get_memory_usage();
        usage > self.config.memory_threshold
    }

    /// Perform memory cleanup
    pub fn cleanup(&mut self) -> Result<usize> {
        let mut freed_bytes = 0;

        if self.config.dynamic_memory {
            // Clean up old checkpoints
            let now = std::time::Instant::now();
            let old_checkpoints: Vec<_> = self
                .checkpoints
                .iter()
                .enumerate()
                .filter(|(_, checkpoint)| now.duration_since(checkpoint.timestamp).as_secs() > 30)
                .map(|(i, _)| i)
                .collect();

            for i in old_checkpoints.into_iter().rev() {
                if let Some(checkpoint) = self.checkpoints.remove(i) {
                    for tensor in &checkpoint.activations {
                        freed_bytes += self.estimate_tensor_size(tensor)?;
                    }
                }
            }

            // Clean up old offloaded tensors
            let old_tensors: Vec<_> = self
                .offloaded_tensors
                .iter()
                .filter(|(_, (_, timestamp))| now.duration_since(*timestamp).as_secs() > 60)
                .map(|(name, _)| name.clone())
                .collect();

            for name in old_tensors {
                if let Some((tensor, _)) = self.offloaded_tensors.remove(&name) {
                    freed_bytes += self.estimate_tensor_size(&tensor)?;
                }
            }

            // Clear recompute cache if memory pressure is high
            if self.get_memory_usage() > self.config.max_memory_usage {
                for tensors in self.recompute_cache.values() {
                    for tensor in tensors {
                        freed_bytes += self.estimate_tensor_size(tensor)?;
                    }
                }
                self.recompute_cache.clear();
            }
        }

        self.update_memory_usage(-(freed_bytes as isize));
        Ok(freed_bytes)
    }

    /// Store tensor for rematerialization
    pub fn store_for_rematerialization(&mut self, key: String, tensors: Vec<Tensor>) -> Result<()> {
        if !self.config.tensor_rematerialization {
            return Ok(());
        }

        let mut total_size = 0;
        for tensor in &tensors {
            total_size += self.estimate_tensor_size(tensor)?;
        }

        // Only store if under memory threshold
        if total_size < self.config.offload_threshold {
            self.recompute_cache.insert(key, tensors);
        }

        Ok(())
    }

    /// Retrieve tensor for rematerialization
    pub fn retrieve_for_rematerialization(&mut self, key: &str) -> Option<Vec<Tensor>> {
        if !self.config.tensor_rematerialization {
            return None;
        }

        self.recompute_cache.remove(key)
    }

    /// Get memory optimization statistics
    pub fn get_stats(&self) -> MemoryOptimizationStats {
        MemoryOptimizationStats {
            current_memory_usage: self.get_memory_usage(),
            checkpoints_count: self.checkpoints.len(),
            offloaded_tensors_count: self.offloaded_tensors.len(),
            recompute_cache_size: self.recompute_cache.len(),
            memory_threshold: self.config.memory_threshold,
            max_memory_usage: self.config.max_memory_usage,
        }
    }
}

/// Statistics for memory optimization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryOptimizationStats {
    pub current_memory_usage: usize,
    pub checkpoints_count: usize,
    pub offloaded_tensors_count: usize,
    pub recompute_cache_size: usize,
    pub memory_threshold: usize,
    pub max_memory_usage: usize,
}

/// Gradient checkpointing wrapper for model layers
pub struct GradientCheckpointWrapper {
    optimizer: MemoryOptimizer,
    layer_index: usize,
}

impl GradientCheckpointWrapper {
    pub fn new(optimizer: MemoryOptimizer, layer_index: usize) -> Self {
        Self {
            optimizer,
            layer_index,
        }
    }

    /// Forward pass with checkpointing
    pub fn forward_with_checkpoint(&mut self, inputs: Vec<Tensor>) -> Result<Vec<Tensor>> {
        // Create checkpoint before forward pass
        self.optimizer.create_checkpoint(self.layer_index, inputs.clone())?;

        // Perform forward pass (simplified)
        let outputs = inputs; // In real implementation, this would be the actual layer forward pass

        Ok(outputs)
    }

    /// Backward pass with checkpointing
    pub fn backward_with_checkpoint(&mut self, grad_outputs: Vec<Tensor>) -> Result<Vec<Tensor>> {
        // Retrieve activations from checkpoint
        if let Some(_activations) = self.optimizer.get_checkpoint_activations(self.layer_index) {
            // Recompute forward pass using checkpointed activations
            // Then compute gradients (simplified)
            Ok(grad_outputs) // In real implementation, this would be the actual gradient computation
        } else {
            Err(anyhow::anyhow!(
                "No checkpoint found for layer {}",
                self.layer_index
            ))
        }
    }
}

/// CPU offloading manager for large tensors
pub struct CPUOffloadManager {
    optimizer: MemoryOptimizer,
    offload_queue: VecDeque<String>,
}

impl CPUOffloadManager {
    pub fn new(optimizer: MemoryOptimizer) -> Self {
        Self {
            optimizer,
            offload_queue: VecDeque::new(),
        }
    }

    /// Schedule tensor for offloading
    pub fn schedule_offload(&mut self, name: String, tensor: Tensor) -> Result<()> {
        self.optimizer.offload_to_cpu(name.clone(), tensor)?;
        self.offload_queue.push_back(name);
        Ok(())
    }

    /// Process offloading queue
    pub fn process_offload_queue(&mut self) -> Result<()> {
        // Process a batch of offloads to avoid blocking
        let batch_size = 10;
        for _ in 0..batch_size {
            if let Some(_name) = self.offload_queue.pop_front() {
                // Offloading is handled in schedule_offload
                // This is where additional processing could be done
            } else {
                break;
            }
        }
        Ok(())
    }

    /// Retrieve tensor from CPU
    pub fn retrieve_tensor(&mut self, name: &str) -> Option<Tensor> {
        self.optimizer.retrieve_from_cpu(name)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_memory_optimizer_creation() {
        let config = MemoryOptimizationConfig::default();
        let optimizer = MemoryOptimizer::new(config);

        assert_eq!(optimizer.get_memory_usage(), 0);
        assert_eq!(optimizer.checkpoints.len(), 0);
        assert_eq!(optimizer.offloaded_tensors.len(), 0);
    }

    #[test]
    fn test_checkpoint_creation() {
        let config = MemoryOptimizationConfig {
            gradient_checkpointing: true,
            ..Default::default()
        };
        let mut optimizer = MemoryOptimizer::new(config);

        // Create mock tensor
        let tensor = Tensor::zeros(&[2, 3]).expect("tensor operation failed");
        let result = optimizer.create_checkpoint(0, vec![tensor]);

        assert!(result.is_ok());
        assert_eq!(optimizer.checkpoints.len(), 1);
    }

    #[test]
    fn test_memory_cleanup() {
        let config = MemoryOptimizationConfig {
            dynamic_memory: true,
            memory_threshold: 1000,
            ..Default::default()
        };
        let mut optimizer = MemoryOptimizer::new(config);

        // Simulate memory usage
        optimizer.update_memory_usage(2000);
        assert!(optimizer.should_cleanup());

        let freed = optimizer.cleanup().expect("operation failed in test");
        assert!(freed == 0); // No actual tensors to free in this test
    }

    #[test]
    fn test_cpu_offloading() {
        let config = MemoryOptimizationConfig {
            cpu_offloading: true,
            offload_threshold: 100,
            ..Default::default()
        };
        let mut optimizer = MemoryOptimizer::new(config);

        let tensor = Tensor::zeros(&[1000, 1000]).expect("tensor operation failed"); // Large tensor
        let result = optimizer.offload_to_cpu("test_tensor".to_string(), tensor);

        assert!(result.is_ok());
        assert_eq!(optimizer.offloaded_tensors.len(), 1);

        let retrieved = optimizer.retrieve_from_cpu("test_tensor");
        assert!(retrieved.is_some());
        assert_eq!(optimizer.offloaded_tensors.len(), 0);
    }

    #[test]
    fn test_gradient_checkpoint_wrapper() {
        let config = MemoryOptimizationConfig {
            gradient_checkpointing: true,
            ..Default::default()
        };
        let optimizer = MemoryOptimizer::new(config);
        let mut wrapper = GradientCheckpointWrapper::new(optimizer, 0);

        let tensor = Tensor::zeros(&[2, 3]).expect("tensor operation failed");
        let result = wrapper.forward_with_checkpoint(vec![tensor]);

        assert!(result.is_ok());
    }

    #[test]
    fn test_memory_stats() {
        let config = MemoryOptimizationConfig::default();
        let optimizer = MemoryOptimizer::new(config);

        let stats = optimizer.get_stats();
        assert_eq!(stats.current_memory_usage, 0);
        assert_eq!(stats.checkpoints_count, 0);
        assert_eq!(stats.offloaded_tensors_count, 0);
    }
}