Struct GradScaler 

pub struct GradScaler { /* private fields */ }

Dynamic gradient scaler for FP16 mixed precision training.

Scales the loss to prevent gradient underflow in FP16, then unscales gradients before the optimizer step. Automatically adjusts the scale based on whether gradients contain inf/nan.

Example

let mut scaler = GradScaler::new();

// Scale the loss in place, then run the backward pass on it.
scaler.scale_loss(&loss, &ctx, &engine);
// ... backward pass fills `gradients` ...

// Unscale and skip the optimizer step if any gradient overflowed.
if scaler.unscale_and_check(&gradients, &ctx, &engine) {
    // optimizer.step(...)
}
scaler.update();

Implementations

impl GradScaler

pub fn new() -> Self

Create a new GradScaler with default settings.

Default: scale=65536, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000

pub fn with_params( init_scale: f32, growth_factor: f32, backoff_factor: f32, growth_interval: u32, ) -> Self

Create a GradScaler with custom parameters.

pub fn scale_factor(&self) -> f32

Get current scale factor.

pub fn scale_loss( &self, loss: &UnifiedTensor, ctx: &WgpuContext, engine: &ComputeEngine, )

Scale the loss value before backward pass.

Scales loss in place (no new tensor is allocated).

pub fn unscale_and_check( &mut self, gradients: &[UnifiedTensor], ctx: &WgpuContext, engine: &ComputeEngine, ) -> bool

Unscale gradients and check for inf/nan.

Returns true if gradients are valid (no overflow), false otherwise. If false, you should skip the optimizer step.

pub fn update(&mut self)

Update the scale factor based on whether overflow occurred.

Call this at the end of each training step.

pub fn had_overflow(&self) -> bool

Check if overflow was detected in the last step.

pub fn set_enabled(&mut self, enabled: bool)

Enable or disable the scaler.

pub fn is_enabled(&self) -> bool

Check if scaler is enabled.

pub fn stats(&self) -> GradScalerStats

Get statistics for logging.

pub fn checkpoint(&self) -> GradScalerCheckpoint

Snapshot the scaler for a checkpoint.

pub fn load_checkpoint(&mut self, checkpoint: &GradScalerCheckpoint)

Restore a snapshot taken by checkpoint.

Trait Implementations

impl Default for GradScaler

fn default() -> Self

Returns the “default value” for a type.