Struct DistributedOptimizer 

pub struct DistributedOptimizer<O: Optimizer> { /* private fields */ }

Distributed optimizer wrapper

This wrapper can be applied to any optimizer to enable distributed training. It handles gradient synchronization across multiple processes before applying the underlying optimizer’s update rule.

Implementations

impl<O: Optimizer> DistributedOptimizer<O>

pub fn new(optimizer: O, config: DistributedConfig) -> OptimizerResult<Self>

Create a new distributed optimizer

pub fn inner(&self) -> &O

Get the underlying optimizer

pub fn inner_mut(&mut self) -> &mut O

Get the underlying optimizer mutably

pub fn config(&self) -> &DistributedConfig

Get the configuration

pub fn set_config(&mut self, config: DistributedConfig)

Update the configuration

pub fn get_communication_stats(&self) -> CommunicationStats

Get communication statistics

pub fn synchronize_gradients(&mut self) -> OptimizerResult<()>

Synchronize gradients across all processes

pub fn create_gradient_buckets( &mut self, parameters: &[Arc<RwLock<Tensor>>], ) -> OptimizerResult<()>

Create gradient buckets for efficient communication

Trait Implementations

impl<O: Optimizer> Optimizer for DistributedOptimizer<O>

fn step(&mut self) -> OptimizerResult<()>

Perform a single optimization step

fn zero_grad(&mut self)

Zero all gradients

fn get_lr(&self) -> Vec<f32>

Get the current learning rate

fn set_lr(&mut self, lr: f32)

Set the learning rate

fn add_param_group( &mut self, params: Vec<Arc<RwLock<Tensor>>>, options: HashMap<String, f32>, )

Add a parameter group

fn state_dict(&self) -> OptimizerResult<OptimizerState>

Get state dict for serialization

fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()>

Load state dict