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