rustytorch/crates/training/rtx-distributed/src/hybrid_parallel.rs

//! Hybrid Parallelism Coordinator
//!
//! This module coordinates multiple parallelism strategies together:
//! - Data Parallelism (DP): Replicate model, partition data
//! - Tensor Parallelism (TP): Partition tensors within layers
//! - Pipeline Parallelism (PP): Partition layers across stages
//! - Sequence Parallelism (SP): Partition along sequence dimension
//!
//! Hybrid parallelism enables training models that are too large for
//! single GPU memory while maximizing throughput.

use crate::error::{DistributedError, Result};
use parking_lot::{Mutex, RwLock};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};

// =============================================================================
// Configuration
// =============================================================================

/// Parallelism dimension type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ParallelDimension {
    /// Data parallelism
    Data,
    /// Tensor parallelism
    Tensor,
    /// Pipeline parallelism
    Pipeline,
    /// Sequence parallelism
    Sequence,
    /// Expert parallelism (for MoE)
    Expert,
}

/// Configuration for hybrid parallelism
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HybridParallelConfig {
    /// Total number of GPUs
    pub total_gpus: usize,
    /// Data parallel degree
    pub dp_degree: usize,
    /// Tensor parallel degree
    pub tp_degree: usize,
    /// Pipeline parallel degree
    pub pp_degree: usize,
    /// Sequence parallel degree (usually equals tp_degree)
    pub sp_degree: usize,
    /// Expert parallel degree (for MoE models)
    pub ep_degree: usize,
    /// Enable sequence parallelism
    pub enable_sequence_parallel: bool,
    /// Enable activation checkpointing
    pub enable_checkpointing: bool,
    /// Number of micro-batches for pipeline
    pub num_micro_batches: usize,
    /// Overlap communication with computation
    pub overlap_comm: bool,
}

impl Default for HybridParallelConfig {
    fn default() -> Self {
        Self {
            total_gpus: 8,
            dp_degree: 2,
            tp_degree: 2,
            pp_degree: 2,
            sp_degree: 1,
            ep_degree: 1,
            enable_sequence_parallel: false,
            enable_checkpointing: true,
            num_micro_batches: 4,
            overlap_comm: true,
        }
    }
}

impl HybridParallelConfig {
    /// Validate configuration
    pub fn validate(&self) -> Result<()> {
        let product = self.dp_degree * self.tp_degree * self.pp_degree;
        if product != self.total_gpus {
            return Err(DistributedError::configuration(format!(
                "DP({}) x TP({}) x PP({}) = {} must equal total_gpus ({})",
                self.dp_degree, self.tp_degree, self.pp_degree, product, self.total_gpus
            )));
        }

        if self.enable_sequence_parallel && self.sp_degree != self.tp_degree {
            return Err(DistributedError::configuration(
                "Sequence parallel degree must equal tensor parallel degree",
            ));
        }

        Ok(())
    }

    /// Create for data parallelism only
    pub fn data_parallel(num_gpus: usize) -> Self {
        Self {
            total_gpus: num_gpus,
            dp_degree: num_gpus,
            tp_degree: 1,
            pp_degree: 1,
            sp_degree: 1,
            ..Default::default()
        }
    }

    /// Create for tensor + data parallelism
    pub fn tensor_data_parallel(tp_degree: usize, dp_degree: usize) -> Self {
        Self {
            total_gpus: tp_degree * dp_degree,
            dp_degree,
            tp_degree,
            pp_degree: 1,
            sp_degree: 1,
            ..Default::default()
        }
    }

    /// Create for 3D parallelism
    pub fn parallelism_3d(dp_degree: usize, tp_degree: usize, pp_degree: usize) -> Self {
        Self {
            total_gpus: dp_degree * tp_degree * pp_degree,
            dp_degree,
            tp_degree,
            pp_degree,
            sp_degree: 1,
            ..Default::default()
        }
    }
}

// =============================================================================
// Process Group Mesh
// =============================================================================

/// Represents a GPU in the parallel mesh
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MeshCoordinate {
    /// Global rank (0 to total_gpus - 1)
    pub global_rank: usize,
    /// Data parallel rank
    pub dp_rank: usize,
    /// Tensor parallel rank
    pub tp_rank: usize,
    /// Pipeline parallel rank (stage)
    pub pp_rank: usize,
    /// Device ID
    pub device_id: i32,
}

impl MeshCoordinate {
    /// Check if this is the first TP rank
    pub fn is_tp_first(&self) -> bool {
        self.tp_rank == 0
    }

    /// Check if this is the last TP rank
    pub fn is_tp_last(&self, tp_degree: usize) -> bool {
        self.tp_rank == tp_degree - 1
    }

    /// Check if this is the first PP stage
    pub fn is_pp_first(&self) -> bool {
        self.pp_rank == 0
    }

    /// Check if this is the last PP stage
    pub fn is_pp_last(&self, pp_degree: usize) -> bool {
        self.pp_rank == pp_degree - 1
    }
}

/// Process group for a subset of GPUs
#[derive(Debug, Clone)]
pub struct ProcessGroupInfo {
    /// Group name
    pub name: String,
    /// Parallelism dimension
    pub dimension: ParallelDimension,
    /// Ranks in this group
    pub ranks: Vec<usize>,
    /// This process's rank within the group
    pub local_rank: usize,
    /// Group size
    pub size: usize,
}

impl ProcessGroupInfo {
    /// Create a new process group info
    pub fn new(
        name: String,
        dimension: ParallelDimension,
        ranks: Vec<usize>,
        local_rank: usize,
    ) -> Self {
        let size = ranks.len();
        Self {
            name,
            dimension,
            ranks,
            local_rank,
            size,
        }
    }

    /// Check if rank is in this group
    pub fn contains(&self, rank: usize) -> bool {
        self.ranks.contains(&rank)
    }
}

/// Manages the GPU mesh for hybrid parallelism
pub struct ProcessGroupMesh {
    /// Configuration
    config: HybridParallelConfig,
    /// This process's global rank
    global_rank: usize,
    /// Mesh coordinate for this process
    coordinate: MeshCoordinate,
    /// Data parallel group
    dp_group: ProcessGroupInfo,
    /// Tensor parallel group
    tp_group: ProcessGroupInfo,
    /// Pipeline parallel group
    pp_group: ProcessGroupInfo,
    /// All process groups by name
    groups: HashMap<String, ProcessGroupInfo>,
}

impl ProcessGroupMesh {
    /// Create a new process group mesh
    pub fn new(config: HybridParallelConfig, global_rank: usize) -> Result<Self> {
        config.validate()?;

        if global_rank >= config.total_gpus {
            return Err(DistributedError::configuration(format!(
                "Global rank {} exceeds total GPUs {}",
                global_rank, config.total_gpus
            )));
        }

        // Calculate mesh coordinates
        // Layout: [DP, TP, PP] with PP varying fastest
        let pp_rank = global_rank % config.pp_degree;
        let tp_rank = (global_rank / config.pp_degree) % config.tp_degree;
        let dp_rank = global_rank / (config.pp_degree * config.tp_degree);

        let coordinate = MeshCoordinate {
            global_rank,
            dp_rank,
            tp_rank,
            pp_rank,
            device_id: global_rank as i32,
        };

        // Build process groups
        let dp_group = Self::build_dp_group(&config, &coordinate);
        let tp_group = Self::build_tp_group(&config, &coordinate);
        let pp_group = Self::build_pp_group(&config, &coordinate);

        let mut groups = HashMap::new();
        groups.insert(dp_group.name.clone(), dp_group.clone());
        groups.insert(tp_group.name.clone(), tp_group.clone());
        groups.insert(pp_group.name.clone(), pp_group.clone());

        Ok(Self {
            config,
            global_rank,
            coordinate,
            dp_group,
            tp_group,
            pp_group,
            groups,
        })
    }

    /// Build data parallel group (same TP and PP ranks)
    fn build_dp_group(config: &HybridParallelConfig, coord: &MeshCoordinate) -> ProcessGroupInfo {
        let mut ranks = Vec::new();
        for dp in 0..config.dp_degree {
            let rank = dp * config.pp_degree * config.tp_degree
                + coord.tp_rank * config.pp_degree
                + coord.pp_rank;
            ranks.push(rank);
        }

        ProcessGroupInfo::new(
            format!("dp_tp{}_pp{}", coord.tp_rank, coord.pp_rank),
            ParallelDimension::Data,
            ranks,
            coord.dp_rank,
        )
    }

    /// Build tensor parallel group (same DP and PP ranks)
    fn build_tp_group(config: &HybridParallelConfig, coord: &MeshCoordinate) -> ProcessGroupInfo {
        let mut ranks = Vec::new();
        for tp in 0..config.tp_degree {
            let rank = coord.dp_rank * config.pp_degree * config.tp_degree
                + tp * config.pp_degree
                + coord.pp_rank;
            ranks.push(rank);
        }

        ProcessGroupInfo::new(
            format!("tp_dp{}_pp{}", coord.dp_rank, coord.pp_rank),
            ParallelDimension::Tensor,
            ranks,
            coord.tp_rank,
        )
    }

    /// Build pipeline parallel group (same DP and TP ranks)
    fn build_pp_group(config: &HybridParallelConfig, coord: &MeshCoordinate) -> ProcessGroupInfo {
        let mut ranks = Vec::new();
        for pp in 0..config.pp_degree {
            let rank = coord.dp_rank * config.pp_degree * config.tp_degree
                + coord.tp_rank * config.pp_degree
                + pp;
            ranks.push(rank);
        }

        ProcessGroupInfo::new(
            format!("pp_dp{}_tp{}", coord.dp_rank, coord.tp_rank),
            ParallelDimension::Pipeline,
            ranks,
            coord.pp_rank,
        )
    }

    /// Get mesh coordinate
    pub fn coordinate(&self) -> &MeshCoordinate {
        &self.coordinate
    }

    /// Get data parallel group
    pub fn dp_group(&self) -> &ProcessGroupInfo {
        &self.dp_group
    }

    /// Get tensor parallel group
    pub fn tp_group(&self) -> &ProcessGroupInfo {
        &self.tp_group
    }

    /// Get pipeline parallel group
    pub fn pp_group(&self) -> &ProcessGroupInfo {
        &self.pp_group
    }

    /// Get group by name
    pub fn get_group(&self, name: &str) -> Option<&ProcessGroupInfo> {
        self.groups.get(name)
    }

    /// Get all group names
    pub fn group_names(&self) -> Vec<String> {
        self.groups.keys().cloned().collect()
    }

    /// Get global rank
    pub fn global_rank(&self) -> usize {
        self.global_rank
    }

    /// Get config
    pub fn config(&self) -> &HybridParallelConfig {
        &self.config
    }

    /// Check if this rank should participate in data parallel AllReduce
    pub fn should_dp_allreduce(&self) -> bool {
        self.config.dp_degree > 1
    }

    /// Check if this rank should participate in tensor parallel AllReduce
    pub fn should_tp_allreduce(&self) -> bool {
        self.config.tp_degree > 1
    }

    /// Get the next rank in pipeline (for send)
    pub fn next_pp_rank(&self) -> Option<usize> {
        if self.coordinate.pp_rank < self.config.pp_degree - 1 {
            Some(self.pp_group.ranks[self.coordinate.pp_rank + 1])
        } else {
            None
        }
    }

    /// Get the previous rank in pipeline (for recv)
    pub fn prev_pp_rank(&self) -> Option<usize> {
        if self.coordinate.pp_rank > 0 {
            Some(self.pp_group.ranks[self.coordinate.pp_rank - 1])
        } else {
            None
        }
    }
}

// =============================================================================
// Hybrid Parallelism Coordinator
// =============================================================================

/// Coordinates hybrid parallelism for distributed training
pub struct HybridParallelCoordinator {
    /// Configuration
    config: HybridParallelConfig,
    /// Process group mesh
    mesh: ProcessGroupMesh,
    /// Current training step
    step: AtomicU64,
    /// Is training active
    active: AtomicBool,
    /// Pending communications
    pending_comms: Mutex<Vec<PendingComm>>,
    /// Statistics
    stats: RwLock<CoordinatorStats>,
}

/// A pending communication
#[derive(Debug, Clone)]
pub struct PendingComm {
    /// Communication ID
    pub id: u64,
    /// Operation type
    pub op: CommOp,
    /// Parallelism dimension
    pub dimension: ParallelDimension,
    /// Source rank
    pub src: usize,
    /// Destination ranks
    pub dst: Vec<usize>,
    /// Data size in bytes
    pub size_bytes: usize,
    /// Creation time
    pub created_at: Instant,
}

/// Communication operation type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommOp {
    /// AllReduce for gradient synchronization
    AllReduce,
    /// AllGather for tensor parallel
    AllGather,
    /// ReduceScatter for FSDP
    ReduceScatter,
    /// Send to next pipeline stage
    Send,
    /// Recv from previous pipeline stage
    Recv,
    /// Broadcast
    Broadcast,
}

impl HybridParallelCoordinator {
    /// Create a new hybrid parallelism coordinator
    pub fn new(config: HybridParallelConfig, global_rank: usize) -> Result<Self> {
        let mesh = ProcessGroupMesh::new(config.clone(), global_rank)?;

        Ok(Self {
            config,
            mesh,
            step: AtomicU64::new(0),
            active: AtomicBool::new(false),
            pending_comms: Mutex::new(Vec::new()),
            stats: RwLock::new(CoordinatorStats::default()),
        })
    }

    /// Start training
    pub fn start(&self) {
        self.active.store(true, Ordering::SeqCst);
    }

    /// Stop training
    pub fn stop(&self) {
        self.active.store(false, Ordering::SeqCst);
    }

    /// Check if active
    pub fn is_active(&self) -> bool {
        self.active.load(Ordering::SeqCst)
    }

    /// Advance to next step
    pub fn next_step(&self) -> u64 {
        self.step.fetch_add(1, Ordering::SeqCst) + 1
    }

    /// Get current step
    pub fn current_step(&self) -> u64 {
        self.step.load(Ordering::SeqCst)
    }

    /// Schedule gradient AllReduce across data parallel group
    pub fn schedule_dp_allreduce(&self, _name: &str, size_bytes: usize) -> Option<u64> {
        if !self.mesh.should_dp_allreduce() {
            return None;
        }

        let group = self.mesh.dp_group();
        let id = self.create_pending_comm(
            CommOp::AllReduce,
            ParallelDimension::Data,
            group.ranks.clone(),
            size_bytes,
        );

        let mut stats = self.stats.write();
        stats.dp_allreduce_count += 1;
        stats.dp_allreduce_bytes += size_bytes;

        Some(id)
    }

    /// Schedule tensor parallel AllGather
    pub fn schedule_tp_allgather(&self, _name: &str, size_bytes: usize) -> Option<u64> {
        if !self.mesh.should_tp_allreduce() {
            return None;
        }

        let group = self.mesh.tp_group();
        let id = self.create_pending_comm(
            CommOp::AllGather,
            ParallelDimension::Tensor,
            group.ranks.clone(),
            size_bytes,
        );

        let mut stats = self.stats.write();
        stats.tp_allgather_count += 1;
        stats.tp_allgather_bytes += size_bytes;

        Some(id)
    }

    /// Schedule tensor parallel ReduceScatter
    pub fn schedule_tp_reduce_scatter(&self, _name: &str, size_bytes: usize) -> Option<u64> {
        if !self.mesh.should_tp_allreduce() {
            return None;
        }

        let group = self.mesh.tp_group();
        let id = self.create_pending_comm(
            CommOp::ReduceScatter,
            ParallelDimension::Tensor,
            group.ranks.clone(),
            size_bytes,
        );

        Some(id)
    }

    /// Schedule pipeline send
    pub fn schedule_pp_send(&self, size_bytes: usize) -> Option<u64> {
        let next_rank = self.mesh.next_pp_rank()?;

        let id = self.create_pending_comm(
            CommOp::Send,
            ParallelDimension::Pipeline,
            vec![next_rank],
            size_bytes,
        );

        let mut stats = self.stats.write();
        stats.pp_send_count += 1;
        stats.pp_send_bytes += size_bytes;

        Some(id)
    }

    /// Schedule pipeline recv
    pub fn schedule_pp_recv(&self, size_bytes: usize) -> Option<u64> {
        let prev_rank = self.mesh.prev_pp_rank()?;

        let id = self.create_pending_comm(
            CommOp::Recv,
            ParallelDimension::Pipeline,
            vec![prev_rank],
            size_bytes,
        );

        let mut stats = self.stats.write();
        stats.pp_recv_count += 1;
        stats.pp_recv_bytes += size_bytes;

        Some(id)
    }

    /// Create a pending communication
    fn create_pending_comm(
        &self,
        op: CommOp,
        dimension: ParallelDimension,
        dst: Vec<usize>,
        size_bytes: usize,
    ) -> u64 {
        static NEXT_ID: AtomicU64 = AtomicU64::new(1);
        let id = NEXT_ID.fetch_add(1, Ordering::SeqCst);

        let comm = PendingComm {
            id,
            op,
            dimension,
            src: self.mesh.global_rank(),
            dst,
            size_bytes,
            created_at: Instant::now(),
        };

        self.pending_comms.lock().push(comm);
        id
    }

    /// Get pending communication count
    pub fn pending_count(&self) -> usize {
        self.pending_comms.lock().len()
    }

    /// Clear pending communications
    pub fn clear_pending(&self) {
        self.pending_comms.lock().clear();
    }

    /// Get mesh
    pub fn mesh(&self) -> &ProcessGroupMesh {
        &self.mesh
    }

    /// Get configuration
    pub fn config(&self) -> &HybridParallelConfig {
        &self.config
    }

    /// Get statistics
    pub fn stats(&self) -> CoordinatorStats {
        self.stats.read().clone()
    }

    /// Get parallelism summary
    pub fn parallelism_summary(&self) -> String {
        format!(
            "DP={} x TP={} x PP={} (Total: {} GPUs)",
            self.config.dp_degree,
            self.config.tp_degree,
            self.config.pp_degree,
            self.config.total_gpus
        )
    }
}

// =============================================================================
// Coordinator Statistics
// =============================================================================

/// Statistics for hybrid parallelism coordinator
#[derive(Debug, Default, Clone)]
pub struct CoordinatorStats {
    /// Data parallel AllReduce count
    pub dp_allreduce_count: usize,
    /// Data parallel AllReduce bytes
    pub dp_allreduce_bytes: usize,
    /// Tensor parallel AllGather count
    pub tp_allgather_count: usize,
    /// Tensor parallel AllGather bytes
    pub tp_allgather_bytes: usize,
    /// Pipeline send count
    pub pp_send_count: usize,
    /// Pipeline send bytes
    pub pp_send_bytes: usize,
    /// Pipeline recv count
    pub pp_recv_count: usize,
    /// Pipeline recv bytes
    pub pp_recv_bytes: usize,
    /// Total communication time
    pub total_comm_time: Duration,
}

impl CoordinatorStats {
    /// Get total communication bytes
    pub fn total_bytes(&self) -> usize {
        self.dp_allreduce_bytes + self.tp_allgather_bytes + self.pp_send_bytes + self.pp_recv_bytes
    }

    /// Get total operation count
    pub fn total_ops(&self) -> usize {
        self.dp_allreduce_count + self.tp_allgather_count + self.pp_send_count + self.pp_recv_count
    }
}

// =============================================================================
// Thread-Safe Wrappers
// =============================================================================

/// Thread-safe shared hybrid parallel coordinator
pub type SharedHybridParallelCoordinator = Arc<HybridParallelCoordinator>;

/// Create a shared hybrid parallel coordinator
pub fn shared_hybrid_coordinator(
    config: HybridParallelConfig,
    global_rank: usize,
) -> Result<SharedHybridParallelCoordinator> {
    Ok(Arc::new(HybridParallelCoordinator::new(
        config,
        global_rank,
    )?))
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_hybrid_config_default() {
        let config = HybridParallelConfig::default();
        assert_eq!(config.total_gpus, 8);
        assert_eq!(config.dp_degree, 2);
        assert_eq!(config.tp_degree, 2);
        assert_eq!(config.pp_degree, 2);
    }

    #[test]
    fn test_config_validate_success() {
        let config = HybridParallelConfig {
            total_gpus: 8,
            dp_degree: 2,
            tp_degree: 2,
            pp_degree: 2,
            ..Default::default()
        };
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_config_validate_failure() {
        let config = HybridParallelConfig {
            total_gpus: 8,
            dp_degree: 2,
            tp_degree: 2,
            pp_degree: 3, // 2 x 2 x 3 = 12 != 8
            ..Default::default()
        };
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_data_parallel_config() {
        let config = HybridParallelConfig::data_parallel(8);
        assert_eq!(config.dp_degree, 8);
        assert_eq!(config.tp_degree, 1);
        assert_eq!(config.pp_degree, 1);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_tensor_data_parallel_config() {
        let config = HybridParallelConfig::tensor_data_parallel(4, 2);
        assert_eq!(config.total_gpus, 8);
        assert_eq!(config.dp_degree, 2);
        assert_eq!(config.tp_degree, 4);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_3d_parallelism_config() {
        let config = HybridParallelConfig::parallelism_3d(2, 2, 2);
        assert_eq!(config.total_gpus, 8);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_mesh_coordinate() {
        let config = HybridParallelConfig::default();
        let mesh = ProcessGroupMesh::new(config, 0).unwrap();
        let coord = mesh.coordinate();

        assert_eq!(coord.global_rank, 0);
        assert_eq!(coord.dp_rank, 0);
        assert_eq!(coord.tp_rank, 0);
        assert_eq!(coord.pp_rank, 0);
    }

    #[test]
    fn test_mesh_coordinate_rank_5() {
        // Layout: [DP, TP, PP] with PP varying fastest
        // Rank 5 = 5 in base [2,2,2]
        // pp = 5 % 2 = 1
        // tp = (5 / 2) % 2 = 0
        // dp = 5 / 4 = 1
        let config = HybridParallelConfig::default();
        let mesh = ProcessGroupMesh::new(config, 5).unwrap();
        let coord = mesh.coordinate();

        assert_eq!(coord.global_rank, 5);
        assert_eq!(coord.pp_rank, 1);
        assert_eq!(coord.tp_rank, 0);
        assert_eq!(coord.dp_rank, 1);
    }

    #[test]
    fn test_dp_group() {
        let config = HybridParallelConfig::default();
        let mesh = ProcessGroupMesh::new(config, 0).unwrap();
        let dp_group = mesh.dp_group();

        assert_eq!(dp_group.dimension, ParallelDimension::Data);
        assert_eq!(dp_group.size, 2); // dp_degree = 2
        assert_eq!(dp_group.local_rank, 0);
    }

    #[test]
    fn test_tp_group() {
        let config = HybridParallelConfig::default();
        let mesh = ProcessGroupMesh::new(config, 0).unwrap();
        let tp_group = mesh.tp_group();

        assert_eq!(tp_group.dimension, ParallelDimension::Tensor);
        assert_eq!(tp_group.size, 2); // tp_degree = 2
    }

    #[test]
    fn test_pp_group() {
        let config = HybridParallelConfig::default();
        let mesh = ProcessGroupMesh::new(config, 0).unwrap();
        let pp_group = mesh.pp_group();

        assert_eq!(pp_group.dimension, ParallelDimension::Pipeline);
        assert_eq!(pp_group.size, 2); // pp_degree = 2
    }

    #[test]
    fn test_pp_next_prev() {
        let config = HybridParallelConfig::default();

        // First stage
        let mesh_0 = ProcessGroupMesh::new(config.clone(), 0).unwrap();
        assert!(mesh_0.next_pp_rank().is_some());
        assert!(mesh_0.prev_pp_rank().is_none());

        // Last stage
        let mesh_1 = ProcessGroupMesh::new(config, 1).unwrap();
        assert!(mesh_1.next_pp_rank().is_none());
        assert!(mesh_1.prev_pp_rank().is_some());
    }

    #[test]
    fn test_should_allreduce() {
        let config = HybridParallelConfig::default();
        let mesh = ProcessGroupMesh::new(config, 0).unwrap();

        assert!(mesh.should_dp_allreduce()); // dp_degree = 2
        assert!(mesh.should_tp_allreduce()); // tp_degree = 2

        let dp_only = HybridParallelConfig::data_parallel(4);
        let mesh_dp = ProcessGroupMesh::new(dp_only, 0).unwrap();
        assert!(mesh_dp.should_dp_allreduce());
        assert!(!mesh_dp.should_tp_allreduce()); // tp_degree = 1
    }

    #[test]
    fn test_coordinator_creation() {
        let config = HybridParallelConfig::default();
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        assert!(!coord.is_active());
        assert_eq!(coord.current_step(), 0);
        assert_eq!(coord.pending_count(), 0);
    }

    #[test]
    fn test_coordinator_start_stop() {
        let config = HybridParallelConfig::default();
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        coord.start();
        assert!(coord.is_active());

        coord.stop();
        assert!(!coord.is_active());
    }

    #[test]
    fn test_coordinator_steps() {
        let config = HybridParallelConfig::default();
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        assert_eq!(coord.current_step(), 0);
        assert_eq!(coord.next_step(), 1);
        assert_eq!(coord.next_step(), 2);
        assert_eq!(coord.current_step(), 2);
    }

    #[test]
    fn test_schedule_dp_allreduce() {
        let config = HybridParallelConfig::default();
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        let id = coord.schedule_dp_allreduce("gradients", 1024);
        assert!(id.is_some());
        assert_eq!(coord.pending_count(), 1);

        let stats = coord.stats();
        assert_eq!(stats.dp_allreduce_count, 1);
        assert_eq!(stats.dp_allreduce_bytes, 1024);
    }

    #[test]
    fn test_schedule_tp_allgather() {
        let config = HybridParallelConfig::default();
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        let id = coord.schedule_tp_allgather("weights", 2048);
        assert!(id.is_some());
        assert_eq!(coord.pending_count(), 1);
    }

    #[test]
    fn test_schedule_pp_send() {
        let config = HybridParallelConfig::default();
        // Rank 0 is first stage, can send
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        let id = coord.schedule_pp_send(512);
        assert!(id.is_some());
    }

    #[test]
    fn test_schedule_pp_recv() {
        let config = HybridParallelConfig::default();
        // Rank 1 is second stage, can recv
        let coord = HybridParallelCoordinator::new(config, 1).unwrap();

        let id = coord.schedule_pp_recv(512);
        assert!(id.is_some());
    }

    #[test]
    fn test_parallelism_summary() {
        let config = HybridParallelConfig::default();
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        let summary = coord.parallelism_summary();
        assert!(summary.contains("DP=2"));
        assert!(summary.contains("TP=2"));
        assert!(summary.contains("PP=2"));
    }

    #[test]
    fn test_clear_pending() {
        let config = HybridParallelConfig::default();
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        coord.schedule_dp_allreduce("g1", 100);
        coord.schedule_dp_allreduce("g2", 200);
        assert_eq!(coord.pending_count(), 2);

        coord.clear_pending();
        assert_eq!(coord.pending_count(), 0);
    }

    #[test]
    fn test_coordinator_stats() {
        let stats = CoordinatorStats {
            dp_allreduce_bytes: 1000,
            tp_allgather_bytes: 500,
            pp_send_bytes: 200,
            pp_recv_bytes: 200,
            dp_allreduce_count: 2,
            tp_allgather_count: 1,
            pp_send_count: 1,
            pp_recv_count: 1,
            ..Default::default()
        };

        assert_eq!(stats.total_bytes(), 1900);
        assert_eq!(stats.total_ops(), 5);
    }

    #[test]
    fn test_shared_coordinator() {
        let config = HybridParallelConfig::default();
        let coord = shared_hybrid_coordinator(config, 0).unwrap();

        assert_eq!(coord.mesh().global_rank(), 0);
    }

    #[test]
    fn test_mesh_coordinate_helpers() {
        let coord = MeshCoordinate {
            global_rank: 0,
            dp_rank: 0,
            tp_rank: 0,
            pp_rank: 0,
            device_id: 0,
        };

        assert!(coord.is_tp_first());
        assert!(coord.is_pp_first());
        assert!(!coord.is_tp_last(2));
        assert!(!coord.is_pp_last(2));
    }

    #[test]
    fn test_process_group_info_contains() {
        let group = ProcessGroupInfo::new(
            "test".to_string(),
            ParallelDimension::Data,
            vec![0, 2, 4, 6],
            0,
        );

        assert!(group.contains(0));
        assert!(group.contains(4));
        assert!(!group.contains(1));
        assert!(!group.contains(3));
    }

    #[test]
    fn test_groups_cover_all_ranks() {
        let config = HybridParallelConfig::default();

        // For each rank, check that all parallel groups are valid
        for rank in 0..8 {
            let mesh = ProcessGroupMesh::new(config.clone(), rank).unwrap();
            assert!(mesh.dp_group().contains(rank));
            assert!(mesh.tp_group().contains(rank));
            assert!(mesh.pp_group().contains(rank));
        }
    }

    #[test]
    fn test_no_dp_allreduce_for_single_dp() {
        let config = HybridParallelConfig::data_parallel(1);
        let coord = HybridParallelCoordinator::new(config, 0).unwrap();

        // With dp_degree=1, no AllReduce should be scheduled
        let id = coord.schedule_dp_allreduce("grad", 100);
        assert!(id.is_none());
    }
}