rustytorch/crates/training/rtx-distributed/src/pipeline_parallel/scheduler.rs

//! Pipeline scheduler for generating execution schedules.

use parking_lot::RwLock;

use super::config::{PipelineConfig, PipelineSchedule};

/// Scheduled operation in the pipeline
#[derive(Debug, Clone)]
pub enum PipelineOp {
    /// Forward pass on a stage
    Forward { stage: usize, micro_batch_id: u64 },
    /// Backward pass on a stage
    Backward { stage: usize, micro_batch_id: u64 },
    /// Send activations to next stage
    SendActivation {
        from_stage: usize,
        to_stage: usize,
        micro_batch_id: u64,
    },
    /// Receive activations from previous stage
    RecvActivation {
        from_stage: usize,
        to_stage: usize,
        micro_batch_id: u64,
    },
    /// Send gradients to previous stage
    SendGradient {
        from_stage: usize,
        to_stage: usize,
        micro_batch_id: u64,
    },
    /// Receive gradients from next stage
    RecvGradient {
        from_stage: usize,
        to_stage: usize,
        micro_batch_id: u64,
    },
    /// Barrier synchronization
    Barrier,
}

/// Pipeline schedule generator
pub struct PipelineScheduler {
    /// Configuration
    config: PipelineConfig,
    /// Generated schedule
    schedule: RwLock<Vec<Vec<PipelineOp>>>,
}

impl PipelineScheduler {
    /// Create a new scheduler
    pub fn new(config: PipelineConfig) -> Self {
        let scheduler = Self {
            config,
            schedule: RwLock::new(Vec::new()),
        };
        scheduler.generate_schedule();
        scheduler
    }

    /// Generate the pipeline schedule
    fn generate_schedule(&self) {
        let schedule = match self.config.schedule {
            PipelineSchedule::GPipe => self.generate_gpipe_schedule(),
            PipelineSchedule::OneFOneBAsync => self.generate_1f1b_schedule(),
            PipelineSchedule::Interleaved => self.generate_interleaved_schedule(),
            PipelineSchedule::Chimera => self.generate_chimera_schedule(),
        };
        *self.schedule.write() = schedule;
    }

    /// Generate GPipe schedule (all forwards, then all backwards)
    fn generate_gpipe_schedule(&self) -> Vec<Vec<PipelineOp>> {
        let num_stages = self.config.num_stages;
        let num_micro_batches = self.config.num_micro_batches;
        let mut schedule = Vec::new();

        // Forward passes
        for mb in 0..num_micro_batches {
            for stage in 0..num_stages {
                let mut step = vec![PipelineOp::Forward {
                    stage,
                    micro_batch_id: mb as u64,
                }];

                if stage < num_stages - 1 {
                    step.push(PipelineOp::SendActivation {
                        from_stage: stage,
                        to_stage: stage + 1,
                        micro_batch_id: mb as u64,
                    });
                }

                schedule.push(step);
            }
        }

        // Backward passes (reverse order)
        for mb in (0..num_micro_batches).rev() {
            for stage in (0..num_stages).rev() {
                let mut step = vec![PipelineOp::Backward {
                    stage,
                    micro_batch_id: mb as u64,
                }];

                if stage > 0 {
                    step.push(PipelineOp::SendGradient {
                        from_stage: stage,
                        to_stage: stage - 1,
                        micro_batch_id: mb as u64,
                    });
                }

                schedule.push(step);
            }
        }

        schedule.push(vec![PipelineOp::Barrier]);
        schedule
    }

    /// Generate 1F1B schedule (one forward, one backward in steady state)
    fn generate_1f1b_schedule(&self) -> Vec<Vec<PipelineOp>> {
        let num_stages = self.config.num_stages;
        let num_micro_batches = self.config.num_micro_batches;
        let mut schedule = Vec::new();

        // Warmup: fill the pipeline with forwards
        for mb in 0..num_stages {
            for stage in 0..=mb.min(num_stages - 1) {
                schedule.push(vec![PipelineOp::Forward {
                    stage,
                    micro_batch_id: mb as u64,
                }]);
            }
        }

        // Steady state: 1F1B
        for mb in num_stages..num_micro_batches {
            // Backward for earlier micro-batch
            let backward_mb = (mb - num_stages) as u64;
            for stage in (0..num_stages).rev() {
                schedule.push(vec![PipelineOp::Backward {
                    stage,
                    micro_batch_id: backward_mb,
                }]);
            }

            // Forward for current micro-batch
            for stage in 0..num_stages {
                schedule.push(vec![PipelineOp::Forward {
                    stage,
                    micro_batch_id: mb as u64,
                }]);
            }
        }

        // Cooldown: drain remaining backwards
        for mb in (num_micro_batches - num_stages)..num_micro_batches {
            for stage in (0..num_stages).rev() {
                schedule.push(vec![PipelineOp::Backward {
                    stage,
                    micro_batch_id: mb as u64,
                }]);
            }
        }

        schedule.push(vec![PipelineOp::Barrier]);
        schedule
    }

    /// Generate interleaved 1F1B schedule (Megatron-LM, arXiv:2104.04473).
    ///
    /// # Virtual stage mapping
    ///
    /// With `p` physical stages and `m` virtual stages per device the full
    /// virtual stage space is `[0, p*m)`.  Physical rank `r` owns virtual
    /// stages `r, r+p, r+2p, …, r+(m-1)*p`.  A micro-batch in its `k`-th
    /// pass through the pipeline traverses virtual stage `k*p + r` on rank `r`.
    ///
    /// # Schedule phases
    ///
    /// 1. **Warmup** – issue `warmup_micros = min((p-1) + (m-1)*p, num_micro)`
    ///    forward ops to fill the pipeline across all virtual stages before any
    ///    backward work begins.
    /// 2. **Steady state** – alternate one backward then one forward (1F1B)
    ///    until all micro-batches have been forwarded.
    /// 3. **Drain** – issue remaining backward ops.
    ///
    /// The theoretical bubble fraction is `(p-1)/(p*m)` vs `(p-1)/p` for
    /// standard 1F1B.
    fn generate_interleaved_schedule(&self) -> Vec<Vec<PipelineOp>> {
        let p = self.config.num_stages;
        // Enforce minimum of 2 virtual stages; mirrors the validate() clamp.
        let m = self.config.num_virtual_stages.max(2);
        let num_micro = self.config.num_micro_batches;
        let rank = self.config.rank;

        // Total virtual stage count across the whole pipeline.
        let total_virtual = p * m;

        // Helper: given a micro-batch index `mb` compute which virtual stage
        // it occupies on this rank.  We cycle through the m virtual stages
        // owned by `rank` in order: r, r+p, r+2p, …  The k-th chunk (0-based)
        // maps to virtual stage `k*p + rank`.
        let virtual_stage = |mb: usize| -> usize { (mb % m) * p + rank };

        let warmup_micros = ((p - 1) + (m - 1) * p).min(num_micro);

        let mut schedule: Vec<Vec<PipelineOp>> = Vec::with_capacity(num_micro * 2 + 1);
        let mut fwd_done: usize = 0;
        let mut bwd_done: usize = 0;

        // ── Phase 1: Warmup ───────────────────────────────────────────────
        for mb in 0..warmup_micros {
            let vs = virtual_stage(mb);
            let mut step = vec![PipelineOp::Forward {
                stage: vs,
                micro_batch_id: mb as u64,
            }];
            if vs + 1 < total_virtual {
                step.push(PipelineOp::SendActivation {
                    from_stage: vs,
                    to_stage: vs + 1,
                    micro_batch_id: mb as u64,
                });
            }
            schedule.push(step);
            fwd_done += 1;
        }

        // ── Phase 2: Steady state 1F1B ────────────────────────────────────
        while fwd_done < num_micro || bwd_done < fwd_done {
            let mut step = Vec::new();

            // Backward first (when one is available to overlap with the next
            // forward send, reducing pipeline bubble).
            if bwd_done < fwd_done {
                let mb = bwd_done;
                let vs = virtual_stage(mb);
                step.push(PipelineOp::Backward {
                    stage: vs,
                    micro_batch_id: mb as u64,
                });
                if vs > 0 {
                    step.push(PipelineOp::SendGradient {
                        from_stage: vs,
                        to_stage: vs - 1,
                        micro_batch_id: mb as u64,
                    });
                }
                bwd_done += 1;
            }

            // Then forward (if still remaining).
            if fwd_done < num_micro {
                let mb = fwd_done;
                let vs = virtual_stage(mb);
                step.push(PipelineOp::Forward {
                    stage: vs,
                    micro_batch_id: mb as u64,
                });
                if vs + 1 < total_virtual {
                    step.push(PipelineOp::SendActivation {
                        from_stage: vs,
                        to_stage: vs + 1,
                        micro_batch_id: mb as u64,
                    });
                }
                fwd_done += 1;
            }

            if step.is_empty() {
                break;
            }
            schedule.push(step);
        }

        // ── Phase 3: Drain remaining backwards ───────────────────────────
        while bwd_done < num_micro {
            let mb = bwd_done;
            let vs = virtual_stage(mb);
            let mut step = vec![PipelineOp::Backward {
                stage: vs,
                micro_batch_id: mb as u64,
            }];
            if vs > 0 {
                step.push(PipelineOp::SendGradient {
                    from_stage: vs,
                    to_stage: vs - 1,
                    micro_batch_id: mb as u64,
                });
            }
            schedule.push(step);
            bwd_done += 1;
        }

        schedule.push(vec![PipelineOp::Barrier]);
        schedule
    }

    /// Generate Chimera schedule (bidirectional)
    fn generate_chimera_schedule(&self) -> Vec<Vec<PipelineOp>> {
        // Simplified - in practice would have bidirectional pipelines
        self.generate_1f1b_schedule()
    }

    /// Get the generated schedule
    pub fn schedule(&self) -> Vec<Vec<PipelineOp>> {
        self.schedule.read().clone()
    }

    /// Get number of steps in schedule
    pub fn num_steps(&self) -> usize {
        self.schedule.read().len()
    }

    /// Theoretical pipeline bubble fraction for the configured schedule.
    ///
    /// | Schedule        | Bubble fraction         |
    /// |-----------------|-------------------------|
    /// | GPipe           | `(p-1) / p`             |
    /// | 1F1B (async)    | `(p-1) / p`             |
    /// | Interleaved     | `(p-1) / (p * m)`       |
    /// | Chimera         | `(p-1) / p`  (approx.)  |
    ///
    /// where `p = num_stages` and `m = num_virtual_stages.max(2)`.
    pub fn bubble_ratio(&self) -> f64 {
        let p = self.config.num_stages as f64;
        match self.config.schedule {
            PipelineSchedule::Interleaved => {
                let m = self.config.num_virtual_stages.max(2) as f64;
                (p - 1.0) / (p * m)
            }
            // GPipe, OneFOneBAsync, Chimera all have the standard bubble.
            _ => (p - 1.0) / p,
        }
    }
}