Initial commit

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

@@ -0,0 +1,193 @@
+//! 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 schedule (reduced bubble)
+    fn generate_interleaved_schedule(&self) -> Vec<Vec<PipelineOp>> {
+        // Simplified interleaved - in practice would have virtual stages
+        self.generate_1f1b_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()
+    }
+}