Struct QLoraTrainer 

pub struct QLoraTrainer { /* private fields */ }

Trainer for QLoRA fine-tuning.

Manages the training loop, gradient computation, and optimizer updates for quantized LoRA training.

Usage

1. Create trainer with config
2. Use var_builder() to create layers that register params in VarMap
3. Call init_optimizer() to set up optimizer with registered params
4. Call training_step() or training_step_lm() for each batch

Implementations

impl QLoraTrainer

pub fn new(config: QLoraTrainingConfig, device: Device) -> QLoraTrainer

Create a new QLoRA trainer.

Arguments

• config - Training configuration
• device - Device for computation

Returns

New trainer instance

pub fn var_builder(&self) -> VarBuilderArgs<'_, Box<dyn SimpleBackend + '_>>

Get a VarBuilder backed by this trainer’s VarMap.

Use this to create QuantizedLinear layers with gradient tracking. Params created through this VarBuilder will be registered in the trainer’s VarMap and trained by the optimizer.

Example

let mut trainer = QLoraTrainer::new(config, device.clone());
let vb = trainer.var_builder();
let layer = QuantizedLinear::from_weight_with_varbuilder(&weight, None, &qlora_config, vb.pp("layer0"))?;
trainer.init_optimizer(&[&layer])?;

pub fn init_optimizer( &mut self, layers: &[&QuantizedLinear], ) -> Result<(), QLoraError>

Initialize the optimizer with trainable parameters.

Creates either a paged or standard AdamW optimizer based on configuration. For paged optimizer, optimizer states are stored on CPU and paged to GPU during updates to reduce VRAM usage.

Important: Layers must be created using var_builder() for standard AdamW, or the optimizer will have no trainable parameters.

Arguments

• layers - The QLoRA layers to train

Errors

Returns error if:

• VarMap is empty (for standard optimizer) - layers weren’t created with var_builder()
• Optimizer initialization fails

Panics

Panics if the VarMap mutex is poisoned.

pub fn state(&self) -> &AdapterTrainingState

Get the current training state.

pub fn current_lr(&self) -> f64

Get the current learning rate.

pub fn global_step(&self) -> usize

Get the current step.

pub fn epoch(&self) -> usize

Get the current epoch.

pub fn training_step( &mut self, layers: &[&QuantizedLinear], input: &Tensor, targets: &Tensor, ) -> Result<f64, QLoraError>

Perform a training step with gradient accumulation.

QLoRA training flow:

1. Forward pass through frozen quantized base + trainable LoRA
2. Compute loss (cross-entropy for LM, MSE for regression)
3. Backward pass - gradients flow only through LoRA weights
4. Accumulate gradients if gradient_accumulation_steps > 1
5. Optimizer step when accumulation complete

Supports both standard AdamW and paged AdamW optimizers.

Arguments

• layers - The QLoRA layers
• input - Input tensor [batch, seq_len, hidden]
• targets - Target tensor (logits or token IDs depending on loss)

Returns

The loss value for this step

Errors

Returns error if forward pass or backward pass fails

Panics

Panics if the VarMap mutex is poisoned.

pub fn training_step_lm( &mut self, layers: &[&QuantizedLinear], input: &Tensor, target_ids: &Tensor, ) -> Result<f64, QLoraError>

Perform training step with cross-entropy loss for language modeling.

Supports both standard AdamW and paged AdamW optimizers.

Arguments

• layers - The QLoRA layers
• input - Input tensor [batch, seq_len, hidden]
• target_ids - Target token IDs [batch, seq_len]

Returns

The cross-entropy loss value

Errors

Returns error if forward pass or loss computation fails

Panics

Panics if the VarMap mutex is poisoned.

pub fn start_epoch(&mut self)

Start a new training epoch.

pub fn should_continue(&self) -> bool

Check if training should continue.

pub fn update_lr(&mut self)

Update learning rate based on schedule.

pub fn config(&self) -> &QLoraTrainingConfig

Get training configuration.

pub fn optimizer_memory_stats(&self) -> Option<(usize, usize)>

Get optimizer memory statistics (CPU bytes, GPU bytes).

pub fn zero_grad(&mut self)

Zero gradients for next accumulation cycle.

Resets the accumulation step counter. Note: In candle, gradients are automatically zeroed when backward_step is called on the optimizer.