Struct Adam4bit 

pub struct Adam4bit { /* private fields */ }

4-bit Adam optimizer with advanced quantization

Implementations

impl Adam4bit

pub fn new( learning_rate: f32, beta1: f32, beta2: f32, epsilon: f32, weight_decay: f32, ) -> Self

Create a new 4-bit Adam optimizer

pub fn with_quantization_config( optimizer_config: Adam4bitOptimizerConfig, quantization_config: AdvancedQuantizationConfig, ) -> Self

Create with custom quantization config

pub fn memory_savings(&self) -> f32

Get memory savings compared to full precision Adam

Trait Implementations

impl Debug for Adam4bit

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

Formats the value using the given formatter.

impl Optimizer for Adam4bit

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.

impl StatefulOptimizer for Adam4bit

type Config = Adam4bitOptimizerConfig

The configuration type for this optimizer.

type State = OptimizerState

The state type used by this optimizer.

fn config(&self) -> &Self::Config

Gets a reference to the optimizer’s configuration.

fn state(&self) -> &Self::State

Gets a reference to the optimizer’s internal state.

fn state_mut(&mut self) -> &mut Self::State

Gets a mutable reference to the optimizer’s internal state.

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

Saves the optimizer state to a dictionary for checkpointing.

fn load_state_dict(&mut self, state: HashMap<String, Tensor>) -> Result<()>

Loads optimizer state from a dictionary during checkpoint restoration.

fn memory_usage(&self) -> StateMemoryStats

Gets memory usage statistics for this optimizer.

fn reset_state(&mut self)

Resets the optimizer state (useful for training restarts).

fn num_parameters(&self) -> usize

Returns the number of parameters being optimized.