diff --git a/CHANGELOG.md b/CHANGELOG.md index 1dce90e95..3feadee3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,11 +4,19 @@ ### Fixed +- **A same-dtype `PrecisionBoundary` leaves no cast in a trace** (`skainet-lang-core`, SKEEP-008 phase A2, #1433): `stage("body", dtype = FP32)` inside an FP32 network called `convert`, which returns its input on the CPU backend but is still recorded by the tracing wrapper, so the exported module carried two `stablehlo.convert` `f32 -> f32`. The boundary now passes a tensor of the target dtype through without calling `convert`. - **`concat` backward with inputs of different sizes** (`skainet-compile-dag`, #1490): the rule split the upstream gradient by one input's size and picked chunk `offset / size`, which only lines up when every input has the same extent along the concat axis; other shapes got another input's slice, or `null`, without an error. Each input now receives its own `narrow` of the upstream at the offset the forward gave it, for any sizes and a negative `dim`. - **`squeeze` backward with no axis** (`skainet-compile-dag`, #1490): `squeeze(x)` removes every size-1 axis, but the backward unsqueezed dim 0 alone, so `x: [1, 3, 1]` got a `[1, 3]` gradient. The gradient is now the upstream reshaped to the input's shape. +### Removed + +- **`dense()`, `stage()` and the mixed-precision scaffold** (`skainet-lang-core`, SKEEP-008 phase A4, #1433): the generic overloads of `NeuralNetworkDsl`, deprecated at `ERROR` level in 0.63.0, are gone; use `dense(outputDimension, dtype = FP16) { }` / `stage(id, dtype = FP16) { }`. `MixedPrecisionModule`, `InternalMixedPrecisionModule` and `MixedPrecisionTensorOps`, deprecated in 0.63.0 and never implemented, are removed, together with `MixedPrecisionException`, which only they could throw. A third-party `NeuralNetworkDsl` implementation drops its three overrides. + ### Added +- **Training through a mixed-precision boundary** (`skainet-lang-core`, `skainet-compile-dag`, SKEEP-008 phase A3, #1433): `TensorOps.convert` is a `@Diff` op. Its backward rule treats the cast as the identity and casts the upstream gradient back to the source dtype, so a gradient entering an FP16 / BF16 region is rounded to that format and the layers behind `dense(dtype = …)` / `stage(id, dtype = …)` train; a cast from or to an integer dtype has no gradient. `MixedPrecisionBackwardTest` checks the rule, a finite-difference gradient through FP16 boundaries, parity with the all-FP32 network, and an SGD run. The inference-only notes are gone; the ops matrix Backward column reads 60/60. +- **SKEEP-008 acceptance tests** (`skainet-lang-core`, `skainet-backend-cpu`, #1433): `repackTo(Dense(4))` is checked bit-exact for Q5_0, Q5_1, Q5_K, Q6_K and `BITNET_B1_58` too, and one FP32 network with both a BF16 `dense` and an FP16 `stage` runs against the FP32 reference. +- **Mixed-precision export test** (`skainet-compile-hlo`, SKEEP-008 phase A2, #1433): `MixedPrecisionExportTest` traces an FP32 network with one BF16 `dense` and one FP16 `stage` and pins the result: one `convert` pair per boundary in the graph, inner layers typed in the narrow dtype, and `stablehlo.convert` `f32 -> bf16`, `bf16 -> f32`, `f32 -> f16`, `f16 -> f32` in the exported module. - **`scaledDotProductAttention` backward** (`skainet-compile-dag`, #1384): the attention adjoint (dQ, dK, dV and the dense mask's gradient) replaces the stub that returned `null`, so attention layers train. It honours `scale`, `causal`, sliding `window` / `rightContext`, `keyLengths` and the dense mask, recomputes the probabilities instead of storing them, and reduces grouped-query heads into their shared K/V. Masked scores get exactly zero gradient. `SdpaBackwardTest` checks every option against finite differences and three cases against PyTorch autograd; `stubbedBackwardOpNames` is empty and the coverage matrix's Backward column is complete. ## [0.63.0] - 2026-11-06 diff --git a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc index 0ce5203ac..1b9164e02 100644 --- a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc +++ b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc @@ -102,8 +102,9 @@ ctx.ops.convert(inner.forward(x, ctx), outerType) // cast back out Three properties follow from using `TensorOps.convert` — an ordinary tensor op — rather than hiding the conversion inside the layer: -* *Free when degenerate.* `convert` returns its input unchanged when the dtypes already - match, so a boundary between equal dtypes costs nothing. +* *Free when degenerate.* A tensor that already has the target dtype is passed through + without a `convert` call, so a boundary between equal dtypes costs nothing and leaves no + cast in a trace or an export. * *Visible in traces and exports.* The tracing wrappers record `convert` like any other op, and the StableHLO converter lowers it to `stablehlo.convert` — a compiled or exported graph shows exactly where precision changes, and a later optimisation pass can fold adjacent @@ -115,14 +116,55 @@ hiding the conversion inside the layer: The dtype is passed as a *value* (`dense(256, dtype = BF16)`), not as a type argument. The earlier API sketched `dense(256)`, but an interface method's type parameter cannot be `reified`, so the implementation had no `KClass` to build the layer's tensors with — those -overloads could never be implemented and are now deprecated. Passing the `DType` object gives +overloads could never be implemented and have been removed. Passing the `DType` object gives the implementation both the instance (for `convert`) and `dtype::class` (for tensor creation). +== What an export looks like + +Tracing a network and exporting it to StableHLO gives one `stablehlo.convert` pair per +boundary, with everything between them typed in the narrow format. For the network above +without its `relu`, so that the BF16 layer is directly followed by the FP16 stage, the +module has this shape (abridged: constants, transposes and bias adds are left out): + +[source,mlir] +---- +%v6 = stablehlo.convert %arg0 : (tensor<1x784xf32>) -> tensor<1x784xbf16> // into the BF16 dense +%v10 = stablehlo.dot_general %v6, %v7, ... -> tensor<1x256xbf16> +%v13 = stablehlo.convert %v12 : (tensor<1x256xbf16>) -> tensor<1x256xf32> // and back out +%v14 = stablehlo.convert %v13 : (tensor<1x256xf32>) -> tensor<1x256xf16> // into the FP16 stage +%v15 = stablehlo.dot_general %v14, %v8, ... -> tensor<1x64xf16> +%v18 = stablehlo.dot_general %v17, %v9, ... -> tensor<1x10xf16> +%v21 = stablehlo.convert %v20 : (tensor<1x10xf16>) -> tensor<1x10xf32> // and back out +---- + +The `bf16 -> f32` cast followed directly by `f32 -> f16` is the price of two adjacent +boundaries with different inner dtypes: the network's own dtype is always restored in +between. `MixedPrecisionExportTest` pins this output. + +== Training through a boundary + +`convert` is differentiable between floating-point dtypes, so a mixed-precision layer or +stage trains like any other. Its backward rule treats the cast as the identity and casts the +gradient back to the source dtype: + +* where the gradient *leaves* a narrow region (the cast back out, seen from the loss side + first) it is rounded to the narrow format, exactly as the forward values were on the way in; +* where it *enters* the outer network again it is widened, which is exact. + +The weights behind a boundary are stored in the narrow dtype, their gradients have that dtype +too, and the optimizer updates them in place. That is a faithful simulation of training in the +narrow format, with its limits: an update smaller than half a unit in the last place of a +weight is lost. Keeping FP32 master weights and scaling the loss to avoid that is automatic +mixed precision, which is out of scope (see below). A cast from or to an integer dtype has no +gradient. + +One thing to know when checking gradients by hand: `input(n)` clears `requiresGrad` on the +tensor it receives unless asked not to (`input(n, requiresGrad = true)`). With the flag +cleared, the first cast has nothing to propagate and the network *input* gets no gradient; the +parameters are not affected. + == What it does not do (yet) -* *Training.* `convert` is not differentiable today, so gradients stop at a boundary: a - mixed-precision layer or stage is inference-only. SKEEP-008 phase A3 makes `convert` a - `@Diff` op, which lifts this. * *Automatic mixed precision.* There is no autocast policy, no FP32 master weights, no loss scaling. The dtype is always an explicit, local choice in the code. * *Integer layers.* A cast to Int8 truncates; it is not quantization. Integer inference goes diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc index df4dc37d3..8ee5a9aaa 100644 --- a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc +++ b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc @@ -72,7 +72,7 @@ Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemen | `TensorOps.lt` | ✖ | — | ❌ | — | — | ✅ | — | `TensorOps.ge` | ✖ | — | ❌ | — | — | ✅ | — | `TensorOps.tril` | ✖ | — | ❌ | ✅ | — | ✅ | — -| `TensorOps.convert` | ✖ | — | ✅ | — | — | ✅ | — +| `TensorOps.convert` | ✖ | — | ✅ | ✅ | — | ✅ | — | `TensorOps.gather` | ✖ | — | ✅ | ✅ | — | ✅ | — | `TensorOps.indexSelect` | ✖ | — | ✅ | ✅ | — | ✅ | — | `TensorOps.exp` | ✖ | — | ✅ | ✅ | — | ✅ | — @@ -83,7 +83,7 @@ Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemen | `TensorOps.scaledDotProductAttention` | ✖ | composite · decompose | ✅ | ✅ | — | ✅ | — | `Similarity.cosineDistance` | ✖ | — | ❌ | — | ✅ | ✅ | ✅ -| *Done* | *1 / 67* | *3 / 67* | *51 / 67* | *59 / 59* | *1 / 67* | *67 / 67* | *1 / 67* +| *Done* | *1 / 67* | *3 / 67* | *51 / 67* | *60 / 60* | *1 / 67* | *67 / 67* | *1 / 67* |=== Per-function detail including notes lives in xref:reference/operators/generated/index.adoc[Operator reference]. diff --git a/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc b/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc index c0fc23960..16e3469d7 100644 --- a/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc +++ b/docs/modules/skeep/pages/008-precision-and-encoding-conversion.adoc @@ -1,12 +1,28 @@ = SKEEP-008: Precision and encoding conversion :description: SKaiNET proposal for the two ways a tensor's numeric form changes after it exists: a per-layer or per-stage dtype in the network DSL (mixed precision), and re-encoding stored weights between dense and packed formats (repackTo and an explicit, lossy requantizeTo). -Status: Draft + +Status: Implemented + Audience: SKaiNET maintainers and contributors; users of the network DSL; loader and kernel authors who consume `TensorEncoding` + Created: 2026-10-21 + Tracking issue: https://github.com/SKaiNET-developers/SKaiNET/issues/1433[#1433] + Origin: the stub inventory behind the operator coverage matrix (https://github.com/SKaiNET-developers/SKaiNET/pull/1370[#1370]). Six `TODO()` stubs in `NetworkBuilder.kt` (`dense()` ×2, `stage()`, in both DSL implementations) and `TensorStorage.repackTo`, which throws for every encoding change, were left without a tracking issue because both change public API semantics. +== Implementation status + +[cols="1,3,2",options="header"] +|=== +| Phase | What shipped | Where + +| B1 | Lossless `repackTo`, lossy `requantizeTo` | https://github.com/SKaiNET-developers/SKaiNET/pull/1470[#1470], 0.63.0 +| A1 | `PrecisionBoundary`, `dense(dtype = …)`, `stage(id, dtype = …)`, deprecations | https://github.com/SKaiNET-developers/SKaiNET/pull/1469[#1469], 0.63.0 +| A2 | Trace and StableHLO export pinned by `MixedPrecisionExportTest`; a same-dtype boundary no longer records a cast | https://github.com/SKaiNET-developers/SKaiNET/pull/1492[#1492] +| A3 | `convert` is `@Diff`; gradient check and training through a boundary in `MixedPrecisionBackwardTest` | https://github.com/SKaiNET-developers/SKaiNET/pull/1492[#1492] +| A4 | Generic `dense()` / `stage()` overloads and the `MixedPrecision*` scaffold removed | https://github.com/SKaiNET-developers/SKaiNET/pull/1492[#1492] +|=== + +The sections below are the proposal as accepted. Where the implementation differs, the text +says so; phase 1 restrictions that A3 lifted are kept for the record. + == Summary SKaiNET has two public entry points that promise to change a tensor's numeric form and do @@ -181,8 +197,9 @@ public class PrecisionBoundary( } ---- -`convert` already returns its input when the dtype matches, so a boundary between equal dtypes -costs nothing. Two consecutive boundaries of the same inner dtype are a cast followed by its +A tensor that already has the target dtype is passed through without calling `convert` (the +tracing wrapper would record the call even though the backend returns its input), so a boundary +between equal dtypes costs nothing and leaves no cast in a trace. Two consecutive boundaries of the same inner dtype are a cast followed by its inverse; a later optimisation pass may fold them, and nothing in this design depends on that. Weights of a `TLayer` dense are created in `TLayer` by the existing initialisers. User-supplied @@ -328,13 +345,19 @@ Each phase is a DARC feature with lane sub-issues, linked from the tracking issu . Should a `reified` extension (`dense(64)`) be offered as sugar over the value form? It reads well, but extensions named like members interact badly with overload resolution and IDE - completion. + completion. + + *Resolved: no.* Only the value form shipped. . Should a stage boundary that ends in the same dtype as the next stage elide the cast pair at - DSL time, or leave that to a compile pass? + DSL time, or leave that to a compile pass? + + *Open.* The DSL elides only the degenerate case (inner dtype equal to the network dtype). Two + adjacent boundaries still export a `narrow -> f32` cast followed by `f32 -> narrow`; folding + that pair is left to a compile pass. . Is `TensorStorage` the right owner of `requantizeTo`, or should it be a free function in - `sk.ainet.lang.tensor.data` next to `Requantizers`, keeping `TensorStorage` a descriptor? + `sk.ainet.lang.tensor.data` next to `Requantizers`, keeping `TensorStorage` a descriptor? + + *Resolved:* it is a member of `TensorStorage`, next to `repackTo`. . Does phase 2 need per-boundary `DTypePolicy` metadata, so that - `DTypeConstraintResolutionPass` can see and validate the boundaries in a compiled graph? + `DTypeConstraintResolutionPass` can see and validate the boundaries in a compiled graph? + + *Open.* A3 did not need it: the boundaries are ordinary `convert` nodes in the graph. == References diff --git a/docs/modules/skeep/pages/index.adoc b/docs/modules/skeep/pages/index.adoc index 05580e0fc..1083211bb 100644 --- a/docs/modules/skeep/pages/index.adoc +++ b/docs/modules/skeep/pages/index.adoc @@ -140,6 +140,6 @@ Every proposal should include: | The browser lane — packed SIMD128 kernels as a sidecar module, quantize-at-load for SafeTensors, and the common-code fixes that keep wasm off the reference kernel | xref:skeep:008-precision-and-encoding-conversion.adoc[SKEEP-008] -| Draft +| Implemented | Precision and encoding conversion — per-layer and per-stage dtype in the network DSL, lossless `repackTo` and lossy `requantizeTo` for stored weights |=== diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt index 8a3e34dc6..13d926946 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt @@ -121,6 +121,40 @@ class MixedPrecisionDslTest { assertTrue(maxDiff > 0f, "FP16 rounding must change the result, or the cast never happened") } + @Test + fun bf16DenseAndFp16StageInOneNetwork_runAndRound() { + // The SKEEP-008 acceptance network: both boundary kinds in one FP32 network. + val mixed = definition { + network(ctx) { + input(2) + dense(3, dtype = BF16, id = "layer1") { + weights { fromArray(w1) } + bias { fromArray(b1) } + } + stage("head", dtype = FP16) { + dense(2, "layer2") { + weights { fromArray(w2) } + bias { fromArray(b2) } + } + } + } + } + + val reference = fp32Reference(ctx).forward(input(), ctx) + val result = mixed.forward(input(), ctx) + + assertEquals(Shape(1, 2), result.shape) + assertEquals(FP32::class, result.dtype, "the last boundary must cast back to the network dtype") + assertEquals(2, mixed.modules.filterIsInstance>().size) + var maxDiff = 0f + for (c in 0 until 2) { + val diff = abs(result.data[0, c] - reference.data[0, c]) + assertTrue(diff <= 5e-2f, "mixed network drifted too far at $c: diff=$diff") + maxDiff = maxOf(maxDiff, diff) + } + assertTrue(maxDiff > 0f, "narrow rounding must change the result, or the casts never happened") + } + @Test fun sameDtypeBoundary_isExact() { // convert returns its input when the dtype matches, so an FP32 stage inside an diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index 2e5222817..a5e205e02 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -106,6 +106,7 @@ public final class sk/ainet/lang/graph/DefaultGradientTape : sk/ainet/lang/graph public fun conv3dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun convTranspose1dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun convTranspose2dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; + public fun convertBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun copy ()Lsk/ainet/tape/ExecutionTape; public fun cosBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun divScalarBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt index 2a02d63c7..668880bef 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt @@ -1156,6 +1156,27 @@ public class DefaultGradientTape( return listOf(upstream.ops.tril(upstream, k)) } + override fun convertBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { + // d(convert)/dx is taken as the identity: the rounding of a narrowing cast is passed + // straight through. The gradient goes back in the source dtype, so a gradient entering + // an FP16 / BF16 region is itself rounded to that format, as the forward values are. + val input = inputs[0] + // A backend may return its input for a same-dtype cast. That call is still recorded, + // with one tensor as both input and output; a gradient for it would be added to itself. + if (session.refOf(input).id == session.refOf(output).id) return listOf(null) + val sourceType = floatDTypeOf(input) ?: return listOf(null) + if (floatDTypeOf(output) == null) return listOf(null) + return listOf(upstream.ops.convert(upstream, sourceType)) + } + + /** The dtype instance of a floating-point tensor that `convert` can target, else `null`. */ + private fun floatDTypeOf(tensor: Tensor): DType? = when (tensor.dtype) { + sk.ainet.lang.types.FP32::class -> sk.ainet.lang.types.FP32 + sk.ainet.lang.types.FP16::class -> sk.ainet.lang.types.FP16 + sk.ainet.lang.types.BF16::class -> sk.ainet.lang.types.BF16 + else -> null + } + override fun gatherBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { // gather(input[vocab,emb], indices, dim=0) -> backward scatter-adds upstream rows into the gathered rows. // Differentiable w.r.t. input only; indices gradient is null. @@ -1475,6 +1496,7 @@ public class DefaultGradientTape( "sin" to ::sinBackward, "cos" to ::cosBackward, "tril" to ::trilBackward, + "convert" to ::convertBackward, "gather" to ::gatherBackward, "indexSelect" to ::indexSelectBackward, "unfold" to ::unfoldBackward, diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/MixedPrecisionBackwardTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/MixedPrecisionBackwardTest.kt new file mode 100644 index 000000000..59482435d --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/MixedPrecisionBackwardTest.kt @@ -0,0 +1,273 @@ +package sk.ainet.exec.autograd + +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue +import sk.ainet.context.Phase +import sk.ainet.exec.tensor.ops.DefaultCpuOps +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.DefaultGradientTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.definition +import sk.ainet.lang.nn.network +import sk.ainet.lang.nn.optim.sgd +import sk.ainet.lang.nn.topology.ModuleParameter +import sk.ainet.lang.nn.topology.collectParams +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.withRequiresGrad +import sk.ainet.lang.trace.GraphSink +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 + +/** + * SKEEP-008 phase A3 (#1433): `convert` is differentiable between floating-point dtypes, so + * gradients cross a [sk.ainet.lang.nn.PrecisionBoundary] and the layers behind it train. + */ +class MixedPrecisionBackwardTest { + + private fun ctx(): DefaultGraphExecutionContext { + val dataFactory = DenseTensorDataFactory() + val graph = DefaultComputeGraph() + return DefaultGraphExecutionContext( + baseOps = DefaultCpuOps(dataFactory), + phase = Phase.TRAIN, + tensorDataFactory = dataFactory, + createTapeFactory = { _ -> DefaultGradientTape(true) }, + computeGraph = graph, + baseSink = GraphSink(graph), + ) + } + + private fun DefaultGraphExecutionContext.floats(shape: Shape, values: FloatArray): Tensor = + fromFloatArray(shape, FP32::class, values) + + private fun values(t: Tensor<*, *>): FloatArray { + val dims = t.shape.dimensions + val idx = IntArray(dims.size) + return FloatArray(t.shape.volume) { + val v = (t.data.get(*idx) as Number).toFloat() + for (d in dims.indices.reversed()) { + if (++idx[d] < dims[d]) break + idx[d] = 0 + } + v + } + } + + /** `sum(out ⊙ weights)`: the gradient of this loss w.r.t. `out` is `weights`. */ + private fun weightedSum(c: DefaultGraphExecutionContext, out: Tensor, weights: FloatArray): Tensor { + val weighted = out.ops.multiply(out, c.floats(out.shape, weights)) + return weighted.ops.sum(weighted) + } + + private fun gradOf( + c: DefaultGraphExecutionContext, + x: Tensor, + loss: DefaultGraphExecutionContext.() -> Tensor, + ): FloatArray { + val (tape, l) = c.record { (this as DefaultGraphExecutionContext).loss() } + (tape as DefaultGradientTape).computeGradients(targets = listOf(l), sources = listOf(x)) + return values(assertNotNull(x.grad, "no gradient reached the source")) + } + + // ── the cast itself ─────────────────────────────────────────────────────────────────────── + + @Test + fun convert_roundTripThroughFp16_passesTheGradientStraightThrough() { + val c = ctx() + val x = c.floats(Shape(2, 2), floatArrayOf(0.1234567f, -1.7654321f, 2.3456789f, 0.0123456f)).withRequiresGrad() + val w = floatArrayOf(1f, 2f, 3f, 4f) // exactly representable in FP16 + val g = gradOf(c, x) { + weightedSum(this, ops.convert(ops.convert(x, FP16), FP32), w) + } + assertContentEquals(w, g) + } + + @Test + fun convert_gradientEnteringANarrowRegion_isRoundedToThatFormat() { + val c = ctx() + val x = c.floats(Shape(1, 2), floatArrayOf(0.5f, -0.25f)).withRequiresGrad() + val w = floatArrayOf(0.1234567f, 1.7654321f) // not representable in BF16 + val g = gradOf(c, x) { + weightedSum(this, ops.convert(ops.convert(x, BF16), FP32), w) + } + for (i in w.indices) { + assertTrue(abs(g[i] - w[i]) <= 1e-2f * abs(w[i]), "gradient $i drifted: ${g[i]} vs ${w[i]}") + } + assertTrue(g.indices.any { g[it] != w[it] }, "the gradient crossed a BF16 tensor and must show its rounding") + } + + @Test + fun convert_toTheSameDtype_doesNotDoubleTheGradient() { + // The CPU backend returns its input here, yet the call is recorded: input and output + // of the recorded op are one tensor, and the rule must not add its gradient to itself. + val c = ctx() + val x = c.floats(Shape(1, 3), floatArrayOf(1f, 2f, 3f)).withRequiresGrad() + val w = floatArrayOf(5f, 6f, 7f) + val g = gradOf(c, x) { weightedSum(this, ops.convert(x, FP32), w) } + assertContentEquals(w, g) + } + + @Test + fun convert_throughAnIntegerDtype_hasNoGradient() { + val c = ctx() + val x = c.floats(Shape(1, 3), floatArrayOf(1.5f, 2.5f, 3.5f)).withRequiresGrad() + val g = gradOf(c, x) { + weightedSum(this, ops.convert(ops.convert(x, Int32), FP32), floatArrayOf(5f, 6f, 7f)) + } + assertContentEquals(floatArrayOf(0f, 0f, 0f), g) + } + + // ── through the DSL boundary ────────────────────────────────────────────────────────────── + + private val w1 = floatArrayOf( + 0.1234567f, -1.7654321f, + 2.3456789f, 0.0123456f, + -0.9876543f, 1.1111111f, + ) + private val b1 = floatArrayOf(0.1111111f, -0.2222222f, 0.3333333f) + private val w2 = floatArrayOf( + 1.9999999f, -0.1111111f, 0.5555555f, + -1.3333333f, 0.7777777f, -0.4444444f, + ) + private val b2 = floatArrayOf(0.0101010f, -0.0202020f) + private val lossWeights = floatArrayOf(1f, -2f) + private val x0 = floatArrayOf(0.3f, -1.1f) + + /** input(2) → dense(3) → dense(2), linear, with the first layer in [first] and the second in a stage of [second]. */ + private fun network(c: DefaultGraphExecutionContext, first: A, second: B): Module = + definition { + network(c) { + // The Input module sets the flag on the tensor it is given; without it the + // first cast has nothing to propagate and the input gets no gradient. + input(2, requiresGrad = true) + dense(3, dtype = first, id = "layer1") { + weights { fromArray(w1) } + bias { fromArray(b1) } + } + stage("head", dtype = second) { + dense(2, "layer2") { + weights { fromArray(w2) } + bias { fromArray(b2) } + } + } + } + } + + private fun lossAt(first: A, second: B, input: FloatArray): Float { + val c = ctx() + val model = network(c, first, second) + val out = model.forward(c.floats(Shape(1, 2), input), c) + return values(weightedSum(c, out, lossWeights))[0] + } + + private class Grads(val input: FloatArray, val params: Map, val paramDtypes: Map) + + private fun grads(first: A, second: B): Grads { + val c = ctx() + val model = network(c, first, second) + val x = c.floats(Shape(1, 2), x0).withRequiresGrad() + val params = model.collectParams() + assertEquals(4, params.size, "two dense layers behind boundaries: ${params.map { it.name }}") + val (tape, loss) = c.record { weightedSum(c, model.forward(x, c), lossWeights) } + @Suppress("UNCHECKED_CAST") + val sources = (listOf(x) + params.map { it.value }) as List> + (tape as DefaultGradientTape).computeGradients(targets = listOf(loss), sources = sources) + return Grads( + input = values(assertNotNull(x.grad)), + params = params.associate { p -> p.name to values(assertNotNull(p.value.grad, "${p.name} has no gradient")) }, + paramDtypes = params.associate { p -> p.name to assertNotNull(p.value.grad).dtype }, + ) + } + + private fun assertClose(expected: FloatArray, actual: FloatArray, relTol: Float, what: String) { + assertEquals(expected.size, actual.size, what) + val scale = expected.maxOf { abs(it) }.coerceAtLeast(1f) + for (i in expected.indices) { + assertTrue( + abs(expected[i] - actual[i]) <= relTol * scale, + "$what[$i]: expected ${expected[i]}, got ${actual[i]} (tolerance ${relTol * scale})", + ) + } + } + + @Test + fun gradientCheck_throughFp16Boundaries_matchesFiniteDifferences() { + // The network is linear, so a central difference is exact up to the rounding the + // boundaries add; a wide step keeps that rounding small against the difference. + val h = 0.5f + val numeric = FloatArray(x0.size) { i -> + val plus = x0.copyOf().also { it[i] += h } + val minus = x0.copyOf().also { it[i] -= h } + (lossAt(FP16, FP16, plus) - lossAt(FP16, FP16, minus)) / (2 * h) + } + assertClose(numeric, grads(FP16, FP16).input, relTol = 1e-2f, what = "dLoss/dInput") + } + + @Test + fun gradients_throughBf16DenseAndFp16Stage_matchTheFp32Network() { + val reference = grads(FP32, FP32) + val mixed = grads(BF16, FP16) + + assertClose(reference.input, mixed.input, relTol = 3e-2f, what = "dLoss/dInput") + for ((name, expected) in reference.params) { + assertClose(expected, assertNotNull(mixed.params[name]), relTol = 3e-2f, what = "dLoss/d$name") + } + assertTrue( + reference.params.any { (name, expected) -> !expected.contentEquals(mixed.params[name]) }, + "narrow gradients must show the rounding, or the casts never happened", + ) + // A parameter's gradient has the parameter's dtype, so an optimizer updates like with like. + assertTrue(mixed.paramDtypes.values.containsAll(listOf(BF16::class, FP16::class)), "${mixed.paramDtypes}") + } + + @Test + fun training_updatesTheLayersBehindABoundary_andLowersTheLoss() { + val c = ctx() + val model = network(c, BF16, FP16) + val x = c.floats(Shape(1, 2), x0) + val target = c.floats(Shape(1, 2), floatArrayOf(0.5f, -0.5f)) + val params = model.collectParams() + val before = params.associate { it.name to values(it.value) } + + fun loss(): Tensor { + val out = model.forward(x, c) + val diff = out.ops.subtract(out, target) + val sq = diff.ops.multiply(diff, diff) + return sq.ops.sum(sq) + } + + val initial = values(loss())[0] + val optimizer = sgd(lr = 0.01) + @Suppress("UNCHECKED_CAST") + params.forEach { optimizer.addParameter(it as ModuleParameter) } + repeat(5) { + val (tape, l) = c.record { loss() } + @Suppress("UNCHECKED_CAST") + (tape as DefaultGradientTape).computeGradients( + targets = listOf(l), + sources = params.map { it.value } as List>, + ) + optimizer.step() + optimizer.zeroGrad() + } + val final = values(loss())[0] + + assertTrue(final < initial, "loss must go down: $initial -> $final") + for (p in params) { + assertTrue( + !before.getValue(p.name).contentEquals(values(p.value)), + "${p.name} sits behind a boundary and must have been updated", + ) + } + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt new file mode 100644 index 000000000..be94a9f59 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/MixedPrecisionExportTest.kt @@ -0,0 +1,154 @@ +package sk.ainet.compile.hlo + +import sk.ainet.compile.opt.passes.DtypeForwardPropagationPass +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.graph.ComputeGraph +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.definition +import sk.ainet.lang.nn.network +import sk.ainet.lang.tape.toComputeGraph +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +/** + * SKEEP-008 phase A2 (#1433): a network with `dense(dtype = …)` / `stage(id, dtype = …)` + * traces with a `convert` node on each side of every [sk.ainet.lang.nn.PrecisionBoundary], + * and exports to StableHLO with one `stablehlo.convert` pair per boundary. + */ +class MixedPrecisionExportTest { + + /** Traces one forward pass of the network [build] makes, on the CPU ops. */ + private fun trace( + build: (DefaultGraphExecutionContext) -> Module, + ): ComputeGraph { + val ctx = DefaultGraphExecutionContext.tape(baseOps = DirectCpuExecutionContext().ops) + val model = build(ctx) + val input = ctx.fromFloatArray( + shape = Shape(1, 4), dtype = FP32::class, + data = floatArrayOf(0.3f, -1.1f, 0.7f, 2.5f), + ) + val inputRefId = ctx.session.refOf(input as Tensor<*, *>).id + val (tape, result) = ctx.record { model.forward(input, ctx) } + assertNotNull(tape, "the forward pass must be recorded") + assertEquals(FP32::class, result?.dtype, "the last boundary must cast back to the network dtype") + return tape.toComputeGraph(synthesizeExternalInputs = true, inputTensorIds = setOf(inputRefId)) + } + + /** FP32 network with one BF16 dense and one FP16 stage: the SKEEP-008 acceptance network. */ + private fun mixedNetwork(ctx: DefaultGraphExecutionContext): Module = definition { + network(ctx) { + input(4) + dense(3, dtype = BF16, id = "bf16Dense") + stage("head", dtype = FP16) { + dense(3, "h1") + dense(2, "h2") + } + } + } + + private fun convertLines(mlir: String): List = + mlir.lines().map { it.trim() }.filter { "stablehlo.convert" in it } + + @Test + fun trace_hasOneConvertPairPerBoundary() { + val graph = trace(::mixedNetwork) + + val casts = graph.getTopologicalOrder() + .filter { it.operation.name == "convert" } + .map { it.inputs.single().dtype to it.outputs.single().dtype } + + assertEquals(4, casts.size, "two boundaries, one cast in and one cast out each: $casts") + assertEquals( + setOf("FP32" to "BF16", "BF16" to "FP32", "FP32" to "FP16", "FP16" to "FP32"), + casts.toSet(), + ) + } + + @Test + fun trace_runsInnerLayersInTheNarrowDtype() { + val graph = trace(::mixedNetwork) + + val matmulDtypes = graph.getTopologicalOrder() + .filter { it.operation.name == "matmul" } + .map { it.outputs.single().dtype } + + // One BF16 dense, then the two dense layers of the FP16 stage. No FP32 matmul: + // every layer of this network sits behind a boundary. + assertEquals(listOf("BF16", "FP16", "FP16"), matmulDtypes) + } + + @Test + fun export_emitsStableHloConvertAtEachBoundary_inOrder() { + val module = StableHloConverterFactory.createExtended().convert(trace(::mixedNetwork), "mixed") + val mlir = module.content + + assertFalse("Unsupported" in mlir, "every op must lower:\n$mlir") + assertFalse("tensor) -> tensor<1x4xbf16>", + "(tensor<1x3xbf16>) -> tensor<1x3xf32>", + "(tensor<1x3xf32>) -> tensor<1x3xf16>", + "(tensor<1x2xf16>) -> tensor<1x2xf32>", + ), + signatures, + "one stablehlo.convert pair per boundary, in execution order:\n$mlir", + ) + + // Inner computation is typed in the narrow dtype; the function signature stays FP32. + assertTrue("-> tensor<1x3xbf16>" in mlir.lines().first { "stablehlo.dot_general" in it }, mlir) + assertEquals( + 2, + mlir.lines().count { "stablehlo.dot_general" in it && it.trim().endsWith("xf16>") }, + "both dense layers of the FP16 stage must be f16:\n$mlir", + ) + assertTrue("(%arg0: tensor<1x4xf32>) -> (tensor<1x2xf32>)" in mlir, mlir) + } + + @Test + fun export_sameDtypeBoundary_emitsNoConvert() { + // convert returns its input when the dtype already matches, so nothing is recorded: + // an FP32 stage inside an FP32 network costs nothing in the exported module either. + val graph = trace { ctx -> + definition { + network(ctx) { + input(4) + stage("body", dtype = FP32) { + dense(3, "h1") + dense(2, "h2") + } + } + } + } + + assertTrue(graph.getTopologicalOrder().none { it.operation.name == "convert" }) + val mlir = StableHloConverterFactory.createExtended().convert(graph, "same").content + assertTrue(convertLines(mlir).isEmpty(), "no cast expected:\n$mlir") + } + + @Test + fun dtypePropagation_keepsTheBoundaries() { + // The pass makes edges dtype-consistent and must treat convert as intentionally + // dtype-changing: a traced mixed-precision graph comes out with the same casts + // and exports to the same module. + val graph = trace(::mixedNetwork) + val propagated = DtypeForwardPropagationPass().apply(graph).graph + + val converter = StableHloConverterFactory.createExtended() + assertEquals( + convertLines(converter.convert(graph, "mixed").content), + convertLines(converter.convert(propagated, "mixed").content), + ) + } +} 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 2974404cb..8136cd139 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 @@ -2591,38 +2591,6 @@ public final class sk/ainet/lang/nn/Conv3d : sk/ainet/lang/nn/Module, sk/ainet/l public final fun outputSize (Lkotlin/Triple;)Lkotlin/Triple; } -public final class sk/ainet/lang/nn/ConversionStats { - public fun ()V - public fun (JJJJJJ)V - public synthetic fun (JJJJJJILkotlin/jvm/internal/DefaultConstructorMarker;)V - public final fun component1 ()J - public final fun component2 ()J - public final fun component3 ()J - public final fun component4 ()J - public final fun component5 ()J - public final fun component6 ()J - public final fun copy (JJJJJJ)Lsk/ainet/lang/nn/ConversionStats; - public static synthetic fun copy$default (Lsk/ainet/lang/nn/ConversionStats;JJJJJJILjava/lang/Object;)Lsk/ainet/lang/nn/ConversionStats; - public fun equals (Ljava/lang/Object;)Z - public final fun getCacheClears ()J - public final fun getCacheHitRatio ()D - public final fun getCacheHits ()J - public final fun getConversionErrors ()J - public final fun getConversionRatio ()D - public final fun getInputConversions ()J - public final fun getOutputConversions ()J - public final fun getTotalForwardCalls ()J - public fun hashCode ()I - public final fun reset ()V - public final fun setCacheClears (J)V - public final fun setCacheHits (J)V - public final fun setConversionErrors (J)V - public final fun setInputConversions (J)V - public final fun setOutputConversions (J)V - public final fun setTotalForwardCalls (J)V - public fun toString ()Ljava/lang/String; -} - public final class sk/ainet/lang/nn/DefaultNeuralNetworkExecutionContext : sk/ainet/lang/nn/NeuralNetworkExecutionContext { public fun ()V public fun (Lsk/ainet/context/Phase;)V @@ -2856,12 +2824,6 @@ public final class sk/ainet/lang/nn/Input : sk/ainet/lang/nn/Module { public final fun getRequiresGrad ()Z } -public abstract class sk/ainet/lang/nn/InternalMixedPrecisionModule : sk/ainet/lang/nn/Module { - public fun (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;)V - public fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; - protected abstract fun forwardImpl (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; -} - public final class sk/ainet/lang/nn/LayerScale : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { public fun (ILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;)V public synthetic fun (ILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -2921,30 +2883,6 @@ public final class sk/ainet/lang/nn/MaxPool2d : sk/ainet/lang/nn/Module { public final fun outputSize (Lkotlin/Pair;)Lkotlin/Pair; } -public final class sk/ainet/lang/nn/MixedPrecisionException : java/lang/RuntimeException { - public fun (Ljava/lang/String;Ljava/lang/Throwable;)V - public synthetic fun (Ljava/lang/String;Ljava/lang/Throwable;ILkotlin/jvm/internal/DefaultConstructorMarker;)V -} - -public abstract class sk/ainet/lang/nn/MixedPrecisionModule : sk/ainet/lang/nn/Module { - public fun (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;)V - public final fun clearConversionCache ()V - protected fun convertInput (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; - protected fun convertOutput (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; - public final fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; - protected abstract fun forwardImpl (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; - public final fun getAutoConvertInput ()Z - public final fun getAutoConvertOutput ()Z - public final fun getCacheSize ()I - protected final fun getConversionOps ()Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps; - public final fun getConversionStats ()Lsk/ainet/lang/nn/ConversionStats; - protected final fun getInputType ()Lsk/ainet/lang/types/DType; - protected final fun getOutputType ()Lsk/ainet/lang/types/DType; - public final fun setAutoConvertInput (Z)V - public final fun setAutoConvertOutput (Z)V - public final fun validatePrecisionConfig ()V -} - public abstract class sk/ainet/lang/nn/Module : sk/ainet/lang/nn/topology/ModuleNode { public fun ()V public fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; @@ -3530,15 +3468,11 @@ public abstract interface class sk/ainet/lang/nn/dsl/NeuralNetworkDsl : sk/ainet public static synthetic fun conv2d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public abstract fun conv3d (ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;)V public static synthetic fun conv3d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V - public abstract fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public abstract fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V public abstract fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V - public abstract fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public abstract fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V - public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V - public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public abstract fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public static synthetic fun flatten$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V @@ -3559,7 +3493,6 @@ public abstract interface class sk/ainet/lang/nn/dsl/NeuralNetworkDsl : sk/ainet public abstract fun sequential (Lkotlin/jvm/functions/Function1;)V public abstract fun softmax (ILjava/lang/String;)V public static synthetic fun softmax$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;ILjava/lang/Object;)V - public abstract fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public abstract fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public abstract fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V public abstract fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V @@ -3576,10 +3509,8 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDsl$DefaultImpls { public static synthetic fun conv2d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun conv2d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun conv3d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V - public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V - public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun flatten$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun groupNorm$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;IIDZLjava/lang/String;ILjava/lang/Object;)V @@ -3605,10 +3536,8 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDslImpl : sk/ainet/lang/nn/ public fun conv2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun conv3d (ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;)V public final fun create ()Lsk/ainet/lang/nn/Module; - public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V - public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public final fun getCurrentShape ()[I @@ -3627,7 +3556,6 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDslImpl : sk/ainet/lang/nn/ public final fun setCurrentShape ([I)V public final fun setLastDimension (I)V public fun softmax (ILjava/lang/String;)V - public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V public fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V @@ -3648,10 +3576,8 @@ public final class sk/ainet/lang/nn/dsl/StageImpl : sk/ainet/lang/nn/dsl/NeuralN public fun conv2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun conv3d (ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;)V public final fun create ()Lsk/ainet/lang/nn/Module; - public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V - public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public final fun getCurrentShape ()[I @@ -3672,7 +3598,6 @@ public final class sk/ainet/lang/nn/dsl/StageImpl : sk/ainet/lang/nn/dsl/NeuralN public final fun setInputDimension (I)V public final fun setLastDimension (I)V public fun softmax (ILjava/lang/String;)V - public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V public fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V @@ -6424,6 +6349,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/DifferentiableTensorOps public abstract fun conv3dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun convTranspose1dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun convTranspose2dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; + public abstract fun convertBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun cosBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun divScalarBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun divideBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; @@ -6673,16 +6599,6 @@ public final class sk/ainet/lang/tensor/ops/MaxPool2dOperation : sk/ainet/lang/t public fun validateInputs (Ljava/util/List;)Lsk/ainet/lang/tensor/ops/ValidationResult; } -public abstract interface class sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps : sk/ainet/lang/tensor/ops/TensorOps { - public abstract fun addMixed (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; - public abstract fun convert (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/tensor/Tensor; -} - -public final class sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps$DefaultImpls { - public static fun matmulWeightTransposed (Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; - public static fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; -} - public final class sk/ainet/lang/tensor/ops/MultiplyOperation : sk/ainet/lang/tensor/ops/BaseOperation { public fun ()V public fun (Ljava/util/Map;)V diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt deleted file mode 100644 index 4e7afd561..000000000 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt +++ /dev/null @@ -1,311 +0,0 @@ -package sk.ainet.lang.nn - -import sk.ainet.context.ExecutionContext -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.tensor.ops.MixedPrecisionTensorOps -import sk.ainet.lang.types.DType -//import sk.ainet.lang.nn.dsl.LayerConfig - -/** - * Abstract base class for modules that support mixed-precision operations. - * This class handles automatic type conversions between different precision types, - * enabling seamless integration of layers with different precision requirements. - * - * The class manages the complexity of precision conversions, allowing derived - * classes to focus on their core functionality while ensuring type safety - * and performance optimization. - * - * @param TInput The precision type for input tensors - * @param TOutput The precision type for output tensors - * @param V The value type corresponding to the DType - * @param inputType The input precision type instance - * @param outputType The output precision type instance - * @param conversionOps Mixed precision tensor operations for handling conversions - * - * Example usage: - * ```kotlin - * class MixedPrecisionLinear( - * inputType: TInput, - * outputType: TOutput, - * conversionOps: MixedPrecisionTensorOps, - * private val inFeatures: Int, - * private val outFeatures: Int - * ) : MixedPrecisionModule(inputType, outputType, conversionOps) { - * - * override fun forwardImpl(input: Tensor): Tensor { - * // Implement layer-specific logic here - * // Input and output conversions are handled automatically - * } - * } - * ``` - */ -@Deprecated( - message = "Never implemented; replaced by PrecisionBoundary (SKEEP-008, #1433). Will be removed.", - replaceWith = ReplaceWith("PrecisionBoundary", "sk.ainet.lang.nn.PrecisionBoundary"), -) -public abstract class MixedPrecisionModule( - protected val inputType: TInput, - protected val outputType: TOutput, - protected val conversionOps: MixedPrecisionTensorOps -) : Module() { - - /** - * Configuration for precision handling within this module. - * Can be used to customize conversion behavior, caching, and validation. - */ - //public var precisionConfig: LayerConfig? = null - - /** - * Whether to enable automatic input conversion. - * When false, input tensors must already be in the expected input type. - */ - public var autoConvertInput: Boolean = true - - /** - * Whether to enable automatic output conversion. - * When false, output tensors will be in the module's internal precision. - */ - public var autoConvertOutput: Boolean = true - - /** - * Cache for converted tensors to avoid repeated conversions. - * Key is the tensor hash, value is the converted tensor. - */ - private val conversionCache = mutableMapOf>() - - /** - * Statistics for monitoring conversion performance. - */ - public val conversionStats: ConversionStats = ConversionStats() - - /** - * Main forward pass implementation that handles precision conversions automatically. - * This method wraps the actual implementation with conversion logic. - */ - final override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = - sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { - conversionStats.totalForwardCalls++ - try { - // Convert input if necessary and requested - val convertedInput = if (autoConvertInput && !isSameType(input, inputType)) { - conversionStats.inputConversions++ - convertInput(input) - } else { - input - } - - // Call the actual implementation - val internalOutput = forwardImpl(convertedInput) - - // Convert output if necessary and requested - val finalOutput = if (autoConvertOutput && !isSameType(internalOutput, inputType)) { - conversionStats.outputConversions++ - @Suppress("UNCHECKED_CAST") - convertOutput(internalOutput) as Tensor - } else { - @Suppress("UNCHECKED_CAST") - internalOutput as Tensor - } - - finalOutput - } catch (e: Exception) { - conversionStats.conversionErrors++ - throw MixedPrecisionException("Error in mixed precision forward pass: ${e.message}", e) - } - } - - /** - * Abstract method that derived classes must implement. - * This method contains the actual layer logic and operates on converted tensors. - * - * @param input Input tensor in the expected input precision - * @return Output tensor in the module's output precision - */ - protected abstract fun forwardImpl(input: Tensor): Tensor - - /** - * Converts input tensor to the module's expected input precision. - * Can be overridden by derived classes for custom conversion logic. - */ - protected open fun convertInput(input: Tensor): Tensor { - val cacheKey = input.hashCode() - - // Check cache first if caching is enabled - /* - if (precisionConfig?.cacheConvertedWeights == true) { - val cached = conversionCache[cacheKey] - if (cached != null) { - conversionStats.cacheHits++ - @Suppress("UNCHECKED_CAST") - return cached as Tensor - } - } - - - */ - // Perform conversion - val converted = conversionOps.convert(input, inputType) - - // Cache result if enabled - /* if (precisionConfig?.cacheConvertedWeights == true) { - conversionCache[cacheKey] = converted - } - - */ - - return converted - } - - /** - * Converts output tensor from the module's internal precision to the expected output precision. - * Can be overridden by derived classes for custom conversion logic. - */ - protected open fun convertOutput(output: Tensor): Tensor { - return conversionOps.convert(output, outputType) - } - - /** - * Validates precision compatibility for the module configuration. - * Should be called during module initialization. - */ - public fun validatePrecisionConfig() { - //precisionConfig?.validate() - - // Additional validation specific to mixed precision modules - if (inputType == outputType && (autoConvertInput || autoConvertOutput)) { - throw IllegalArgumentException( - "Input and output types are the same (${inputType::class.simpleName}), " + - "but conversion is enabled. Consider disabling auto-conversion for better performance." - ) - } - } - - /** - * Clears the conversion cache to free memory. - * Should be called periodically in long-running applications. - */ - public fun clearConversionCache() { - conversionCache.clear() - conversionStats.cacheClears++ - } - - /** - * Gets the current cache size for monitoring purposes. - */ - public fun getCacheSize(): Int = conversionCache.size - - /** - * Helper method to check if two tensors have the same precision type. - */ - private fun isSameType(tensor: Tensor<*, V>, targetType: DType): Boolean { - // This is a simplified check - in practice, this would need to be more sophisticated - return tensor.toString().contains(targetType::class.simpleName ?: "Unknown") - } -} - -/** - * Specialized mixed precision module for cases where input and output types are the same - * but internal computation uses a different precision. - * - * This is commonly used for layers where weights are stored in low precision (e.g., INT8) - * but computations are performed in higher precision (e.g., FP16) for accuracy. - * - * @param TIO The input/output precision type - * @param TInternal The internal computation precision type - * @param V The value type corresponding to the DType - */ -@Deprecated( - message = "Never implemented; replaced by PrecisionBoundary (SKEEP-008, #1433). Will be removed.", - replaceWith = ReplaceWith("PrecisionBoundary", "sk.ainet.lang.nn.PrecisionBoundary"), -) -public abstract class InternalMixedPrecisionModule( - private val ioType: TIO, - private val internalType: TInternal, - private val conversionOps: MixedPrecisionTensorOps -) : Module() { - - /** - * Forward implementation that ensures output is converted back to input type. - */ - override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = - sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { - // Convert input to internal precision if needed - val internalInput = if (ioType != internalType) { - conversionOps.convert(input, internalType) - } else { - @Suppress("UNCHECKED_CAST") - input as Tensor - } - - // Perform internal computation - val internalOutput = forwardImpl(internalInput) - - // Convert back to IO precision if needed - if (internalType != ioType) { - conversionOps.convert(internalOutput, ioType) - } else { - @Suppress("UNCHECKED_CAST") - internalOutput as Tensor - } - } - - /** - * Abstract method for internal computation in the internal precision. - */ - protected abstract fun forwardImpl(input: Tensor): Tensor -} - -/** - * Statistics class for tracking mixed precision conversion performance. - */ -public data class ConversionStats( - var totalForwardCalls: Long = 0, - var inputConversions: Long = 0, - var outputConversions: Long = 0, - var conversionErrors: Long = 0, - var cacheHits: Long = 0, - var cacheClears: Long = 0 -) { - /** - * Gets the cache hit ratio as a percentage. - */ - public fun getCacheHitRatio(): Double { - val totalCacheableOperations = inputConversions + outputConversions - return if (totalCacheableOperations > 0) { - (cacheHits.toDouble() / totalCacheableOperations) * 100.0 - } else { - 0.0 - } - } - - /** - * Gets the conversion ratio (conversions per forward call). - */ - public fun getConversionRatio(): Double { - return if (totalForwardCalls > 0) { - (inputConversions + outputConversions).toDouble() / totalForwardCalls - } else { - 0.0 - } - } - - /** - * Resets all statistics. - */ - public fun reset() { - totalForwardCalls = 0 - inputConversions = 0 - outputConversions = 0 - conversionErrors = 0 - cacheHits = 0 - cacheClears = 0 - } -} - -/** - * Exception thrown when mixed precision operations fail. - */ -public class MixedPrecisionException( - message: String, - cause: Throwable? = null -) : RuntimeException(message, cause) \ No newline at end of file diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt index fa83e6cca..b875f69fc 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt @@ -9,13 +9,17 @@ import sk.ainet.lang.types.DType * the input is cast to [innerType], [inner] runs, and the result is cast back to * [outerType] (SKEEP-008, #1433). * - * The casts use [sk.ainet.lang.tensor.ops.TensorOps.convert], which returns its input - * unchanged when the dtypes already match, so a boundary between equal dtypes costs - * nothing. `convert` is recorded by the tracing wrappers and lowers to - * `stablehlo.convert`, so the boundary stays visible in traces and exports. + * The casts use [sk.ainet.lang.tensor.ops.TensorOps.convert]. A tensor that already has + * the target dtype is passed through without calling it, so a boundary between equal + * dtypes costs nothing and leaves no cast in a trace. Otherwise `convert` is recorded by + * the tracing wrappers and lowers to `stablehlo.convert`, so the boundary stays visible + * in traces and exports. * - * `convert` is not differentiable (SKEEP-008 phase A3 lifts this), so gradients stop at - * the boundary: the wrapped layers do not train. + * `convert` is differentiable between floating-point dtypes, so gradients cross the + * boundary and the wrapped layers train. The gradient is cast like the values: it is + * rounded to [innerType] where it enters the wrapped module and widened back to + * [outerType] where it leaves. There is no FP32 master copy of the weights and no loss + * scaling; that is automatic mixed precision, which this module does not do. */ public class PrecisionBoundary( override val name: String, @@ -31,7 +35,17 @@ public class PrecisionBoundary( get() = listOf(inner as Module) override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { - val x = ctx.ops.convert(input, innerType) - return ctx.ops.convert(inner.forward(x, ctx), outerType) + val x = cast(input, innerType, ctx) + return cast(inner.forward(x, ctx), outerType, ctx) } + + // A tracing wrapper records every `convert` call, also the one that returns its input, + // so the same-dtype case is skipped here rather than left to the backend. + @Suppress("UNCHECKED_CAST") + private fun cast( + tensor: Tensor, + target: TTo, + ctx: ExecutionContext, + ): Tensor = + if (tensor.dtype == target::class) tensor as Tensor else ctx.ops.convert(tensor, target) } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt index fb910c4d0..a0b79ac09 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt @@ -169,8 +169,9 @@ public interface NeuralNetworkDsl : NetworkDslItem { /** * Creates a dense layer whose weights and computation use [dtype] while the network stays * in its own dtype: the layer's input is cast to [dtype] and its output cast back - * (SKEEP-008, #1433). The cast is `TensorOps.convert`, which is not differentiable yet, - * so a mixed-precision layer is inference-only: gradients stop at its boundary. + * (SKEEP-008, #1433). The cast is `TensorOps.convert`, which is differentiable between + * floating-point dtypes, so the layer trains: gradients cross the boundary and are + * rounded to [dtype] on the way in, as the forward values are. * * @param outputDimension The number of neurons/output features * @param dtype The dtype the layer's weights and computation use (e.g. `FP16`, `BF16`) @@ -184,27 +185,6 @@ public interface NeuralNetworkDsl : NetworkDslItem { content: DENSE.() -> Unit = {}, ) - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - public fun dense( - outputDimension: Int, - id: String = "", - content: DENSE.() -> Unit = {} - ): Module - - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - public fun dense( - id: String = "", - content: DENSE.() -> Unit = {} - ): Module - /** * Applies an activation function as a separate layer. * @@ -453,8 +433,9 @@ public interface NeuralNetworkDsl : NetworkDslItem { /** * Creates a stage whose layers all use [dtype] while the network stays in its own dtype: * the stage's input is cast to [dtype] and its output cast back (SKEEP-008, #1433). - * The cast is `TensorOps.convert`, which is not differentiable yet, so a mixed-precision - * stage is inference-only: gradients stop at its boundary. + * The cast is `TensorOps.convert`, which is differentiable between floating-point dtypes, + * so the stage trains: gradients cross the boundary and are rounded to [dtype] on the way + * in, as the forward values are. * * @param id Identifier for the stage * @param dtype The dtype all layers in this stage use (e.g. `FP16`, `BF16`) @@ -465,16 +446,6 @@ public interface NeuralNetworkDsl : NetworkDslItem { dtype: TStage, content: NeuralNetworkDsl.() -> Unit, ) - - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - public fun stage( - id: String, - content: NeuralNetworkDsl.() -> Unit - ): Module } public interface WandBTensorValueContext { @@ -1387,46 +1358,6 @@ public class StageImpl( ) } - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - override fun dense( - outputDimension: Int, - id: String, - content: DENSE.() -> Unit - ): Module = error( - "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + - "dense(outputDimension, dtype = FP16, id) { ... }" - ) - - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - override fun dense( - id: String, - content: DENSE.() -> Unit - ): Module = error( - "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + - "dense(outputDimension, dtype = FP16, id) { ... }" - ) - - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - override fun stage( - id: String, - content: NeuralNetworkDsl.() -> Unit - ): Module = error( - "Mixed-precision stage was redesigned (SKEEP-008): pass the dtype as a value, " + - "stage(id, dtype = FP16) { ... }" - ) - override fun batchNorm( numFeatures: Int, eps: Double, @@ -1843,46 +1774,6 @@ public class NeuralNetworkDslImpl( ) } - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - override fun dense( - outputDimension: Int, - id: String, - content: DENSE.() -> Unit - ): Module = error( - "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + - "dense(outputDimension, dtype = FP16, id) { ... }" - ) - - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - override fun dense( - id: String, - content: DENSE.() -> Unit - ): Module = error( - "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + - "dense(outputDimension, dtype = FP16, id) { ... }" - ) - - @Deprecated( - message = "The dtype type argument cannot be implemented (no KClass at runtime, " + - "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }", - level = DeprecationLevel.ERROR, - ) - override fun stage( - id: String, - content: NeuralNetworkDsl.() -> Unit - ): Module = error( - "Mixed-precision stage was redesigned (SKEEP-008): pass the dtype as a value, " + - "stage(id, dtype = FP16) { ... }" - ) - override fun batchNorm( numFeatures: Int, eps: Double, diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt deleted file mode 100644 index f29a90fc3..000000000 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt +++ /dev/null @@ -1,20 +0,0 @@ -package sk.ainet.lang.tensor.ops - -import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.types.DType - -@Deprecated( - message = "Never implemented: TensorOps.convert already covers the cast, and the DSL's " + - "mixed-precision layers use PrecisionBoundary (SKEEP-008, #1433). Will be removed.", -) -public interface MixedPrecisionTensorOps : TensorOps { - override fun convert( - tensor: Tensor, - targetType: TTo - ): Tensor - - public fun addMixed( - a: Tensor, - b: Tensor - ): Tensor -} \ No newline at end of file diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index ad78125b8..570963ea1 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -348,7 +348,15 @@ public interface TensorOps { @Diff public fun tril(tensor: Tensor, k: Int = 0): Tensor - // Type conversion operations. Non-differentiable by design (dtype cast). + /** + * Casts [tensor] to [targetType]. Narrowing to FP16 / BF16 rounds to the format. + * + * Differentiable between floating-point dtypes (SKEEP-008 phase A3): the backward pass + * casts the upstream gradient back to the dtype of [tensor], treating the rounding as + * the identity (a straight-through estimate). A cast from or to an integer dtype has no + * gradient. + */ + @Diff public fun convert( tensor: Tensor, targetType: TTo diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/RepackRequantizeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/RepackRequantizeTest.kt index 85e81fb4d..91374f1ae 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/RepackRequantizeTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/RepackRequantizeTest.kt @@ -81,6 +81,35 @@ class RepackRequantizeTest { assertContentEquals(dequantized, floatsOf(dense)) } + @Test + fun repackTo_encodingsWithoutAQuantizer_toDense_reproduceDequantizedValues_bitExact() { + // No quantizer produces these, so the blocks are synthetic: every byte below 0x3C keeps + // the FP16 / FP32 scale fields finite. 512 elements is two K-quant super-blocks. + val blockShape = Shape(2, 256) + val encodings = listOf( + TensorEncoding.Q5_0, TensorEncoding.Q5_1, TensorEncoding.Q5_K, TensorEncoding.Q6_K, + TensorEncoding.BITNET_B1_58, + ) + for (encoding in encodings) { + val byteCount = checkNotNull(encoding.physicalBytes(blockShape.volume.toLong())).toInt() + val bytes = ByteArray(byteCount) { ((it * 37 + 11) % 60).toByte() } + val source = TensorStorage( + shape = blockShape, + dtype = FP32, + encoding = encoding, + buffer = BufferHandleFactory.owned(bytes), + ) + val dequantized = floatsOf(source) + assertTrue(dequantized.any { it != 0f }, "$encoding: the synthetic blocks must decode to something") + + val dense = source.repackTo(TensorEncoding.Dense(4)) + + assertEquals(TensorEncoding.Dense(4), dense.encoding, "$encoding") + assertEquals(blockShape, dense.shape, "$encoding") + assertContentEquals(dequantized, floatsOf(dense), "$encoding") + } + } + @Test fun repackTo_narrowDense_widensExactly() { // Dense(2) FP16 bytes -> Dense(4) FP32; every FP16 value is exactly representable. diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOpsTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeCompatibilityTest.kt similarity index 79% rename from skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOpsTest.kt rename to skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeCompatibilityTest.kt index 39b1bdec0..dbd802d63 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOpsTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeCompatibilityTest.kt @@ -1,30 +1,15 @@ -package sk.ainet.lang.tensor.ops +package sk.ainet.lang.types import kotlin.test.Test import kotlin.test.assertTrue import kotlin.test.assertNotNull -import sk.ainet.lang.types.* import kotlin.test.assertEquals /** - * Unit tests for MixedPrecisionTensorOps conversion methods. - * - * This test suite covers Task 1.4.1 requirements for testing TensorOps conversion methods - * and ensures that all conversion methods work correctly with different precision types. - * - * Note: This is a simplified test that validates the interface exists and compiles. - * Full functional testing would require more complex tensor implementations. + * DType compatibility, promotion and precision-validation rules: which dtypes mix, what + * they promote to, and what the validation helpers report. */ -class MixedPrecisionTensorOpsTest { - - @Test - fun testMixedPrecisionTensorOpsInterfaceExists() { - // Test that the MixedPrecisionTensorOps interface exists - // This is a compile-time validation - if this compiles, the interface exists - val interfaceName = MixedPrecisionTensorOps::class.simpleName - assertNotNull(interfaceName, "MixedPrecisionTensorOps interface should exist") - assertEquals(interfaceName, "MixedPrecisionTensorOps", "Interface should have correct name") - } +class DTypeCompatibilityTest { @Test fun testDTypeCompatibilityMethods() {