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 to ensure all ranks start from identical parameters (broadcast from rank 0).

pub fn sync_gradients(&self)

Synchronizes gradients across all processes. Should be called after the backward pass. All-reduces gradients so every rank gets the average gradient across all ranks.

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. Sets the training mode.

fn zero_grad(&self)

Zeros all gradients of parameters.

fn to_device(&self, device: Device)

Moves all parameters to the specified device.

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

Returns the module name for debugging.