Struct ZeROState 

pub struct ZeROState {
    pub step: usize,
    pub optimizer_states: HashMap<String, HashMap<String, Tensor>>,
    pub gradient_partitions: HashMap<String, GradientBuffer>,
    pub parameter_partitions: HashMap<String, ParameterPartition>,
    pub communication_buffers: HashMap<String, Tensor>,
}

ZeRO optimizer state management

Fields

step: usize

Current step number

optimizer_states: HashMap<String, HashMap<String, Tensor>>

Partitioned optimizer states per parameter group

gradient_partitions: HashMap<String, GradientBuffer>

Partitioned gradients (for Stage 2+)

parameter_partitions: HashMap<String, ParameterPartition>

Partitioned parameters (for Stage 3)

communication_buffers: HashMap<String, Tensor>

Communication buffers for all-gather operations

Implementations

impl ZeROState

pub fn new() -> Self

pub fn zero_grad(&mut self)

Reset gradients for next iteration

pub fn step(&mut self)

Increment step counter

pub fn memory_usage(&self) -> HashMap<String, usize>

Get memory usage statistics

Trait Implementations

impl Clone for ZeROState

fn clone(&self) -> ZeROState

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ZeROState

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for ZeROState

fn default() -> Self

Returns the “default value” for a type.