#[derive(Debug, Clone, PartialEq, Default)]
pub enum CompositionMode {
#[default]
None,
DataParallel { shards: usize },
ModelParallel { stages: usize },
Pipeline { depth: usize, overlap: bool },
Batch { batch_size: usize, prefetch: usize },
}
impl CompositionMode {
pub fn data_parallel(shards: usize) -> Self {
CompositionMode::DataParallel { shards }
}
pub fn batch(size: usize) -> Self {
CompositionMode::Batch {
batch_size: size,
prefetch: 2,
}
}
pub fn pipeline(depth: usize) -> Self {
CompositionMode::Pipeline {
depth,
overlap: true,
}
}
}