Module Mixed_prec.Loss_scaler

Dynamic loss-scale state, host-managed: the scale (and its reciprocal) live in tiny device-resident tensors embedded in the compiled routines, and are overwritten via Context.set_values when the schedule changes the scale — no recompilation. The schedule is torch-AMP's: multiply the scale by backoff_factor whenever a step produced non-finite gradients (the step is skipped), and by growth_factor after growth_interval consecutive good steps.

type t = private {
  1. scale : Ocannl_tensor.Tensor.t;
    (*

    Pass as Train.grad_update ~loss_scale.

    *)
  2. unscale : Ocannl_tensor.Tensor.t;
    (*

    The reciprocal; pass as Train.sgd_update ~grad_unscale.

    *)
  3. mutable scale_val : float;
  4. mutable good_steps : int;
  5. growth_factor : float;
  6. backoff_factor : float;
  7. growth_interval : int;
}

Read-only outside this module: the schedule below is the only writer, and the device copies of the scale follow it only through set_scale.

val create : ?init_scale:float -> ?growth_factor:float -> ?backoff_factor:float -> ?growth_interval:int -> unit -> t

Defaults: init_scale = 65536., growth_factor = 2., backoff_factor = 0.5, growth_interval = 200.

val scale_value : t -> float

The current scale, as the host last set it.

val set_scale : t -> Context.t -> float -> Context.t

Sets the scale to v and writes it (and its reciprocal) to the device tensors.

val update : t -> Context.t -> grads_finite:bool -> Context.t

One step of the schedule: backoff on non-finite gradients, growth after growth_interval consecutive good steps.

val update_n : t -> Context.t -> grads_finite:bool -> steps:int -> Context.t

Like update, crediting steps steps at once — for the fused gated recipe (gated_step), whose host only samples the checksum every check_interval steps: a finite sample credits the whole window toward growth, a non-finite one backs off once. Growth consumes growth_interval from the credited count and keeps the remainder (growing once per full interval contained in it), so the average growth cadence matches the per-step schedule even when the sampling interval does not divide growth_interval.

val save_state : t -> string

The scaler's host state -- the current scale, the good-step count toward the next growth, and the schedule's parameters, floats in exact hexadecimal -- as one line, for a resumable checkpoint to carry in Checkpoint.progress's extra under a key of the loop's choosing (gh-ocannl-1726). The device copies of the scale are inputs the routines only read, so Checkpoint.training_state does not hold them, and a scaler re-created at a resume would restart from init_scale with no good steps: the resumed run would skip and grow on other steps than the saved one.

val restore_state : t -> Context.t -> string -> Context.t

restore_state t ctx s puts the state save_state encoded in s back into t, a scaler created afresh by the resumed program, and writes the restored scale (and its reciprocal) to the device tensors, as set_scale does. Raises Failure when s is not such an encoding, or was saved under other growth_factor, backoff_factor or growth_interval values than t's (naming both): the schedule belongs to the program, and adopting or overriding it silently would not continue the saved run.