Struct DistributedDataParallel 

pub struct DistributedDataParallel { /* private fields */ }

Distributed Data Parallel model wrapper for multi-node training

Implementations

impl DistributedDataParallel

pub fn new( model: Box<dyn Model<f32>>, device: Device, process_group: Arc<CommunicationGroup>, comm_runtime: Arc<RwLock<CommunicationRuntime>>, config: DDPConfig, ) -> Result<Self>

Create new DistributedDataParallel model wrapper

pub fn sync_gradients(&mut self) -> Result<()>

Perform gradient synchronization using all-reduce

pub fn process_group(&self) -> &CommunicationGroup

Get process group information

pub fn local_rank(&self) -> usize

Get local rank within process group

pub fn world_size(&self) -> usize

Get world size (total number of processes)

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

Set bucket size for gradient communication

Trait Implementations

impl Model<f32> for DistributedDataParallel

fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>>

fn parameters(&self) -> Vec<&Tensor<f32>>

fn parameters_mut(&mut self) -> Vec<&mut Tensor<f32>>

fn set_training(&mut self, training: bool)

fn zero_grad(&mut self)

fn as_any(&self) -> &dyn Any

Provide access to Any for downcasting

fn extract_features(&self, input: &Tensor<T>) -> Result<Option<Vec<Tensor<T>>>>

Extract intermediate features for knowledge distillation Returns None if the model doesn’t support feature extraction