From e86a4d7c74f4e4fd945270a12429882c0303e6a4 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Wed, 2 Sep 2026 11:22:38 +0200 Subject: [PATCH] feat(io-safetensors): sharded-index ParametersLoader riding the shared materializer (#1246) Extract the per-tensor dtype dispatch, byte/dequant helpers, and DTypePolicy mappers from SafeTensorsParametersLoader into an internal SafeTensorsMaterializer (pure refactor; primitive-typed signature so it serves both StreamingSafeTensorInfo and ShardedTensorInfo), then add ShardedSafeTensorsParametersLoader consuming model.safetensors.index.json via StreamingShardedSafeTensorsReader.openFromIndex. The sharded loader adds a fail-fast dtype pre-scan (one aggregated error before any tensor is delivered, mirroring the GGUF loader's #919 contract) and a tensorFilter hook so family-side skip policy (size guards, name allowlists) stays out of the engine while all dtype/policy handling stays in it. withPolicy has signature parity with the single-file factory. Tests: commonTest policy-routing + pre-scan suite; jvmTest 2-shard fixture (cross-shard name-sorted delivery, KEEP_NATIVE vs DEQUANT arms, tensorFilter, IncompleteShard vs allowPartial, fail-fast-before-first- delivery) and a 1-shard-index vs single-file parity test guarding the extraction refactor. Co-Authored-By: Claude Fable 5 --- .../io/safetensors/SafeTensorsMaterializer.kt | 332 ++++++++++++++++ .../SafeTensorsParametersLoader.kt | 286 +------------- .../ShardedSafeTensorsParametersLoader.kt | 148 ++++++++ ...edSafeTensorsParametersLoaderPolicyTest.kt | 132 +++++++ ...ardedSafeTensorsParametersLoaderJvmTest.kt | 359 ++++++++++++++++++ 5 files changed, 989 insertions(+), 268 deletions(-) create mode 100644 skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt create mode 100644 skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt create mode 100644 skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderPolicyTest.kt create mode 100644 skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderJvmTest.kt diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt new file mode 100644 index 000000000..d4b971161 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsMaterializer.kt @@ -0,0 +1,332 @@ +package sk.ainet.io.safetensors + +import sk.ainet.context.ExecutionContext +import sk.ainet.io.model.DataType +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 +import sk.ainet.lang.types.Int8 +import kotlin.math.pow +import kotlin.reflect.KClass + +/** + * Shared per-tensor materialization for SafeTensors loaders. + * + * Owns the dtype-dispatch (raw little-endian bytes → typed [Tensor]) and the + * narrow-float policy handling used by both [SafeTensorsParametersLoader] + * (single file) and [ShardedSafeTensorsParametersLoader] (index + shards). + * + * The signature is deliberately primitive-typed (name/dataType/shape/bytes) + * rather than taking a tensor-info object: the single-file reader surfaces + * [StreamingSafeTensorInfo] while the sharded reader surfaces + * [ShardedTensorInfo], and the two are unrelated types. + */ +internal object SafeTensorsMaterializer { + + /** + * Materialize one tensor from its raw on-disk bytes. + * + * Conversion rules (identical to the historical + * [SafeTensorsParametersLoader] behavior): + * - F32/F64 → FP32 (F64 downcast with warning) + * - F16 → FP32 dequant, or native [Fp16DenseTensorData] under KEEP_NATIVE + * - BF16 → FP32 dequant, or native [Bf16DenseTensorData] under KEEP_NATIVE + * - I32/I64 → Int32 (I64 downcast with warning) + * - I8/U8/I16/U16/U32/U64/BOOL/UNKNOWN → Int8 raw bytes + * + * Each arm `require`s the matching requested [dtype] and throws otherwise. + */ + @Suppress("UNCHECKED_CAST") + fun materialize( + ctx: ExecutionContext, + dtype: KClass, + name: String, + dataType: DataType, + rawDtype: String, + shape: Shape, + bytes: ByteArray, + bf16Policy: Bf16LoadPolicy, + fp16Policy: NarrowFloatLoadPolicy, + ): Tensor = when (dataType) { + DataType.FLOAT32 -> { + require(dtype == FP32::class) { + "SafeTensors F32 tensor '$name' requires FP32 dtype, got ${dtype.simpleName}" + } + val floats = bytesToFloatArray(bytes) + // Wrap the decoded array (zero-copy) — it was freshly allocated by bytesToFloatArray + ctx.wrapFloatArray(shape, dtype, floats) as Tensor + } + + DataType.FLOAT64 -> { + require(dtype == FP32::class) { + "SafeTensors F64 tensor '$name' requires FP32 dtype (downcast), got ${dtype.simpleName}" + } + println("WARNING: Downcasting F64 tensor '$name' to F32") + val doubles = bytesToDoubleArray(bytes) + val floats = FloatArray(doubles.size) { doubles[it].toFloat() } + ctx.wrapFloatArray(shape, dtype, floats) as Tensor + } + + DataType.FLOAT16 -> { + require(dtype == FP32::class) { + "SafeTensors F16 tensor '$name' requires FP32 dtype, got ${dtype.simpleName}" + } + when (fp16Policy) { + NarrowFloatLoadPolicy.DEQUANT_TO_FP32 -> { + val floats = dequantF16(bytes) + ctx.wrapFloatArray(shape, dtype, floats) as Tensor + } + NarrowFloatLoadPolicy.KEEP_NATIVE -> { + // Mirrors the BF16 arm below: wrap the on-disk F16 bytes directly. + // dtype stays FP32 from the consumer's POV (the tensor data decodes + // on read); the storage type is what a narrow-float matmul dispatch + // pattern-matches on. + val fp16Data = Fp16DenseTensorData(shape, bytes) + ctx.fromData(fp16Data as TensorData, dtype) + } + } + } + + DataType.BFLOAT16 -> { + require(dtype == FP32::class) { + "SafeTensors BF16 tensor '$name' requires FP32 dtype, got ${dtype.simpleName}" + } + when (bf16Policy) { + Bf16LoadPolicy.DEQUANT_TO_FP32 -> { + val floats = dequantBF16(bytes) + ctx.wrapFloatArray(shape, dtype, floats) as Tensor + } + Bf16LoadPolicy.KEEP_NATIVE -> { + // Wrap the on-disk BF16 bytes directly. dtype stays FP32 from + // the consumer's POV (Bf16TensorData : TensorData + // decodes on read); the storage type is what the matmul + // dispatch will pattern-match on to pick the BF16 SPI kernel. + val bf16Data = Bf16DenseTensorData(shape, bytes) + ctx.fromData(bf16Data as TensorData, dtype) + } + } + } + + DataType.INT32 -> { + require(dtype == Int32::class) { + "SafeTensors I32 tensor '$name' requires Int32 dtype, got ${dtype.simpleName}" + } + val ints = bytesToIntArray(bytes) + ctx.wrapIntArray(shape, dtype, ints) as Tensor + } + + DataType.INT64 -> { + require(dtype == Int32::class) { + "SafeTensors I64 tensor '$name' requires Int32 dtype (downcast), got ${dtype.simpleName}" + } + println("WARNING: Downcasting I64 tensor '$name' to I32") + val longs = bytesToLongArray(bytes) + val ints = IntArray(longs.size) { longs[it].toInt() } + ctx.wrapIntArray(shape, dtype, ints) as Tensor + } + + DataType.INT8 -> { + require(dtype == Int8::class) { + "SafeTensors I8 tensor '$name' requires Int8 dtype, got ${dtype.simpleName}" + } + ctx.fromByteArray(shape, dtype, bytes) as Tensor + } + + DataType.UINT8 -> { + require(dtype == Int8::class) { + "SafeTensors U8 tensor '$name' requires Int8 dtype, got ${dtype.simpleName}" + } + // U8 stored as signed bytes (reinterpret) + ctx.fromByteArray(shape, dtype, bytes) as Tensor + } + + DataType.INT16, DataType.UINT16, + DataType.UINT32, DataType.UINT64 -> { + // Store as raw bytes for now + require(dtype == Int8::class) { + "SafeTensors $rawDtype tensor '$name' requires Int8 dtype (raw bytes), got ${dtype.simpleName}" + } + ctx.fromByteArray(shape, dtype, bytes) as Tensor + } + + DataType.BOOL -> { + require(dtype == Int8::class) { + "SafeTensors BOOL tensor '$name' requires Int8 dtype, got ${dtype.simpleName}" + } + ctx.fromByteArray(shape, dtype, bytes) as Tensor + } + + DataType.UNKNOWN -> { + println("WARNING: Unknown dtype '$rawDtype' for tensor '$name'. Storing as raw bytes.") + require(dtype == Int8::class) { + "Unknown SafeTensors dtype requires Int8 dtype for raw bytes storage" + } + ctx.fromByteArray(shape, dtype, bytes) as Tensor + } + + else -> { + error("Unsupported SafeTensors dtype: $dataType for tensor '$name'") + } + } + + /** + * The dtype the requested [dtype] KClass must be for a tensor of + * [dataType] to materialize, or `null` when [materialize] accepts it + * under any policy. Used by fail-fast pre-scans to reject a load + * before any tensor is delivered. + */ + fun requiredDType(dataType: DataType): KClass = when (dataType) { + DataType.FLOAT32, DataType.FLOAT64, DataType.FLOAT16, DataType.BFLOAT16 -> FP32::class + DataType.INT32, DataType.INT64 -> Int32::class + else -> Int8::class + } + + // ========== Byte Conversion Helpers ========== + + internal fun bytesToFloatArray(bytes: ByteArray): FloatArray { + val out = FloatArray(bytes.size / 4) + for (i in out.indices) { + val offset = i * 4 + val bits = (bytes[offset].toInt() and 0xFF) or + ((bytes[offset + 1].toInt() and 0xFF) shl 8) or + ((bytes[offset + 2].toInt() and 0xFF) shl 16) or + ((bytes[offset + 3].toInt() and 0xFF) shl 24) + out[i] = Float.fromBits(bits) + } + return out + } + + internal fun bytesToDoubleArray(bytes: ByteArray): DoubleArray { + val out = DoubleArray(bytes.size / 8) + for (i in out.indices) { + val offset = i * 8 + val bits = (bytes[offset].toLong() and 0xFF) or + ((bytes[offset + 1].toLong() and 0xFF) shl 8) or + ((bytes[offset + 2].toLong() and 0xFF) shl 16) or + ((bytes[offset + 3].toLong() and 0xFF) shl 24) or + ((bytes[offset + 4].toLong() and 0xFF) shl 32) or + ((bytes[offset + 5].toLong() and 0xFF) shl 40) or + ((bytes[offset + 6].toLong() and 0xFF) shl 48) or + ((bytes[offset + 7].toLong() and 0xFF) shl 56) + out[i] = Double.fromBits(bits) + } + return out + } + + internal fun bytesToIntArray(bytes: ByteArray): IntArray { + val out = IntArray(bytes.size / 4) + for (i in out.indices) { + val offset = i * 4 + out[i] = (bytes[offset].toInt() and 0xFF) or + ((bytes[offset + 1].toInt() and 0xFF) shl 8) or + ((bytes[offset + 2].toInt() and 0xFF) shl 16) or + ((bytes[offset + 3].toInt() and 0xFF) shl 24) + } + return out + } + + internal fun bytesToLongArray(bytes: ByteArray): LongArray { + val out = LongArray(bytes.size / 8) + for (i in out.indices) { + val offset = i * 8 + out[i] = (bytes[offset].toLong() and 0xFF) or + ((bytes[offset + 1].toLong() and 0xFF) shl 8) or + ((bytes[offset + 2].toLong() and 0xFF) shl 16) or + ((bytes[offset + 3].toLong() and 0xFF) shl 24) or + ((bytes[offset + 4].toLong() and 0xFF) shl 32) or + ((bytes[offset + 5].toLong() and 0xFF) shl 40) or + ((bytes[offset + 6].toLong() and 0xFF) shl 48) or + ((bytes[offset + 7].toLong() and 0xFF) shl 56) + } + return out + } + + // ========== Dequantization Helpers ========== + + internal fun dequantF16(bytes: ByteArray): FloatArray { + val out = FloatArray(bytes.size / 2) + for (i in out.indices) { + val offset = i * 2 + val half = (bytes[offset].toInt() and 0xFF) or + ((bytes[offset + 1].toInt() and 0xFF) shl 8) + out[i] = halfToFloat(half) + } + return out + } + + internal fun dequantBF16(bytes: ByteArray): FloatArray { + val out = FloatArray(bytes.size / 2) + for (i in out.indices) { + val offset = i * 2 + val bf16Low = bytes[offset].toInt() and 0xFF + val bf16High = bytes[offset + 1].toInt() and 0xFF + // BF16 is just the upper 16 bits of F32 + val bits = (bf16High shl 24) or (bf16Low shl 16) + out[i] = Float.fromBits(bits) + } + return out + } + + private fun halfToFloat(hbits: Int): Float { + val mant = hbits and 0x03FF + val exp = hbits and 0x7C00 + val sign = hbits and 0x8000 + return when (exp) { + 0 -> { + // Subnormal + val v = (mant.toFloat() / 1024.0f) * (2.0f).pow(-14) + if (sign != 0) -v else v + } + 0x7C00 -> { + // Inf/NaN + val v = if (mant == 0) Float.POSITIVE_INFINITY else Float.NaN + if (sign != 0) -v else v + } + else -> { + // Normal + val v = (1.0f + mant.toFloat() / 1024.0f) * (2.0f).pow((exp shr 10) - 15) + if (sign != 0) -v else v + } + } + } + + // ========== Policy Mapping ========== + + internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = + mapPolicyToNarrow(policy, BF16) + + internal fun mapPolicyToFp16(policy: DTypePolicy): NarrowFloatLoadPolicy = + mapPolicyToNarrow(policy, FP16) + + /** + * Resolve [policy] for one narrow-float source format. A tensor is kept native only when + * the policy names *that* format — `Require(BF16)` must not keep F16 tensors packed, and + * vice versa, since neither can be converted to the other without a lossy re-encode. + */ + private fun mapPolicyToNarrow(policy: DTypePolicy, native: DType): NarrowFloatLoadPolicy = + when (policy) { + DTypePolicy.Any -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + is DTypePolicy.Require -> when (policy.target) { + native -> NarrowFloatLoadPolicy.KEEP_NATIVE + // The other narrow format, or FP32: this format still widens. + BF16, FP16, FP32 -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + else -> throw IllegalArgumentException( + "SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " + + "the loader produces FP32 / BF16 / FP16 / Int32 / Int8 tensors depending on " + + "source dtype; it cannot fabricate ${policy.target.name} from arbitrary sources.", + ) + } + is DTypePolicy.Prefer -> if (policy.target == native) NarrowFloatLoadPolicy.KEEP_NATIVE + else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + is DTypePolicy.OneOf -> if (native in policy.allowed) NarrowFloatLoadPolicy.KEEP_NATIVE + else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + } +} 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 ef89d4e98..22dc1905f 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 @@ -3,20 +3,10 @@ package sk.ainet.io.safetensors import sk.ainet.context.ExecutionContext import sk.ainet.io.ParametersLoader import sk.ainet.io.RandomAccessSource -import sk.ainet.io.model.DataType import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor -import sk.ainet.lang.tensor.data.Bf16DenseTensorData -import sk.ainet.lang.tensor.data.Fp16DenseTensorData -import sk.ainet.lang.tensor.data.TensorData -import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType import sk.ainet.lang.types.DTypePolicy -import sk.ainet.lang.types.FP16 -import sk.ainet.lang.types.FP32 -import sk.ainet.lang.types.Int32 -import sk.ainet.lang.types.Int8 -import kotlin.math.pow import kotlin.reflect.KClass /** @@ -25,7 +15,7 @@ import kotlin.reflect.KClass * Uses [StreamingSafeTensorsReader] for memory-efficient loading - only * parses the header (~1KB-1MB) and loads tensors on-demand. * - * Supported conversions: + * Supported conversions (see [SafeTensorsMaterializer]): * - F32/F64 tensors -> FP32 (F64 downcast with warning) * - I32/I64 tensors -> Int32 (I64 downcast with warning) * - I8/U8 tensors -> Int8 @@ -36,6 +26,10 @@ import kotlin.reflect.KClass * into TensorData, avoiding a second allocation. The raw-byte decode step * (little-endian bytes → typed array) is still necessary. * + * Single-file only: for HF checkpoints shipped as `model.safetensors.index.json` + * plus `model-NNNNN-of-NNNNN.safetensors` shards, use + * [ShardedSafeTensorsParametersLoader] (#1246). + * * @param sourceProvider Factory providing RandomAccessSource to the SafeTensors file * @param onProgress Optional progress callback (current, total, tensorName) * @param bf16Policy How to handle `BFLOAT16` tensors. Default is @@ -65,128 +59,17 @@ class SafeTensorsParametersLoader( val bytes = reader.loadTensorData(tensorInfo) val shape = Shape(*tensorInfo.shape.map { it.toInt() }.toIntArray()) - @Suppress("UNCHECKED_CAST") - val tensor: Tensor = when (tensorInfo.dataType) { - DataType.FLOAT32 -> { - require(dtype == FP32::class) { - "SafeTensors F32 tensor '${tensorInfo.name}' requires FP32 dtype, got ${dtype.simpleName}" - } - val floats = bytesToFloatArray(bytes) - // Wrap the decoded array (zero-copy) — it was freshly allocated by bytesToFloatArray - ctx.wrapFloatArray(shape, dtype, floats) as Tensor - } - - DataType.FLOAT64 -> { - require(dtype == FP32::class) { - "SafeTensors F64 tensor '${tensorInfo.name}' requires FP32 dtype (downcast), got ${dtype.simpleName}" - } - println("WARNING: Downcasting F64 tensor '${tensorInfo.name}' to F32") - val doubles = bytesToDoubleArray(bytes) - val floats = FloatArray(doubles.size) { doubles[it].toFloat() } - ctx.wrapFloatArray(shape, dtype, floats) as Tensor - } - - DataType.FLOAT16 -> { - require(dtype == FP32::class) { - "SafeTensors F16 tensor '${tensorInfo.name}' requires FP32 dtype, got ${dtype.simpleName}" - } - when (fp16Policy) { - NarrowFloatLoadPolicy.DEQUANT_TO_FP32 -> { - val floats = dequantF16(bytes) - ctx.wrapFloatArray(shape, dtype, floats) as Tensor - } - NarrowFloatLoadPolicy.KEEP_NATIVE -> { - // Mirrors the BF16 arm below: wrap the on-disk F16 bytes directly. - // dtype stays FP32 from the consumer's POV (the tensor data decodes - // on read); the storage type is what a narrow-float matmul dispatch - // pattern-matches on. - val fp16Data = Fp16DenseTensorData(shape, bytes) - ctx.fromData(fp16Data as TensorData, dtype) - } - } - } - - DataType.BFLOAT16 -> { - require(dtype == FP32::class) { - "SafeTensors BF16 tensor '${tensorInfo.name}' requires FP32 dtype, got ${dtype.simpleName}" - } - when (bf16Policy) { - Bf16LoadPolicy.DEQUANT_TO_FP32 -> { - val floats = dequantBF16(bytes) - ctx.wrapFloatArray(shape, dtype, floats) as Tensor - } - Bf16LoadPolicy.KEEP_NATIVE -> { - // Wrap the on-disk BF16 bytes directly. dtype stays FP32 from - // the consumer's POV (Bf16TensorData : TensorData - // decodes on read); the storage type is what the matmul - // dispatch will pattern-match on to pick the BF16 SPI kernel. - val bf16Data = Bf16DenseTensorData(shape, bytes) - ctx.fromData(bf16Data as TensorData, dtype) - } - } - } - - DataType.INT32 -> { - require(dtype == Int32::class) { - "SafeTensors I32 tensor '${tensorInfo.name}' requires Int32 dtype, got ${dtype.simpleName}" - } - val ints = bytesToIntArray(bytes) - ctx.wrapIntArray(shape, dtype, ints) as Tensor - } - - DataType.INT64 -> { - require(dtype == Int32::class) { - "SafeTensors I64 tensor '${tensorInfo.name}' requires Int32 dtype (downcast), got ${dtype.simpleName}" - } - println("WARNING: Downcasting I64 tensor '${tensorInfo.name}' to I32") - val longs = bytesToLongArray(bytes) - val ints = IntArray(longs.size) { longs[it].toInt() } - ctx.wrapIntArray(shape, dtype, ints) as Tensor - } - - DataType.INT8 -> { - require(dtype == Int8::class) { - "SafeTensors I8 tensor '${tensorInfo.name}' requires Int8 dtype, got ${dtype.simpleName}" - } - ctx.fromByteArray(shape, dtype, bytes) as Tensor - } - - DataType.UINT8 -> { - require(dtype == Int8::class) { - "SafeTensors U8 tensor '${tensorInfo.name}' requires Int8 dtype, got ${dtype.simpleName}" - } - // U8 stored as signed bytes (reinterpret) - ctx.fromByteArray(shape, dtype, bytes) as Tensor - } - - DataType.INT16, DataType.UINT16, - DataType.UINT32, DataType.UINT64 -> { - // Store as raw bytes for now - require(dtype == Int8::class) { - "SafeTensors ${tensorInfo.dtype} tensor '${tensorInfo.name}' requires Int8 dtype (raw bytes), got ${dtype.simpleName}" - } - ctx.fromByteArray(shape, dtype, bytes) as Tensor - } - - DataType.BOOL -> { - require(dtype == Int8::class) { - "SafeTensors BOOL tensor '${tensorInfo.name}' requires Int8 dtype, got ${dtype.simpleName}" - } - ctx.fromByteArray(shape, dtype, bytes) as Tensor - } - - DataType.UNKNOWN -> { - println("WARNING: Unknown dtype '${tensorInfo.dtype}' for tensor '${tensorInfo.name}'. Storing as raw bytes.") - require(dtype == Int8::class) { - "Unknown SafeTensors dtype requires Int8 dtype for raw bytes storage" - } - ctx.fromByteArray(shape, dtype, bytes) as Tensor - } - - else -> { - error("Unsupported SafeTensors dtype: ${tensorInfo.dataType} for tensor '${tensorInfo.name}'") - } - } + val tensor: Tensor = SafeTensorsMaterializer.materialize( + ctx = ctx, + dtype = dtype, + name = tensorInfo.name, + dataType = tensorInfo.dataType, + rawDtype = tensorInfo.dtype, + shape = shape, + bytes = bytes, + bf16Policy = bf16Policy, + fp16Policy = fp16Policy, + ) onTensorLoaded(tensorInfo.name, tensor) current++ @@ -195,115 +78,6 @@ class SafeTensorsParametersLoader( } } - // ========== Byte Conversion Helpers ========== - - private fun bytesToFloatArray(bytes: ByteArray): FloatArray { - val out = FloatArray(bytes.size / 4) - for (i in out.indices) { - val offset = i * 4 - val bits = (bytes[offset].toInt() and 0xFF) or - ((bytes[offset + 1].toInt() and 0xFF) shl 8) or - ((bytes[offset + 2].toInt() and 0xFF) shl 16) or - ((bytes[offset + 3].toInt() and 0xFF) shl 24) - out[i] = Float.fromBits(bits) - } - return out - } - - private fun bytesToDoubleArray(bytes: ByteArray): DoubleArray { - val out = DoubleArray(bytes.size / 8) - for (i in out.indices) { - val offset = i * 8 - val bits = (bytes[offset].toLong() and 0xFF) or - ((bytes[offset + 1].toLong() and 0xFF) shl 8) or - ((bytes[offset + 2].toLong() and 0xFF) shl 16) or - ((bytes[offset + 3].toLong() and 0xFF) shl 24) or - ((bytes[offset + 4].toLong() and 0xFF) shl 32) or - ((bytes[offset + 5].toLong() and 0xFF) shl 40) or - ((bytes[offset + 6].toLong() and 0xFF) shl 48) or - ((bytes[offset + 7].toLong() and 0xFF) shl 56) - out[i] = Double.fromBits(bits) - } - return out - } - - private fun bytesToIntArray(bytes: ByteArray): IntArray { - val out = IntArray(bytes.size / 4) - for (i in out.indices) { - val offset = i * 4 - out[i] = (bytes[offset].toInt() and 0xFF) or - ((bytes[offset + 1].toInt() and 0xFF) shl 8) or - ((bytes[offset + 2].toInt() and 0xFF) shl 16) or - ((bytes[offset + 3].toInt() and 0xFF) shl 24) - } - return out - } - - private fun bytesToLongArray(bytes: ByteArray): LongArray { - val out = LongArray(bytes.size / 8) - for (i in out.indices) { - val offset = i * 8 - out[i] = (bytes[offset].toLong() and 0xFF) or - ((bytes[offset + 1].toLong() and 0xFF) shl 8) or - ((bytes[offset + 2].toLong() and 0xFF) shl 16) or - ((bytes[offset + 3].toLong() and 0xFF) shl 24) or - ((bytes[offset + 4].toLong() and 0xFF) shl 32) or - ((bytes[offset + 5].toLong() and 0xFF) shl 40) or - ((bytes[offset + 6].toLong() and 0xFF) shl 48) or - ((bytes[offset + 7].toLong() and 0xFF) shl 56) - } - return out - } - - // ========== Dequantization Helpers ========== - - private fun dequantF16(bytes: ByteArray): FloatArray { - val out = FloatArray(bytes.size / 2) - for (i in out.indices) { - val offset = i * 2 - val half = (bytes[offset].toInt() and 0xFF) or - ((bytes[offset + 1].toInt() and 0xFF) shl 8) - out[i] = halfToFloat(half) - } - return out - } - - private fun dequantBF16(bytes: ByteArray): FloatArray { - val out = FloatArray(bytes.size / 2) - for (i in out.indices) { - val offset = i * 2 - val bf16Low = bytes[offset].toInt() and 0xFF - val bf16High = bytes[offset + 1].toInt() and 0xFF - // BF16 is just the upper 16 bits of F32 - val bits = (bf16High shl 24) or (bf16Low shl 16) - out[i] = Float.fromBits(bits) - } - return out - } - - private fun halfToFloat(hbits: Int): Float { - val mant = hbits and 0x03FF - val exp = hbits and 0x7C00 - val sign = hbits and 0x8000 - return when (exp) { - 0 -> { - // Subnormal - val v = (mant.toFloat() / 1024.0f) * (2.0f).pow(-14) - if (sign != 0) -v else v - } - 0x7C00 -> { - // Inf/NaN - val v = if (mant == 0) Float.POSITIVE_INFINITY else Float.NaN - if (sign != 0) -v else v - } - else -> { - // Normal - val v = (1.0f + mant.toFloat() / 1024.0f) * (2.0f).pow((exp shr 10) - 15) - if (sign != 0) -v else v - } - } - } - companion object { /** @@ -343,33 +117,9 @@ class SafeTensorsParametersLoader( ) internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = - mapPolicyToNarrow(policy, BF16) + SafeTensorsMaterializer.mapPolicyToBf16(policy) internal fun mapPolicyToFp16(policy: DTypePolicy): NarrowFloatLoadPolicy = - mapPolicyToNarrow(policy, FP16) - - /** - * Resolve [policy] for one narrow-float source format. A tensor is kept native only when - * the policy names *that* format — `Require(BF16)` must not keep F16 tensors packed, and - * vice versa, since neither can be converted to the other without a lossy re-encode. - */ - private fun mapPolicyToNarrow(policy: DTypePolicy, native: DType): NarrowFloatLoadPolicy = - when (policy) { - DTypePolicy.Any -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 - is DTypePolicy.Require -> when (policy.target) { - native -> NarrowFloatLoadPolicy.KEEP_NATIVE - // The other narrow format, or FP32: this format still widens. - BF16, FP16, FP32 -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 - else -> throw IllegalArgumentException( - "SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " + - "the loader produces FP32 / BF16 / FP16 / Int32 / Int8 tensors depending on " + - "source dtype; it cannot fabricate ${policy.target.name} from arbitrary sources.", - ) - } - is DTypePolicy.Prefer -> if (policy.target == native) NarrowFloatLoadPolicy.KEEP_NATIVE - else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 - is DTypePolicy.OneOf -> if (native in policy.allowed) NarrowFloatLoadPolicy.KEEP_NATIVE - else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 - } + SafeTensorsMaterializer.mapPolicyToFp16(policy) } } diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt new file mode 100644 index 000000000..57e776b53 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoader.kt @@ -0,0 +1,148 @@ +package sk.ainet.io.safetensors + +import sk.ainet.context.ExecutionContext +import sk.ainet.io.ParametersLoader +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.DTypePolicy +import kotlin.reflect.KClass + +/** + * [ParametersLoader] for sharded SafeTensors checkpoints — the + * `model.safetensors.index.json` + `model-NNNNN-of-NNNNN.safetensors` layout + * HuggingFace uses for multi-file models (#1246). + * + * Sharded counterpart of [SafeTensorsParametersLoader]: the per-tensor + * materialization and narrow-float policy handling are shared verbatim via + * [SafeTensorsMaterializer], so both loaders produce identical tensors for + * identical bytes. + * + * Reading rides [StreamingShardedSafeTensorsReader.openFromIndex]: every + * shard file is opened eagerly and its handle stays open until the load + * completes (the reader is closed when `load` returns). Tensors are + * delivered in name-sorted order, regardless of which shard holds them. + * + * Unlike the single-file loader, unsupported-dtype failures are raised + * *before* any tensor is delivered: a pre-scan over the (filtered) index + * aggregates every tensor whose SafeTensors dtype cannot materialize into + * the requested [DType] and throws one [IllegalArgumentException] listing + * them all — a 40-shard load should not die on tensor 900 (mirrors the + * GGUF loader's fail-fast contract, #919). + * + * Platform note: this loader resolves shard files by path + * (`openRandomAccessSource`), which yields no sources on js/wasm/ + * androidNativeArm32. A provider-per-shard factory + * (`(shardFilename) -> RandomAccessSource`) for those platforms is a + * planned follow-up and needs an + * `openFromParsedIndex(index, shardProvider)` companion on + * [StreamingShardedSafeTensorsReader]. + * + * @param indexPath Path to `model.safetensors.index.json`. + * @param onProgress Optional progress callback (current, total, tensorName); + * `total` reflects the filtered tensor count. + * @param bf16Policy How to handle `BFLOAT16` tensors (see + * [SafeTensorsParametersLoader]). + * @param fp16Policy How to handle `FLOAT16` tensors. + * @param allowPartial If true, missing shards are tolerated and only tensors + * from present shards are delivered; if false (default), a missing shard + * throws [SafeTensorsShardException.IncompleteShard]. + * @param tensorFilter Optional predicate over [ShardedTensorInfo]; tensors + * for which it returns false are neither materialized nor delivered, and + * are exempt from the fail-fast dtype pre-scan. This is the hook for + * family-side skip policy (size guards, name allowlists) — dtype/policy + * handling stays engine-side. `null` loads everything. + */ +public class ShardedSafeTensorsParametersLoader( + private val indexPath: String, + 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 allowPartial: Boolean = false, + private val tensorFilter: ((ShardedTensorInfo) -> Boolean)? = null, +) : ParametersLoader { + + override suspend fun load( + ctx: ExecutionContext, + dtype: KClass, + onTensorLoaded: (String, Tensor) -> Unit + ) { + StreamingShardedSafeTensorsReader.openFromIndex(indexPath, allowPartial).use { reader -> + val tensors = tensorFilter?.let { filter -> reader.tensors.filter(filter) } + ?: reader.tensors + + failFastOnUnsupportedTensorTypes(tensors, dtype) + + val total = tensors.size.toLong() + var current = 0L + + for (tensorInfo in tensors) { + val bytes = reader.loadTensorData(tensorInfo) + val shape = Shape(*tensorInfo.shape.map { it.toInt() }.toIntArray()) + + val tensor: Tensor = SafeTensorsMaterializer.materialize( + ctx = ctx, + dtype = dtype, + name = tensorInfo.name, + dataType = tensorInfo.dataType, + rawDtype = tensorInfo.dtype, + shape = shape, + bytes = bytes, + bf16Policy = bf16Policy, + fp16Policy = fp16Policy, + ) + + onTensorLoaded(tensorInfo.name, tensor) + current++ + onProgress(current, total, tensorInfo.name) + } + } + } + + public companion object { + + /** + * Constructs a [ShardedSafeTensorsParametersLoader] from a generalised + * [DTypePolicy]. Signature parity with + * [SafeTensorsParametersLoader.withPolicy]; the policy → behaviour + * mapping is identical (documented there). + */ + public fun withPolicy( + indexPath: String, + policy: DTypePolicy, + onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, + allowPartial: Boolean = false, + tensorFilter: ((ShardedTensorInfo) -> Boolean)? = null, + ): ShardedSafeTensorsParametersLoader = ShardedSafeTensorsParametersLoader( + indexPath = indexPath, + onProgress = onProgress, + bf16Policy = SafeTensorsMaterializer.mapPolicyToBf16(policy), + fp16Policy = SafeTensorsMaterializer.mapPolicyToFp16(policy), + allowPartial = allowPartial, + tensorFilter = tensorFilter, + ) + + /** + * Pre-scan [tensors] and throw one aggregated error naming every + * tensor whose dtype cannot materialize into [dtype], before any + * tensor is delivered. + */ + internal fun failFastOnUnsupportedTensorTypes( + tensors: List, + dtype: KClass, + ) { + val mismatches = tensors.filter { SafeTensorsMaterializer.requiredDType(it.dataType) != dtype } + if (mismatches.isNotEmpty()) { + val listing = mismatches.joinToString(separator = "\n") { + " - '${it.name}' (${it.dtype}, shard ${it.shardLocation}) requires " + + "${SafeTensorsMaterializer.requiredDType(it.dataType).simpleName}" + } + throw IllegalArgumentException( + "ShardedSafeTensorsParametersLoader: ${mismatches.size} of ${tensors.size} tensors " + + "cannot materialize as ${dtype.simpleName}; no tensors were delivered. " + + "Use tensorFilter to exclude them or load with the required dtype:\n$listing", + ) + } + } + } +} diff --git a/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderPolicyTest.kt b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderPolicyTest.kt new file mode 100644 index 000000000..c5d5d5e10 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderPolicyTest.kt @@ -0,0 +1,132 @@ +package sk.ainet.io.safetensors + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue +import sk.ainet.io.model.DataType +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 +import sk.ainet.lang.types.Int8 +import sk.ainet.lang.types.Ternary + +/** + * Unit tests for [ShardedSafeTensorsParametersLoader]'s policy routing and + * fail-fast pre-scan. Mirrors the strategy documented in + * [SafeTensorsParametersLoaderPolicyTest]: `withPolicy` is a thin wrapper + * over [SafeTensorsMaterializer]'s mappers plus the constructor, so testing + * the mappers and the pre-scan covers the routing logic without a real + * multi-shard fixture (that lives in the jvmTest suite). + */ +class ShardedSafeTensorsParametersLoaderPolicyTest { + + // ---- Policy mapping parity with the single-file loader ---- + + @Test + fun sharded_and_single_file_loaders_share_one_policy_mapper() { + // Both loaders must route DTypePolicy through the same mapper: spot-check + // the four policy shapes for BF16 and FP16 arms. + for (policy in listOf( + DTypePolicy.Any, + DTypePolicy.Require(BF16), + DTypePolicy.Require(FP32), + DTypePolicy.Prefer(BF16), + DTypePolicy.OneOf(setOf(BF16, FP32)), + )) { + assertEquals( + SafeTensorsParametersLoader.mapPolicyToBf16(policy), + SafeTensorsMaterializer.mapPolicyToBf16(policy), + "BF16 mapping diverged for $policy", + ) + assertEquals( + SafeTensorsParametersLoader.mapPolicyToFp16(policy), + SafeTensorsMaterializer.mapPolicyToFp16(policy), + "FP16 mapping diverged for $policy", + ) + } + } + + @Test + fun require_unsatisfiable_dtype_throws() { + assertFailsWith { + SafeTensorsMaterializer.mapPolicyToBf16(DTypePolicy.Require(Ternary)) + } + } + + // ---- Fail-fast pre-scan ---- + + private fun info(name: String, dtype: String, dataType: DataType) = ShardedTensorInfo( + base = StreamingSafeTensorInfo( + name = name, + dtype = dtype, + dataType = dataType, + shape = listOf(2L, 2L), + elementCount = 4L, + dataOffsetStart = 0L, + dataOffsetEnd = 16L, + sizeInBytes = 16L, + absoluteDataOffset = 8L, + ), + shardFilename = "model-00001-of-00002.safetensors", + shardIndex = 1, + totalShards = 2, + ) + + @Test + fun pre_scan_passes_for_matching_dtypes() { + ShardedSafeTensorsParametersLoader.failFastOnUnsupportedTensorTypes( + listOf( + info("a.weight", "F32", DataType.FLOAT32), + info("b.weight", "BF16", DataType.BFLOAT16), + info("c.weight", "F16", DataType.FLOAT16), + ), + FP32::class, + ) + } + + @Test + fun pre_scan_aggregates_all_mismatches_in_one_error() { + val error = assertFailsWith { + ShardedSafeTensorsParametersLoader.failFastOnUnsupportedTensorTypes( + listOf( + info("ok.weight", "F32", DataType.FLOAT32), + info("ids.tokens", "I64", DataType.INT64), + info("mask.bits", "BOOL", DataType.BOOL), + ), + FP32::class, + ) + } + val message = error.message ?: "" + assertTrue("ids.tokens" in message, "expected first mismatch listed: $message") + assertTrue("mask.bits" in message, "expected second mismatch listed: $message") + assertTrue("2 of 3" in message, "expected aggregate count: $message") + assertTrue("ok.weight" !in message, "matching tensor must not be listed: $message") + } + + @Test + fun pre_scan_names_the_required_dtype_per_tensor() { + val error = assertFailsWith { + ShardedSafeTensorsParametersLoader.failFastOnUnsupportedTensorTypes( + listOf(info("ids.tokens", "I32", DataType.INT32)), + Int8::class, + ) + } + assertTrue("Int32" in (error.message ?: "")) + } + + @Test + fun required_dtype_mapping_matches_materializer_arms() { + assertEquals(FP32::class, SafeTensorsMaterializer.requiredDType(DataType.FLOAT32)) + assertEquals(FP32::class, SafeTensorsMaterializer.requiredDType(DataType.FLOAT64)) + assertEquals(FP32::class, SafeTensorsMaterializer.requiredDType(DataType.FLOAT16)) + assertEquals(FP32::class, SafeTensorsMaterializer.requiredDType(DataType.BFLOAT16)) + assertEquals(Int32::class, SafeTensorsMaterializer.requiredDType(DataType.INT32)) + assertEquals(Int32::class, SafeTensorsMaterializer.requiredDType(DataType.INT64)) + assertEquals(Int8::class, SafeTensorsMaterializer.requiredDType(DataType.INT8)) + assertEquals(Int8::class, SafeTensorsMaterializer.requiredDType(DataType.UINT8)) + assertEquals(Int8::class, SafeTensorsMaterializer.requiredDType(DataType.BOOL)) + assertEquals(Int8::class, SafeTensorsMaterializer.requiredDType(DataType.UNKNOWN)) + } +} diff --git a/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderJvmTest.kt b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderJvmTest.kt new file mode 100644 index 000000000..8dc18ad20 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/ShardedSafeTensorsParametersLoaderJvmTest.kt @@ -0,0 +1,359 @@ +package sk.ainet.io.safetensors + +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import java.nio.file.Path +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue +import kotlinx.coroutines.runBlocking +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.DTypePolicy +import sk.ainet.lang.types.FP32 + +/** + * End-to-end coverage of [ShardedSafeTensorsParametersLoader] against a + * genuine 2-shard fixture (`model-0000X-of-00002.safetensors` + index), + * plus the 1-shard-index vs single-file parity test that guards the + * [SafeTensorsMaterializer] extraction refactor. + * + * Fixture layout: + * - shard 1: `alpha.weight` (F32 2×2), `bravo.weight` (BF16 4) + * - shard 2: `charlie.weight` (F16 4), `zulu.weight` (F32 3) + * Name-sorted delivery order is therefore alpha, bravo, charlie, zulu — + * interleaving would differ if delivery were shard-ordered only. + */ +class ShardedSafeTensorsParametersLoaderJvmTest { + + private val bf16AbsTol = 1e-2f + + private val alphaValues = floatArrayOf(1f, 2f, 3f, 4f) + private val bravoValues = floatArrayOf(0.5f, -1.0f, 2.5f, -64.0f) + private val charlieValues = floatArrayOf(1.0f, -2.0f, 0.25f, 8.0f) // exactly representable in F16 + private val zuluValues = floatArrayOf(-7f, 0f, 42f) + + // ---- fixture builders ---- + + private fun fp32ToBf16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bf16 = (values[i].toRawBits() ushr 16) and 0xFFFF + out[i * 2] = (bf16 and 0xFF).toByte() + out[i * 2 + 1] = ((bf16 ushr 8) and 0xFF).toByte() + } + return out + } + + /** Normal-range-only FP32 → F16 encoder; test values must be exactly representable. */ + private fun fp32ToF16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = values[i].toRawBits() + val sign = (bits ushr 16) and 0x8000 + val half = if (bits and 0x7FFFFFFF == 0) { + sign // ±0 + } else { + val exp32 = (bits ushr 23) and 0xFF + val mant = bits and 0x7FFFFF + val expH = exp32 - 127 + 15 + check(expH in 1..30) { "test value ${values[i]} not a normal F16" } + check(mant and 0x1FFF == 0) { "test value ${values[i]} not exact in F16" } + sign or (expH shl 10) or (mant ushr 13) + } + out[i * 2] = (half and 0xFF).toByte() + out[i * 2 + 1] = ((half ushr 8) and 0xFF).toByte() + } + return out + } + + private fun fp32Bytes(values: FloatArray): ByteArray { + val buf = ByteBuffer.allocate(values.size * 4).order(ByteOrder.LITTLE_ENDIAN) + values.forEach { buf.putFloat(it) } + return buf.array() + } + + private fun writeShard(path: Path, entries: List, ByteArray>>>) { + var offset = 0L + val headerJson = entries.joinToString(prefix = "{", postfix = "}", separator = ",") { (name, dtype, shapeAndBytes) -> + val (shape, bytes) = shapeAndBytes + val start = offset + offset += bytes.size + "\"$name\": {\"dtype\": \"$dtype\", \"shape\": [${shape.joinToString(", ")}], " + + "\"data_offsets\": [$start, $offset]}" + } + val headerBytes = headerJson.toByteArray(Charsets.UTF_8) + path.toFile().outputStream().use { out -> + out.write( + ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(headerBytes.size.toLong()).array(), + ) + out.write(headerBytes) + entries.forEach { (_, _, shapeAndBytes) -> out.write(shapeAndBytes.second) } + } + } + + /** + * Write the standard 2-shard fixture into [dir]; returns the index path. + * Set [omitShard2] to leave shard 2 missing on disk (still referenced by the index). + */ + private fun writeTwoShardFixture(dir: Path, omitShard2: Boolean = false): String { + val shard1 = "model-00001-of-00002.safetensors" + val shard2 = "model-00002-of-00002.safetensors" + writeShard( + dir.resolve(shard1), + listOf( + Triple("alpha.weight", "F32", listOf(2L, 2L) to fp32Bytes(alphaValues)), + Triple("bravo.weight", "BF16", listOf(4L) to fp32ToBf16Bytes(bravoValues)), + ), + ) + if (!omitShard2) { + writeShard( + dir.resolve(shard2), + listOf( + Triple("charlie.weight", "F16", listOf(4L) to fp32ToF16Bytes(charlieValues)), + Triple("zulu.weight", "F32", listOf(3L) to fp32Bytes(zuluValues)), + ), + ) + } + val totalSize = Files.size(dir.resolve(shard1)) + + (if (omitShard2) 0L else Files.size(dir.resolve(shard2))) + val indexPath = dir.resolve("model.safetensors.index.json") + Files.writeString( + indexPath, + """ + { + "metadata": {"total_size": $totalSize}, + "weight_map": { + "alpha.weight": "$shard1", + "bravo.weight": "$shard1", + "charlie.weight": "$shard2", + "zulu.weight": "$shard2" + } + } + """.trimIndent(), + ) + return indexPath.toString() + } + + private fun loadAll(loader: ShardedSafeTensorsParametersLoader): Pair, Map>> = + runBlocking { + val ctx = DirectCpuExecutionContext.create() + val order = mutableListOf() + val out = mutableMapOf>() + loader.load(ctx, FP32::class) { name, tensor -> + order.add(name) + out[name] = tensor + } + order to out + } + + private fun assertClose(expected: FloatArray, actual: FloatArray, tol: Float, label: String) { + assertEquals(expected.size, actual.size, "$label size") + for (i in expected.indices) { + assertTrue( + abs(expected[i] - actual[i]) <= tol, + "$label mismatch at $i: expected=${expected[i]} actual=${actual[i]}", + ) + } + } + + // ---- tests ---- + + @Test + fun `delivers all tensors across both shards name-sorted with exact values`() { + val dir = Files.createTempDirectory("sharded-loader-") + try { + val indexPath = writeTwoShardFixture(dir) + val (order, tensors) = loadAll(ShardedSafeTensorsParametersLoader(indexPath)) + + assertEquals(listOf("alpha.weight", "bravo.weight", "charlie.weight", "zulu.weight"), order) + assertContentEquals(alphaValues, tensors["alpha.weight"]!!.data.copyToFloatArray(), "alpha exact") + assertContentEquals(zuluValues, tensors["zulu.weight"]!!.data.copyToFloatArray(), "zulu exact") + assertClose(bravoValues, tensors["bravo.weight"]!!.data.copyToFloatArray(), bf16AbsTol, "bravo") + // F16 fixture values are exactly representable — zero tolerance. + assertContentEquals(charlieValues, tensors["charlie.weight"]!!.data.copyToFloatArray(), "charlie exact") + assertEquals(listOf(2, 2), tensors["alpha.weight"]!!.shape.dimensions.toList()) + } finally { + dir.toFile().deleteRecursively() + } + } + + @Test + fun `default policy dequants narrow floats and keep_native preserves storage types`() { + val dir = Files.createTempDirectory("sharded-loader-policy-") + try { + val indexPath = writeTwoShardFixture(dir) + + val (_, dequanted) = loadAll(ShardedSafeTensorsParametersLoader(indexPath)) + assertTrue(dequanted["bravo.weight"]!!.data is FloatArrayTensorData<*>, "default BF16 → FloatArray") + assertTrue(dequanted["charlie.weight"]!!.data is FloatArrayTensorData<*>, "default F16 → FloatArray") + + val (_, kept) = loadAll( + ShardedSafeTensorsParametersLoader( + indexPath = indexPath, + bf16Policy = Bf16LoadPolicy.KEEP_NATIVE, + fp16Policy = NarrowFloatLoadPolicy.KEEP_NATIVE, + ), + ) + assertTrue(kept["bravo.weight"]!!.data is Bf16DenseTensorData, "KEEP_NATIVE BF16 storage") + assertTrue(kept["charlie.weight"]!!.data is Fp16DenseTensorData, "KEEP_NATIVE F16 storage") + assertTrue(kept["alpha.weight"]!!.data is FloatArrayTensorData<*>, "F32 unaffected by policies") + } finally { + dir.toFile().deleteRecursively() + } + } + + @Test + fun `withPolicy routes Require BF16 to keep_native for bf16 only`() { + val dir = Files.createTempDirectory("sharded-loader-withpolicy-") + try { + val indexPath = writeTwoShardFixture(dir) + val (_, tensors) = loadAll( + ShardedSafeTensorsParametersLoader.withPolicy(indexPath, DTypePolicy.Require(BF16)), + ) + assertTrue(tensors["bravo.weight"]!!.data is Bf16DenseTensorData, "Require(BF16) keeps BF16 native") + assertTrue( + tensors["charlie.weight"]!!.data is FloatArrayTensorData<*>, + "Require(BF16) must not keep F16 native", + ) + } finally { + dir.toFile().deleteRecursively() + } + } + + @Test + fun `tensorFilter excludes tensors from delivery and progress total`() { + val dir = Files.createTempDirectory("sharded-loader-filter-") + try { + val indexPath = writeTwoShardFixture(dir) + val totals = mutableSetOf() + val loader = ShardedSafeTensorsParametersLoader( + indexPath = indexPath, + onProgress = { _, total, _ -> totals.add(total) }, + tensorFilter = { it.name != "bravo.weight" }, + ) + val (order, tensors) = loadAll(loader) + assertEquals(listOf("alpha.weight", "charlie.weight", "zulu.weight"), order) + assertTrue("bravo.weight" !in tensors) + assertEquals(setOf(3L), totals, "progress total must reflect the filtered count") + } finally { + dir.toFile().deleteRecursively() + } + } + + @Test + fun `missing shard throws IncompleteShard unless allowPartial delivers shard-1 tensors only`() { + val dir = Files.createTempDirectory("sharded-loader-partial-") + try { + val indexPath = writeTwoShardFixture(dir, omitShard2 = true) + + assertFailsWith { + loadAll(ShardedSafeTensorsParametersLoader(indexPath)) + } + + val (order, tensors) = loadAll( + ShardedSafeTensorsParametersLoader(indexPath, allowPartial = true), + ) + assertEquals(listOf("alpha.weight", "bravo.weight"), order) + assertContentEquals(alphaValues, tensors["alpha.weight"]!!.data.copyToFloatArray()) + } finally { + dir.toFile().deleteRecursively() + } + } + + @Test + fun `fail-fast dtype mismatch throws before any tensor is delivered`() { + val dir = Files.createTempDirectory("sharded-loader-failfast-") + try { + val shard = "model-00001-of-00001.safetensors" + val intBytes = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putInt(7).putInt(9).array() + writeShard( + dir.resolve(shard), + listOf( + Triple("good.weight", "F32", listOf(2L) to fp32Bytes(floatArrayOf(1f, 2f))), + Triple("ids.tokens", "I32", listOf(2L) to intBytes), + ), + ) + val indexPath = dir.resolve("model.safetensors.index.json") + Files.writeString( + indexPath, + """{"metadata":{"total_size":${Files.size(dir.resolve(shard))}},""" + + """"weight_map":{"good.weight":"$shard","ids.tokens":"$shard"}}""", + ) + + var delivered = 0 + val error = assertFailsWith { + runBlocking { + val ctx = DirectCpuExecutionContext.create() + ShardedSafeTensorsParametersLoader(indexPath.toString()) + .load(ctx, FP32::class) { _, _ -> delivered++ } + } + } + assertEquals(0, delivered, "fail-fast must fire before the first onTensorLoaded") + assertTrue("ids.tokens" in (error.message ?: ""), "error must name the offending tensor") + } finally { + dir.toFile().deleteRecursively() + } + } + + @Test + fun `one-shard index parity with single-file loader`() { + // Guards the SafeTensorsMaterializer extraction: identical bytes through + // either loader must yield identical tensors (values and storage types). + val dir = Files.createTempDirectory("sharded-loader-parity-") + try { + val shard = "model.safetensors" + writeShard( + dir.resolve(shard), + listOf( + Triple("alpha.weight", "F32", listOf(2L, 2L) to fp32Bytes(alphaValues)), + Triple("bravo.weight", "BF16", listOf(4L) to fp32ToBf16Bytes(bravoValues)), + Triple("charlie.weight", "F16", listOf(4L) to fp32ToF16Bytes(charlieValues)), + ), + ) + val indexPath = dir.resolve("model.safetensors.index.json") + Files.writeString( + indexPath, + """{"metadata":{"total_size":${Files.size(dir.resolve(shard))}},""" + + """"weight_map":{"alpha.weight":"$shard","bravo.weight":"$shard","charlie.weight":"$shard"}}""", + ) + + val singleFile: Map> = runBlocking { + val ctx = DirectCpuExecutionContext.create() + val out = mutableMapOf>() + SafeTensorsParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(File(dir.toFile(), shard)) }, + ).load(ctx, FP32::class) { name, tensor -> out[name] = tensor } + out + } + val (_, sharded) = loadAll(ShardedSafeTensorsParametersLoader(indexPath.toString())) + + assertEquals(singleFile.keys, sharded.keys) + for ((name, single) in singleFile) { + val fromSharded = sharded[name]!! + assertEquals( + single.data::class, fromSharded.data::class, + "storage type parity for '$name'", + ) + assertContentEquals( + single.data.copyToFloatArray(), fromSharded.data.copyToFloatArray(), + "value parity for '$name'", + ) + } + } finally { + dir.toFile().deleteRecursively() + } + } +}