rlx_tensor/
schedule.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Learning-rate schedules for [`Func::train_step_at`](crate::Func::train_step_at).
6//!
7//! A schedule maps a global step to a learning rate; `train_step_at` calls
8//! `Optimizer::set_lr` with that value before each step. Available with the
9//! `optim` feature.
10
11/// Common learning-rate schedules.
12#[derive(Clone, Copy, Debug, PartialEq)]
13pub enum LrSchedule {
14    /// Fixed learning rate.
15    Constant(f32),
16    /// Multiply `base` by `gamma` every `step_size` steps (staircase).
17    Step {
18        base: f32,
19        step_size: usize,
20        gamma: f32,
21    },
22    /// Cosine decay from `base` to `min` over `total` steps (then holds `min`).
23    Cosine { base: f32, min: f32, total: usize },
24    /// Linear warmup from 0 to `base` over `warmup` steps, then holds `base`.
25    Warmup { base: f32, warmup: usize },
26    /// Linear warmup to `base` over `warmup`, then cosine decay to `min` by
27    /// `total` — the standard transformer schedule.
28    WarmupCosine {
29        base: f32,
30        min: f32,
31        warmup: usize,
32        total: usize,
33    },
34}
35
36impl LrSchedule {
37    /// Learning rate at `step` (0-based).
38    pub fn lr_at(&self, step: usize) -> f32 {
39        match *self {
40            LrSchedule::Constant(lr) => lr,
41            LrSchedule::Step {
42                base,
43                step_size,
44                gamma,
45            } => {
46                let k = step.checked_div(step_size).unwrap_or(0);
47                base * gamma.powi(k as i32)
48            }
49            LrSchedule::Cosine { base, min, total } => cosine(base, min, step, total),
50            LrSchedule::Warmup { base, warmup } => {
51                if warmup > 0 && step < warmup {
52                    base * (step as f32 + 1.0) / warmup as f32
53                } else {
54                    base
55                }
56            }
57            LrSchedule::WarmupCosine {
58                base,
59                min,
60                warmup,
61                total,
62            } => {
63                if warmup > 0 && step < warmup {
64                    base * (step as f32 + 1.0) / warmup as f32
65                } else {
66                    cosine(base, min, step - warmup, total.saturating_sub(warmup))
67                }
68            }
69        }
70    }
71}
72
73/// Cosine decay from `base` to `min` over `total` steps; holds `min` after.
74fn cosine(base: f32, min: f32, step: usize, total: usize) -> f32 {
75    if total == 0 {
76        return min;
77    }
78    let t = (step.min(total)) as f32 / total as f32;
79    let factor = 0.5 * (1.0 + (std::f32::consts::PI * t).cos());
80    min + (base - min) * factor
81}
rlx_tensor/schedule.rs

rlx_tensor/
schedule.rs
