Struct CPUOffloadedOptimizer 

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

A wrapper that enables CPU offloading for any optimizer.

Implementations

impl<T: Optimizer + StatefulOptimizer> CPUOffloadedOptimizer<T>

pub fn new(base_optimizer: T, config: CPUOffloadConfig) -> Self

Creates a new CPU-offloaded optimizer wrapper.

pub fn with_default_config(base_optimizer: T) -> Self

Creates a CPU-offloaded optimizer with default configuration.

pub fn get_memory_stats(&self) -> &CPUOffloadStats

Get the current memory statistics.

pub fn get_memory_savings_bytes(&self) -> usize

Get the total memory savings (GPU memory freed).

pub fn get_memory_savings_percent(&self) -> f32

Get the memory savings as a percentage of total optimizer memory.

pub fn prefetch_states(&mut self, keys: &[String], device: &str) -> Result<()>

Prefetch tensors that will be needed soon.

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

Clean up GPU cache of states that won’t be needed soon.

pub fn get_config(&self) -> &CPUOffloadConfig

Get the configuration.

pub fn set_config(&mut self, config: CPUOffloadConfig)

Update the configuration.

Trait Implementations

impl<T: Optimizer> Optimizer for CPUOffloadedOptimizer<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<(), TrustformersError>

Accumulates gradients for gradient accumulation.

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

Applies accumulated gradients after gradient accumulation.