From 7e9b33dd9dd7bfab16540e2109ed171f48e17583 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 2 Sep 2026 18:10:11 +0200 Subject: [PATCH] feat(io-safetensors): tensorFilter on the single-file SafeTensorsParametersLoader (#1256) Parity with ShardedSafeTensorsParametersLoader: an optional predicate over the file's tensor headers, applied before any tensor is read or delivered and before the progress total is computed; threaded through withPolicy. Lets a family load selectively from a checkpoint that also carries tensors the requested dtype cannot accept, instead of failing on the first unmapped one. Unblocks the Voxtral and llm-core legacy-path collapses downstream. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 10 +++ .../SafeTensorsParametersLoader.kt | 13 ++- .../SafeTensorsParametersLoaderFilterTest.kt | 87 +++++++++++++++++++ 3 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFilterTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 530b5dc6..6cb3afb0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,16 @@ ## [Unreleased] +### Added + +- **`tensorFilter` on the single-file `SafeTensorsParametersLoader`** + ([#1256](https://github.com/SKaiNET-developers/SKaiNET/issues/1256)): parity with the sharded + loader — an optional predicate over the tensor headers; filtered-out tensors are neither read + nor delivered and do not count toward progress. Lets a family load selectively from a + checkpoint that also carries tensors the requested dtype cannot accept (int64 index tables, + custom quantized payloads) instead of failing on the first unmapped one. Threaded through + `withPolicy`. + ## [0.53.0] - 2026-09-02 Headline: **the export pipeline emits billion-parameter models.** Tracing a 4.5B-parameter diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt index 22dc1905..c1dd64c6 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt @@ -37,12 +37,21 @@ import kotlin.reflect.KClass * existing consumers. Flip to [Bf16LoadPolicy.KEEP_NATIVE] to keep * weights in their on-disk BF16 layout and let the matmul dispatch * route to a vectorised BF16 kernel. + * @param tensorFilter Optional predicate over the file's tensor headers; a + * tensor for which it returns `false` is neither read nor delivered, and + * does not count toward [onProgress]'s total. `null` (default) loads every + * tensor. This is the family-side hook for name allowlists and size guards + * — parity with [ShardedSafeTensorsParametersLoader]'s filter (#1256): a + * checkpoint that also carries tensors the requested dtype cannot accept + * (int64 index tables, custom quantized payloads) can be loaded selectively + * instead of failing on the first unmapped tensor. */ class SafeTensorsParametersLoader( private val sourceProvider: () -> RandomAccessSource, private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, private val bf16Policy: Bf16LoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, private val fp16Policy: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + private val tensorFilter: ((StreamingSafeTensorInfo) -> Boolean)? = null, ) : ParametersLoader { override suspend fun load( @@ -51,7 +60,7 @@ class SafeTensorsParametersLoader( onTensorLoaded: (String, Tensor) -> Unit ) { StreamingSafeTensorsReader.open(sourceProvider()).use { reader -> - val tensors = reader.tensors + val tensors = tensorFilter?.let { keep -> reader.tensors.filter(keep) } ?: reader.tensors val total = tensors.size.toLong() var current = 0L @@ -109,11 +118,13 @@ class SafeTensorsParametersLoader( sourceProvider: () -> RandomAccessSource, policy: DTypePolicy, onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, + tensorFilter: ((StreamingSafeTensorInfo) -> Boolean)? = null, ): SafeTensorsParametersLoader = SafeTensorsParametersLoader( sourceProvider = sourceProvider, onProgress = onProgress, bf16Policy = mapPolicyToBf16(policy), fp16Policy = mapPolicyToFp16(policy), + tensorFilter = tensorFilter, ) internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = diff --git a/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFilterTest.kt b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFilterTest.kt new file mode 100644 index 00000000..e872af9e --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFilterTest.kt @@ -0,0 +1,87 @@ +package sk.ainet.io.safetensors + +import kotlinx.coroutines.runBlocking +import org.junit.Test +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP32 +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1256: `tensorFilter` on the single-file loader — parity with the sharded + * loader. A checkpoint that also carries tensors the requested dtype cannot + * accept (here an INT64 index table next to FP32 weights) must be loadable + * selectively; without the filter the per-arm `require` still fails as before. + */ +class SafeTensorsParametersLoaderFilterTest { + + /** One F32 `weight` [2,2] and one I64 `positions` [3] in a single file. */ + private fun mixedFile(): File { + val weight = floatArrayOf(1f, 2f, 3f, 4f) + val positions = longArrayOf(0L, 1L, 2L) + val weightBytes = ByteBuffer.allocate(weight.size * 4).order(ByteOrder.LITTLE_ENDIAN) + .apply { weight.forEach { putFloat(it) } }.array() + val posBytes = ByteBuffer.allocate(positions.size * 8).order(ByteOrder.LITTLE_ENDIAN) + .apply { positions.forEach { putLong(it) } }.array() + val header = """{"weight":{"dtype":"F32","shape":[2,2],"data_offsets":[0,${weightBytes.size}]},""" + + """"positions":{"dtype":"I64","shape":[3],"data_offsets":[${weightBytes.size},${weightBytes.size + posBytes.size}]}}""" + val headerBytes = header.toByteArray(Charsets.UTF_8) + val file = Files.createTempFile("filter_test", ".safetensors").toFile().also { it.deleteOnExit() } + file.outputStream().use { os -> + os.write(ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN).putLong(headerBytes.size.toLong()).array()) + os.write(headerBytes) + os.write(weightBytes) + os.write(posBytes) + } + return file + } + + @Test + fun filter_skips_unwanted_tensors_and_progress_counts_only_the_kept_ones() = runBlocking { + val file = mixedFile() + val delivered = linkedMapOf>() + val totals = mutableSetOf() + val loader = SafeTensorsParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file) }, + onProgress = { _, total, _ -> totals += total }, + tensorFilter = { it.name == "weight" }, + ) + loader.load(DirectCpuExecutionContext(), FP32::class) { name, t -> delivered[name] = t } + + assertEquals(setOf("weight"), delivered.keys) + assertEquals(setOf(1L), totals, "progress total must reflect the filtered count") + val data = delivered.getValue("weight").data as FloatArrayTensorData<*> + assertTrue(floatArrayOf(1f, 2f, 3f, 4f).contentEquals(data.buffer)) + } + + @Test + fun without_filter_the_unaccepted_dtype_still_fails_as_before() = runBlocking { + val file = mixedFile() + val loader = SafeTensorsParametersLoader(sourceProvider = { JvmRandomAccessSource.open(file) }) + assertFailsWith { + loader.load(DirectCpuExecutionContext(), FP32::class) { _, _ -> } + } + Unit + } + + @Test + fun withPolicy_threads_the_filter() = runBlocking { + val file = mixedFile() + val delivered = mutableListOf() + SafeTensorsParametersLoader.withPolicy( + sourceProvider = { JvmRandomAccessSource.open(file) }, + policy = DTypePolicy.Any, + tensorFilter = { it.name != "positions" }, + ).load(DirectCpuExecutionContext(), FP32::class) { name, _ -> delivered += name } + assertEquals(listOf("weight"), delivered) + } +}