Struct PagedAdamWState 

pub struct PagedAdamWState {
    pub exp_avg: HashMap<String, Tensor>,
    pub exp_avg_sq: HashMap<String, Tensor>,
    pub steps: HashMap<String, usize>,
    pub page_size: usize,
    pub max_gpu_memory: usize,
    pub current_gpu_usage: usize,
    /* private fields */
}

Paged optimizer state for CPU offloading.

Stores optimizer states (momentum, variance) on CPU and pages them to GPU as needed during parameter updates. This enables training large models on limited VRAM by trading off memory for compute.

Matches Python QLoRA’s --optim paged_adamw_32bit behavior.

Fields

exp_avg: HashMap<String, Tensor>

First moment estimates (CPU tensors, paged to GPU on demand).

exp_avg_sq: HashMap<String, Tensor>

Second moment estimates (CPU tensors, paged to GPU on demand).

steps: HashMap<String, usize>

Step counts per parameter.

page_size: usize

Page size in bytes.

max_gpu_memory: usize

Maximum GPU memory for optimizer states (0 = unlimited).

current_gpu_usage: usize

Current GPU memory usage (bytes).

Implementations

impl PagedAdamWState

pub fn new(page_size: usize, max_gpu_memory: usize) -> PagedAdamWState

Create new paged optimizer state.

pub fn init_param( &mut self, name: &str, shape: &[usize], _device: &Device, ) -> Result<(), QLoraError>

Initialize state for a parameter.

Errors

Returns error if tensor creation fails.

pub fn page_to_device( &mut self, name: &str, device: &Device, ) -> Result<(Tensor, Tensor), QLoraError>

Page state to GPU for update, returns (exp_avg, exp_avg_sq) on target device.

Updates LRU tracking and GPU memory usage.

Errors

Returns error if device transfer fails.

pub fn page_to_cpu( &mut self, name: &str, exp_avg: &Tensor, exp_avg_sq: &Tensor, ) -> Result<(), QLoraError>

Page state back to CPU after update.

Updates GPU memory tracking.

Errors

Returns error if device transfer fails.

pub fn increment_step(&mut self, name: &str)

Increment step count for a parameter.

pub fn get_step(&self, name: &str) -> usize

Get step count for a parameter.

pub fn is_gpu_resident(&self, name: &str) -> bool

Check if a parameter’s optimizer state is currently GPU-resident.

pub fn gpu_resident_count(&self) -> usize

Get the number of parameters currently GPU-resident.

Trait Implementations

impl Debug for PagedAdamWState

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

Formats the value using the given formatter.