Struct ZeRoOptimizer 

pub struct ZeRoOptimizer { /* private fields */ }

ZeRO optimizer state

Implementations

impl ZeRoOptimizer

pub fn new(config: ZeRoConfig, learning_rate: f32) -> Self

Create new ZeRO optimizer

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

Partition parameters across ranks

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

Partition gradients (Stage 2+)

pub fn reduce_scatter_gradients(&mut self) -> Result<(), String>

Reduce-scatter gradients across ranks

pub fn all_gather_parameters(&self) -> Result<HashMap<String, Tensor>, String>

All-gather parameters (Stage 3)

pub fn step(&mut self) -> Result<(), String>

Optimizer step with ZeRO

pub fn offload_to_cpu(&mut self, name: &str) -> Result<(), String>

Offload to CPU

pub fn load_from_cpu(&mut self, name: &str) -> Result<(), String>

Load from CPU

pub fn memory_savings_ratio(&self) -> f32

Get memory savings ratio

pub fn get_stats(&self) -> ZeRoStats

Get statistics