Struct AdamW4bit 

pub struct AdamW4bit { /* private fields */ }

4-bit AdamW: Adam4bit with decoupled weight decay.

Identical quantized state (NF4 momentum and variance) and identical adaptive step; the only difference is that λ·w is subtracted from the parameter directly rather than folded into the gradient, which is what makes AdamW’s decay independent of the gradient magnitude.

Implementations

impl AdamW4bit

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

Creates a 4-bit AdamW optimizer.

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

Creates a 4-bit AdamW optimizer with a custom quantization configuration.

pub fn memory_savings(&self) -> f32

Memory saved relative to full-precision AdamW, measured from the live buffers.

Trait Implementations

impl Debug for AdamW4bit

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

Formats the value using the given formatter.

impl Optimizer for AdamW4bit

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 AdamW4bit

type Config = <Adam4bit as StatefulOptimizer>::Config

The configuration type for this optimizer.

type State = <Adam4bit as StatefulOptimizer>::State

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.

fn save_state(&self, path: &Path) -> Result<()>

Saves the optimizer state to path.

fn load_state(&mut self, path: &Path) -> Result<()>

Loads the optimizer state written by Self::save_state.