From e52e7b4d301376be372095ffea729db2e853b5f8 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Fri, 2 Oct 2026 23:32:53 +0200 Subject: [PATCH] Add LossScale for mixed precision --- .../scala/deepwit/optimizer/LossScale.scala | 61 ++++++++++ .../deepwit/optimizer/LossScaleSuite.scala | 102 ++++++++++++++++ .../mixedPrecision/MixedPrecision.scala | 111 ++++++++++++++++++ 3 files changed, 274 insertions(+) create mode 100644 core/src/main/scala/deepwit/optimizer/LossScale.scala create mode 100644 core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala create mode 100644 examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala diff --git a/core/src/main/scala/deepwit/optimizer/LossScale.scala b/core/src/main/scala/deepwit/optimizer/LossScale.scala new file mode 100644 index 0000000..87df54a --- /dev/null +++ b/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)) diff --git a/core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala b/core/src/test/scala/deepwit/optimizer/LossScaleSuite.scala new file mode 100644 index 0000000..ff3af5f --- /dev/null +++ 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) diff --git a/examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala b/examples/src/main/scala/deepwit/examples/mixedPrecision/MixedPrecision.scala new file mode 100644 index 0000000..da290a4 --- /dev/null +++ b/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)