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 = {
  1. scale : Ocannl_tensor.Operation.DSL_modules.Tensor.t;
    (*

    Pass as Train.grad_update ~loss_scale.

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

    The reciprocal; pass as Train.sgd_update ~grad_unscale.

    *)
  3. mutable scale_val : Base.float;
  4. mutable good_steps : Base.int;
  5. growth_factor : Base.float;
  6. backoff_factor : Base.float;
  7. growth_interval : Base.int;
}
val host_scalar : l:Base.string -> Base.float -> Ocannl_tensor.Operation.DSL_modules.Tensor.t
val create : ?init_scale:Base.float -> ?growth_factor:Base.float -> ?backoff_factor:Base.float -> ?growth_interval:Base.int -> unit -> t
val scale_value : t -> Base.float
val set_scale : t -> Context.t -> Base.float -> Context.t
val update : t -> Context.t -> grads_finite:bool -> Context.t
val update_n : t -> Context.t -> grads_finite:bool -> steps:Base__Int.t -> 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.