trustformers_training/

parallelism_3d.rs

use crate::distributed::ProcessGroup;
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};
use trustformers_core::tensor::Tensor;
use trustformers_core::traits::Model;

/// 3D Parallelism Configuration
///
/// Combines data parallelism (DP), model parallelism (MP), and pipeline parallelism (PP)
/// to efficiently scale transformer training across multiple GPUs and nodes.
///
/// Key concepts:
/// - Data Parallelism: Each process has a full copy of the model and trains on different data
/// - Model Parallelism: Model parameters are split across processes within a layer
/// - Pipeline Parallelism: Model layers are split across processes, enabling pipeline execution
///
/// The total number of processes = dp_size * mp_size * pp_size
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParallelismConfig {
    /// Data parallel group size
    pub dp_size: usize,
    /// Model parallel group size
    pub mp_size: usize,
    /// Pipeline parallel group size
    pub pp_size: usize,
    /// Number of micro-batches for pipeline parallelism
    pub num_micro_batches: usize,
    /// Whether to use gradient accumulation
    pub gradient_accumulation: bool,
    /// Number of gradient accumulation steps
    pub accumulation_steps: usize,
    /// Whether to use activation checkpointing
    pub activation_checkpointing: bool,
    /// Communication backend preference
    pub comm_backend: CommBackend,
    /// Pipeline scheduling strategy
    pub pipeline_schedule: PipelineSchedule,
    /// Memory optimization level
    pub memory_optimization: MemoryOptimization,
}

impl Default for ParallelismConfig {
    fn default() -> Self {
        Self {
            dp_size: 1,
            mp_size: 1,
            pp_size: 1,
            num_micro_batches: 4,
            gradient_accumulation: true,
            accumulation_steps: 1,
            activation_checkpointing: true,
            comm_backend: CommBackend::NCCL,
            pipeline_schedule: PipelineSchedule::GPipe,
            memory_optimization: MemoryOptimization::Medium,
        }
    }
}

/// Communication backend options
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CommBackend {
    NCCL,
    Gloo,
    MPI,
    InfiniBand,
}

/// Pipeline scheduling strategies
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum PipelineSchedule {
    /// Google's GPipe scheduling
    GPipe,
    /// PipeDream scheduling
    PipeDream,
    /// PipeDream-2BW (bidirectional weight updates)
    PipeDream2BW,
    /// Interleaved 1F1B (One Forward One Backward)
    Interleaved1F1B,
    /// Adaptive scheduling based on communication patterns
    Adaptive,
}

/// Memory optimization levels
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MemoryOptimization {
    None,
    Low,
    Medium,
    High,
    Extreme,
}

/// 3D Parallelism Coordinator
///
/// Manages the coordination between different parallelism strategies
/// and handles communication between process groups.
#[allow(dead_code)]
pub struct Parallelism3D {
    config: ParallelismConfig,
    #[allow(dead_code)]
    global_rank: usize,
    world_size: usize,

    // Process group ranks and mappings
    dp_rank: usize,
    mp_rank: usize,
    pp_rank: usize,

    // Process groups for different parallelism types
    dp_group: Arc<dyn ProcessGroup>,
    mp_group: Arc<dyn ProcessGroup>,
    pp_group: Arc<dyn ProcessGroup>,

    // Pipeline state management
    pipeline_state: Arc<RwLock<PipelineState>>,

    // Communication statistics
    comm_stats: Arc<Mutex<CommunicationStats>>,

    // Memory management
    memory_manager: Arc<Mutex<MemoryManager>>,
}

/// Pipeline execution state
#[derive(Debug, Default)]
#[allow(dead_code)]
struct PipelineState {
    #[allow(dead_code)]
    current_micro_batch: usize,
    forward_passes_completed: usize,
    backward_passes_completed: usize,
    pipeline_bubbles: usize,
    stage_timings: HashMap<usize, Duration>,
    communication_overhead: Duration,
}

/// Communication statistics for 3D parallelism
#[derive(Debug, Default)]
#[allow(dead_code)]
struct CommunicationStats {
    dp_all_reduce_time: Duration,
    mp_all_reduce_time: Duration,
    pp_send_recv_time: Duration,
    #[allow(dead_code)]
    total_bytes_communicated: u64,
    communication_efficiency: f32,
    bandwidth_utilization: f32,
}

/// Memory management for 3D parallelism
#[derive(Debug)]
#[allow(dead_code)]
struct MemoryManager {
    #[allow(dead_code)]
    activation_memory_pool: HashMap<String, Vec<Tensor>>,
    gradient_memory_pool: HashMap<String, Vec<Tensor>>,
    peak_memory_usage: u64,
    current_memory_usage: u64,
    memory_optimization_level: MemoryOptimization,
    checkpointed_activations: HashMap<String, Vec<Tensor>>,
}

impl Default for MemoryManager {
    fn default() -> Self {
        Self {
            activation_memory_pool: HashMap::new(),
            gradient_memory_pool: HashMap::new(),
            peak_memory_usage: 0,
            current_memory_usage: 0,
            memory_optimization_level: MemoryOptimization::Medium,
            checkpointed_activations: HashMap::new(),
        }
    }
}

impl Parallelism3D {
    /// Create a new 3D parallelism coordinator
    pub fn new(
        config: ParallelismConfig,
        global_rank: usize,
        world_size: usize,
        dp_group: Arc<dyn ProcessGroup>,
        mp_group: Arc<dyn ProcessGroup>,
        pp_group: Arc<dyn ProcessGroup>,
    ) -> Result<Self> {
        // Validate configuration
        if config.dp_size * config.mp_size * config.pp_size != world_size {
            return Err(anyhow!(
                "Invalid parallelism configuration: dp_size ({}) * mp_size ({}) * pp_size ({}) != world_size ({})",
                config.dp_size, config.mp_size, config.pp_size, world_size
            ));
        }

        // Calculate local ranks for each parallelism type
        let dp_rank = global_rank / (config.mp_size * config.pp_size);
        let mp_rank = (global_rank / config.pp_size) % config.mp_size;
        let pp_rank = global_rank % config.pp_size;

        let memory_manager = MemoryManager {
            memory_optimization_level: config.memory_optimization.clone(),
            ..Default::default()
        };

        Ok(Self {
            config,
            global_rank,
            world_size,
            dp_rank,
            mp_rank,
            pp_rank,
            dp_group,
            mp_group,
            pp_group,
            pipeline_state: Arc::new(RwLock::new(PipelineState::default())),
            comm_stats: Arc::new(Mutex::new(CommunicationStats::default())),
            memory_manager: Arc::new(Mutex::new(memory_manager)),
        })
    }

    /// Execute forward pass with 3D parallelism
    pub fn forward_pass<M: Model>(
        &self,
        model: &M,
        inputs: &[Tensor],
        micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        let _start_time = Instant::now();

        // Handle different pipeline scheduling strategies
        match self.config.pipeline_schedule {
            PipelineSchedule::GPipe => self.forward_gpipe(model, inputs, micro_batch_id),
            PipelineSchedule::PipeDream => self.forward_pipedream(model, inputs, micro_batch_id),
            PipelineSchedule::PipeDream2BW => {
                self.forward_pipedream_2bw(model, inputs, micro_batch_id)
            },
            PipelineSchedule::Interleaved1F1B => {
                self.forward_interleaved_1f1b(model, inputs, micro_batch_id)
            },
            PipelineSchedule::Adaptive => self.forward_adaptive(model, inputs, micro_batch_id),
        }
    }

    /// Execute backward pass with 3D parallelism
    pub fn backward_pass<M: Model>(
        &self,
        model: &mut M,
        gradients: &[Tensor],
        micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        let _start_time = Instant::now();

        // Handle different pipeline scheduling strategies for backward pass
        match self.config.pipeline_schedule {
            PipelineSchedule::GPipe => self.backward_gpipe(model, gradients, micro_batch_id),
            PipelineSchedule::PipeDream => {
                self.backward_pipedream(model, gradients, micro_batch_id)
            },
            PipelineSchedule::PipeDream2BW => {
                self.backward_pipedream_2bw(model, gradients, micro_batch_id)
            },
            PipelineSchedule::Interleaved1F1B => {
                self.backward_interleaved_1f1b(model, gradients, micro_batch_id)
            },
            PipelineSchedule::Adaptive => self.backward_adaptive(model, gradients, micro_batch_id),
        }
    }

    /// Synchronize gradients across all parallelism dimensions
    pub fn synchronize_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
        let start_time = Instant::now();

        // Step 1: Reduce-scatter within model parallel group
        if self.config.mp_size > 1 {
            self.mp_reduce_scatter_gradients(gradients)?;
        }

        // Step 2: All-reduce within data parallel group
        if self.config.dp_size > 1 {
            self.dp_all_reduce_gradients(gradients)?;
        }

        // Step 3: All-gather within model parallel group
        if self.config.mp_size > 1 {
            self.mp_all_gather_gradients(gradients)?;
        }

        // Update communication statistics
        let mut stats = self.comm_stats.lock().expect("lock should not be poisoned");
        stats.dp_all_reduce_time += start_time.elapsed();

        Ok(())
    }

    /// Optimize memory usage based on configuration
    pub fn optimize_memory(&self, tensors: &mut [Tensor]) -> Result<()> {
        let memory_manager = self.memory_manager.lock().expect("lock should not be poisoned");

        match memory_manager.memory_optimization_level {
            MemoryOptimization::None => {
                // No optimization
                Ok(())
            },
            MemoryOptimization::Low => {
                // Basic activation checkpointing
                self.apply_activation_checkpointing(tensors, 4)
            },
            MemoryOptimization::Medium => {
                // Activation checkpointing + gradient compression
                self.apply_activation_checkpointing(tensors, 2)?;
                self.apply_gradient_compression(tensors)
            },
            MemoryOptimization::High => {
                // All optimizations + CPU offloading
                self.apply_activation_checkpointing(tensors, 1)?;
                self.apply_gradient_compression(tensors)?;
                self.apply_cpu_offloading(tensors)
            },
            MemoryOptimization::Extreme => {
                // ZeRO-style optimization
                self.apply_zero_optimization(tensors)
            },
        }
    }

    /// Handle pipeline bubble optimization
    pub fn optimize_pipeline_bubbles(&self) -> Result<()> {
        let state = self.pipeline_state.write().expect("lock should not be poisoned");

        // Analyze pipeline timing patterns
        let total_stages = self.config.pp_size;
        let avg_stage_time = state.stage_timings.values().sum::<Duration>() / total_stages as u32;

        // Identify bottleneck stages
        let mut bottleneck_stages = Vec::new();
        for (stage, timing) in &state.stage_timings {
            if *timing > avg_stage_time * 2 {
                bottleneck_stages.push(*stage);
            }
        }

        // Apply bubble reduction strategies
        if !bottleneck_stages.is_empty() {
            self.apply_load_balancing(&bottleneck_stages)?;
        }

        // Track bubble statistics
        let pipeline_efficiency = 1.0
            - (state.pipeline_bubbles as f32
                / (state.forward_passes_completed + state.backward_passes_completed) as f32);

        if pipeline_efficiency < 0.8 {
            self.adjust_micro_batch_size()?;
        }

        Ok(())
    }

    /// Get comprehensive 3D parallelism statistics
    pub fn get_statistics(&self) -> Result<Parallelism3DStats> {
        let pipeline_state = self.pipeline_state.read().expect("lock should not be poisoned");
        let comm_stats = self.comm_stats.lock().expect("lock should not be poisoned");
        let memory_manager = self.memory_manager.lock().expect("lock should not be poisoned");

        Ok(Parallelism3DStats {
            dp_rank: self.dp_rank,
            mp_rank: self.mp_rank,
            pp_rank: self.pp_rank,
            pipeline_efficiency: self.calculate_pipeline_efficiency(&pipeline_state),
            communication_efficiency: comm_stats.communication_efficiency,
            memory_efficiency: self.calculate_memory_efficiency(&memory_manager),
            total_communication_time: comm_stats.dp_all_reduce_time
                + comm_stats.mp_all_reduce_time
                + comm_stats.pp_send_recv_time,
            peak_memory_usage: memory_manager.peak_memory_usage,
            pipeline_bubbles: pipeline_state.pipeline_bubbles,
            micro_batches_processed: pipeline_state.forward_passes_completed,
        })
    }

    // Private implementation methods for different pipeline strategies

    fn forward_gpipe<M: Model>(
        &self,
        model: &M,
        inputs: &[Tensor],
        _micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        // GPipe: Sequential forward passes, then sequential backward passes

        if self.pp_rank == 0 {
            // First stage: process input
            let outputs = self.process_pipeline_stage(model, inputs, 0)?;

            // Send to next stage
            if self.config.pp_size > 1 {
                self.send_to_next_stage(&outputs)?;
            }

            Ok(outputs)
        } else {
            // Intermediate/final stages: receive from previous, process, send to next
            let received_inputs = self.receive_from_previous_stage()?;
            let outputs = self.process_pipeline_stage(model, &received_inputs, self.pp_rank)?;

            if self.pp_rank < self.config.pp_size - 1 {
                self.send_to_next_stage(&outputs)?;
            }

            Ok(outputs)
        }
    }

    fn forward_pipedream<M: Model>(
        &self,
        model: &M,
        inputs: &[Tensor],
        micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        // PipeDream: Interleaved forward and backward passes
        // Implementation would be more complex, involving asynchronous execution
        self.forward_gpipe(model, inputs, micro_batch_id) // Simplified for now
    }

    fn forward_pipedream_2bw<M: Model>(
        &self,
        model: &M,
        inputs: &[Tensor],
        micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        // PipeDream-2BW: Bidirectional weight updates
        self.forward_gpipe(model, inputs, micro_batch_id) // Simplified for now
    }

    fn forward_interleaved_1f1b<M: Model>(
        &self,
        model: &M,
        inputs: &[Tensor],
        micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        // Interleaved 1F1B: One forward, one backward pattern
        self.forward_gpipe(model, inputs, micro_batch_id) // Simplified for now
    }

    fn forward_adaptive<M: Model>(
        &self,
        model: &M,
        inputs: &[Tensor],
        micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        // Adaptive scheduling based on runtime characteristics

        // Choose strategy based on current performance metrics
        let stats = self.comm_stats.lock().expect("lock should not be poisoned");
        let communication_time_ratio = stats.pp_send_recv_time.as_millis() as f32
            / (stats.dp_all_reduce_time.as_millis() + stats.mp_all_reduce_time.as_millis() + 1)
                as f32;

        if communication_time_ratio > 2.0 {
            // High communication overhead, use GPipe
            self.forward_gpipe(model, inputs, micro_batch_id)
        } else {
            // Low communication overhead, use interleaved
            self.forward_interleaved_1f1b(model, inputs, micro_batch_id)
        }
    }

    // Backward pass implementations (similar pattern)
    fn backward_gpipe<M: Model>(
        &self,
        _model: &mut M,
        gradients: &[Tensor],
        _micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        // Implement GPipe backward pass
        Ok(gradients.to_vec()) // Simplified
    }

    fn backward_pipedream<M: Model>(
        &self,
        _model: &mut M,
        gradients: &[Tensor],
        _micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        Ok(gradients.to_vec()) // Simplified
    }

    fn backward_pipedream_2bw<M: Model>(
        &self,
        _model: &mut M,
        gradients: &[Tensor],
        _micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        Ok(gradients.to_vec()) // Simplified
    }

    fn backward_interleaved_1f1b<M: Model>(
        &self,
        _model: &mut M,
        gradients: &[Tensor],
        _micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        Ok(gradients.to_vec()) // Simplified
    }

    fn backward_adaptive<M: Model>(
        &self,
        _model: &mut M,
        gradients: &[Tensor],
        _micro_batch_id: usize,
    ) -> Result<Vec<Tensor>> {
        Ok(gradients.to_vec()) // Simplified
    }

    // Communication methods
    fn mp_reduce_scatter_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
        // Reduce-scatter operation within model parallel group
        // This distributes the reduction computation across MP ranks
        for _tensor in gradients.iter_mut() {
            // Simplified: would implement actual reduce-scatter
        }
        Ok(())
    }

    fn dp_all_reduce_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
        // All-reduce operation within data parallel group
        self.dp_group.all_reduce(gradients)?;

        // Average gradients by DP group size
        for _tensor in gradients.iter_mut() {
            // tensor.div_scalar(self.config.dp_size as f32)?;
        }

        Ok(())
    }

    fn mp_all_gather_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
        // All-gather operation within model parallel group
        // This assembles the full gradient tensors from scattered pieces
        for _tensor in gradients.iter_mut() {
            // Simplified: would implement actual all-gather
        }
        Ok(())
    }

    fn send_to_next_stage(&self, _tensors: &[Tensor]) -> Result<()> {
        // Send tensors to next pipeline stage
        // In a real implementation, this would use the PP process group
        Ok(())
    }

    fn receive_from_previous_stage(&self) -> Result<Vec<Tensor>> {
        // Receive tensors from previous pipeline stage
        // In a real implementation, this would use the PP process group
        Ok(vec![Tensor::zeros(&[1])?]) // Placeholder
    }

    fn process_pipeline_stage<M: Model>(
        &self,
        _model: &M,
        inputs: &[Tensor],
        _stage: usize,
    ) -> Result<Vec<Tensor>> {
        // Process inputs through a specific pipeline stage
        // This would involve running a subset of model layers
        Ok(inputs.to_vec()) // Simplified
    }

    // Memory optimization methods
    fn apply_activation_checkpointing(
        &self,
        tensors: &mut [Tensor],
        checkpoint_ratio: usize,
    ) -> Result<()> {
        let mut memory_manager = self.memory_manager.lock().expect("lock should not be poisoned");

        // Save every Nth activation for checkpointing
        for (i, tensor) in tensors.iter().enumerate() {
            if i % checkpoint_ratio == 0 {
                memory_manager
                    .checkpointed_activations
                    .entry(format!("checkpoint_{}", i))
                    .or_default()
                    .push(tensor.clone());
            }
        }

        Ok(())
    }

    fn apply_gradient_compression(&self, tensors: &mut [Tensor]) -> Result<()> {
        // Apply gradient compression techniques
        for _tensor in tensors.iter_mut() {
            // Simplified: could implement various compression schemes
            // - Quantization
            // - Sparsification
            // - Low-rank approximation
        }
        Ok(())
    }

    fn apply_cpu_offloading(&self, _tensors: &mut [Tensor]) -> Result<()> {
        // Offload tensors to CPU memory when not actively used
        Ok(())
    }

    fn apply_zero_optimization(&self, _tensors: &mut [Tensor]) -> Result<()> {
        // Apply ZeRO-style optimizer state partitioning
        Ok(())
    }

    // Performance optimization methods
    fn apply_load_balancing(&self, _bottleneck_stages: &[usize]) -> Result<()> {
        // Implement dynamic load balancing for pipeline stages
        Ok(())
    }

    fn adjust_micro_batch_size(&self) -> Result<()> {
        // Dynamically adjust micro-batch size to reduce pipeline bubbles
        Ok(())
    }

    // Statistics calculation methods
    fn calculate_pipeline_efficiency(&self, state: &PipelineState) -> f32 {
        if state.forward_passes_completed == 0 {
            return 0.0;
        }

        let total_passes = state.forward_passes_completed + state.backward_passes_completed;
        1.0 - (state.pipeline_bubbles as f32 / total_passes as f32)
    }

    fn calculate_memory_efficiency(&self, memory_manager: &MemoryManager) -> f32 {
        if memory_manager.peak_memory_usage == 0 {
            return 1.0;
        }

        memory_manager.current_memory_usage as f32 / memory_manager.peak_memory_usage as f32
    }
}

/// Comprehensive statistics for 3D parallelism
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Parallelism3DStats {
    pub dp_rank: usize,
    pub mp_rank: usize,
    pub pp_rank: usize,
    pub pipeline_efficiency: f32,
    pub communication_efficiency: f32,
    pub memory_efficiency: f32,
    pub total_communication_time: Duration,
    pub peak_memory_usage: u64,
    pub pipeline_bubbles: usize,
    pub micro_batches_processed: usize,
}

/// Manager for coordinating 3D parallelism across training
pub struct Parallelism3DManager {
    coordinators: HashMap<String, Arc<Parallelism3D>>,
    global_config: ParallelismConfig,
    performance_tracker: Arc<Mutex<PerformanceTracker>>,
}

#[derive(Debug, Default)]
#[allow(dead_code)]
struct PerformanceTracker {
    #[allow(dead_code)]
    iteration_times: Vec<Duration>,
    communication_times: Vec<Duration>,
    memory_usage_samples: Vec<u64>,
    efficiency_scores: Vec<f32>,
}

impl Parallelism3DManager {
    /// Create a new 3D parallelism manager
    pub fn new(config: ParallelismConfig) -> Self {
        Self {
            coordinators: HashMap::new(),
            global_config: config,
            performance_tracker: Arc::new(Mutex::new(PerformanceTracker::default())),
        }
    }

    /// Register a new 3D parallelism coordinator
    pub fn register_coordinator(
        &mut self,
        name: String,
        coordinator: Arc<Parallelism3D>,
    ) -> Result<()> {
        self.coordinators.insert(name, coordinator);
        Ok(())
    }

    /// Get aggregate statistics across all coordinators
    pub fn get_aggregate_stats(&self) -> Result<AggregateParallelismStats> {
        let mut aggregate = AggregateParallelismStats::default();

        for coordinator in self.coordinators.values() {
            let stats = coordinator.get_statistics()?;
            aggregate.total_pipeline_efficiency += stats.pipeline_efficiency;
            aggregate.total_communication_time += stats.total_communication_time;
            aggregate.total_memory_usage += stats.peak_memory_usage;
            aggregate.total_micro_batches += stats.micro_batches_processed;
        }

        if !self.coordinators.is_empty() {
            aggregate.average_pipeline_efficiency =
                aggregate.total_pipeline_efficiency / self.coordinators.len() as f32;
        }

        aggregate.num_coordinators = self.coordinators.len();

        Ok(aggregate)
    }

    /// Optimize configuration based on performance metrics
    pub fn optimize_configuration(&mut self) -> Result<ParallelismConfig> {
        let tracker = self.performance_tracker.lock().expect("lock should not be poisoned");

        if tracker.efficiency_scores.is_empty() {
            return Ok(self.global_config.clone());
        }

        let avg_efficiency =
            tracker.efficiency_scores.iter().sum::<f32>() / tracker.efficiency_scores.len() as f32;
        let mut optimized_config = self.global_config.clone();

        // Adjust micro-batch size based on efficiency
        if avg_efficiency < 0.8 {
            optimized_config.num_micro_batches = (optimized_config.num_micro_batches * 2).min(16);
        } else if avg_efficiency > 0.95 {
            optimized_config.num_micro_batches = (optimized_config.num_micro_batches / 2).max(1);
        }

        // Adjust memory optimization based on usage patterns
        let avg_memory_usage = tracker.memory_usage_samples.iter().sum::<u64>()
            / tracker.memory_usage_samples.len() as u64;
        if avg_memory_usage > (0.9 * (32u64 * 1024 * 1024 * 1024) as f64) as u64 {
            // 32GB threshold
            optimized_config.memory_optimization = MemoryOptimization::High;
        }

        Ok(optimized_config)
    }
}

/// Aggregate statistics across multiple 3D parallelism coordinators
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct AggregateParallelismStats {
    pub num_coordinators: usize,
    pub total_pipeline_efficiency: f32,
    pub average_pipeline_efficiency: f32,
    pub total_communication_time: Duration,
    pub total_memory_usage: u64,
    pub total_micro_batches: usize,
}

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

    #[test]
    fn test_parallelism_config_validation() {
        let config = ParallelismConfig {
            dp_size: 2,
            mp_size: 2,
            pp_size: 2,
            ..Default::default()
        };

        let world_size = 8; // 2 * 2 * 2
        assert_eq!(config.dp_size * config.mp_size * config.pp_size, world_size);
    }

    #[test]
    fn test_rank_calculation() {
        let config = ParallelismConfig {
            dp_size: 2,
            mp_size: 2,
            pp_size: 2,
            ..Default::default()
        };

        let global_rank = 5;
        let _world_size = 8;

        let dp_rank = global_rank / (config.mp_size * config.pp_size);
        let mp_rank = (global_rank / config.pp_size) % config.mp_size;
        let pp_rank = global_rank % config.pp_size;

        assert_eq!(dp_rank, 1);
        assert_eq!(mp_rank, 0);
        assert_eq!(pp_rank, 1);
    }

    #[test]
    fn test_3d_parallelism_creation() {
        let config = ParallelismConfig {
            dp_size: 2,
            mp_size: 1,
            pp_size: 1,
            ..Default::default()
        };

        let dp_group = Arc::new(SimulatedProcessGroup::new(0, 2));
        let mp_group = Arc::new(SimulatedProcessGroup::new(0, 1));
        let pp_group = Arc::new(SimulatedProcessGroup::new(0, 1));

        let parallelism = Parallelism3D::new(config, 0, 2, dp_group, mp_group, pp_group);

        assert!(parallelism.is_ok());
        let p = parallelism.expect("operation failed in test");
        assert_eq!(p.dp_rank, 0);
        assert_eq!(p.mp_rank, 0);
        assert_eq!(p.pp_rank, 0);
    }

    #[test]
    fn test_memory_optimization_levels() {
        use MemoryOptimization::*;

        let levels = vec![None, Low, Medium, High, Extreme];

        for level in levels {
            let config = ParallelismConfig {
                memory_optimization: level,
                ..Default::default()
            };

            // Should be able to serialize/deserialize
            let json = serde_json::to_string(&config).expect("JSON serialization failed");
            let deserialized: ParallelismConfig =
                serde_json::from_str(&json).expect("JSON deserialization failed");

            assert!(matches!(deserialized.memory_optimization, _));
        }
    }

    #[test]
    fn test_pipeline_schedule_types() {
        use PipelineSchedule::*;

        let schedules = vec![GPipe, PipeDream, PipeDream2BW, Interleaved1F1B, Adaptive];

        for schedule in schedules {
            let config = ParallelismConfig {
                pipeline_schedule: schedule,
                ..Default::default()
            };

            // Should be able to serialize/deserialize
            let json = serde_json::to_string(&config).expect("JSON serialization failed");
            let deserialized: ParallelismConfig =
                serde_json::from_str(&json).expect("JSON deserialization failed");

            assert!(matches!(deserialized.pipeline_schedule, _));
        }
    }

    #[test]
    fn test_3d_parallelism_manager() {
        let config = ParallelismConfig::default();
        let mut manager = Parallelism3DManager::new(config);

        // Test configuration optimization with empty data
        let optimized_config = manager.optimize_configuration();
        assert!(optimized_config.is_ok());

        // Test aggregate stats with no coordinators
        let stats = manager.get_aggregate_stats();
        assert!(stats.is_ok());

        let stats = stats.expect("operation failed in test");
        assert_eq!(stats.num_coordinators, 0);
        assert_eq!(stats.average_pipeline_efficiency, 0.0);
    }

    #[test]
    fn test_config_serialization() {
        let config = ParallelismConfig {
            dp_size: 4,
            mp_size: 2,
            pp_size: 8,
            num_micro_batches: 16,
            gradient_accumulation: true,
            accumulation_steps: 4,
            activation_checkpointing: true,
            comm_backend: CommBackend::NCCL,
            pipeline_schedule: PipelineSchedule::Interleaved1F1B,
            memory_optimization: MemoryOptimization::High,
        };

        let json = serde_json::to_string(&config).expect("JSON serialization failed");
        let deserialized: ParallelismConfig =
            serde_json::from_str(&json).expect("JSON deserialization failed");

        assert_eq!(config.dp_size, deserialized.dp_size);
        assert_eq!(config.mp_size, deserialized.mp_size);
        assert_eq!(config.pp_size, deserialized.pp_size);
        assert_eq!(config.num_micro_batches, deserialized.num_micro_batches);
        assert_eq!(
            config.gradient_accumulation,
            deserialized.gradient_accumulation
        );
        assert_eq!(config.accumulation_steps, deserialized.accumulation_steps);
        assert_eq!(
            config.activation_checkpointing,
            deserialized.activation_checkpointing
        );
    }
}