Module Train.Lr_schedule

Host-side learning-rate schedules: pure functions from the step number to a float, fed to the device via scheduled_learning_rate (or any host-written scalar). All schedules start with a linear warmup over warmup_steps steps (base_lr * (step+1) / warmup_steps; no warmup when 0) and decay toward base_lr *. final_frac at total_steps. Steps beyond total_steps clamp to the final value.

type kind =
  1. | Constant
    (*

    base_lr after warmup; final_frac is ignored.

    *)
  2. | Cosine
    (*

    Half-cosine from base_lr down to base_lr *. final_frac.

    *)
  3. | Linear
    (*

    Straight line from base_lr down to base_lr *. final_frac.

    *)
  4. | Wsd of {
    1. decay_frac : Base.float;
    }
    (*

    Warmup-stable-decay (arXiv:2405.18392): hold base_lr until the final decay_frac fraction of total_steps (llm.c uses 0.2), then decay as 1 - sqrt(ratio).

    *)
type t = {
  1. kind : kind;
  2. base_lr : Base.float;
  3. warmup_steps : Base.int;
  4. total_steps : Base.int;
  5. final_frac : Base.float;
    (*

    The final learning rate as a fraction of base_lr.

    *)
}
val learning_rate : t -> step:Base__Int.t -> Base__Float.t

The learning rate at step (0-based).