Struct MemoryBandwidthOptimizer 

pub struct MemoryBandwidthOptimizer { /* private fields */ }

Memory-bandwidth co-optimization

Implementations

impl MemoryBandwidthOptimizer

pub fn new( optimizer: Box<dyn Optimizer>, memory_threshold: usize, bandwidth_threshold: f64, base_batch_size: usize, ) -> Self

Create a new memory-bandwidth co-optimizer

pub fn adjust_batch_size( &mut self, memory_usage: usize, bandwidth_usage: f64, ) -> usize

Adjust batch size based on memory and bandwidth usage

pub fn get_utilization(&self) -> ResourceUtilization

Get current resource utilization

Trait Implementations

impl Optimizer for MemoryBandwidthOptimizer

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<(), TrustformersError>

Accumulates gradients for gradient accumulation.

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

Applies accumulated gradients after gradient accumulation.