Struct ModelParallelContext 

pub struct ModelParallelContext { /* private fields */ }

Model parallel context managing distributed execution

Implementations

impl ModelParallelContext

pub fn new(config: ModelParallelConfig) -> Result<Self>

Initialize model parallel context

pub fn rank(&self) -> usize

Get current process rank

pub fn world_size(&self) -> usize

Get total world size

pub fn partition_tensor( &self, tensor: &Tensor, split_dim: usize, ) -> Result<DistributedTensor>

Partition a tensor across devices

pub fn all_gather(&self, distributed: &DistributedTensor) -> Result<Tensor>

Gather distributed tensor to full tensor

pub fn reduce_scatter( &self, tensor: &Tensor, split_dim: usize, ) -> Result<Tensor>

Reduce scattered tensor across devices

pub fn all_reduce(&self, tensor: &mut Tensor) -> Result<()>

All-reduce operation for gradient synchronization

pub fn broadcast(&self, tensor: &mut Tensor, root: usize) -> Result<()>

Broadcast tensor from root rank to all other ranks