Struct DataParallel 

pub struct DataParallel { /* private fields */ }

Data parallel training

Replicates the model across GPUs and splits data batches. Gradients are averaged across all GPUs after backward pass.

Implementations

impl DataParallel

pub fn new(config: DistributedConfig, device_ids: Vec<usize>) -> Self

pub fn split_batch(&self, batch: &Tensor) -> Vec<Tensor>

Split batch across devices

pub fn all_reduce_gradients( &self, gradients: &HashMap<String, Tensor>, ) -> HashMap<String, Tensor>

All-reduce gradients across devices

pub fn broadcast_parameters( &self, parameters: &HashMap<String, Tensor>, ) -> HashMap<String, Tensor>

Broadcast model parameters from rank 0 to all ranks