Struct Trainer 

pub struct Trainer {
    pub model: Lfm2Model,
    pub optimizer: Box<dyn Optimizer>,
    pub ctx: Arc<WgpuContext>,
    pub engine: Arc<ComputeEngine>,
    pub clip_grad_norm: Option<f32>,
    pub grad_accum_steps: usize,
    pub step_in_accum: usize,
    pub scaler: Option<GradScaler>,
    pub distributed_grad_sync: Option<FileGradSyncConfig>,
    pub lr_scheduler: Option<Box<dyn LRScheduler>>,
    pub length_bucket: usize,
    /* private fields */
}

Runs optimization steps on a model.

Fields

model: Lfm2Model

The model being trained.

optimizer: Box<dyn Optimizer>

Optimizer applied to the trainable parameters.

ctx: Arc<WgpuContext>

GPU context.

engine: Arc<ComputeEngine>

Kernel engine.

clip_grad_norm: Option<f32>

Clip the global gradient norm to this value before each step.

grad_accum_steps: usize

Micro-batches (sequences) per optimizer step.

step_in_accum: usize

Micro-batches accumulated so far in the current window.

scaler: Option<GradScaler>

Dynamic loss scaler (fp16-style overflow protection).

distributed_grad_sync: Option<FileGradSyncConfig>

Multi-process gradient averaging through a shared directory.

lr_scheduler: Option<Box<dyn LRScheduler>>

Learning-rate schedule, advanced once per optimizer step (before it is read, so step k uses the schedule’s rate at k).

length_bucket: usize

Sequences are right-padded to a multiple of this many tokens (with the padding excluded from the loss). Causal attention keeps real tokens from seeing the padding, so results are unchanged, but GPU buffer sizes repeat and can be reused. 1 disables padding.

Implementations

impl Trainer

pub fn new( model: Lfm2Model, optimizer: Box<dyn Optimizer>, ctx: Arc<WgpuContext>, engine: Arc<ComputeEngine>, ) -> Self

A trainer with no clipping, accumulation, scaler or schedule.

pub fn with_length_bucket(self, tokens: usize) -> Self

Pad sequences to a multiple of tokens (see length_bucket).

pub fn with_clip_grad_norm(self, max_norm: f32) -> Self

Clip the global gradient L2 norm to max_norm.

pub fn with_grad_accum_steps(self, steps: usize) -> Self

Accumulate gradients over steps sequences per optimizer step. Each sequence’s loss is divided by steps, so the update matches the mean over the window.

pub fn with_grad_scaler(self, scaler: GradScaler) -> Self

Scale the loss before backward and skip steps whose gradients overflow.

pub fn with_file_grad_sync(self, cfg: FileGradSyncConfig) -> Self

Average gradients across processes before each optimizer step.

pub fn with_lr_scheduler(self, scheduler: Box<dyn LRScheduler>) -> Self

Set the learning rate from scheduler before each optimizer step.

pub fn optimizer_steps(&self) -> usize

Optimizer steps taken so far.

pub fn set_optimizer_steps(&mut self, steps: usize)

Restore the optimizer-step counter (and advance the LR schedule to match) when resuming from a checkpoint.

pub fn train_step( &mut self, input_ids: &[u32], target_ids: &[u32], train_params: &[Tensor], ) -> Result<StepStats, TrainError>

Train on one sequence with every target supervised.

pub fn train_step_masked( &mut self, input_ids: &[u32], target_ids: &[u32], loss_mask: Option<&[u8]>, train_params: &[Tensor], ) -> Result<StepStats, TrainError>

Train on one sequence, optionally ignoring target positions via loss_mask. loss_mask is aligned with target_ids (1 = include, 0 = ignore).

Gradients accumulate into train_params; every grad_accum_steps calls they are (optionally) synced, clipped and applied, then cleared.

pub fn train_batch( &mut self, batch: &Batch, train_params: &[Tensor], ) -> Result<Vec<StepStats>, TrainError>

Train on every sequence of batch (one micro-batch per sequence, padding stripped). Sequences without supervised tokens are skipped.