Struct ZeROOptimizer 

pub struct ZeROOptimizer<T: Optimizer> { /* private fields */ }

Main ZeRO optimizer that wraps an underlying optimizer

Implementations

impl<T: Optimizer> ZeROOptimizer<T>

pub fn new( base_optimizer: T, config: ZeROConfig, mp_context: Arc<ModelParallelContext>, ) -> Result<Self>

Create a new ZeRO optimizer

pub fn register_parameters( &mut self, parameters: HashMap<String, Tensor>, ) -> Result<()>

Register parameters with ZeRO optimizer

pub fn update_gradients( &mut self, gradients: HashMap<String, Tensor>, ) -> Result<()>

Update gradients for ZeRO optimization

pub fn gather_parameters( &mut self, parameter_names: &[String], ) -> Result<HashMap<String, Tensor>>

Gather parameters for forward pass (Stage 3 only)

pub fn release_parameters(&mut self, parameter_names: &[String]) -> Result<()>

Release gathered parameters to save memory (Stage 3 only)

pub fn get_memory_stats(&self) -> &ZeROMemoryStats

Get memory statistics

pub fn check_memory_usage(&self) -> bool

Check if memory usage exceeds threshold

pub fn get_stage(&self) -> ZeROStage

Get current ZeRO stage

pub fn base_optimizer(&self) -> &T

Get the underlying base optimizer

pub fn base_optimizer_mut(&mut self) -> &mut T

Get mutable reference to base optimizer

pub fn mp_context(&self) -> &Arc<ModelParallelContext>

Get model parallel context

pub fn optimizer_step(&mut self) -> Result<()>

Perform optimizer step with ZeRO optimizations

Trait Implementations

impl<T: Optimizer> Optimizer for ZeROOptimizer<T>

fn update(&mut self, parameter: &mut Tensor, grad: &Tensor) -> Result<()>

Updates a parameter based on its gradient.

fn zero_grad(&mut self)

Clears all accumulated gradients.

fn step(&mut self)

Performs a single optimization step.

fn get_lr(&self) -> f32

Gets the current learning rate.

fn set_lr(&mut self, lr: f32)

Sets a new learning rate.

fn accumulate_grad( &mut self, parameter: &mut Tensor, grad: &Tensor, ) -> Result<()>

Accumulates gradients for gradient accumulation.

fn apply_accumulated_grads(&mut self, accumulation_steps: usize) -> Result<()>

Applies accumulated gradients after gradient accumulation.