Struct DistributedDataParallel 

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

Distributed Data Parallel wrapper for models

Implementations

impl<M: Module> DistributedDataParallel<M>

pub fn new( module: M, process_group: Arc<ProcessGroup>, device_ids: Vec<usize>, output_device: Option<usize>, broadcast_buffers: bool, bucket_cap_mb: f32, ) -> TorshResult<Self>

Create a new DDP wrapper

pub fn new_with_bucket_config( module: M, process_group: Arc<ProcessGroup>, device_ids: Vec<usize>, output_device: Option<usize>, broadcast_buffers: bool, bucket_config: BucketConfig, ) -> TorshResult<Self>

Create a new DDP wrapper with custom bucket configuration

pub fn new_with_configs( module: M, process_group: Arc<ProcessGroup>, device_ids: Vec<usize>, output_device: Option<usize>, broadcast_buffers: bool, bucket_config: BucketConfig, overlap_config: OverlapConfig, ) -> TorshResult<Self>

Create a new DDP wrapper with custom configurations

pub async fn sync_gradients(&mut self) -> TorshResult<()>

Synchronize gradients across all processes

pub fn register_gradient_hooks(&self) -> TorshResult<()>

Register gradient synchronization hooks This should be called during the backward pass to automatically sync gradients

pub fn register_gradient_async( &self, param_name: &str, gradient: Tensor, ) -> TorshResult<()>

Register a gradient for asynchronous synchronization (overlap mode) This should be called when a gradient becomes available during backward pass

pub fn check_unused_parameters(&self) -> TorshResult<Vec<String>>

Check for unused parameters and issue warnings

pub fn start_iteration(&self) -> TorshResult<()>

Start a new iteration (reset unused parameter tracking)

pub fn get_overlap_stats(&self) -> HashMap<String, Value>

Get overlap computation statistics

pub fn has_gradients(&self) -> bool

Check if any parameters have gradients

pub fn zero_grad(&mut self) -> TorshResult<()>

Zero all gradients

pub fn get_sync_stats(&self) -> GradientSyncStats

Get gradient synchronization statistics

pub fn set_bucketing_enabled(&mut self, enabled: bool) -> TorshResult<()>

Enable/disable gradient bucketing at runtime

pub fn get_bucket_info(&self) -> Vec<BucketInfo>

Get bucket information for debugging

pub async fn check_gradient_consistency(&self) -> TorshResult<bool>

Perform a gradient consistency check across all processes This is useful for debugging distributed training issues

Trait Implementations

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

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

Forward pass through the module

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

Get all parameters in the module (non-recursive)

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

Get named parameters (non-recursive)

fn training(&self) -> bool

Check if in training mode

fn train(&mut self)

Set training mode

fn eval(&mut self)

Set evaluation mode

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

Move module to device

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

Get all parameters recursively including submodules

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

Get all named parameters recursively with module prefixes

fn set_training(&mut self, _training: bool)

Set training mode (internal implementation)

fn load_state_dict( &mut self, state_dict: &HashMap<String, Tensor>, strict: bool, ) -> Result<(), TorshError>

Load state dictionary into the module

fn load_state_dict_strict( &mut self, state_dict: &HashMap<String, Tensor>, ) -> Result<(), TorshError>

Load state dictionary with default strict=true

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

Save state dictionary from the module

fn name(&self) -> Option<&str>

Get the module name (optional, for debugging and serialization)

fn buffers(&self) -> Vec<Arc<RwLock<RawRwLock, Tensor>>>

Get all buffers (non-trainable parameters)

fn named_buffers(&self) -> HashMap<String, Arc<RwLock<RawRwLock, Tensor>>>

Get named buffers

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

Get all direct child modules

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

Get all direct child modules with names

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

where Self: Sized,

Get all modules recursively (depth-first traversal)

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

where Self: Sized,

Get all modules recursively with hierarchical names

fn zero_grad(&mut self)

Zero all gradients recursively

fn num_parameters(&self) -> usize

Count total number of parameters

fn num_trainable_parameters(&self) -> usize

Count trainable parameters

fn memory_usage(&self) -> usize

Get memory usage estimate in bytes

fn freeze(&mut self)

Freeze all parameters (set requires_grad = false)

fn unfreeze(&mut self)

Unfreeze all parameters (set requires_grad = true)

fn extra_repr(&self) -> String

Get string representation

fn register_hook( &mut self, _hook_type: HookType, _callback: Box<dyn Fn(&dyn Module, &Tensor, Option<&Tensor>) -> Result<(), TorshError> + Send + Sync>, ) -> Option<HookHandle>

Register a hook for this module (default implementation does nothing)

fn remove_hook(&mut self, _hook_type: HookType, _handle: HookHandle) -> bool

Remove a hook by handle (default implementation does nothing)

fn execute_hooks( &self, _hook_type: HookType, _input: &Tensor, _output: Option<&Tensor>, ) -> Result<(), TorshError>

Execute hooks of a specific type (default implementation does nothing)

fn forward_with_hooks(&self, input: &Tensor) -> Result<Tensor, TorshError>

Forward pass with hooks support

fn has_hooks(&self, _hook_type: HookType) -> bool

Check if module has hooks registered

fn call(&self, input: &Tensor) -> Result<Tensor, TorshError>

Convenient method to call forward and handle common patterns

fn apply(&self, input: &Tensor) -> Result<Tensor, TorshError>

Apply the module to input (alias for forward)

fn has_parameters(&self) -> bool

Check if the module has any parameters

fn has_children(&self) -> bool

Check if the module has any child modules

fn parameter_count(&self) -> usize

Get parameter count (convenience method)

fn trainable_parameter_count(&self) -> usize

Get trainable parameter count (convenience method)

fn memory_usage_mb(&self) -> f64

Get memory usage in MB (convenience method)

fn toggle_training(&mut self)

Toggle training mode (convenience method)

fn eval_mode(&self) -> bool

Check if module is in evaluation mode

fn sequential_forward( modules: &[&dyn Module], input: Tensor, ) -> Result<Tensor, TorshError>

where Self: Sized,

Sequential forward pass through multiple modules

fn batch_forward(&self, inputs: &[Tensor]) -> Result<Vec<Tensor>, TorshError>

Apply module multiple times with different inputs (batch processing)

fn conditional_forward( &self, input: &Tensor, condition: bool, ) -> Result<Tensor, TorshError>

Forward with condition - only apply if condition is true

fn residual_forward(&self, input: &Tensor) -> Result<Tensor, TorshError>

Forward with residual connection

fn module_info(&self) -> ModuleInfo

Get detailed module information for debugging

fn check_training_readiness(&self) -> Result<(), TorshError>

Check if module is ready for training

fn parameter_names_matching(&self, pattern: &str) -> Vec<String>

Get parameter names matching a pattern

fn parameters_by_type(&self, param_type: &str) -> HashMap<String, Parameter>

Get parameters by layer type (e.g., “weight”, “bias”)

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

Clone module parameters (for creating copies or checkpoints)

fn diagnose(&self) -> ModuleDiagnostics

Quick diagnostic check of module health