Struct MixedPrecisionTrainer 

pub struct MixedPrecisionTrainer<M, O, B: Backend> { /* private fields */ }

Mixed-precision training: reduced-precision forward/backward with FP32 master weights.

Why mixed precision?

• FP16/BF16 is 2× faster on GPUs with tensor cores (V100, A100, H100)
• Half-precision uses half the memory for activations, enabling larger batches
• FP32 master weights prevent precision loss during gradient updates

How it works:

1. Inputs and targets are cast to compute_dtype (F16 or BF16) before forward
2. The forward pass runs with reduced-precision activations
3. Dynamic loss scaling prevents gradient underflow in half-precision:
   • Loss is multiplied by a scale factor before backward
   • Gradients are divided by the same factor after
   • If overflow (NaN/Inf) is detected, the step is skipped and scale reduces
4. Gradients are cast back to FP32 and applied to FP32 master weights

Compute dtype options:

• DType::F16: 16-bit IEEE float, range ±65504, good for most training
• DType::BF16: bfloat16, same range as F32 but less precision, preferred when available
• DType::F32: Standard precision (disables casting, only does loss scaling)

Example

let model = Linear::<CpuBackend>::new(784, 10, true, DType::F32, &CpuDevice)?;
let optimizer = Adam::new(model.parameters(), 1e-3);
let mut trainer = MixedPrecisionTrainer::new(
    model, optimizer, DType::F16, Default::default(),
);

for (input, target) in data {
    let metrics = trainer.train_step(&input, &target, mse_loss)?;
    println!("loss={:.4}, scale={}", metrics.loss, metrics.loss_scale);
}

Implementations

impl<M, O, B> MixedPrecisionTrainer<M, O, B>

where M: Module<B>, O: Optimizer<B>, B: Backend,

pub fn new( model: M, optimizer: O, compute_dtype: DType, config: LossScaleConfig, ) -> Self

Create a new mixed-precision trainer.

The model and optimizer should use FP32 parameters. compute_dtype sets the precision for forward/backward (F16, BF16, or F32).

pub fn model(&self) -> &M

Reference to the model.

pub fn model_mut(&mut self) -> &mut M

Mutable reference to the model.

pub fn optimizer(&self) -> &O

Reference to the optimizer.

pub fn loss_scale(&self) -> f64

Current loss scale.

pub fn compute_dtype(&self) -> DType

The compute dtype (F16, BF16, or F32).

pub fn skipped_steps(&self) -> u64

Total number of skipped steps.

pub fn train_step<F>( &mut self, input: &Tensor<B>, target: &Tensor<B>, loss_fn: F, ) -> Result<MixedPrecisionMetrics>

where F: Fn(&Tensor<B>, &Tensor<B>) -> Result<Tensor<B>>,

Perform one mixed-precision training step.

The input and target are cast to compute_dtype for the forward pass. Dynamic loss scaling is applied to prevent gradient underflow. Gradients are cast back to FP32 and applied to FP32 master weights.

Arguments

• input: input tensor (any dtype, will be cast to compute_dtype)
• target: target tensor (any dtype, will be cast to compute_dtype)
• loss_fn: function computing scalar loss from (prediction, target)

Returns

MixedPrecisionMetrics with loss value and scaling info.