Struct PagedAdamW 

pub struct PagedAdamW { /* private fields */ }

Paged AdamW optimizer with CPU offloading.

Implements AdamW with optimizer state paging to CPU memory, matching Python’s paged_adamw_32bit from bitsandbytes.

Memory Behavior

• Optimizer states (exp_avg, exp_avg_sq) stored on CPU
• States paged to GPU only during parameter update
• Enables training 7B+ models on 24GB GPUs with QLoRA

Implementations

impl PagedAdamW

pub fn new( lr: f64, weight_decay: f64, page_size: usize, max_gpu_memory: usize, ) -> PagedAdamW

Create a new paged AdamW optimizer.

Arguments

• lr - Learning rate
• weight_decay - Weight decay coefficient
• page_size - Page size in bytes for CPU offloading
• max_gpu_memory - Maximum GPU memory for optimizer states (0 = unlimited)

pub fn with_betas(self, beta1: f64, beta2: f64) -> PagedAdamW

Create with custom betas.

pub fn init(&mut self, params: &[(String, Tensor)]) -> Result<(), QLoraError>

Initialize optimizer state for parameters.

Errors

Returns error if state initialization fails.

pub fn set_lr(&mut self, lr: f64)

Set learning rate.

pub fn lr(&self) -> f64

Get current learning rate.

pub fn step_param( &mut self, name: &str, param: &mut Tensor, grad: &Tensor, ) -> Result<(), QLoraError>

Perform optimizer step for a single parameter.

Implements AdamW update with CPU paging:

m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
v_t = β₂ * v_{t-1} + (1 - β₂) * g_t²
m̂_t = m_t / (1 - β₁^t)
v̂_t = v_t / (1 - β₂^t)
θ_t = θ_{t-1} - lr * (m̂_t / (√v̂_t + ε) + λ * θ_{t-1})

Errors

Returns error if tensor operations fail.

pub fn memory_stats(&self) -> (usize, usize)

Get memory usage statistics.