Struct DDP 

pub struct DDP<M: Module<T>, T: Float> { /* private fields */ }

Distributed Data Parallel module wrapper.

Wraps an inner Module and provides [sync_gradients] to allreduce parameter gradients across all ranks. Parameters are grouped into buckets (default 25 MB) for efficient communication.

let ddp = DDP::new(model, backend);

loop {
    let output = ddp.module().forward(&input)?;
    let loss = criterion.forward(&output, &target)?;
    ferrotorch_core::backward(&loss)?;
    ddp.sync_gradients()?;
    optimizer.step()?;
    optimizer.zero_grad()?;
}

Implementations

impl<M: Module<T>, T: Float> DDP<M, T>

pub fn new(module: M, backend: Arc<dyn Backend>) -> Self

Wrap a module for distributed data-parallel training.

Parameters are assigned to ~25 MB gradient buckets in reverse order (matching PyTorch’s convention — backward computes gradients in reverse parameter order, so the first bucket fills first).

pub fn with_bucket_size( module: M, backend: Arc<dyn Backend>, bucket_size_bytes: usize, ) -> Self

Wrap a module with a custom bucket size (in bytes).

pub fn module(&self) -> &M

Immutable access to the inner module (for forward pass, etc.).

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

Mutable access to the inner module (for train/eval mode, etc.).

pub fn into_inner(self) -> M

Consume the DDP wrapper and return the inner module.

pub fn backend(&self) -> &Arc<dyn Backend>

The backend used for communication.

pub fn sync_gradients(&self) -> FerrotorchResult<()>

Allreduce parameter gradients across ranks using gradient bucketing.

Parameters are grouped into ~25 MB buckets. Each bucket is allreduced independently as a single flat buffer. This enables future overlapped communication where the first bucket can start transferring while backward is still computing later gradients.

Call this after backward() and before optimizer.step().

pub fn overlapped_sync_gradients(&self) -> FerrotorchResult<()>

Allreduce parameter gradients with bucket-level parallelism.

Like [sync_gradients], but processes all buckets concurrently using std::thread::scope. Each bucket’s allreduce runs in its own thread, overlapping communication across buckets. All threads complete before this method returns.

This provides communication/computation overlap when backward and sync run on different threads, and communication overlap across buckets even in the synchronous case.

pub fn broadcast_parameters(&mut self, root: usize) -> FerrotorchResult<()>

Broadcast model parameters from root rank to all other ranks.

Ensures all ranks start with identical weights. Call once before the training loop begins.

Warning

This replaces the Parameter objects in the module. Any optimizer that holds references to the old parameters must be re-initialized after calling this method, otherwise optimizer state (momentum, adaptive learning rates, etc.) will refer to stale parameters.

Trait Implementations

impl<M: Module<T>, T: Float> Module<T> for DDP<M, T>

fn forward(&self, input: &Tensor<T>) -> FerrotorchResult<Tensor<T>>

Forward pass. Takes input tensor, returns output tensor.

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

Iterate over all learnable parameters.

fn parameters_mut(&mut self) -> Vec<&mut Parameter<T>>

Iterate over all learnable parameters mutably.

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

Named parameters for state dict serialization.

fn train(&mut self)

Set training mode. Affects dropout, batchnorm, etc.

fn eval(&mut self)

Set evaluation mode.

fn is_training(&self) -> bool

Whether the module is in training mode.

fn to_device(&mut self, device: Device) -> Result<(), FerrotorchError>

Move all parameters to a device.

fn state_dict(&self) -> HashMap<String, Tensor<T>>

Export parameters as a state dict.

fn load_state_dict( &mut self, state: &HashMap<String, Tensor<T>>, strict: bool, ) -> Result<(), FerrotorchError>

Load parameters from a state dict.