Struct PipelineParallelCoordinator 

pub struct PipelineParallelCoordinator { /* private fields */ }

Pipeline parallelism coordinator.

Implementations

impl PipelineParallelCoordinator

pub fn new( config: DistributedConfig, backend: Arc<RwLock<dyn CommunicationBackend>>, num_stages: usize, ) -> Self

Create a new pipeline parallel coordinator.

pub fn set_micro_batch_size(&mut self, size: usize)

Set micro-batch size for pipeline parallelism.

pub fn stage_for_rank(&self, rank: usize) -> usize

Get the stage assignment for this rank.

pub fn send_activations(&self, stage: usize) -> Result<(), ExecutorError>

Send activations to next stage.

pub fn recv_activations(&self, stage: usize) -> Result<(), ExecutorError>

Receive activations from previous stage.

pub fn send_gradients(&self, stage: usize) -> Result<(), ExecutorError>

Send gradients to previous stage.

pub fn recv_gradients(&self, stage: usize) -> Result<(), ExecutorError>

Receive gradients from next stage.

pub fn num_stages(&self) -> usize

Get the number of stages in the pipeline.

pub fn micro_batch_size(&self) -> usize

Get the micro-batch size.

pub fn config(&self) -> &DistributedConfig

Get the configuration.