Add LossScale for mixed precision

core/src/main/scala/deepwit/optimizer/LossScale.scala

@@ -0,0 +1,61 @@
+package deepwit.optimizer
+
+import dimwit.*
+import dimwit.TreeOf
+import dimwit.TreeOf.mapLeaves
+import dimwit.TreeOf.ops.**!
+
+/** Dynamic loss scaling, for training in Float16 with Float32 master weights.
+  *
+  * Float16 rounds values below about 6e-8 to zero, and small gradients often are. Multiplying the loss by
+  * `scale` multiplies every gradient by it as well, which keeps them in range. Too large a scale overflows
+  * to inf instead, so a step whose gradients are not all finite is skipped, and [[next]] adapts the scale.
+  *
+  * BFloat16 has the range of Float32 and needs none of this.
+  *
+  * {{{
+  * val grads = state.lossScale.unscaled(Autodiff.grad(p => state.lossScale.scaled(cost(p)))(state.params))
+  * val (params, optState) = optimizer.update(grads, state.params, state.optState)
+  * val finite = allFinite(grads)
+  * TrainState(
+  *   select(finite, params, state.params),
+  *   select(finite, optState, state.optState),
+  *   state.lossScale.next(finite)
+  * )
+  * }}}
+  *
+  * @param finiteSteps Steps in a row with finite gradients since the scale last changed.
+  */
+case class LossScale(scale: Tensor0[Float32], finiteSteps: Tensor0[Int32]):
+
+  def scaled(loss: Tensor0[Float32]): Tensor0[Float32] = loss * scale
+
+  def unscaled[H: TensorTree, V](scaledGrads: Grad[H])(using TreeOf[H, V])(using IsFloating[V]): Grad[H] =
+    Grad(scaledGrads.value **! (1f / scale).asFloat(VType[V]))
+
+  /** Halves the scale after an overflow, and doubles it after `growthInterval` finite steps in a row.
+    *
+    * @param minScale If the scale gets stuck here, the forward pass overflows, which no scale can fix.
+    */
+  def next(
+      gradsFinite: Tensor0[Bool],
+      growthInterval: Int = 2000,
+      growthFactor: Float = 2f,
+      backoffFactor: Float = 0.5f,
+      minScale: Float = 1f
+  ): LossScale =
+    val counted = LossScale(scale, finiteSteps + 1)
+    val grown = LossScale(scale * growthFactor, 0)
+    val backedOff = LossScale(maximum(scale * backoffFactor, minScale), 0)
+    val afterFiniteStep = select(counted.finiteSteps >= growthInterval, grown, counted)
+    select(gradsFinite, afterFiniteStep, backedOff)
+
+object LossScale:
+  def initial(scale: Float = 65536f): LossScale = LossScale(Tensor0(scale), Tensor0(0))
+
+def allFinite[H: TensorTree, V](grads: Grad[H])(using TreeOf[H, V])(using IsFloating[V]): Tensor0[Bool] =
+  grads.value.mapLeaves([T <: Tuple] => (labels: Labels[T]) ?=> (grad: Tensor[T, V]) => grad.isfinite.all).reduce(_ and _)
+
+/** `ifTrue` if `condition` holds, else `ifFalse`, leaf by leaf. Unlike an `if`, this works inside `jit`. */
+def select[P: TensorTree](condition: Tensor0[Bool], ifTrue: P, ifFalse: P): P =
+  TensorTree[P].zipMap(ifTrue, ifFalse, [T <: Tuple, V] => (labels: Labels[T]) ?=> (a: Tensor[T, V], b: Tensor[T, V]) => where_!(condition, a, b))

core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala

@@ -0,0 +1,102 @@
+package deepwit.optimizer
+
+import deepwit.*
+import deepwit.base.AffineLayer
+import dimwit.*
+import dimwit.optimizer.Adam
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.funspec.AnyFunSpec
+
+class LossScaleSuite extends AnyFunSpec with Matchers:
+
+  private def vector(values: Float*): Tensor1[A, Float32] =
+    Tensor(Shape1(Axis[A] -> values.size)).fromArray(values.toArray)
+
+  private val params = AffineLayer.Params(
+    weight = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 2)).fromArray(Array(1f, 2f, 3f, 4f)),
+    bias = Tensor(Shape1(Axis[B] -> 2)).fromArray(Array(5f, 6f))
+  )
+
+  private def cost(p: AffineLayer.Params[A, B, Float32]): Tensor0[Float32] =
+    p.weight.pow(2).sum + p.bias.sum
+
+  describe("unscaled"):
+
+    it("recovers the gradients of the unscaled loss"):
+      val lossScale = LossScale.initial(1024f)
+      val expected = Autodiff.grad(cost)(params).value
+      val grads = lossScale.unscaled(Autodiff.grad((p: AffineLayer.Params[A, B, Float32]) => lossScale.scaled(cost(p)))(params)).value
+      grads.weight should approxEqual(expected.weight, 1e-5f)
+      grads.bias should approxEqual(expected.bias, 1e-5f)
+
+  describe("allFinite"):
+
+    it("holds for finite gradients"):
+      allFinite(Grad(vector(1f, -2f))).item shouldBe true
+
+    it("fails on inf and NaN"):
+      allFinite(Grad(vector(1f, Float.PositiveInfinity))).item shouldBe false
+      allFinite(Grad(vector(Float.NaN, 1f))).item shouldBe false
+
+  describe("next"):
+
+    it("halves the scale and resets the count after an overflow"):
+      val next = LossScale(Tensor0(1024f), Tensor0(7)).next(Tensor0(false))
+      next.scale.item shouldBe 512f
+      next.finiteSteps.item shouldBe 0
+
+    it("counts finite steps and doubles the scale after growthInterval of them"):
+      val afterTwo = (1 to 2).foldLeft(LossScale.initial(1024f))((lossScale, _) => lossScale.next(Tensor0(true), growthInterval = 3))
+      afterTwo.scale.item shouldBe 1024f
+      afterTwo.finiteSteps.item shouldBe 2
+
+      val afterThree = afterTwo.next(Tensor0(true), growthInterval = 3)
+      afterThree.scale.item shouldBe 2048f
+      afterThree.finiteSteps.item shouldBe 0
+
+    it("does not back off below minScale"):
+      LossScale(Tensor0(1f), Tensor0(0)).next(Tensor0(false)).scale.item shouldBe 1f
+
+  describe("select"):
+
+    it("keeps or replaces every leaf, whatever its shape, inside jit"):
+      val optimizer = Adam(Tensor0(0.1f))
+      val optState = optimizer.init(params)
+      val (newParams, newOptState) = optimizer.update(Autodiff.grad(cost)(params), params, optState)
+
+      val step = jit: (finite: Tensor0[Bool]) =>
+        (select(finite, newParams, params), select(finite, newOptState, optState))
+
+      val (skippedParams, skippedOptState) = step(Tensor0(false))
+      skippedParams.weight should approxEqual(params.weight, 0f)
+      skippedParams.bias should approxEqual(params.bias, 0f)
+      skippedOptState.momentums.weight should approxEqual(optState.momentums.weight, 0f)
+      skippedOptState.beta1t.item shouldBe optState.beta1t.item
+
+      val (takenParams, takenOptState) = step(Tensor0(true))
+      takenParams.weight should approxEqual(newParams.weight, 0f)
+      takenOptState.beta1t.item shouldBe newOptState.beta1t.item
+
+  describe("in Float16"):
+
+    val tiny = Tensor0(1e-4f).asFloat(VType[Float16])
+
+    // The gradient is tiny · tiny = 1e-8, below Float16's smallest value of about 6e-8.
+    def tinyGradientCost(w: Tensor0[Float32]): Tensor0[Float32] =
+      (w.asFloat(VType[Float16]) * tiny * tiny).asFloat(VType[Float32])
+
+    def gradWith(lossScale: LossScale): Grad[Tensor0[Float32]] =
+      lossScale.unscaled(Autodiff.grad((w: Tensor0[Float32]) => lossScale.scaled(tinyGradientCost(w)))(Tensor0(1f)))
+
+    it("loses a small gradient without scaling"):
+      Autodiff.grad(tinyGradientCost)(Tensor0(1f)).value.item shouldBe 0f
+
+    it("keeps it with scaling, after backing off from an overflow"):
+      // The loss's own gradient of 1, scaled by 2^16, exceeds Float16's largest value of 65504.
+      val initial = LossScale.initial()
+      val overflowed = gradWith(initial)
+      allFinite(overflowed).item shouldBe false
+
+      val grad = gradWith(initial.next(allFinite(overflowed)))
+      allFinite(grad).item shouldBe true
+      grad.value.item shouldBe (1e-8f +- 1e-10f)

examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala

@@ -0,0 +1,111 @@
+package deepwit.examples.mixedPrecision
+
+import dimwit.*
+import dimwit.optimizer.{Adam, AdamState}
+import dimwit.stats.Uniform
+import dimwit.TreeOf.ops.asFloats
+
+import deepwit.training.after
+import deepwit.activation.gelu
+import deepwit.base.{AffineFormLayer, AffineLayer}
+import deepwit.loss.SquaredError
+import deepwit.optimizer.{LossScale, allFinite, select}
+
+trait Sample derives Label
+trait Feature derives Label
+trait Embedding derives Label
+
+class MLP[V: IsFloating](params: MLP.Params[V]) extends (Tensor1[Feature, V] => Tensor0[V]):
+
+  private val layer1 = AffineLayer(params.layer1)
+  private val layer2 = AffineLayer(params.layer2)
+  private val output = AffineFormLayer(params.output)
+
+  def apply(x: Tensor1[Feature, V]): Tensor0[V] = output(gelu(layer2(gelu(layer1(x)))))
+
+object MLP:
+
+  case class Params[V](
+      layer1: AffineLayer.Params[Feature, Embedding, V],
+      layer2: AffineLayer.Params[Embedding, Embedding, V],
+      output: AffineFormLayer.Params[Embedding, V]
+  )
+
+  object Params:
+
+    def init(hiddenSize: Int, key: Key): Params[Float32] =
+      val (layer1Key, layer2Key, outputKey) = key.splitToTuple(3)
+      val hiddenExtent = Axis[Embedding] -> hiddenSize
+      Params(
+        layer1 = AffineLayer.Params.init(Axis[Feature] -> 1, hiddenExtent, layer1Key),
+        layer2 = AffineLayer.Params.init(hiddenExtent, hiddenExtent, layer2Key),
+        output = AffineFormLayer.Params.init(hiddenExtent, outputKey)
+      )
+
+case class TrainState(
+    params: MLP.Params[Float32],
+    optimizerState: AdamState[MLP.Params[Float32]],
+    lossScale: LossScale,
+    skippedSteps: Tensor0[Int32]
+)
+
+/** Fits a curve of amplitude 1e-3 in Float32, in Float16, and in Float16 with a [[LossScale]].
+  *
+  * The small amplitude makes the gradients small. Float16 rounds them to zero and barely learns; with the
+  * loss scale it matches Float32. Predicting zero everywhere would give an RMSE of about 0.71 amplitudes.
+  */
+@main
+def mixedPrecision(): Unit =
+
+  val numIterations = 3_000
+  val numSamples = 256
+  val hiddenSize = 32
+  val learningRate = 3e-3f
+  val amplitude = 1e-3f
+
+  val (dataKey, initKey) = Key(42).split2()
+  val shape = Shape(Axis[Sample] -> numSamples, Axis[Feature] -> 1)
+  val xs = Uniform(Tensor(shape).fill(-1f), Tensor(shape).fill(1f)).sample(dataKey)
+  val ys = (xs.slice(Axis[Feature].at(0)) *! (2f * math.Pi.toFloat)).sin *! amplitude
+
+  /** Runs the model in `vtype` and the loss in Float32. The master weights stay Float32, and so do their gradients. */
+  def costIn[V: IsFloating](vtype: VType[V])(params: MLP.Params[Float32]): Tensor0[Float32] =
+    val model = MLP(params.asFloats(vtype))
+    zipvmap(Axis[Sample])(xs.asFloat(vtype), ys): (x, y) =>
+      val prediction = model(x) * amplitude
+      SquaredError(y, prediction.asFloat(VType[Float32]))
+    .mean
+
+  val optimizer = Adam(learningRate)
+
+  def plainStep[V: IsFloating](vtype: VType[V])(state: TrainState): TrainState =
+    val grads = Autodiff.grad(costIn(vtype))(state.params)
+    val (params, optimizerState) = optimizer.update(grads, state.params, state.optimizerState)
+    state.copy(params = params, optimizerState = optimizerState)
+
+  def scaledStep(state: TrainState): TrainState =
+    val lossScale = state.lossScale
+    val grads = lossScale.unscaled(Autodiff.grad((params: MLP.Params[Float32]) => lossScale.scaled(costIn(VType[Float16])(params)))(state.params))
+    val (params, optimizerState) = optimizer.update(grads, state.params, state.optimizerState)
+    val finite = allFinite(grads)
+    TrainState(
+      select(finite, params, state.params),
+      select(finite, optimizerState, state.optimizerState),
+      lossScale.next(finite, growthInterval = 200),
+      where(finite, state.skippedSteps, state.skippedSteps + 1)
+    )
+
+  val initialState =
+    val params = MLP.Params.init(hiddenSize, initKey)
+    TrainState(params, optimizer.init(params), LossScale.initial(), Tensor0(0))
+
+  def run(name: String, step: TrainState => TrainState, scaled: Boolean = false): Unit =
+    val jitStep = jitDonatingUnsafe(step)
+    val state = Iterator.iterate(initialState)(jitStep).after(numIterations)
+    val rmse = math.sqrt(costIn(VType[Float32])(state.params).item) / amplitude
+    val scaleInfo = if scaled then f"   final scale: 2^${math.log(state.lossScale.scale.item) / math.log(2)}%.0f   skipped steps: ${state.skippedSteps.item}" else ""
+    println(f"$name%-20s RMSE / amplitude: $rmse%.4f$scaleInfo")
+
+  run("Float32", plainStep(VType[Float32]))
+  run("Float16", plainStep(VType[Float16]))
+  run("Float16 + LossScale", scaledStep, scaled = true)