Struct DistributedDataParallel 

pub struct DistributedDataParallel<M: Module> { /* private fields */ }

Wrapper that enables distributed data parallel training.

DDP replicates the model across multiple processes and synchronizes gradients during the backward pass.

Implementations

impl<M: Module> DistributedDataParallel<M>

pub fn new(module: M, process_group: ProcessGroup) -> Self

Creates a new DDP wrapper.

pub fn broadcast_buffers(self, broadcast: bool) -> Self

Sets whether to broadcast buffers from rank 0.

pub fn gradient_as_bucket_view(self, bucket_view: bool) -> Self

Sets whether to use gradient bucketing.

pub fn module(&self) -> &M

Returns a reference to the underlying module.

pub fn module_mut(&mut self) -> &mut M

Returns a mutable reference to the underlying module.

pub fn process_group(&self) -> &ProcessGroup

Returns the process group.

pub fn sync_parameters(&mut self)

Synchronizes model parameters across all processes. Should be called once at the start of training.

pub fn sync_gradients(&self)

Synchronizes gradients across all processes. Should be called after the backward pass.

pub fn forward(&self, input: &Variable) -> Variable

Performs forward pass with gradient synchronization.

Trait Implementations

impl<M: Module> Module for DistributedDataParallel<M>

fn forward(&self, input: &Variable) -> Variable

Performs the forward pass.

fn parameters(&self) -> Vec<Parameter>

Returns all parameters of this module.

fn train(&mut self)

Sets the module to training mode.

fn eval(&mut self)

Sets the module to evaluation mode.

fn is_training(&self) -> bool

Returns whether the module is in training mode.

fn named_parameters(&self) -> HashMap<String, Parameter>

Returns named parameters of this module.

fn num_parameters(&self) -> usize

Returns the number of trainable parameters.

fn set_training(&mut self, _training: bool)

Sets the training mode.

fn zero_grad(&self)

Zeros all gradients of parameters.

fn name(&self) -> &'static str

Returns the module name for debugging.