## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/realizar/src/parallel.rs

    Line
    Count
    Source
    1
    //! Multi-GPU and Distributed Inference
    2
    //!
    3
    //! Per spec §10: Implements parallelism strategies for 70B+ model inference.
    4
    //! Reference: [11] Shoeybi et al. (2019) "Megatron-LM: Training Multi-Billion Parameter LMs"
    5
    //!
    6
    //! ## Parallelism Strategies
    7
    //!
    8
    //! | Strategy | Description | Use Case | Scaling |
    9
    //! |----------|-------------|----------|---------|
    10
    //! | Tensor Parallel (TP) | Split tensors across GPUs | Within node | 2-8 GPUs |
    11
    //! | Pipeline Parallel (PP) | Split layers across GPUs | Across nodes | 2-64 GPUs |
    12
    //! | Data Parallel (DP) | Replicate model, split batches | High throughput | Any |
    13
    //!
    14
    //! ## Performance Target
    15
    //!
    16
    //! Per spec §1.3: >85% scaling efficiency for 2-8 GPUs (Amdahl's law measurement)
    17
    18
    // Module-level clippy allows
    19
    #![allow(clippy::must_use_candidate)]
    20
    #![allow(clippy::return_self_not_must_use)]
    21
    #![allow(clippy::missing_errors_doc)]
    22
    23
    use serde::{Deserialize, Serialize};
    24
    use std::collections::HashMap;
    25
    use std::sync::Arc;
    26
    use thiserror::Error;
    27
    28
    /// Error type for parallelism operations
    29
    #[derive(Debug, Error)]
    30
    pub enum ParallelError {
    31
        /// Invalid rank
    32
        #[error("Invalid rank {rank} for world size {world_size}")]
    33
        InvalidRank {
    34
            /// The invalid rank value
    35
            rank: usize,
    36
            /// The total world size
    37
            world_size: usize,
    38
        },
    39
    40
        /// Invalid world size
    41
        #[error("Invalid world size: {0}")]
    42
        InvalidWorldSize(usize),
    43
    44
        /// Communication error
    45
        #[error("Communication error: {0}")]
    46
        CommunicationError(String),
    47
    48
        /// Tensor shape mismatch
    49
        #[error("Tensor shape mismatch: expected {expected:?}, got {got:?}")]
    50
        ShapeMismatch {
    51
            /// Expected shape
    52
            expected: Vec<usize>,
    53
            /// Actual shape
    54
            got: Vec<usize>,
    55
        },
    56
    57
        /// Pipeline stage error
    58
        #[error("Pipeline stage error: {0}")]
    59
        PipelineError(String),
    60
    61
        /// Not initialized
    62
        #[error("Parallel context not initialized")]
    63
        NotInitialized,
    64
    }
    65
    66
    /// Reduce operation for collective communications
    67
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
    68
    pub enum ReduceOp {
    69
        /// Sum all values
    70
        Sum,
    71
        /// Take maximum
    72
        Max,
    73
        /// Take minimum
    74
        Min,
    75
        /// Average (Sum / world_size)
    76
        Avg,
    77
    }
    78
    79
    /// Parallel configuration
    80
    #[derive(Debug, Clone, Serialize, Deserialize)]
    81
    pub struct ParallelConfig {
    82
        /// Tensor parallel size (within node)
    83
        pub tp_size: usize,
    84
        /// Pipeline parallel size (across nodes)
    85
        pub pp_size: usize,
    86
        /// Data parallel size (batch distribution)
    87
        pub dp_size: usize,
    88
        /// Current global rank
    89
        pub rank: usize,
    90
        /// Total world size
    91
        pub world_size: usize,
    92
    }
    93
    94
    impl ParallelConfig {
    95
        /// Create a new parallel configuration
    96
        ///
    97
        /// # Arguments
    98
        ///
    99
        /// * `tp_size` - Tensor parallel size (typically 2, 4, or 8)
    100
        /// * `pp_size` - Pipeline parallel size (number of stages)
    101
        /// * `dp_size` - Data parallel size (batch replication)
    102
        /// * `rank` - Current process rank
    103
    11
        pub fn new(
    104
    11
            tp_size: usize,
    105
    11
            pp_size: usize,
    106
    11
            dp_size: usize,
    107
    11
            rank: usize,
    108
    11
        ) -> Result<Self, ParallelError> {
    109
    11
            let world_size = tp_size * pp_size * dp_size;
    110
    111
    11
            if world_size == 0 {
    112
    4
                return Err(ParallelError::InvalidWorldSize(0));
    113
    7
            }
    114
    115
    7
            if rank >= world_size {
    116
    1
                return Err(ParallelError::InvalidRank { rank, world_size });
    117
    6
            }
    118
    119
    6
            Ok(Self {
    120
    6
                tp_size,
    121
    6
                pp_size,
    122
    6
                dp_size,
    123
    6
                rank,
    124
    6
                world_size,
    125
    6
            })
    126
    11
        }
    127
    128
        /// Create single-GPU configuration (no parallelism)
    129
    4
        pub fn single() -> Self {
    130
    4
            Self {
    131
    4
                tp_size: 1,
    132
    4
                pp_size: 1,
    133
    4
                dp_size: 1,
    134
    4
                rank: 0,
    135
    4
                world_size: 1,
    136
    4
            }
    137
    4
        }
    138
    139
        /// Get tensor parallel rank within TP group
    140
    4
        pub fn tp_rank(&self) -> usize {
    141
    4
            self.rank % self.tp_size
    142
    4
        }
    143
    144
        /// Get pipeline parallel stage
    145
    4
        pub fn pp_stage(&self) -> usize {
    146
    4
            (self.rank / self.tp_size) % self.pp_size
    147
    4
        }
    148
    149
        /// Get data parallel rank
    150
    1
        pub fn dp_rank(&self) -> usize {
    151
    1
            self.rank / (self.tp_size * self.pp_size)
    152
    1
        }
    153
    154
        /// Check if this is the first TP rank
    155
    1
        pub fn is_tp_first(&self) -> bool {
    156
    1
            self.tp_rank() == 0
    157
    1
        }
    158
    159
        /// Check if this is the last TP rank
    160
    1
        pub fn is_tp_last(&self) -> bool {
    161
    1
            self.tp_rank() == self.tp_size - 1
    162
    1
        }
    163
    164
        /// Check if this is the first PP stage
    165
    1
        pub fn is_pp_first(&self) -> bool {
    166
    1
            self.pp_stage() == 0
    167
    1
        }
    168
    169
        /// Check if this is the last PP stage
    170
    1
        pub fn is_pp_last(&self) -> bool {
    171
    1
            self.pp_stage() == self.pp_size - 1
    172
    1
        }
    173
    }
    174
    175
    impl Default for ParallelConfig {
    176
    0
        fn default() -> Self {
    177
    0
            Self::single()
    178
    0
        }
    179
    }
    180
    181
    /// Mock tensor for parallelism testing
    182
    /// In production, this would be replaced with trueno::Tensor
    183
    #[derive(Debug, Clone)]
    184
    pub struct ParallelTensor {
    185
        /// Shape of the tensor
    186
        pub shape: Vec<usize>,
    187
        /// Data (f32 for simplicity)
    188
        pub data: Vec<f32>,
    189
    }
    190
    191
    impl ParallelTensor {
    192
        /// Create a new tensor
    193
    15
        pub fn new(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, ParallelError> {
    194
    15
            let expected_size: usize = shape.iter().product();
    195
    15
            if data.len() != expected_size {
    196
    0
                return Err(ParallelError::ShapeMismatch {
    197
    0
                    expected: vec![expected_size],
    198
    0
                    got: vec![data.len()],
    199
    0
                });
    200
    15
            }
    201
    15
            Ok(Self { shape, data })
    202
    15
        }
    203
    204
        /// Create a zeros tensor
    205
    1
        pub fn zeros(shape: Vec<usize>) -> Self {
    206
    1
            let size: usize = shape.iter().product();
    207
    1
            Self {
    208
    1
                shape,
    209
    1
                data: vec![0.0; size],
    210
    1
            }
    211
    1
        }
    212
    213
        /// Get a narrow slice along a dimension
    214
    4
        pub fn narrow(&self, dim: usize, start: usize, length: usize) -> Result<Self, ParallelError> {
    215
    4
            if dim >= self.shape.len() {
    216
    0
                return Err(ParallelError::ShapeMismatch {
    217
    0
                    expected: vec![dim],
    218
    0
                    got: self.shape.clone(),
    219
    0
                });
    220
    4
            }
    221
    222
            // For 2D tensors (matrices), implement proper narrowing
    223
    4
            if self.shape.len() == 2 {
    224
    4
                let rows = self.shape[0];
    225
    4
                let cols = self.shape[1];
    226
    227
    4
                if dim == 0 {
    228
                    // Narrow rows
    229
    3
                    let mut new_data = Vec::with_capacity(length * cols);
    230
    8
                    for row in start3..(start + length)3 {
    231
    8
                        let row_start = row * cols;
    232
    8
                        new_data.extend_from_slice(&self.data[row_start..row_start + cols]);
    233
    8
                    }
    234
    3
                    let new_shape = vec![length, cols];
    235
    3
                    return Ok(Self {
    236
    3
                        shape: new_shape,
    237
    3
                        data: new_data,
    238
    3
                    });
    239
    1
                }
    240
                // Narrow columns
    241
    1
                let mut new_data = Vec::with_capacity(rows * length);
    242
    2
                for row in 0..rows1 {
    243
    2
                    let row_start = row * cols;
    244
    2
                    new_data
    245
    2
                        .extend_from_slice(&self.data[row_start + start..row_start + start + length]);
    246
    2
                }
    247
    1
                let new_shape = vec![rows, length];
    248
    1
                return Ok(Self {
    249
    1
                    shape: new_shape,
    250
    1
                    data: new_data,
    251
    1
                });
    252
    0
            }
    253
    254
            // For 1D tensors
    255
    0
            if self.shape.len() == 1 {
    256
    0
                let new_data = self.data[start..start + length].to_vec();
    257
    0
                return Ok(Self {
    258
    0
                    shape: vec![length],
    259
    0
                    data: new_data,
    260
    0
                });
    261
    0
            }
    262
    263
            // Fallback: simplified implementation
    264
    0
            let new_data = self.data[start..start + length].to_vec();
    265
    0
            let mut new_shape = self.shape.clone();
    266
    0
            new_shape[dim] = length;
    267
    0
            Ok(Self {
    268
    0
                shape: new_shape,
    269
    0
                data: new_data,
    270
    0
            })
    271
    4
        }
    272
    273
        /// Transpose for 2D tensors
    274
    3
        pub fn transpose(&self) -> Result<Self, ParallelError> {
    275
    3
            if self.shape.len() != 2 {
    276
    0
                return Err(ParallelError::ShapeMismatch {
    277
    0
                    expected: vec![2],
    278
    0
                    got: vec![self.shape.len()],
    279
    0
                });
    280
    3
            }
    281
    282
    3
            let rows = self.shape[0];
    283
    3
            let cols = self.shape[1];
    284
    3
            let mut new_data = vec![0.0; rows * cols];
    285
    286
    8
            for i in 0..rows3 {
    287
    38
                for j in 0..cols8 {
    288
    38
                    new_data[j * rows + i] = self.data[i * cols + j];
    289
    38
                }
    290
            }
    291
    292
    3
            Ok(Self {
    293
    3
                shape: vec![cols, rows],
    294
    3
                data: new_data,
    295
    3
            })
    296
    3
        }
    297
    298
        /// Matrix multiplication (simplified)
    299
    3
        pub fn matmul(&self, other: &Self) -> Result<Self, ParallelError> {
    300
    3
            if self.shape.len() != 2 || other.shape.len() != 2 {
    301
    0
                return Err(ParallelError::ShapeMismatch {
    302
    0
                    expected: vec![2, 2],
    303
    0
                    got: vec![self.shape.len(), other.shape.len()],
    304
    0
                });
    305
    3
            }
    306
    307
    3
            let m = self.shape[0];
    308
    3
            let k = self.shape[1];
    309
    3
            let n = other.shape[1];
    310
    311
    3
            if k != other.shape[0] {
    312
    0
                return Err(ParallelError::ShapeMismatch {
    313
    0
                    expected: vec![k],
    314
    0
                    got: vec![other.shape[0]],
    315
    0
                });
    316
    3
            }
    317
    318
    3
            let mut result = vec![0.0; m * n];
    319
    320
    3
            for i in 0..m {
    321
    8
                for j in 0..n3 {
    322
    8
                    let mut sum = 0.0;
    323
    36
                    for l in 0..k8 {
    324
    36
                        sum += self.data[i * k + l] * other.data[l * n + j];
    325
    36
                    }
    326
    8
                    result[i * n + j] = sum;
    327
                }
    328
            }
    329
    330
    3
            Ok(Self {
    331
    3
                shape: vec![m, n],
    332
    3
                data: result,
    333
    3
            })
    334
    3
        }
    335
    336
        /// Add another tensor element-wise
    337
    1
        pub fn add(&self, other: &Self) -> Result<Self, ParallelError> {
    338
    1
            if self.shape != other.shape {
    339
    0
                return Err(ParallelError::ShapeMismatch {
    340
    0
                    expected: self.shape.clone(),
    341
    0
                    got: other.shape.clone(),
    342
    0
                });
    343
    1
            }
    344
    345
    1
            let data: Vec<f32> = self
    346
    1
                .data
    347
    1
                .iter()
    348
    1
                .zip(&other.data)
    349
    4
                .map1(|(a, b)| a + b)
    350
    1
                .collect();
    351
    1
            Ok(Self {
    352
    1
                shape: self.shape.clone(),
    353
    1
                data,
    354
    1
            })
    355
    1
        }
    356
    357
        /// Sum all elements
    358
    1
        pub fn sum(&self) -> f32 {
    359
    1
            self.data.iter().sum()
    360
    1
        }
    361
    362
        /// Number of elements
    363
    1
        pub fn numel(&self) -> usize {
    364
    1
            self.data.len()
    365
    1
        }
    366
    }
    367
    368
    /// Mock communicator for collective operations
    369
    /// In production, this would use NCCL or MPI
    370
    #[derive(Debug, Clone)]
    371
    pub struct Communicator {
    372
        /// World size
    373
        world_size: usize,
    374
        /// Current rank
    375
        rank: usize,
    376
        /// test buffers for testing
    377
        #[allow(dead_code)]
    378
        buffers: Arc<std::sync::RwLock<HashMap<usize, Vec<f32>>>>,
    379
    }
    380
    381
    impl Communicator {
    382
        /// Create a new communicator
    383
    15
        pub fn new(world_size: usize, rank: usize) -> Result<Self, ParallelError> {
    384
    15
            if rank >= world_size {
    385
    2
                return Err(ParallelError::InvalidRank { rank, world_size });
    386
    13
            }
    387
    13
            Ok(Self {
    388
    13
                world_size,
    389
    13
                rank,
    390
    13
                buffers: Arc::new(std::sync::RwLock::new(HashMap::new())),
    391
    13
            })
    392
    15
        }
    393
    394
        /// All-reduce operation
    395
    3
        pub fn all_reduce(
    396
    3
            &self,
    397
    3
            tensor: &ParallelTensor,
    398
    3
            op: ReduceOp,
    399
    3
        ) -> Result<ParallelTensor, ParallelError> {
    400
            // In a real implementation, this would use NCCL
    401
            // For testing, we simulate single-process behavior
    402
    3
            match op {
    403
                ReduceOp::Sum => {
    404
                    // Single process: multiply by world_size to simulate sum from all ranks
    405
    2
                    let data: Vec<f32> = tensor
    406
    2
                        .data
    407
    2
                        .iter()
    408
    4
                        .map2(|x| x * self.world_size as f32)
    409
    2
                        .collect();
    410
    2
                    Ok(ParallelTensor {
    411
    2
                        shape: tensor.shape.clone(),
    412
    2
                        data,
    413
    2
                    })
    414
                },
    415
                ReduceOp::Avg => {
    416
                    // Average: no change in single process (sum / world_size = value)
    417
    1
                    Ok(tensor.clone())
    418
                },
    419
                ReduceOp::Max | ReduceOp::Min => {
    420
                    // Single process: return as-is
    421
    0
                    Ok(tensor.clone())
    422
                },
    423
            }
    424
    3
        }
    425
    426
        /// All-gather operation
    427
    1
        pub fn all_gather(&self, tensor: &ParallelTensor) -> Result<ParallelTensor, ParallelError> {
    428
            // Simulate all-gather by replicating data world_size times
    429
    1
            let mut data = Vec::with_capacity(tensor.data.len() * self.world_size);
    430
    2
            for _ in 0..self.world_size1 {
    431
    2
                data.extend_from_slice(&tensor.data);
    432
    2
            }
    433
    434
    1
            let mut new_shape = tensor.shape.clone();
    435
    1
            if !new_shape.is_empty() {
    436
    1
                new_shape[0] *= self.world_size;
    437
    1
            }0
    438
    439
    1
            Ok(ParallelTensor {
    440
    1
                shape: new_shape,
    441
    1
                data,
    442
    1
            })
    443
    1
        }
    444
    445
        /// Reduce-scatter operation
    446
    0
        pub fn reduce_scatter(
    447
    0
            &self,
    448
    0
            tensor: &ParallelTensor,
    449
    0
            op: ReduceOp,
    450
    0
        ) -> Result<ParallelTensor, ParallelError> {
    451
            // Reduce then scatter: each rank gets 1/world_size of the result
    452
    0
            let chunk_size = tensor.data.len() / self.world_size;
    453
    0
            let start = self.rank * chunk_size;
    454
    0
            let end = start + chunk_size;
    455
    456
    0
            let chunk_data: Vec<f32> = match op {
    457
    0
                ReduceOp::Sum => tensor.data[start..end]
    458
    0
                    .iter()
    459
    0
                    .map(|x| x * self.world_size as f32)
    460
    0
                    .collect(),
    461
    0
                ReduceOp::Avg | ReduceOp::Max | ReduceOp::Min => tensor.data[start..end].to_vec(),
    462
            };
    463
    464
    0
            let mut new_shape = tensor.shape.clone();
    465
    0
            if !new_shape.is_empty() {
    466
    0
                new_shape[0] /= self.world_size;
    467
    0
            }
    468
    469
    0
            Ok(ParallelTensor {
    470
    0
                shape: new_shape,
    471
    0
                data: chunk_data,
    472
    0
            })
    473
    0
        }
    474
    475
        /// Barrier synchronization
    476
    1
        pub fn barrier(&self) -> Result<(), ParallelError> {
    477
            // In real implementation, this would synchronize all processes
    478
    1
            Ok(())
    479
    1
        }
    480
    481
        /// Get world size
    482
    1
        pub fn world_size(&self) -> usize {
    483
    1
            self.world_size
    484
    1
        }
    485
    486
        /// Get rank
    487
    1
        pub fn rank(&self) -> usize {
    488
    1
            self.rank
    489
    1
        }
    490
    }
    491
    492
    /// Tensor Parallelism for multi-GPU inference
    493
    /// Reference: [11] Megatron-LM tensor parallelism
    494
    #[derive(Debug)]
    495
    pub struct TensorParallel {
    496
        /// Number of tensor parallel ranks
    497
        tp_size: usize,
    498
        /// Current rank within TP group
    499
        rank: usize,
    500
        /// Communication group
    501
        comm: Communicator,
    502
    }
    503
    504
    impl TensorParallel {
    505
        /// Create a new tensor parallel context
    506
    10
        pub fn new(tp_size: usize, rank: usize) -> Result<Self, ParallelError> {
    507
    10
            if tp_size == 0 {
    508
    1
                return Err(ParallelError::InvalidWorldSize(0));
    509
    9
            }
    510
    9
            if rank >= tp_size {
    511
    2
                return Err(ParallelError::InvalidRank {
    512
    2
                    rank,
    513
    2
                    world_size: tp_size,
    514
    2
                });
    515
    7
            }
    516
    517
    7
            let comm = Communicator::new(tp_size, rank)?0;
    518
    519
    7
            Ok(Self {
    520
    7
                tp_size,
    521
    7
                rank,
    522
    7
                comm,
    523
    7
            })
    524
    10
        }
    525
    526
        /// Get chunk size for weight sharding
    527
    5
        pub fn chunk_size(&self, total_size: usize) -> usize {
    528
    5
            total_size / self.tp_size
    529
    5
        }
    530
    531
        /// Column-parallel linear (for MLP first layer, attention QKV)
    532
        ///
    533
        /// Each rank holds weight[:, rank*chunk:(rank+1)*chunk]
    534
        /// No communication needed as outputs are independent
    535
    1
        pub fn column_parallel_linear(
    536
    1
            &self,
    537
    1
            input: &ParallelTensor,
    538
    1
            weight: &ParallelTensor,
    539
    1
            bias: Option<&ParallelTensor>,
    540
    1
        ) -> Result<ParallelTensor, ParallelError> {
    541
            // Get local weight slice
    542
    1
            let output_dim = weight.shape[0];
    543
    1
            let chunk = self.chunk_size(output_dim);
    544
    1
            let local_weight = weight.narrow(0, self.rank * chunk, chunk)?0;
    545
    546
            // Transpose weight for matmul: (out_chunk, in) -> (in, out_chunk)
    547
    1
            let weight_t = local_weight.transpose()?0;
    548
    549
            // Local matmul: (batch, in) @ (in, out_chunk) -> (batch, out_chunk)
    550
    1
            let mut local_output = input.matmul(&weight_t)?0;
    551
    552
            // Add local bias if present
    553
    1
            if let Some(b0) = bias {
    554
    0
                let local_bias = b.narrow(0, self.rank * chunk, chunk)?;
    555
                // Broadcast bias addition
    556
    0
                let bias_expanded = ParallelTensor {
    557
    0
                    shape: local_output.shape.clone(),
    558
    0
                    data: local_output
    559
    0
                        .data
    560
    0
                        .iter()
    561
    0
                        .enumerate()
    562
    0
                        .map(|(i, v)| v + local_bias.data[i % local_bias.data.len()])
    563
    0
                        .collect(),
    564
                };
    565
    0
                local_output = bias_expanded;
    566
    1
            }
    567
    568
    1
            Ok(local_output)
    569
    1
        }
    570
    571
        /// Row-parallel linear (for MLP second layer, attention output)
    572
        ///
    573
        /// Each rank holds weight[rank*chunk:(rank+1)*chunk, :]
    574
        /// Requires all-reduce to sum partial results
    575
    1
        pub fn row_parallel_linear(
    576
    1
            &self,
    577
    1
            input: &ParallelTensor,
    578
    1
            weight: &ParallelTensor,
    579
    1
            bias: Option<&ParallelTensor>,
    580
    1
        ) -> Result<ParallelTensor, ParallelError> {
    581
            // Get local weight slice (rows)
    582
    1
            let input_dim = weight.shape[0];
    583
    1
            let chunk = self.chunk_size(input_dim);
    584
    1
            let local_weight = weight.narrow(0, self.rank * chunk, chunk)?0;
    585
    586
            // Transpose for matmul
    587
    1
            let weight_t = local_weight.transpose()?0;
    588
    589
            // Local matmul
    590
    1
            let local_output = input.matmul(&weight_t)?0;
    591
    592
            // All-reduce to sum partial results
    593
    1
            let mut output = self.comm.all_reduce(&local_output, ReduceOp::Sum)?0;
    594
    595
            // Add bias only on rank 0 to avoid double counting
    596
    1
            if self.rank == 0 {
    597
    1
                if let Some(b0) = bias {
    598
    0
                    output = output.add(b)?;
    599
    1
                }
    600
    0
            }
    601
    602
    1
            Ok(output)
    603
    1
        }
    604
    605
        /// Get TP rank
    606
    1
        pub fn rank(&self) -> usize {
    607
    1
            self.rank
    608
    1
        }
    609
    610
        /// Get TP size
    611
    2
        pub fn tp_size(&self) -> usize {
    612
    2
            self.tp_size
    613
    2
        }
    614
    }
    615
    616
    /// Pipeline stage info
    617
    #[derive(Debug, Clone, Serialize, Deserialize)]
    618
    pub struct PipelineStage {
    619
        /// Stage index
    620
        pub index: usize,
    621
        /// Start layer index
    622
        pub start_layer: usize,
    623
        /// End layer index (exclusive)
    624
        pub end_layer: usize,
    625
        /// Number of layers in this stage
    626
        pub num_layers: usize,
    627
    }
    628
    629
    /// Pipeline Parallelism for multi-node inference
    630
    /// Reference: [11] GPipe-style pipeline parallelism
    631
    #[derive(Debug)]
    632
    pub struct PipelineParallel {
    633
        /// Number of pipeline stages
    634
        pp_size: usize,
    635
        /// Current stage
    636
        stage: usize,
    637
        /// Stage info
    638
        stage_info: PipelineStage,
    639
        /// Micro-batch size for pipelining
    640
        micro_batch_size: usize,
    641
        /// Stats
    642
        stats: PipelineStats,
    643
    }
    644
    645
    /// Pipeline execution statistics
    646
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    647
    pub struct PipelineStats {
    648
        /// Total micro-batches processed
    649
        pub micro_batches_processed: u64,
    650
        /// Total pipeline bubbles (idle time)
    651
        pub bubble_time_ms: f64,
    652
        /// Average stage latency
    653
        pub avg_stage_latency_ms: f64,
    654
        /// Total forward passes
    655
        pub forward_passes: u64,
    656
    }
    657
    658
    impl PipelineParallel {
    659
        /// Create a new pipeline parallel context
    660
        ///
    661
        /// # Arguments
    662
        ///
    663
        /// * `pp_size` - Number of pipeline stages
    664
        /// * `stage` - Current stage index (0 to pp_size-1)
    665
        /// * `total_layers` - Total number of layers to distribute
    666
        /// * `micro_batch_size` - Size of micro-batches for pipelining
    667
    16
        pub fn new(
    668
    16
            pp_size: usize,
    669
    16
            stage: usize,
    670
    16
            total_layers: usize,
    671
    16
            micro_batch_size: usize,
    672
    16
        ) -> Result<Self, ParallelError> {
    673
    16
            if pp_size == 0 {
    674
    0
                return Err(ParallelError::InvalidWorldSize(0));
    675
    16
            }
    676
    16
            if stage >= pp_size {
    677
    0
                return Err(ParallelError::InvalidRank {
    678
    0
                    rank: stage,
    679
    0
                    world_size: pp_size,
    680
    0
                });
    681
    16
            }
    682
    683
            // Distribute layers evenly across stages
    684
    16
            let layers_per_stage = total_layers / pp_size;
    685
    16
            let extra_layers = total_layers % pp_size;
    686
    687
            // Earlier stages get extra layers if uneven
    688
    16
            let start_layer = stage * layers_per_stage + stage.min(extra_layers);
    689
    16
            let num_layers = layers_per_stage + usize::from(stage < extra_layers);
    690
    16
            let end_layer = start_layer + num_layers;
    691
    692
    16
            let stage_info = PipelineStage {
    693
    16
                index: stage,
    694
    16
                start_layer,
    695
    16
                end_layer,
    696
    16
                num_layers,
    697
    16
            };
    698
    699
    16
            Ok(Self {
    700
    16
                pp_size,
    701
    16
                stage,
    702
    16
                stage_info,
    703
    16
                micro_batch_size,
    704
    16
                stats: PipelineStats::default(),
    705
    16
            })
    706
    16
        }
    707
    708
        /// Get stage info
    709
    5
        pub fn stage_info(&self) -> &PipelineStage {
    710
    5
            &self.stage_info
    711
    5
        }
    712
    713
        /// Get micro-batch size
    714
    2
        pub fn micro_batch_size(&self) -> usize {
    715
    2
            self.micro_batch_size
    716
    2
        }
    717
    718
        /// Check if this is the first stage
    719
    5
        pub fn is_first_stage(&self) -> bool {
    720
    5
            self.stage == 0
    721
    5
        }
    722
    723
        /// Check if this is the last stage
    724
    5
        pub fn is_last_stage(&self) -> bool {
    725
    5
            self.stage == self.pp_size - 1
    726
    5
        }
    727
    728
        /// Get number of stages
    729
    2
        pub fn num_stages(&self) -> usize {
    730
    2
            self.pp_size
    731
    2
        }
    732
    733
        /// Get current stage index
    734
    1
        pub fn stage(&self) -> usize {
    735
    1
            self.stage
    736
    1
        }
    737
    738
        /// Calculate theoretical bubble ratio (idle time fraction)
    739
        /// Bubble ratio = (pp_size - 1) / (pp_size + num_microbatches - 1)
    740
    1
        pub fn bubble_ratio(&self, num_microbatches: usize) -> f32 {
    741
    1
            if num_microbatches == 0 {
    742
    0
                return 1.0;
    743
    1
            }
    744
    1
            (self.pp_size - 1) as f32 / (self.pp_size + num_microbatches - 1) as f32
    745
    1
        }
    746
    747
        /// Get statistics
    748
    1
        pub fn stats(&self) -> &PipelineStats {
    749
    1
            &self.stats
    750
    1
        }
    751
    752
        /// Record a micro-batch processed
    753
    2
        pub fn record_micro_batch(&mut self, stage_latency_ms: f64) {
    754
    2
            self.stats.micro_batches_processed += 1;
    755
    2
            self.stats.forward_passes += 1;
    756
    757
            // Update running average
    758
    2
            let n = self.stats.micro_batches_processed as f64;
    759
    2
            self.stats.avg_stage_latency_ms =
    760
    2
                (self.stats.avg_stage_latency_ms * (n - 1.0) + stage_latency_ms) / n;
    761
    2
        }
    762
    }
    763
    764
    /// ZeRO-Inference memory offload
    765
    /// Reference: [10] Microsoft DeepSpeed ZeRO-Inference
    766
    #[derive(Debug, Clone, Serialize, Deserialize)]
    767
    #[allow(clippy::struct_excessive_bools)] // Config struct - bools are appropriate
    768
    pub struct ZeroOffload {
    769
        /// Offload optimizer states to CPU
    770
        pub offload_optimizer: bool,
    771
        /// Offload parameters to CPU
    772
        pub offload_params: bool,
    773
        /// Offload activations to CPU
    774
        pub offload_activations: bool,
    775
        /// Pin memory for faster CPU-GPU transfer
    776
        pub pin_memory: bool,
    777
        /// Overlap compute and communication
    778
        pub overlap_comm: bool,
    779
    }
    780
    781
    impl Default for ZeroOffload {
    782
    7
        fn default() -> Self {
    783
    7
            Self {
    784
    7
                offload_optimizer: true,
    785
    7
                offload_params: false,
    786
    7
                offload_activations: false,
    787
    7
                pin_memory: true,
    788
    7
                overlap_comm: true,
    789
    7
            }
    790
    7
        }
    791
    }
    792
    793
    impl ZeroOffload {
    794
        /// Create inference-optimized config (offload everything)
    795
    4
        pub fn inference() -> Self {
    796
    4
            Self {
    797
    4
                offload_optimizer: false, // No optimizer in inference
    798
    4
                offload_params: true,
    799
    4
                offload_activations: true,
    800
    4
                pin_memory: true,
    801
    4
                overlap_comm: true,
    802
    4
            }
    803
    4
        }
    804
    805
        /// Estimate memory savings ratio
    806
    4
        pub fn memory_savings_ratio(&self) -> f32 {
    807
    4
            let mut ratio = 1.0;
    808
    4
            if self.offload_params {
    809
    2
                ratio *= 0.5; // Params on CPU
    810
    2
            }
    811
    4
            if self.offload_activations {
    812
    2
                ratio *= 0.7; // Activations on CPU
    813
    2
            }
    814
    4
            1.0 - ratio
    815
    4
        }
    816
    }
    817
    818
    /// Distributed inference context combining all parallelism strategies
    819
    #[derive(Debug)]
    820
    pub struct DistributedContext {
    821
        /// Parallel configuration
    822
        config: ParallelConfig,
    823
        /// Tensor parallelism (if enabled)
    824
        tensor_parallel: Option<TensorParallel>,
    825
        /// Pipeline parallelism (if enabled)
    826
        pipeline_parallel: Option<PipelineParallel>,
    827
        /// ZeRO offload settings
    828
        zero_offload: ZeroOffload,
    829
        /// Initialized flag
    830
        initialized: bool,
    831
    }
    832
    833
    impl DistributedContext {
    834
        /// Create a new distributed context
    835
    4
        pub fn new(config: ParallelConfig) -> Result<Self, ParallelError> {
    836
    4
            let tensor_parallel = if config.tp_size > 1 {
    837
    1
                Some(TensorParallel::new(config.tp_size, config.tp_rank())?0)
    838
            } else {
    839
    3
                None
    840
            };
    841
    842
            // Note: Pipeline parallel requires layer count, initialized separately
    843
    4
            let pipeline_parallel = None;
    844
    845
    4
            Ok(Self {
    846
    4
                config,
    847
    4
                tensor_parallel,
    848
    4
                pipeline_parallel,
    849
    4
                zero_offload: ZeroOffload::default(),
    850
    4
                initialized: true,
    851
    4
            })
    852
    4
        }
    853
    854
        /// Initialize pipeline parallelism
    855
    1
        pub fn init_pipeline(
    856
    1
            &mut self,
    857
    1
            total_layers: usize,
    858
    1
            micro_batch_size: usize,
    859
    1
        ) -> Result<(), ParallelError> {
    860
    1
            if self.config.pp_size > 1 {
    861
    1
                self.pipeline_parallel = Some(PipelineParallel::new(
    862
    1
                    self.config.pp_size,
    863
    1
                    self.config.pp_stage(),
    864
    1
                    total_layers,
    865
    1
                    micro_batch_size,
    866
    0
                )?);
    867
    0
            }
    868
    1
            Ok(())
    869
    1
        }
    870
    871
        /// Set ZeRO offload configuration
    872
    1
        pub fn set_zero_offload(&mut self, zero: ZeroOffload) {
    873
    1
            self.zero_offload = zero;
    874
    1
        }
    875
    876
        /// Get parallel configuration
    877
    0
        pub fn config(&self) -> &ParallelConfig {
    878
    0
            &self.config
    879
    0
        }
    880
    881
        /// Get tensor parallel context
    882
    3
        pub fn tensor_parallel(&self) -> Option<&TensorParallel> {
    883
    3
            self.tensor_parallel.as_ref()
    884
    3
        }
    885
    886
        /// Get pipeline parallel context
    887
    3
        pub fn pipeline_parallel(&self) -> Option<&PipelineParallel> {
    888
    3
            self.pipeline_parallel.as_ref()
    889
    3
        }
    890
    891
        /// Get mutable pipeline parallel context
    892
    0
        pub fn pipeline_parallel_mut(&mut self) -> Option<&mut PipelineParallel> {
    893
    0
            self.pipeline_parallel.as_mut()
    894
    0
        }
    895
    896
        /// Get ZeRO offload config
    897
    1
        pub fn zero_offload(&self) -> &ZeroOffload {
    898
    1
            &self.zero_offload
    899
    1
        }
    900
    901
        /// Check if distributed execution is enabled
    902
    2
        pub fn is_distributed(&self) -> bool {
    903
    2
            self.config.world_size > 1
    904
    2
        }
    905
    906
        /// Check if initialized
    907
    1
        pub fn is_initialized(&self) -> bool {
    908
    1
            self.initialized
    909
    1
        }
    910
    }
    911
    912
    #[cfg(test)]
    913
    mod tests {
    914
        use super::*;
    915
    916
        // =========================================================================
    917
        // ParallelConfig Tests
    918
        // =========================================================================
    919
    920
        #[test]
    921
    1
        fn test_parallel_config_new() {
    922
    1
            let config = ParallelConfig::new(2, 2, 2, 0).expect("test");
    923
    1
            assert_eq!(config.tp_size, 2);
    924
    1
            assert_eq!(config.pp_size, 2);
    925
    1
            assert_eq!(config.dp_size, 2);
    926
    1
            assert_eq!(config.world_size, 8);
    927
    1
            assert_eq!(config.rank, 0);
    928
    1
        }
    929
    930
        #[test]
    931
    1
        fn test_parallel_config_single() {
    932
    1
            let config = ParallelConfig::single();
    933
    1
            assert_eq!(config.tp_size, 1);
    934
    1
            assert_eq!(config.pp_size, 1);
    935
    1
            assert_eq!(config.dp_size, 1);
    936
    1
            assert_eq!(config.world_size, 1);
    937
    1
            assert_eq!(config.rank, 0);
    938
    1
        }
    939
    940
        #[test]
    941
    1
        fn test_parallel_config_invalid_rank() {
    942
    1
            let result = ParallelConfig::new(2, 2, 2, 100);
    943
    1
            assert!(result.is_err());
    944
    1
        }
    945
    946
        #[test]
    947
    1
        fn test_parallel_config_invalid_world_size() {
    948
    1
            let result = ParallelConfig::new(0, 0, 0, 0);
    949
    1
            assert!(result.is_err());
    950
    1
        }
    951
    952
        #[test]
    953
    1
        fn test_parallel_config_ranks() {
    954
            // World size = 2 * 2 * 2 = 8
    955
            // Rank 5: tp_rank=1, pp_stage=0, dp_rank=1
    956
    1
            let config = ParallelConfig::new(2, 2, 2, 5).expect("test");
    957
    1
            assert_eq!(config.tp_rank(), 1);
    958
    1
            assert_eq!(config.pp_stage(), 0);
    959
    1
            assert_eq!(config.dp_rank(), 1);
    960
    1
        }
    961
    962
        #[test]
    963
    1
        fn test_parallel_config_first_last_checks() {
    964
    1
            let config = ParallelConfig::new(2, 2, 1, 0).expect("test");
    965
    1
            assert!(config.is_tp_first());
    966
    1
            assert!(!config.is_tp_last());
    967
    1
            assert!(config.is_pp_first());
    968
    1
            assert!(!config.is_pp_last());
    969
    1
        }
    970
    971
        // =========================================================================
    972
        // ParallelTensor Tests
    973
        // =========================================================================
    974
    975
        #[test]
    976
    1
        fn test_parallel_tensor_new() {
    977
    1
            let tensor =
    978
    1
                ParallelTensor::new(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("test");
    979
    1
            assert_eq!(tensor.shape, vec![2, 3]);
    980
    1
            assert_eq!(tensor.numel(), 6);
    981
    1
        }
    982
    983
        #[test]
    984
    1
        fn test_parallel_tensor_zeros() {
    985
    1
            let tensor = ParallelTensor::zeros(vec![2, 3]);
    986
    1
            assert_eq!(tensor.sum(), 0.0);
    987
    1
        }
    988
    989
        #[test]
    990
    1
        fn test_parallel_tensor_narrow_rows() {
    991
    1
            let tensor = ParallelTensor::new(vec![4, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
    992
    1
                .expect("test");
    993
    1
            let narrowed = tensor.narrow(0, 1, 2).expect("test");
    994
    1
            assert_eq!(narrowed.shape, vec![2, 2]);
    995
    1
            assert_eq!(narrowed.data, vec![3.0, 4.0, 5.0, 6.0]);
    996
    1
        }
    997
    998
        #[test]
    999
    1
        fn test_parallel_tensor_narrow_cols() {
    1000
    1
            let tensor = ParallelTensor::new(vec![2, 4], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
    1001
    1
                .expect("test");
    1002
    1
            let narrowed = tensor.narrow(1, 1, 2).expect("test");
    1003
    1
            assert_eq!(narrowed.shape, vec![2, 2]);
    1004
    1
            assert_eq!(narrowed.data, vec![2.0, 3.0, 6.0, 7.0]);
    1005
    1
        }
    1006
    1007
        #[test]
    1008
    1
        fn test_parallel_tensor_transpose() {
    1009
    1
            let tensor =
    1010
    1
                ParallelTensor::new(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("test");
    1011
    1
            let transposed = tensor.transpose().expect("test");
    1012
    1
            assert_eq!(transposed.shape, vec![3, 2]);
    1013
    1
            assert_eq!(transposed.data, vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
    1014
    1
        }
    1015
    1016
        #[test]
    1017
    1
        fn test_parallel_tensor_matmul() {
    1018
            // [1, 2] @ [[1, 2], [3, 4]] = [7, 10]
    1019
    1
            let a = ParallelTensor::new(vec![1, 2], vec![1.0, 2.0]).expect("test");
    1020
    1
            let b = ParallelTensor::new(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0]).expect("test");
    1021
    1
            let c = a.matmul(&b).expect("test");
    1022
    1
            assert_eq!(c.shape, vec![1, 2]);
    1023
    1
            assert_eq!(c.data, vec![7.0, 10.0]);
    1024
    1
        }
    1025
    1026
        #[test]
    1027
    1
        fn test_parallel_tensor_add() {
    1028
    1
            let a = ParallelTensor::new(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0]).expect("test");
    1029
    1
            let b = ParallelTensor::new(vec![2, 2], vec![5.0, 6.0, 7.0, 8.0]).expect("test");
    1030
    1
            let c = a.add(&b).expect("test");
    1031
    1
            assert_eq!(c.data, vec![6.0, 8.0, 10.0, 12.0]);
    1032
    1
        }
    1033
    1034
        // =========================================================================
    1035
        // Communicator Tests
    1036
        // =========================================================================
    1037
    1038
        #[test]
    1039
    1
        fn test_communicator_new() {
    1040
    1
            let comm = Communicator::new(4, 2).expect("test");
    1041
    1
            assert_eq!(comm.world_size(), 4);
    1042
    1
            assert_eq!(comm.rank(), 2);
    1043
    1
        }
    1044
    1045
        #[test]
    1046
    1
        fn test_communicator_invalid_rank() {
    1047
    1
            let result = Communicator::new(4, 10);
    1048
    1
            assert!(result.is_err());
    1049
    1
        }
    1050
    1051
        #[test]
    1052
    1
        fn test_communicator_all_reduce_sum() {
    1053
    1
            let comm = Communicator::new(4, 0).expect("test");
    1054
    1
            let tensor = ParallelTensor::new(vec![2], vec![1.0, 2.0]).expect("test");
    1055
    1
            let result = comm.all_reduce(&tensor, ReduceOp::Sum).expect("test");
    1056
            // test: multiply by world_size
    1057
    1
            assert_eq!(result.data, vec![4.0, 8.0]);
    1058
    1
        }
    1059
    1060
        #[test]
    1061
    1
        fn test_communicator_all_reduce_avg() {
    1062
    1
            let comm = Communicator::new(4, 0).expect("test");
    1063
    1
            let tensor = ParallelTensor::new(vec![2], vec![1.0, 2.0]).expect("test");
    1064
    1
            let result = comm.all_reduce(&tensor, ReduceOp::Avg).expect("test");
    1065
    1
            assert_eq!(result.data, vec![1.0, 2.0]);
    1066
    1
        }
    1067
    1068
        #[test]
    1069
    1
        fn test_communicator_all_gather() {
    1070
    1
            let comm = Communicator::new(2, 0).expect("test");
    1071
    1
            let tensor = ParallelTensor::new(vec![2], vec![1.0, 2.0]).expect("test");
    1072
    1
            let result = comm.all_gather(&tensor).expect("test");
    1073
    1
            assert_eq!(result.shape, vec![4]);
    1074
    1
            assert_eq!(result.data, vec![1.0, 2.0, 1.0, 2.0]);
    1075
    1
        }
    1076
    1077
        #[test]
    1078
    1
        fn test_communicator_barrier() {
    1079
    1
            let comm = Communicator::new(4, 0).expect("test");
    1080
    1
            assert!(comm.barrier().is_ok());
    1081
    1
        }
    1082
    1083
        // =========================================================================
    1084
        // TensorParallel Tests
    1085
        // =========================================================================
    1086
    1087
        #[test]
    1088
    1
        fn test_tensor_parallel_new() {
    1089
    1
            let tp = TensorParallel::new(4, 2).expect("test");
    1090
    1
            assert_eq!(tp.tp_size(), 4);
    1091
    1
            assert_eq!(tp.rank(), 2);
    1092
    1
        }
    1093
    1094
        #[test]
    1095
    1
        fn test_tensor_parallel_invalid_rank() {
    1096
    1
            let result = TensorParallel::new(4, 10);
    1097
    1
            assert!(result.is_err());
    1098
    1
        }
    1099
    1100
        #[test]
    1101
    1
        fn test_tensor_parallel_chunk_size() {
    1102
    1
            let tp = TensorParallel::new(4, 0).expect("test");
    1103
    1
            assert_eq!(tp.chunk_size(100), 25);
    1104
    1
            assert_eq!(tp.chunk_size(16), 4);
    1105
    1
        }
    1106
    1107
        #[test]
    1108
    1
        fn test_tensor_parallel_column_linear() {
    1109
    1
            let tp = TensorParallel::new(2, 0).expect("test");
    1110
    1111
            // Input: (1, 4), Weight: (8, 4) split to (4, 4) per rank
    1112
    1
            let input = ParallelTensor::new(vec![1, 4], vec![1.0, 1.0, 1.0, 1.0]).expect("test");
    1113
    1
            let weight =
    1114
    32
                ParallelTensor::new1(vec!1[8, 4], (0..32)1.map1(|i| i as f32).collect1()).expect1("test"1);
    1115
    1116
    1
            let output = tp
    1117
    1
                .column_parallel_linear(&input, &weight, None)
    1118
    1
                .expect("test");
    1119
            // Output should be (1, 4) - chunk of full output
    1120
    1
            assert_eq!(output.shape, vec![1, 4]);
    1121
    1
        }
    1122
    1123
        #[test]
    1124
    1
        fn test_tensor_parallel_row_linear() {
    1125
    1
            let tp = TensorParallel::new(2, 0).expect("test");
    1126
    1127
            // Row parallel: Weight (4, 8) split to (2, 8) per rank
    1128
            // After transpose: (8, 2)
    1129
            // Input needs to be (batch, 8) to matmul with (8, 2) -> output (batch, 2)
    1130
    1
            let input = ParallelTensor::new(vec![1, 8], vec![1.0; 8]).expect("test");
    1131
    1
            let weight =
    1132
    32
                ParallelTensor::new1(vec!1[4, 8], (0..32)1.map1(|i| i as f32).collect1()).expect1("test"1);
    1133
    1134
    1
            let output = tp.row_parallel_linear(&input, &weight, None).expect("test");
    1135
            // Output shape after row parallel
    1136
    1
            assert!(!output.data.is_empty());
    1137
            // After all-reduce, output shape is (1, 2)
    1138
    1
            assert_eq!(output.shape[0], 1);
    1139
    1
        }
    1140
    1141
        // =========================================================================
    1142
        // PipelineParallel Tests
    1143
        // =========================================================================
    1144
    1145
        #[test]
    1146
    1
        fn test_pipeline_parallel_new() {
    1147
    1
            let pp = PipelineParallel::new(4, 1, 24, 4).expect("test");
    1148
    1
            assert_eq!(pp.num_stages(), 4);
    1149
    1
            assert_eq!(pp.stage(), 1);
    1150
    1
            assert_eq!(pp.micro_batch_size(), 4);
    1151
    1
        }
    1152
    1153
        #[test]
    1154
    1
        fn test_pipeline_parallel_layer_distribution() {
    1155
            // 24 layers across 4 stages = 6 layers each
    1156
    1
            let pp = PipelineParallel::new(4, 0, 24, 4).expect("test");
    1157
    1
            let info = pp.stage_info();
    1158
    1
            assert_eq!(info.start_layer, 0);
    1159
    1
            assert_eq!(info.end_layer, 6);
    1160
    1
            assert_eq!(info.num_layers, 6);
    1161
    1162
    1
            let pp2 = PipelineParallel::new(4, 3, 24, 4).expect("test");
    1163
    1
            let info2 = pp2.stage_info();
    1164
    1
            assert_eq!(info2.start_layer, 18);
    1165
    1
            assert_eq!(info2.end_layer, 24);
    1166
    1
        }
    1167
    1168
        #[test]
    1169
    1
        fn test_pipeline_parallel_uneven_layers() {
    1170
            // 25 layers across 4 stages: 7, 6, 6, 6
    1171
    1
            let pp = PipelineParallel::new(4, 0, 25, 4).expect("test");
    1172
    1
            assert_eq!(pp.stage_info().num_layers, 7);
    1173
    1174
    1
            let pp1 = PipelineParallel::new(4, 1, 25, 4).expect("test");
    1175
    1
            assert_eq!(pp1.stage_info().num_layers, 6);
    1176
    1
        }
    1177
    1178
        #[test]
    1179
    1
        fn test_pipeline_parallel_first_last() {
    1180
    1
            let first = PipelineParallel::new(4, 0, 24, 4).expect("test");
    1181
    1
            assert!(first.is_first_stage());
    1182
    1
            assert!(!first.is_last_stage());
    1183
    1184
    1
            let last = PipelineParallel::new(4, 3, 24, 4).expect("test");
    1185
    1
            assert!(!last.is_first_stage());
    1186
    1
            assert!(last.is_last_stage());
    1187
    1
        }
    1188
    1189
        #[test]
    1190
    1
        fn test_pipeline_parallel_bubble_ratio() {
    1191
    1
            let pp = PipelineParallel::new(4, 0, 24, 4).expect("test");
    1192
            // Bubble = (4-1) / (4 + 8 - 1) = 3/11 ≈ 0.27
    1193
    1
            let ratio = pp.bubble_ratio(8);
    1194
    1
            assert!(ratio > 0.2 && ratio < 0.4);
    1195
    1
        }
    1196
    1197
        #[test]
    1198
    1
        fn test_pipeline_parallel_stats() {
    1199
    1
            let mut pp = PipelineParallel::new(4, 0, 24, 4).expect("test");
    1200
    1
            pp.record_micro_batch(10.0);
    1201
    1
            pp.record_micro_batch(12.0);
    1202
    1203
    1
            let stats = pp.stats();
    1204
    1
            assert_eq!(stats.micro_batches_processed, 2);
    1205
    1
            assert_eq!(stats.forward_passes, 2);
    1206
    1
            assert!((stats.avg_stage_latency_ms - 11.0).abs() < 0.1);
    1207
    1
        }
    1208
    1209
        // =========================================================================
    1210
        // ZeroOffload Tests
    1211
        // =========================================================================
    1212
    1213
        #[test]
    1214
    1
        fn test_zero_offload_default() {
    1215
    1
            let zero = ZeroOffload::default();
    1216
    1
            assert!(zero.offload_optimizer);
    1217
    1
            assert!(!zero.offload_params);
    1218
    1
            assert!(zero.pin_memory);
    1219
    1
        }
    1220
    1221
        #[test]
    1222
    1
        fn test_zero_offload_inference() {
    1223
    1
            let zero = ZeroOffload::inference();
    1224
    1
            assert!(!zero.offload_optimizer);
    1225
    1
            assert!(zero.offload_params);
    1226
    1
            assert!(zero.offload_activations);
    1227
    1
        }
    1228
    1229
        #[test]
    1230
    1
        fn test_zero_offload_memory_savings() {
    1231
    1
            let zero = ZeroOffload::default();
    1232
    1
            let savings = zero.memory_savings_ratio();
    1233
    1
            assert!(savings >= 0.0 && savings <= 1.0);
    1234
    1235
    1
            let zero_inference = ZeroOffload::inference();
    1236
    1
            let savings_inference = zero_inference.memory_savings_ratio();
    1237
    1
            assert!(savings_inference > savings);
    1238
    1
        }
    1239
    1240
        // =========================================================================
    1241
        // DistributedContext Tests
    1242
        // =========================================================================
    1243
    1244
        #[test]
    1245
    1
        fn test_distributed_context_single() {
    1246
    1
            let config = ParallelConfig::single();
    1247
    1
            let ctx = DistributedContext::new(config).expect("test");
    1248
    1249
    1
            assert!(!ctx.is_distributed());
    1250
    1
            assert!(ctx.is_initialized());
    1251
    1
            assert!(ctx.tensor_parallel().is_none());
    1252
    1
            assert!(ctx.pipeline_parallel().is_none());
    1253
    1
        }
    1254
    1255
        #[test]
    1256
    1
        fn test_distributed_context_with_tp() {
    1257
    1
            let config = ParallelConfig::new(4, 1, 1, 0).expect("test");
    1258
    1
            let ctx = DistributedContext::new(config).expect("test");
    1259
    1260
    1
            assert!(ctx.is_distributed());
    1261
    1
            assert!(ctx.tensor_parallel().is_some());
    1262
    1
            assert_eq!(ctx.tensor_parallel().expect("test").tp_size(), 4);
    1263
    1
        }
    1264
    1265
        #[test]
    1266
    1
        fn test_distributed_context_init_pipeline() {
    1267
    1
            let config = ParallelConfig::new(1, 4, 1, 0).expect("test");
    1268
    1
            let mut ctx = DistributedContext::new(config).expect("test");
    1269
    1270
    1
            ctx.init_pipeline(24, 4).expect("test");
    1271
    1
            assert!(ctx.pipeline_parallel().is_some());
    1272
    1
            assert_eq!(ctx.pipeline_parallel().expect("test").num_stages(), 4);
    1273
    1
        }
    1274
    1275
        #[test]
    1276
    1
        fn test_distributed_context_zero_offload() {
    1277
    1
            let config = ParallelConfig::single();
    1278
    1
            let mut ctx = DistributedContext::new(config).expect("test");
    1279
    1280
    1
            ctx.set_zero_offload(ZeroOffload::inference());
    1281
    1
            assert!(ctx.zero_offload().offload_params);
    1282
    1
        }
    1283
    1284
        // =========================================================================
    1285
        // ReduceOp Tests
    1286
        // =========================================================================
    1287
    1288
        #[test]
    1289
    1
        fn test_reduce_op_serialization() {
    1290
    1
            let op = ReduceOp::Sum;
    1291
    1
            let json = serde_json::to_string(&op).expect("test");
    1292
    1
            let deserialized: ReduceOp = serde_json::from_str(&json).expect("test");
    1293
    1
            assert_eq!(op, deserialized);
    1294
    1
        }
    1295
    1296
        // =========================================================================
    1297
        // Error Tests
    1298
        // =========================================================================
    1299
    1300
        #[test]
    1301
    1
        fn test_parallel_error_display() {
    1302
    1
            let err = ParallelError::InvalidRank {
    1303
    1
                rank: 10,
    1304
    1
                world_size: 4,
    1305
    1
            };
    1306
    1
            assert!(err.to_string().contains("10"));
    1307
    1
            assert!(err.to_string().contains("4"));
    1308
    1309
    1
            let err2 = ParallelError::CommunicationError("timeout".to_string());
    1310
    1
            assert!(err2.to_string().contains("timeout"));
    1311
    1
        }
    1312
    1313
        // =========================================================================
    1314
        // Extended Coverage Tests: ParallelConfig
    1315
        // =========================================================================
    1316
    1317
        #[test]
    1318
    1
        fn test_parallel_config_world_size_calculation_ext_cov() {
    1319
    1
            let config = ParallelConfig::new(2, 2, 2, 0).expect("test");
    1320
    1
            assert_eq!(config.world_size, 8);
    1321
    1
        }
    1322
    1323
        #[test]
    1324
    1
        fn test_parallel_config_single_debug_ext_cov() {
    1325
    1
            let config = ParallelConfig::single();
    1326
    1
            let debug_str = format!("{:?}", config);
    1327
    1
            assert!(debug_str.contains("tp_size"));
    1328
    1
            assert!(debug_str.contains("pp_size"));
    1329
    1
        }
    1330
    1331
        #[test]
    1332
    1
        fn test_parallel_config_invalid_zero_tp_ext_cov() {
    1333
    1
            let result = ParallelConfig::new(0, 1, 1, 0);
    1334
    1
            assert!(result.is_err());
    1335
    1
        }
    1336
    1337
        #[test]
    1338
    1
        fn test_parallel_config_invalid_zero_pp_ext_cov() {
    1339
    1
            let result = ParallelConfig::new(1, 0, 1, 0);
    1340
    1
            assert!(result.is_err());
    1341
    1
        }
    1342
    1343
        #[test]
    1344
    1
        fn test_parallel_config_invalid_zero_dp_ext_cov() {
    1345
    1
            let result = ParallelConfig::new(1, 1, 0, 0);
    1346
    1
            assert!(result.is_err());
    1347
    1
        }
    1348
    1349
        // =========================================================================
    1350
        // Extended Coverage Tests: ReduceOp
    1351
        // =========================================================================
    1352
    1353
        #[test]
    1354
    1
        fn test_reduce_op_all_variants_ext_cov() {
    1355
    1
            let ops = [ReduceOp::Sum, ReduceOp::Min, ReduceOp::Max, ReduceOp::Avg];
    1356
    5
            for op4 in ops {
    1357
    4
                let json = serde_json::to_string(&op).expect("serialize");
    1358
    4
                let _: ReduceOp = serde_json::from_str(&json).expect("deserialize");
    1359
    4
            }
    1360
    1
        }
    1361
    1362
        #[test]
    1363
    1
        fn test_reduce_op_clone_ext_cov() {
    1364
    1
            let op = ReduceOp::Max;
    1365
    1
            let cloned = op;
    1366
    1
            assert_eq!(op, cloned);
    1367
    1
        }
    1368
    1369
        #[test]
    1370
    1
        fn test_reduce_op_debug_ext_cov() {
    1371
    1
            let op = ReduceOp::Avg;
    1372
    1
            let debug_str = format!("{:?}", op);
    1373
    1
            assert!(debug_str.contains("Avg"));
    1374
    1
        }
    1375
    1376
        // =========================================================================
    1377
        // Extended Coverage Tests: ParallelError
    1378
        // =========================================================================
    1379
    1380
        #[test]
    1381
    1
        fn test_parallel_error_all_variants_ext_cov() {
    1382
    1
            let errors: [ParallelError; 5] = [
    1383
    1
                ParallelError::InvalidRank {
    1384
    1
                    rank: 5,
    1385
    1
                    world_size: 4,
    1386
    1
                },
    1387
    1
                ParallelError::InvalidWorldSize(0),
    1388
    1
                ParallelError::CommunicationError("timeout".to_string()),
    1389
    1
                ParallelError::ShapeMismatch {
    1390
    1
                    expected: vec![2, 3],
    1391
    1
                    got: vec![3, 2],
    1392
    1
                },
    1393
    1
                ParallelError::PipelineError("stage error".to_string()),
    1394
    1
            ];
    1395
    6
            for err5 in errors {
    1396
    5
                let _ = err.to_string();
    1397
    5
            }
    1398
    1
        }
    1399
    1400
        #[test]
    1401
    1
        fn test_parallel_error_shape_mismatch_ext_cov() {
    1402
    1
            let err = ParallelError::ShapeMismatch {
    1403
    1
                expected: vec![10, 20],
    1404
    1
                got: vec![20, 10],
    1405
    1
            };
    1406
    1
            let msg = err.to_string();
    1407
    1
            assert!(msg.contains("10") || msg.contains("20")0);
    1408
    1
        }
    1409
    1410
        #[test]
    1411
    1
        fn test_parallel_error_debug_ext_cov() {
    1412
    1
            let err = ParallelError::NotInitialized;
    1413
    1
            let debug_str = format!("{:?}", err);
    1414
    1
            assert!(debug_str.contains("NotInitialized"));
    1415
    1
        }
    1416
    1417
        // =========================================================================
    1418
        // Extended Coverage Tests: ZeroOffload
    1419
        // =========================================================================
    1420
    1421
        #[test]
    1422
    1
        fn test_zero_offload_clone_ext_cov() {
    1423
    1
            let zero = ZeroOffload::inference();
    1424
    1
            let cloned = zero.clone();
    1425
    1
            assert_eq!(zero.offload_params, cloned.offload_params);
    1426
    1
            assert_eq!(zero.offload_activations, cloned.offload_activations);
    1427
    1
        }
    1428
    1429
        #[test]
    1430
    1
        fn test_zero_offload_debug_ext_cov() {
    1431
    1
            let zero = ZeroOffload::default();
    1432
    1
            let debug_str = format!("{:?}", zero);
    1433
    1
            assert!(debug_str.contains("offload_optimizer"));
    1434
    1
            assert!(debug_str.contains("pin_memory"));
    1435
    1
        }
    1436
    1437
        #[test]
    1438
    1
        fn test_zero_offload_memory_savings_extremes_ext_cov() {
    1439
            // Test with all offload options enabled
    1440
    1
            let full_offload = ZeroOffload {
    1441
    1
                offload_optimizer: true,
    1442
    1
                offload_params: true,
    1443
    1
                offload_activations: true,
    1444
    1
                pin_memory: true,
    1445
    1
                overlap_comm: true,
    1446
    1
            };
    1447
    1
            let savings = full_offload.memory_savings_ratio();
    1448
    1
            assert!(savings >= 0.0);
    1449
    1
            assert!(savings <= 1.0);
    1450
    1451
            // Test with no offload
    1452
    1
            let no_offload = ZeroOffload {
    1453
    1
                offload_optimizer: false,
    1454
    1
                offload_params: false,
    1455
    1
                offload_activations: false,
    1456
    1
                pin_memory: false,
    1457
    1
                overlap_comm: false,
    1458
    1
            };
    1459
    1
            let no_savings = no_offload.memory_savings_ratio();
    1460
    1
            assert!(no_savings >= 0.0);
    1461
    1
            assert!(no_savings < savings);
    1462
    1
        }
    1463
    1464
        // =========================================================================
    1465
        // Extended Coverage Tests: PipelineStats
    1466
        // =========================================================================
    1467
    1468
        #[test]
    1469
    1
        fn test_pipeline_stats_clone_debug_ext_cov() {
    1470
    1
            let stats = PipelineStats {
    1471
    1
                micro_batches_processed: 100,
    1472
    1
                forward_passes: 100,
    1473
    1
                bubble_time_ms: 5.0,
    1474
    1
                avg_stage_latency_ms: 10.5,
    1475
    1
            };
    1476
    1
            let cloned = stats.clone();
    1477
    1
            assert_eq!(
    1478
                stats.micro_batches_processed,
    1479
                cloned.micro_batches_processed
    1480
            );
    1481
    1482
    1
            let debug_str = format!("{:?}", stats);
    1483
    1
            assert!(debug_str.contains("micro_batches_processed"));
    1484
    1
            assert!(debug_str.contains("bubble_time_ms"));
    1485
    1
        }
    1486
    1487
        // =========================================================================
    1488
        // Extended Coverage Tests: TensorParallel
    1489
        // =========================================================================
    1490
    1491
        #[test]
    1492
    1
        fn test_tensor_parallel_chunk_size_ext_cov() {
    1493
    1
            let tp = TensorParallel::new(4, 0).expect("test");
    1494
    1
            let chunk = tp.chunk_size(1000);
    1495
    1
            assert_eq!(chunk, 250); // 1000 / 4 = 250
    1496
    1
        }
    1497
    1498
        #[test]
    1499
    1
        fn test_tensor_parallel_debug_ext_cov() {
    1500
    1
            let tp = TensorParallel::new(8, 2).expect("test");
    1501
    1
            let debug_str = format!("{:?}", tp);
    1502
    1
            assert!(debug_str.contains("tp_size"));
    1503
    1
            assert!(debug_str.contains("rank"));
    1504
    1
        }
    1505
    1506
        #[test]
    1507
    1
        fn test_tensor_parallel_invalid_rank_ext_cov() {
    1508
    1
            let result = TensorParallel::new(4, 10);
    1509
    1
            assert!(result.is_err());
    1510
    1
        }
    1511
    1512
        #[test]
    1513
    1
        fn test_tensor_parallel_invalid_size_ext_cov() {
    1514
    1
            let result = TensorParallel::new(0, 0);
    1515
    1
            assert!(result.is_err());
    1516
    1
        }
    1517
    1518
        // =========================================================================
    1519
        // Extended Coverage Tests: Communicator
    1520
        // =========================================================================
    1521
    1522
        #[test]
    1523
    1
        fn test_communicator_debug_ext_cov() {
    1524
    1
            let comm = Communicator::new(4, 0).expect("test");
    1525
    1
            let debug_str = format!("{:?}", comm);
    1526
    1
            assert!(debug_str.contains("world_size"));
    1527
    1
            assert!(debug_str.contains("rank"));
    1528
    1
        }
    1529
    1530
        #[test]
    1531
    1
        fn test_communicator_invalid_rank_ext_cov() {
    1532
    1
            let result = Communicator::new(4, 10);
    1533
    1
            assert!(result.is_err());
    1534
    1
        }
    1535
    1536
        // =========================================================================
    1537
        // Extended Coverage Tests: PipelineParallel
    1538
        // =========================================================================
    1539
    1540
        #[test]
    1541
    1
        fn test_pipeline_parallel_stage_info_ext_cov() {
    1542
            // PipelineParallel::new(pp_size, stage, total_layers, micro_batch_size)
    1543
    1
            let pp = PipelineParallel::new(4, 0, 24, 4).expect("test");
    1544
    1
            let info = pp.stage_info();
    1545
    1
            assert_eq!(info.start_layer, 0);
    1546
    1
            assert_eq!(info.num_layers, 6); // 24 / 4 = 6 layers per stage
    1547
    1
        }
    1548
    1549
        #[test]
    1550
    1
        fn test_pipeline_parallel_debug_ext_cov() {
    1551
    1
            let pp = PipelineParallel::new(4, 0, 24, 4).expect("test");
    1552
    1
            let debug_str = format!("{:?}", pp);
    1553
    1
            assert!(debug_str.contains("pp_size"));
    1554
    1
            assert!(debug_str.contains("stage"));
    1555
    1
        }
    1556
    1557
        #[test]
    1558
    1
        fn test_pipeline_parallel_micro_batch_size_ext_cov() {
    1559
    1
            let pp = PipelineParallel::new(4, 0, 24, 8).expect("test");
    1560
    1
            assert_eq!(pp.micro_batch_size(), 8);
    1561
    1
        }
    1562
    1563
        #[test]
    1564
    1
        fn test_pipeline_parallel_first_last_stage_ext_cov() {
    1565
    1
            let first = PipelineParallel::new(4, 0, 24, 4).expect("test");
    1566
    1
            let last = PipelineParallel::new(4, 3, 24, 4).expect("test");
    1567
    1
            let middle = PipelineParallel::new(4, 1, 24, 4).expect("test");
    1568
    1569
    1
            assert!(first.is_first_stage());
    1570
    1
            assert!(!first.is_last_stage());
    1571
    1572
    1
            assert!(!last.is_first_stage());
    1573
    1
            assert!(last.is_last_stage());
    1574
    1575
    1
            assert!(!middle.is_first_stage());
    1576
    1
            assert!(!middle.is_last_stage());
    1577
    1
        }
    1578
    }