From dfc861a1815d16de69dbe11d224c7b7fcaf4ad10 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Tue, 25 Aug 2026 16:45:35 +0200 Subject: [PATCH] feat(io): report every weight the loader re-encodes on the way in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1117. Slice 4 of #1109. "Why is this model bigger than the file" had no answer except whoever remembered which policy was set. Now each conversion the loader performs emits a TraceEvent.AdapterInserted naming the tensor and the sizes either side; a weight that arrives in the form it will be used in emits nothing. Reusing AdapterInserted rather than adding an event type, because its own documentation already describes this — "the dispatcher inserted a conversion (dequantize, requantize, gather)" — and a loader-side re-encode is the same thing at a different moment. It also means the Perfetto, JFR and android.os.Trace exporters need no changes to show it, which is the third acceptance criterion satisfied by not writing code. The event gained one defaulted field, bytesBefore, and a bytesDelta derived from it: a conversion's cost is the difference between two sizes and one number cannot express it. Perfetto and JFR now carry both. Three conversions emit, not one. The issue asked for the ones the resolver causes, which today means form-driven dequantization alone — but the loader performs two widenings nobody ever asked for, and those are the ones most worth seeing. Ternary tensors widen to FP32 whatever the policy, because packed ternary storage does not exist yet (#1033): roughly 20×, and no flag on this loader turns it off. Narrow floats double unless keepF16Native or keepBf16Native is set. Reporting only the requested conversion would answer the easy half of the question. The sink defaults to NoopTraceSink and every emission is guarded on isEnabled, so a caller who asked for no trace builds no Formats and allocates nothing. Gate: scripts/pr-gate.sh — all legs passed. Co-Authored-By: Claude Fable 5 --- .../io/gguf/StreamingGgufParametersLoader.kt | 76 ++++++++++ .../sk/ainet/io/gguf/WeightFormTraceTest.kt | 137 ++++++++++++++++++ .../api/jvm/skainet-lang-core.api | 11 +- .../memory/trace/PerfettoTraceExporter.kt | 7 +- .../sk/ainet/lang/memory/trace/TraceEvent.kt | 16 +- .../ainet/lang/memory/trace/JfrTraceSink.kt | 4 +- 6 files changed, 243 insertions(+), 8 deletions(-) create mode 100644 skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormTraceTest.kt diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index 387d359dc..b84293597 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -4,6 +4,12 @@ import sk.ainet.context.ExecutionContext import sk.ainet.io.ParametersLoader import sk.ainet.io.RandomAccessSource import sk.ainet.io.gguf.dequant.DequantOps +import sk.ainet.lang.memory.Format +import sk.ainet.lang.memory.ScopeKind +import sk.ainet.lang.memory.trace.NoopTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.memory.trace.TraceSink +import sk.ainet.lang.tensor.TensorId import sk.ainet.lang.memory.plan.EncodingRequest import sk.ainet.lang.memory.plan.WeightByteOrder import sk.ainet.lang.memory.plan.WeightForm @@ -115,8 +121,52 @@ public class StreamingGgufParametersLoader( * silently resolved, so nobody loses a setting they thought they had. */ private val weightForm: WeightForm? = null, + /** + * Where conversions are reported (#1117). + * + * A weight that arrives in the form it will be used in costs nothing and says nothing. One that + * is re-encoded on the way in — dequantized because no kernel can feed its encoding, widened + * because it is a narrow float, or widened because ternary packed storage does not exist yet + * (#1033) — emits a `TraceEvent.AdapterInserted` naming the tensor and the sizes either side. + * That is the difference between "why is this model 3 GB" being answerable and being folklore. + * + * Defaults to [NoopTraceSink]: nothing is recorded and nothing is allocated. + */ + private val traceSink: TraceSink = NoopTraceSink, ) : ParametersLoader { + /** + * Report that [tensorName] was re-encoded from [from] to [to] on the way in. + * + * Guarded on [TraceSink.isEnabled] so the common case builds no `Format`s and allocates + * nothing. The tensor is named by parsing its GGUF name as a [TensorId] — `blk.0.attn_q.weight` + * is already dotted, and renders back as itself. + */ + private fun traceConversion( + kind: String, + tensorName: String, + from: Format, + to: Format, + bytesBefore: Long, + bytesAfter: Long, + ) { + if (!traceSink.isEnabled) return + traceSink.emit( + TraceEvent.AdapterInserted( + kind = kind, + from = from, + to = to, + bytes = bytesAfter, + target = runCatching { TensorId.parse(tensorName) }.getOrNull(), + scope = ScopeKind.MODEL, + bytesBefore = bytesBefore, + ), + ) + } + + /** The dense FP32 size of [elements] — what every widening in this loader converts to. */ + private fun denseFp32Bytes(elements: Long): Long = elements * 4 + /** The three axes as one value: [weightForm] if given, otherwise what the three parameters say. */ private val form: WeightForm = weightForm ?: WeightForm( encoding = when (quantPolicy) { @@ -259,6 +309,10 @@ public class StreamingGgufParametersLoader( ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) } else { // Loader-owned widened array — zero-copy wrap (#782). + traceConversion( + "widen-f16", tensorInfo.name, Format.dense(FP16), Format.dense(FP32), + rawBytes.size.toLong(), denseFp32Bytes(tensorInfo.nElements), + ) ctx.wrapFloatArray(shape, dtype, dequantF16(rawBytes)) as Tensor } else -> null @@ -271,6 +325,10 @@ public class StreamingGgufParametersLoader( ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) } else { // Loader-owned widened array — zero-copy wrap (#782). + traceConversion( + "widen-bf16", tensorInfo.name, Format.dense(BF16), Format.dense(FP32), + rawBytes.size.toLong(), denseFp32Bytes(tensorInfo.nElements), + ) ctx.wrapFloatArray(shape, dtype, dequantBF16(rawBytes)) as Tensor } else -> null @@ -331,6 +389,14 @@ public class StreamingGgufParametersLoader( (dtype == FP32::class || dtype == FP16::class) ) { val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt()) + traceConversion( + kind = "dequantize-on-load", + tensorName = tensorInfo.name, + from = ggufFormat(tensorInfo.tensorType, rawBytes.size.toLong()), + to = Format.dense(FP32), + bytesBefore = rawBytes.size.toLong(), + bytesAfter = denseFp32Bytes(tensorInfo.nElements), + ) return ctx.wrapFloatArray(shape, dtype, dest) as Tensor } if (tensorInfo.tensorType == GGMLQuantizationType.TQ1_0 || tensorInfo.tensorType == GGMLQuantizationType.TQ2_0) { @@ -342,6 +408,16 @@ public class StreamingGgufParametersLoader( "load as FP32, so the requested dtype $dtype is not supported" } val dest = DequantOps.dequantFromBytes(rawBytes, tensorInfo.tensorType, tensorInfo.nElements.toInt()) + // Worth seeing precisely because no policy asked for it: a 1.6-bit weight arriving as + // FP32 is a ~20× widening that no flag on this loader can currently turn off (#1033). + traceConversion( + kind = "widen-ternary-no-packed-storage", + tensorName = tensorInfo.name, + from = ggufFormat(tensorInfo.tensorType, rawBytes.size.toLong()), + to = Format.dense(FP32), + bytesBefore = rawBytes.size.toLong(), + bytesAfter = denseFp32Bytes(tensorInfo.nElements), + ) @Suppress("UNCHECKED_CAST") return ctx.wrapFloatArray(shape, dtype, dest) as Tensor } diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormTraceTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormTraceTest.kt new file mode 100644 index 000000000..c57f570bc --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormTraceTest.kt @@ -0,0 +1,137 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.EncodingRequest +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import java.io.File +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1117: a weight that is re-encoded on the way in says so; one that is not stays silent. + * + * The question this makes answerable is "why is this model bigger than the file". Before, the + * answer lived in whoever remembered which policy was set. + */ +@OptIn(ExperimentalMemoryApi::class) +class WeightFormTraceTest { + + private fun file(): File = SyntheticGguf.write( + SyntheticGguf.tensor("blk.0.attn_q.weight", GGMLQuantizationType.Q4_K, elements = 1024), + SyntheticGguf.tensor("blk.0.ffn_up.weight", GGMLQuantizationType.Q8_0, elements = 1024), + SyntheticGguf.tensor("token_embd.weight", GGMLQuantizationType.F32, elements = 512), + ) + + private fun load(f: File, form: WeightForm?, sink: RecordingTraceSink) { + val ctx = DefaultDataExecutionContext() + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(f) }, + weightForm = form, + traceSink = sink, + ).load(ctx, FP32::class) { _, _ -> } + } + } + + private fun conversions(sink: RecordingTraceSink): List = + sink.events().filterIsInstance() + + @Test + fun `a weight held as stored converts nothing and says nothing`() { + val f = file() + try { + val sink = RecordingTraceSink() + load(f, WeightForm.AS_STORED_ON_HEAP, sink) + assertTrue( + conversions(sink).isEmpty(), + "the common case must stay silent, got ${conversions(sink).map { it.kind }}", + ) + } finally { + f.delete() + } + } + + @Test + fun `a dequantized weight reports the tensor and both sizes`() { + val f = file() + try { + val sink = RecordingTraceSink() + load(f, WeightForm(encoding = EncodingRequest.DequantizeTo(FP32)), sink) + + val events = conversions(sink) + assertEquals(2, events.size, "one per quantized weight; F32 needs no conversion: ${events.map { it.kind }}") + + val q4k = events.single { it.from.encoding == TensorEncoding.Q4_K } + assertEquals("dequantize-on-load", q4k.kind) + assertEquals("blk.0.attn_q.weight", q4k.target?.canonical, "the event names the tensor") + assertEquals(1024L * 4, q4k.bytes, "afterwards: dense FP32") + assertTrue(q4k.bytesBefore < q4k.bytes, "before: the packed bytes, ${q4k.bytesBefore}") + assertEquals(q4k.bytes - q4k.bytesBefore, q4k.bytesDelta, "the delta is what it cost") + assertTrue(q4k.bytesDelta > 0, "a dequantization only ever adds") + } finally { + f.delete() + } + } + + @Test + fun `the widening nobody asked for is the one most worth seeing`() { + // Ternary tensors widen to FP32 whatever the policy, because packed ternary storage does + // not exist yet (#1033). No flag reveals that; the trace does. + val f = SyntheticGguf.write( + SyntheticGguf.tensor("blk.0.ffn_down.weight", GGMLQuantizationType.TQ2_0, elements = 512), + ) + try { + val sink = RecordingTraceSink() + load(f, WeightForm.AS_STORED_ON_HEAP, sink) + + val event = conversions(sink).single() + assertEquals("widen-ternary-no-packed-storage", event.kind) + assertEquals(512L * 4, event.bytes) + assertTrue( + event.bytes > event.bytesBefore * 5, + "a ~2-bit weight as FP32 is a large multiple: ${event.bytesBefore} → ${event.bytes}", + ) + } finally { + f.delete() + } + } + + @Test + fun `a narrow float widened at load is reported too`() { + val f = SyntheticGguf.write(SyntheticGguf.tensor("w", GGMLQuantizationType.F16, elements = 256)) + try { + val sink = RecordingTraceSink() + load(f, WeightForm.AS_STORED_ON_HEAP, sink) + + val event = conversions(sink).single() + assertEquals("widen-f16", event.kind) + assertEquals(256L * 2, event.bytesBefore) + assertEquals(256L * 4, event.bytes, "F16 to FP32 doubles it") + } finally { + f.delete() + } + } + + @Test + fun `the default loader records nothing at all`() { + // NoopTraceSink is the default, so a caller who never asked for a trace pays for none. + val f = file() + try { + val ctx = DefaultDataExecutionContext() + runBlocking { + StreamingGgufParametersLoader(sourceProvider = { JvmRandomAccessSource.open(f) }) + .load(ctx, FP32::class) { _, _ -> } + } + } finally { + f.delete() + } + } +} 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 2c542fba2..869f5e484 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 @@ -2042,8 +2042,8 @@ public abstract interface class sk/ainet/lang/memory/trace/TraceEvent { } public final class sk/ainet/lang/memory/trace/TraceEvent$AdapterInserted : sk/ainet/lang/memory/trace/TraceEvent { - public fun (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;J)V - public synthetic fun (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JJ)V + public synthetic fun (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JJILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; public final fun component2 ()Lsk/ainet/lang/memory/Format; public final fun component3 ()Lsk/ainet/lang/memory/Format; @@ -2051,10 +2051,13 @@ public final class sk/ainet/lang/memory/trace/TraceEvent$AdapterInserted : sk/ai public final fun component5 ()Lsk/ainet/lang/tensor/TensorId; public final fun component6 ()Lsk/ainet/lang/memory/ScopeKind; public final fun component7 ()J - public final fun copy (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;J)Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted; - public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted;Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted; + public final fun component8 ()J + public final fun copy (Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JJ)Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted; + public static synthetic fun copy$default (Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted;Ljava/lang/String;Lsk/ainet/lang/memory/Format;Lsk/ainet/lang/memory/Format;JLsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/ScopeKind;JJILjava/lang/Object;)Lsk/ainet/lang/memory/trace/TraceEvent$AdapterInserted; public fun equals (Ljava/lang/Object;)Z public final fun getBytes ()J + public final fun getBytesBefore ()J + public final fun getBytesDelta ()J public final fun getFrom ()Lsk/ainet/lang/memory/Format; public final fun getKind ()Ljava/lang/String; public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt index de673c74d..891b5030e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/PerfettoTraceExporter.kt @@ -64,7 +64,12 @@ public object PerfettoTraceExporter { is TraceEvent.AdapterInserted -> emit( instant( "adapter:${e.kind}", "adapter", ts, MAIN_TID, - mapOf("from" to e.from.toString(), "to" to e.to.toString(), "bytes" to e.bytes.toString(), "target" to (e.target?.canonical ?: "—")), + mapOf( + "from" to e.from.toString(), "to" to e.to.toString(), + "bytes" to e.bytes.toString(), "bytesBefore" to e.bytesBefore.toString(), + "bytesDelta" to e.bytesDelta.toString(), + "target" to (e.target?.canonical ?: "—"), + ), ), ) is TraceEvent.Allocation -> { diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt index e7ecf0747..9d3821e10 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/trace/TraceEvent.kt @@ -36,7 +36,14 @@ public sealed interface TraceEvent { override val timeNanos: Long = TraceClock.nowNanos(), ) : TraceEvent - /** The dispatcher inserted a conversion (dequantize, requantize, gather) — always visible (§5.1). */ + /** + * A conversion was inserted (dequantize, requantize, relayout, gather) — always visible (§5.1). + * + * Emitted by the dispatcher when it adapts an operand mid-forward, and by the loader when a + * resolved `WeightForm` re-encodes a weight on the way in (#1109/#1117). [bytes] is the size + * afterwards and [bytesBefore] the size before, so "why is this model 3 GB" has an answer in + * the trace: the two differ by exactly what the conversion cost. + */ public data class AdapterInserted( val kind: String, val from: Format, @@ -44,8 +51,13 @@ public sealed interface TraceEvent { val bytes: Long, val target: TensorId? = null, val scope: ScopeKind = ScopeKind.FORWARD, + /** Size before the conversion; defaults to [bytes] for conversions that do not change size. */ + val bytesBefore: Long = bytes, override val timeNanos: Long = TraceClock.nowNanos(), - ) : TraceEvent + ) : TraceEvent { + /** What this conversion added (or, when negative, saved). */ + public val bytesDelta: Long get() = bytes - bytesBefore + } /** A storage was allocated. [site] is the allocation site in debug mode, [origin] the TensorId it backs. */ public data class Allocation( diff --git a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/trace/JfrTraceSink.kt b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/trace/JfrTraceSink.kt index f4272395e..226ba094d 100644 --- a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/trace/JfrTraceSink.kt +++ b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/trace/JfrTraceSink.kt @@ -50,6 +50,7 @@ public class JfrTraceSink : TraceSink { @Label("From") var from: String? = null @Label("To") var to: String? = null @Label("Bytes") var bytes: Long = 0 + @Label("Bytes before") var bytesBefore: Long = 0 @Label("Target") var target: String? = null } @@ -72,7 +73,8 @@ public class JfrTraceSink : TraceSink { storageId = event.storageId; scope = event.scope.name; bytes = event.bytes; freed = true }.commit() is TraceEvent.AdapterInserted -> AdapterEvent().apply { - kind = event.kind; from = event.from.toString(); to = event.to.toString(); bytes = event.bytes; target = event.target?.canonical + kind = event.kind; from = event.from.toString(); to = event.to.toString() + bytes = event.bytes; bytesBefore = event.bytesBefore; target = event.target?.canonical }.commit() // PhaseBegin is covered by PhaseEnd's duration; counters/plans have no JFR analogue yet. else -> Unit