Struct GradientBuffer 

pub struct GradientBuffer {
    pub name: String,
    pub local_gradient: Tensor,
    pub accumulated_gradient: Option<Tensor>,
    pub accumulation_steps: usize,
    pub partition_info: PartitionInfo,
}

Gradient buffer for ZeRO Stage 2+

Fields

name: String

Buffer name

local_gradient: Tensor

Local gradient shard

accumulated_gradient: Option<Tensor>

Accumulated gradients

accumulation_steps: usize

Number of accumulated steps

partition_info: PartitionInfo

Partition metadata

Implementations

impl GradientBuffer

pub fn new( name: String, local_gradient: Tensor, partition_info: PartitionInfo, ) -> Self

pub fn zero(&mut self)

Zero the gradient buffer

pub fn accumulate(&mut self, gradient: &Tensor) -> Result<()>

Accumulate gradient

pub fn get_accumulated(&self) -> Option<Tensor>

Get the accumulated gradient (averaged if needed)

pub fn memory_usage(&self) -> usize

Get memory usage of this buffer

Trait Implementations

impl Clone for GradientBuffer

fn clone(&self) -> GradientBuffer

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for GradientBuffer

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

Formats the value using the given formatter.