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 and buffers to a device.

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

Export parameters and buffers as a state dict (torch parity).

fn buffers(&self) -> Vec<&Buffer<T>>

Iterate over all non-trainable buffers (e.g. running mean / variance in BatchNorm). Default returns empty — concrete modules with buffers override.

fn buffers_mut(&mut self) -> Vec<&mut Buffer<T>>

Mutable iteration over all buffers. Default returns empty.

fn named_buffers(&self) -> Vec<(String, &Buffer<T>)>

Named buffers (dot-separated paths for nested modules). Default returns empty.

fn as_any(&self) -> Option<&(dyn Any + 'static)>

Downcast hook for type-erased buffer-loader dispatch. (#984)

fn children(&self) -> Vec<&dyn Module<T>>

Direct child modules. Default returns empty (leaf module).

fn named_children(&self) -> Vec<(String, &dyn Module<T>)>

Direct child modules with their attribute names. Default returns empty.

fn modules(&self) -> Vec<&dyn Module<T>>

where Self: Sized,

All modules in this subtree, depth-first (self first, then each child’s descendants in order).

fn descendants_dyn(&self) -> Vec<&dyn Module<T>>

All strict descendants of self in depth-first order. Object-safe.

fn named_modules(&self) -> Vec<(String, &dyn Module<T>)>

where Self: Sized,

All modules in this subtree with dot-separated path names. The root is named ""; children paths are joined with ..

fn named_descendants_dyn(&self) -> Vec<(String, &dyn Module<T>)>

Strict descendants with dot-paths. Object-safe.

fn with_forward_hook( self, hook: Box<dyn Fn(&Tensor<T>, &Tensor<T>) + Sync + Send>, ) -> (HookedModule<Self, T>, HookHandle)

where Self: Sized,

Wrap this module in a HookedModule and register a forward hook. Returns the wrapper paired with a HookHandle that can be used to remove the hook later. The wrapper implements Module<T> itself, so it slots into any place the original module did. Mirrors torch.nn.Module.register_forward_hook.

fn with_forward_pre_hook( self, hook: Box<dyn Fn(&Tensor<T>) -> Result<Tensor<T>, FerrotorchError> + Sync + Send>, ) -> (HookedModule<Self, T>, HookHandle)

where Self: Sized,

Wrap this module in a HookedModule and register a forward pre-hook. See Self::with_forward_hook. Mirrors torch.nn.Module.register_forward_pre_hook.

fn with_backward_hook( self, hook: Box<dyn Fn(&Tensor<T>, &Tensor<T>) + Sync + Send>, ) -> (HookedModule<Self, T>, HookHandle)

where Self: Sized,

Wrap this module in a HookedModule and register a backward hook. See Self::with_forward_hook. Mirrors torch.nn.Module.register_backward_hook.

fn zero_grad(&self) -> Result<(), FerrotorchError>

Set the gradient of every parameter to None.

fn requires_grad_(&mut self, requires_grad: bool)

Toggle requires_grad on every parameter (freeze / unfreeze the module). Mirrors torch.nn.Module.requires_grad_.

fn apply_to_parameters(&mut self, f: &mut dyn FnMut(&mut Parameter<T>))

Apply a function to every parameter in this module. Mirrors torch.nn.Module.apply for the parameter case (true apply recurses over all submodules; the recursive form requires &mut dyn Module which conflicts with this trait’s &mut self borrow).

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

Load parameters from a state dict.