Struct PipelineParallel 

pub struct PipelineParallel<B: Backend> { /* private fields */ }

Pipeline-parallel executor using GPipe-style micro-batching.

Splits a model into sequential stages and processes micro-batches through the pipeline. This increases throughput by overlapping computation across stages.

Example

let stage0 = PipelineStage::new(0)
    .add_layer(Box::new(Linear::new(784, 256, true, DType::F32, &dev)?));
let stage1 = PipelineStage::new(1)
    .add_layer(Box::new(Linear::new(256, 10, true, DType::F32, &dev)?));

let pipeline = PipelineParallel::new(vec![stage0, stage1], 4);
let output = pipeline.forward(&input)?;

Implementations

impl<B: Backend> PipelineParallel<B>

pub fn new(stages: Vec<PipelineStage<B>>, num_micro_batches: usize) -> Self

Create a pipeline with the given stages and micro-batch count.

pub fn forward(&self, x: &Tensor<B>) -> Result<Tensor<B>>

Full forward pass through all stages.

Splits the input into num_micro_batches micro-batches, runs each through the pipeline sequentially, and concatenates the results.

In a multi-device setting, stages would run on different devices with inter-device transfers between stages.

pub fn parameters(&self) -> Vec<Tensor<B>>

Collect all parameters from all stages (for optimizer).

pub fn num_stages(&self) -> usize

Number of stages.

pub fn stage(&self, idx: usize) -> Option<&PipelineStage<B>>

Get a reference to a specific stage.