From 655980690e3e8c7e77c86db2a8a3b9266c707423 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 23 Jul 2026 13:38:20 +0200 Subject: [PATCH] feat(optim): mutable learning rate and warmup+cosine LR schedule AdamOptimizer and SgdOptimizer expose lr as a public var so training loops can adjust it between steps without recreating the optimizer and losing the moment estimates. A stateless LrSchedule fun interface plus a linearWarmupCosineDecay factory cover the GPT-style pretraining recipe: linear ramp from initialLr to peakLr over the warmup steps, then cosine decay to minLr. Closes #865 --- .../sk/ainet/exec/optim/SgdOptimizerTest.kt | 24 ++++++++ .../api/android/skainet-lang-core.api | 13 ++++ .../api/jvm/skainet-lang-core.api | 13 ++++ .../kotlin/sk/ainet/lang/nn/optim/Adam.kt | 7 ++- .../sk/ainet/lang/nn/optim/LrSchedule.kt | 56 +++++++++++++++++ .../kotlin/sk/ainet/lang/nn/optim/SGD.kt | 6 +- .../sk/ainet/lang/nn/optim/LrScheduleTest.kt | 60 +++++++++++++++++++ 7 files changed, 177 insertions(+), 2 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/LrSchedule.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/optim/LrScheduleTest.kt diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/optim/SgdOptimizerTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/optim/SgdOptimizerTest.kt index 65a98006..e898bba7 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/optim/SgdOptimizerTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/optim/SgdOptimizerTest.kt @@ -86,4 +86,28 @@ class SgdOptimizerTest { opt.step() assertEquals(3.9025f, w.value.data[0, 0], 1e-3f) } + + @Test + fun lr_can_be_rescheduled_between_steps() { + val ctx = DirectCpuExecutionContext(phase = Phase.TRAIN) + val w = param1x1(10f) + + val opt = sk.ainet.lang.nn.optim.SgdOptimizer(lr = 0.1) + opt.addParameter(w) + + // Step 1 at lr = 0.1: w = 10 - 0.1 * 2 = 9.8 + val g1 = ctx.full(Shape(1, 1), FP32::class, 2.0) + w.value.accumulateGrad(g1 as sk.ainet.lang.tensor.Tensor) + opt.step() + opt.zeroGrad() + assertEquals(9.8f, w.value.data[0, 0], 1e-6f) + + // Reschedule to lr = 0.5 (as a warmup/decay schedule would): + // w = 9.8 - 0.5 * 2 = 8.8 + opt.lr = 0.5 + val g2 = ctx.full(Shape(1, 1), FP32::class, 2.0) + w.value.accumulateGrad(g2 as sk.ainet.lang.tensor.Tensor) + opt.step() + assertEquals(8.8f, w.value.data[0, 0], 1e-6f) + } } diff --git a/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api index c9c152e6..79359b30 100644 --- a/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api @@ -1897,11 +1897,22 @@ public final class sk/ainet/lang/nn/optim/AdamOptimizer : sk/ainet/lang/nn/optim public synthetic fun (DDDDDZZILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V + public final fun getLr ()D public final fun reset ()V + public final fun setLr (D)V public fun step ()V public fun zeroGrad ()V } +public abstract interface class sk/ainet/lang/nn/optim/LrSchedule { + public abstract fun lrAt (I)D +} + +public final class sk/ainet/lang/nn/optim/LrScheduleKt { + public static final fun linearWarmupCosineDecay (IIDDD)Lsk/ainet/lang/nn/optim/LrSchedule; + public static synthetic fun linearWarmupCosineDecay$default (IIDDDILjava/lang/Object;)Lsk/ainet/lang/nn/optim/LrSchedule; +} + public abstract interface class sk/ainet/lang/nn/optim/Optimizer { public abstract fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V public abstract fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V @@ -1928,6 +1939,8 @@ public final class sk/ainet/lang/nn/optim/SgdOptimizer : sk/ainet/lang/nn/optim/ public synthetic fun (DDDILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V + public final fun getLr ()D + public final fun setLr (D)V public fun step ()V public fun zeroGrad ()V } diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 4c7cbad8..5db976e1 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -2125,11 +2125,22 @@ public final class sk/ainet/lang/nn/optim/AdamOptimizer : sk/ainet/lang/nn/optim public synthetic fun (DDDDDZZILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V + public final fun getLr ()D public final fun reset ()V + public final fun setLr (D)V public fun step ()V public fun zeroGrad ()V } +public abstract interface class sk/ainet/lang/nn/optim/LrSchedule { + public abstract fun lrAt (I)D +} + +public final class sk/ainet/lang/nn/optim/LrScheduleKt { + public static final fun linearWarmupCosineDecay (IIDDD)Lsk/ainet/lang/nn/optim/LrSchedule; + public static synthetic fun linearWarmupCosineDecay$default (IIDDDILjava/lang/Object;)Lsk/ainet/lang/nn/optim/LrSchedule; +} + public abstract interface class sk/ainet/lang/nn/optim/Optimizer { public abstract fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V public abstract fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V @@ -2156,6 +2167,8 @@ public final class sk/ainet/lang/nn/optim/SgdOptimizer : sk/ainet/lang/nn/optim/ public synthetic fun (DDDILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V + public final fun getLr ()D + public final fun setLr (D)V public fun step ()V public fun zeroGrad ()V } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/Adam.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/Adam.kt index 707228cb..78a8210d 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/Adam.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/Adam.kt @@ -33,7 +33,12 @@ import kotlin.math.pow * @param amsgrad If true, uses the AMSGrad variant that maintains the maximum of all v_t (default: false) */ public class AdamOptimizer @kotlin.jvm.JvmOverloads constructor( - private val lr: Double = 0.001, + /** + * Learning rate. Mutable so learning-rate schedules (warmup, cosine + * decay, …) can adjust it between steps without recreating the optimizer + * and losing the moment estimates. See [LrSchedule]. + */ + public var lr: Double = 0.001, private val beta1: Double = 0.9, private val beta2: Double = 0.999, private val epsilon: Double = 1e-8, diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/LrSchedule.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/LrSchedule.kt new file mode 100644 index 00000000..79a0526e --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/LrSchedule.kt @@ -0,0 +1,56 @@ +package sk.ainet.lang.nn.optim + +import kotlin.math.PI +import kotlin.math.cos + +/** + * A learning-rate schedule: maps a global training step to a learning rate. + * + * Schedules are stateless and decoupled from optimizers — a training loop + * assigns the value to the optimizer's mutable `lr` before each step: + * + * ```kotlin + * val schedule = linearWarmupCosineDecay(totalSteps, warmupSteps, peakLr = 5e-4) + * // inside the loop: + * optimizer.lr = schedule.lrAt(globalStep) + * ``` + */ +public fun interface LrSchedule { + /** The learning rate to use for [step] (0-based, `0 <= step < totalSteps`). */ + public fun lrAt(step: Int): Double +} + +/** + * Linear warmup followed by cosine decay — the schedule used by GPT-style + * pretraining recipes. + * + * - steps `0 ..< warmupSteps`: linear ramp from [initialLr] to [peakLr] + * - steps `warmupSteps ..< totalSteps`: cosine decay from [peakLr] to [minLr] + * + * @param totalSteps total number of optimization steps + * @param warmupSteps steps spent ramping up (must be in `1 ..< totalSteps`) + * @param peakLr learning rate at the end of the warmup + * @param initialLr learning rate at step 0 + * @param minLr learning rate approached at the end of training + */ +public fun linearWarmupCosineDecay( + totalSteps: Int, + warmupSteps: Int, + peakLr: Double, + initialLr: Double = 3e-5, + minLr: Double = 1e-6, +): LrSchedule { + require(totalSteps > 0) { "totalSteps must be positive, was $totalSteps" } + require(warmupSteps in 1 until totalSteps) { + "warmupSteps must be in 1 until totalSteps=$totalSteps, was $warmupSteps" + } + return LrSchedule { step -> + require(step in 0 until totalSteps) { "step $step out of range [0, $totalSteps)" } + if (step < warmupSteps) { + initialLr + (peakLr - initialLr) * step / warmupSteps + } else { + val progress = (step - warmupSteps).toDouble() / (totalSteps - warmupSteps) + minLr + (peakLr - minLr) * 0.5 * (1 + cos(PI * progress)) + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/SGD.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/SGD.kt index 6fd2b421..88e8429a 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/SGD.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/SGD.kt @@ -9,7 +9,11 @@ import sk.ainet.lang.types.DType * Stochastic Gradient Descent optimizer with optional momentum and weight decay. */ public class SgdOptimizer @kotlin.jvm.JvmOverloads constructor( - private val lr: Double, + /** + * Learning rate. Mutable so learning-rate schedules can adjust it between + * steps. See [LrSchedule]. + */ + public var lr: Double, private val momentum: Double = 0.0, private val weightDecay: Double = 0.0, ) : Optimizer { diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/optim/LrScheduleTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/optim/LrScheduleTest.kt new file mode 100644 index 00000000..30b94fbd --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/optim/LrScheduleTest.kt @@ -0,0 +1,60 @@ +package sk.ainet.lang.nn.optim + +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +class LrScheduleTest { + + private val total = 100 + private val warmup = 20 + private val peak = 1e-3 + private val initial = 3e-5 + private val min = 1e-6 + + private val schedule = linearWarmupCosineDecay(total, warmup, peak, initial, min) + + @Test + fun warmup_starts_at_initialLr_and_ramps_linearly() { + assertEquals(initial, schedule.lrAt(0), 1e-12) + val lrs = (0 until warmup).map { schedule.lrAt(it) } + lrs.zipWithNext().forEach { (a, b) -> assertTrue(b > a, "warmup must strictly increase") } + // Linear: constant increments + val increments = lrs.zipWithNext().map { (a, b) -> b - a } + increments.zipWithNext().forEach { (a, b) -> assertEquals(a, b, 1e-12) } + } + + @Test + fun peak_is_reached_at_the_end_of_warmup() { + assertEquals(peak, schedule.lrAt(warmup), 1e-12) + } + + @Test + fun decay_is_monotonic_and_approaches_minLr() { + val lrs = (warmup until total).map { schedule.lrAt(it) } + lrs.zipWithNext().forEach { (a, b) -> assertTrue(b <= a, "decay must not increase") } + assertTrue(abs(lrs.last() - min) < peak * 0.01, "final lr ${lrs.last()} should approach $min") + } + + @Test + fun invalid_arguments_are_rejected() { + assertFailsWith { linearWarmupCosineDecay(0, 1, peak) } + assertFailsWith { linearWarmupCosineDecay(10, 0, peak) } + assertFailsWith { linearWarmupCosineDecay(10, 10, peak) } + assertFailsWith { schedule.lrAt(-1) } + assertFailsWith { schedule.lrAt(total) } + } + + @Test + fun optimizer_lr_is_assignable_between_steps() { + val adam = AdamOptimizer(lr = 0.001) + adam.lr = schedule.lrAt(warmup) + assertEquals(peak, adam.lr, 1e-12) + + val sgd = SgdOptimizer(lr = 0.1) + sgd.lr = 0.05 + assertEquals(0.05, sgd.lr, 1e-12) + } +}