diff --git a/BUILD.bazel b/BUILD.bazel index 77c5090a..ed616a39 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -349,7 +349,10 @@ cc_library( cc_library( name = "weights", srcs = ["gemma/weights.cc"], - hdrs = ["gemma/weights.h"], + hdrs = [ + "gemma/weights.h", + "gemma/weights_internal.h", + ], deps = [ ":configs", ":gemma_args", @@ -379,6 +382,23 @@ cc_test( ], ) +cc_test( + name = "weights_test", + srcs = ["gemma/weights_test.cc"], + deps = [ + ":gemma_args", + ":mat", + ":threading_context", + ":weights", + "//compression:compress", + "//compression:types", + "//io", + "//io:blob_store", + "@googletest//:gtest_main", # buildcleaner: keep + "@highway//:hwy", + ], +) + # For building all tests in one command, so we can test several. test_suite( name = "ops_tests", diff --git a/CMakeLists.txt b/CMakeLists.txt index 059145fe..f10c6be8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -206,6 +206,7 @@ set(SOURCES gemma/vit.h gemma/weights.cc gemma/weights.h + gemma/weights_internal.h io/blob_store.cc io/blob_store.h io/fields.cc @@ -383,6 +384,7 @@ set(GEMMA_TEST_FILES deepseek/deepseek_test.cc gemma/gemma_args_test.cc gemma/tensor_info_test.cc + gemma/weights_test.cc io/blob_store_test.cc io/fields_test.cc ops/bench_matmul.cc diff --git a/gemma/gemma_args.h b/gemma/gemma_args.h index 8ba8f0d7..08b68490 100644 --- a/gemma/gemma_args.h +++ b/gemma/gemma_args.h @@ -53,6 +53,7 @@ struct LoaderArgs : public ArgsBase { Path weights; // weights file location Tristate map; Tristate to_bf16; + Tristate sfp_embedding; Tristate wrapping; template @@ -65,6 +66,13 @@ struct LoaderArgs : public ArgsBase { "Enable memory-mapping? -1 = auto, 0 = no, 1 = yes."); visitor(to_bf16, "to_bf16", Tristate::kDefault, "Convert weights to bf16? -1 = auto, 0 = no, 1 = yes."); + visitor(sfp_embedding, "sfp_embedding", Tristate::kDefault, + "Compress the embedding table to SFP while loading? This halves\n" + " its footprint and, for models with tied input/output\n" + " embeddings, the bandwidth of the per-token logits MatMul,\n" + " in exchange for 8-bit precision on that tensor. Enabling this\n" + " disables automatic mapping; explicit --map=1 still wins.\n" + " -1 = auto (currently off), 0 = no, 1 = yes."); visitor(wrapping, "wrapping", Tristate::kDefault, "Enable prompt wrapping? Specify 0 for pre-2025 format PT models."); } diff --git a/gemma/weights.cc b/gemma/weights.cc index 535f8a88..c699ac6a 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -29,19 +29,24 @@ #include "gemma/configs.h" #include "gemma/gemma_args.h" #include "gemma/model_store.h" +#include "gemma/weights_internal.h" +#include "hwy/base.h" +#include "hwy/highway.h" +#include "hwy/profiler.h" #include "io/blob_store.h" #include "util/mat.h" #include "util/threading_context.h" #include "util/zones.h" -#include "hwy/base.h" -#include "hwy/highway.h" -#include "hwy/profiler.h" // TODO: move into foreach_target #include "compression/compress-inl.h" namespace gcpp { +namespace hn = hwy::HWY_NAMESPACE; + +using weights_internal::TensorToRead; + static std::mutex g_mat_owners_mutex; // Copies att_weights from `attn_vec_einsum_w`. @@ -306,7 +311,7 @@ static void HWY_MAYBE_UNUSED InitAttWeightsI8( hwy::AlignedFreeUniquePtr att_weights_tmp = hwy::AllocateAligned(model_dim * heads * qkv_dim); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, attn_vec_einsum_w.Span(), 0, attn_vec_einsum_w_tmp.get(), model_dim * heads * qkv_dim); @@ -370,7 +375,7 @@ static void HWY_MAYBE_UNUSED SplitW1I8(const LayerConfig& layer_config, hwy::AlignedFreeUniquePtr w_tmp = hwy::AllocateAligned(total_size); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, gating_einsum_w.Span(), 0, w_tmp.get(), total_size); @@ -427,7 +432,7 @@ static void HWY_MAYBE_UNUSED SplitAttW1I8(const LayerConfig& layer_config, hwy::AlignedFreeUniquePtr w_tmp = hwy::AllocateAligned(w1_size); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, qkv_einsum_w.Span(), 0, w_tmp.get(), w1_size); @@ -466,7 +471,7 @@ static void HWY_MAYBE_UNUSED SplitAttW1I8(const LayerConfig& layer_config, hwy::AlignedFreeUniquePtr w_tmp = hwy::AllocateAligned(total_size); - const hwy::HWY_NAMESPACE::ScalableTag df; + const hn::ScalableTag df; HWY_NAMESPACE::DecompressAndZeroPad(df, qkv_einsum_w.Span(), 0, w_tmp.get(), total_size); @@ -657,13 +662,20 @@ std::vector WeightsPtrs::AddTensorDataToWriter( } // Decides whether to read or map based on heuristics and user override. -static WeightsPtrs::Mode ChooseMode(uint64_t file_bytes, - const LoaderArgs& loader, - const InferenceArgs& inference, - const Allocator& allocator) { +WeightsPtrs::Mode weights_internal::ChooseMode(uint64_t file_bytes, + const LoaderArgs& loader, + const InferenceArgs& inference, + const Allocator& allocator) { Tristate to_bf16 = loader.to_bf16; Tristate map = loader.map; + // An explicit request to convert the embedding requires owned memory. Do + // not let the automatic mapping heuristic override that request. An + // explicit --map=1 is still honored and diagnosed in ReadFromBlobs. + if (loader.sfp_embedding == Tristate::kTrue && map == Tristate::kDefault) { + map = Tristate::kFalse; + } + // Disable mapping if not padded to the base page size. if (file_bytes % allocator.BasePageBytes() != 0) { if (map == Tristate::kTrue) { // Only complain if explicitly requested. @@ -706,18 +718,6 @@ static WeightsPtrs::Mode ChooseMode(uint64_t file_bytes, : WeightsPtrs::Mode::kRead; } -struct TensorToRead { - MatPtr* mat; - BlobRange range; - // Some tensors opt out of padding via kPacked flags. - MatPadding padding; - - // only for kReadBF16 - bool keep_type = false; - Type prev_type; - size_t prev_packed_bytes = 0; -}; - // Allocates multiple in parallel and binds to NUMA nodes. static void AllocateAndBindAll(std::vector& tensors, const WeightsPtrs::Mode mode, @@ -735,8 +735,16 @@ static void AllocateAndBindAll(std::vector& tensors, tensor.prev_type = mat.GetType(); tensor.prev_packed_bytes = mat.PackedBytes(); - // We only care about MatMul inputs; skip F32 or small tensors. - if (tensor.prev_type == Type::kF32 || mat.Rows() < 1024) { + // Only worthwhile from 16/32-bit types; the others are already <= 8 + // bits, and NUQ is smaller than SFP. + if (tensor.to_sfp && tensor.prev_type != Type::kF32 && + tensor.prev_type != Type::kBF16) { + tensor.to_sfp = false; + } + if (tensor.to_sfp) { + mat.SetType(Type::kSFP); + // We only care about MatMul inputs; skip F32 or small tensors. + } else if (tensor.prev_type == Type::kF32 || mat.Rows() < 1024) { tensor.keep_type = true; tensor.padding = MatPadding::kPacked; // single I/O for simplicity } else if (mode == WeightsPtrs::Mode::kReadBF16) { @@ -771,7 +779,7 @@ static void MapAll(const std::vector& tensors, template static void DecompressToBF16(MatPtr& mat, const hwy::AlignedFreeUniquePtr& buf) { - hwy::HWY_NAMESPACE::ScalableTag dbf; + hn::ScalableTag dbf; const size_t cols = mat.Cols(); const size_t num_packed = CompressedArrayElements(mat.Extents().Area()); @@ -836,6 +844,188 @@ static void ReadAllToBF16(const std::vector& tensors, }); } +// Tensors flagged `to_sfp`, in any mode that reads rather than maps: + +// Number of rows to read per parallel task. Rows are grouped so that the reads +// are large enough to be efficient, but small enough that the transient buffers +// are negligible: staging the whole tensor would defeat the purpose of +// compressing it. +static size_t RowsPerChunk(size_t row_bytes) { + constexpr size_t kTargetBytes = 4 * 1024 * 1024; + return HWY_MAX(size_t{1}, kTargetBytes / HWY_MAX(size_t{1}, row_bytes)); +} + +// Holds the per-task buffers for one chunk of `rows_per_chunk` rows, so that +// both passes below can reuse the same code. +template +class Chunk { + public: + Chunk(size_t cols, size_t rows_per_chunk) + : cols_(cols), rows_per_chunk_(rows_per_chunk) { + const hn::ScalableTag df; + const size_t NF = hn::Lanes(df); + buf_ = hwy::AllocateAligned(rows_per_chunk * cols * sizeof(T)); + // `DecompressAndZeroPad` writes whole vectors, and `compress-inl.h` + // requires up to two of them beyond the requested count. + raw_ = hwy::AllocateAligned( + hwy::RoundUpTo(rows_per_chunk * cols, NF) + 4 * NF); + HWY_ASSERT(buf_ && raw_); + } + + // Reads rows `[begin, end)` from the file, where the tensor is stored as + // type `T` and packed, hence `mat.Stride()` does not apply. Returns the + // decompressed values, ignoring `mat.Scale()`. + float* Decompress(const TensorToRead& tensor, const BlobReader& reader, + size_t begin, size_t end) { + HWY_DASSERT(end - begin <= rows_per_chunk_); + const size_t row_bytes = cols_ * sizeof(T); + const size_t num = (end - begin) * cols_; + HWY_ASSERT(reader.file().Read(tensor.range.offset + begin * row_bytes, + num * sizeof(T), buf_.get())); + + // Rows are contiguous in both source and destination, hence decompress the + // entire chunk at once: this is faster, and prevents the zero padding from + // overwriting the start of the next row. + const hn::ScalableTag df; + const PackedSpan packed{HWY_RCAST_ALIGNED(T*, buf_.get()), num}; + HWY_NAMESPACE::DecompressAndZeroPad(df, packed, 0, raw_.get(), num); + return raw_.get(); + } + + private: + size_t cols_; + size_t rows_per_chunk_; + hwy::AlignedFreeUniquePtr buf_; + hwy::AlignedFreeUniquePtr raw_; +}; + +// Returns the largest magnitude in the tensor, ignoring `mat.Scale()`. +template +static float MaxAbs(const TensorToRead& tensor, const BlobReader& reader, + size_t rows_per_chunk, size_t num_chunks, + ThreadingContext& ctx) { + const MatPtr& mat = *tensor.mat; + const size_t rows = mat.Rows(); + const size_t cols = mat.Cols(); + // Indexed by chunk rather than by worker, so that we need not know how many + // workers `ParallelFor` will use. + std::vector chunk_max(num_chunks, 0.0f); + + ParallelFor(Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, + Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { + GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); + const size_t begin = chunk * rows_per_chunk; + const size_t end = HWY_MIN(begin + rows_per_chunk, rows); + Chunk buffers(cols, rows_per_chunk); + const float* raw = + buffers.Decompress(tensor, reader, begin, end); + + using DF = hn::ScalableTag; + using VF = hn::Vec; + const DF df; + const size_t NF = hn::Lanes(df); + const size_t num = (end - begin) * cols; + VF max0 = hn::Zero(df); + VF max1 = hn::Zero(df); + size_t i = 0; + for (; i + 2 * NF <= num; i += 2 * NF) { + max0 = hn::Max(max0, hn::Abs(hn::LoadU(df, raw + i))); + max1 = hn::Max(max1, hn::Abs(hn::LoadU(df, raw + i + NF))); + } + if (i + NF <= num) { + max0 = hn::Max(max0, hn::Abs(hn::LoadU(df, raw + i))); + i += NF; + } + float maxabs = hn::ReduceMax(df, hn::Max(max0, max1)); + for (; i < num; ++i) { + maxabs = HWY_MAX(maxabs, hwy::ScalarAbs(raw[i])); + } + chunk_max[chunk] = maxabs; + }); + + float maxabs = 0.0f; + for (const float m : chunk_max) maxabs = HWY_MAX(maxabs, m); + return maxabs; +} + +// Reads the tensor as stored in the file (type `T`) and compresses it into +// `mat`, whose type `AllocateAndBindAll` already changed to `Type::kSFP`. +// The file is read twice because `SfpStream` encodes a limited range of +// magnitudes, hence we may need a per-tensor scale, which requires knowing the +// largest magnitude before encoding anything. The second read is typically +// served from the OS cache. +template +static void CompressToSFP(const TensorToRead& tensor, const BlobReader& reader, + ThreadingContext& ctx) { + MatPtr& mat = *tensor.mat; + const size_t rows = mat.Rows(); + const size_t cols = mat.Cols(); + const size_t rows_per_chunk = RowsPerChunk(cols * sizeof(T)); + const size_t num_chunks = hwy::DivCeil(rows, rows_per_chunk); + + const float prev_scale = mat.Scale(); + const float maxabs = + MaxAbs(tensor, reader, rows_per_chunk, num_chunks, ctx) * prev_scale; + const float scale = + (maxabs <= SfpStream::kMax) ? 1.0f : maxabs / SfpStream::kMax; + const float mul = prev_scale / scale; + + ParallelFor(Parallelism::kFlat, num_chunks, ctx, /*cluster_idx=*/0, + Callers::kReadAllToSFP, [&](uint64_t chunk, size_t thread) { + GCPP_ZONE(ctx, thread, Zones::kStartupWeightsReadAllToSFP); + const size_t begin = chunk * rows_per_chunk; + const size_t end = HWY_MIN(begin + rows_per_chunk, rows); + Chunk buffers(cols, rows_per_chunk); + float* raw = buffers.Decompress(tensor, reader, begin, end); + + if (mul != 1.0f) { + for (size_t i = 0; i < (end - begin) * cols; ++i) { + // Clamp because rounding may still exceed `kMax`. + const float magn = + HWY_MIN(SfpStream::kMax, hwy::ScalarAbs(raw[i] * mul)); + raw[i] = hwy::ScalarCopySign(magn, raw[i]); + } + } + + // Row by row because destination rows are padded, whereas `raw` + // is not. This is safe because `SfpStream` is a per-value + // encoding. + CompressPerThread tls; // unused by SFP, which is stateless + for (size_t r = begin; r < end; ++r) { + const PackedSpan row{ + HWY_RCAST_ALIGNED(SfpStream*, mat.RowBytes(r)), cols}; + HWY_NAMESPACE::Compress(raw + (r - begin) * cols, cols, tls, + row, + /*packed_ofs=*/0); + } + }); + + mat.SetScale(scale); +} + +void weights_internal::ReadAllToSFP(const std::vector& tensors, + const BlobReader& reader, + ThreadingContext& ctx) { + PROFILER_ZONE("Startup.Weights.ReadAllToSFP"); + // Usually a single (large) tensor, hence parallelize within, not across. + for (const TensorToRead& tensor : tensors) { + // CompressToSFP derives read sizes from the tensor shape. Ensure those + // reads cannot cross the blob boundary if metadata is inconsistent. + HWY_ASSERT_M(tensor.range.bytes == tensor.prev_packed_bytes, + tensor.mat->Name()); + switch (tensor.prev_type) { + case Type::kF32: + CompressToSFP(tensor, reader, ctx); + break; + case Type::kBF16: + CompressToSFP(tensor, reader, ctx); + break; + default: + HWY_ABORT("Unsupported type %s", TypeName(tensor.prev_type)); + } + } +} + // Mode == kRead: static std::vector MakeBatches( @@ -934,13 +1124,20 @@ static MapPtr MapOrReadAll(std::vector& tensors, AllocateAndBindAll(tensors, *mode, mat_owners, ctx); } + // `MakeBatches` and `ReadAllToBF16` read into the destination rows, hence + // tensors that require a compression pass are handled separately. + std::vector to_sfp, rest; + for (const TensorToRead& tensor : tensors) { + (tensor.to_sfp ? to_sfp : rest).push_back(tensor); + } + if (!to_sfp.empty()) weights_internal::ReadAllToSFP(to_sfp, reader, ctx); + if (*mode == WeightsPtrs::Mode::kReadBF16) { - ReadAllToBF16(tensors, reader, ctx); + ReadAllToBF16(rest, reader, ctx); return MapPtr(); } - const std::vector batches = - MakeBatches(tensors, reader.file_bytes()); + const std::vector batches = MakeBatches(rest, reader.file_bytes()); ReadBatches(reader, batches, ctx); return MapPtr(); } @@ -973,7 +1170,22 @@ WeightsPtrs::Mode WeightsPtrs::ReadFromBlobs(const ModelStore& model, HWY_ABORT("Tensor %s is required but not found in file.", t.mat.Name()); }); - Mode mode = ChooseMode(reader.file_bytes(), loader, inference, ctx.allocator); + Mode mode = weights_internal::ChooseMode(reader.file_bytes(), loader, + inference, ctx.allocator); + + // Compressing the input embedding to SFP halves its footprint. For models + // with tied input/output embeddings, it also halves the weight bandwidth of + // the per-token logits MatMul. + if (loader.sfp_embedding == Tristate::kTrue) { + if (mode == Mode::kMap) { + HWY_WARN("Cannot have sfp_embedding && map, ignoring sfp_embedding."); + } else { + for (TensorToRead& tensor : tensors) { + if (tensor.mat == &embedder_input_embedding) tensor.to_sfp = true; + } + } + } + mapped_ = MapOrReadAll(tensors, reader, &mode, mat_owners, ctx); { diff --git a/gemma/weights_internal.h b/gemma/weights_internal.h new file mode 100644 index 00000000..1d84ba6f --- /dev/null +++ b/gemma/weights_internal.h @@ -0,0 +1,61 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_INTERNAL_H_ +#define THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_INTERNAL_H_ + +#include +#include + +#include + +#include "compression/types.h" +#include "gemma/gemma_args.h" +#include "gemma/weights.h" +#include "io/blob_store.h" +#include "util/allocator.h" +#include "util/mat.h" +#include "util/threading_context.h" + +namespace gcpp { +namespace weights_internal { + +// Describes one tensor whose file bytes will be mapped, read, or converted. +// Kept in this internal header so loader behavior can be tested without a full +// model-sized weights file. +struct TensorToRead { + MatPtr* mat; + BlobRange range; + MatPadding padding; + + // Only for kReadBF16. + bool keep_type = false; + // Convert to Type::kSFP while loading. + bool to_sfp = false; + Type prev_type; + size_t prev_packed_bytes = 0; +}; + +WeightsPtrs::Mode ChooseMode(uint64_t file_bytes, const LoaderArgs& loader, + const InferenceArgs& inference, + const Allocator& allocator); + +void ReadAllToSFP(const std::vector& tensors, + const BlobReader& reader, ThreadingContext& ctx); + +} // namespace weights_internal +} // namespace gcpp + +#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_INTERNAL_H_ diff --git a/gemma/weights_test.cc b/gemma/weights_test.cc new file mode 100644 index 00000000..78e810f1 --- /dev/null +++ b/gemma/weights_test.cc @@ -0,0 +1,221 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#include +#include +#include +#include + +#include "compression/compress-inl.h" +#include "compression/types.h" +#include "gemma/gemma_args.h" +#include "gemma/weights_internal.h" +#include "gtest/gtest.h" +#include "hwy/base.h" +#include "io/blob_store.h" +#include "io/io.h" +#include "util/basics.h" +#include "util/mat.h" +#include "util/threading_context.h" + +namespace gcpp { +namespace { + +class TemporaryBlob { + public: + TemporaryBlob() { + const int fd = mkstemp(path_); + HWY_ASSERT(fd >= 0); + HWY_ASSERT(close(fd) == 0); + } + + ~TemporaryBlob() { unlink(path_); } + + Path path() const { return Path(path_); } + + private: + char path_[sizeof("/tmp/weights_test.sbs-XXXXXX")] = + "/tmp/weights_test.sbs-XXXXXX"; +}; + +static ThreadingContext MakeContext() { + ThreadingArgs args; + args.max_threads = 2; + args.pin = Tristate::kFalse; + args.bind = Tristate::kFalse; + return ThreadingContext(args); +} + +template +static float SourceValue(const T value) { + return hwy::ConvertScalarTo(value); +} + +template +static T MakeSourceValue(const float value) { + return hwy::ConvertScalarTo(value); +} + +template +static void TestSFPConversion(const Type source_type, const float prev_scale, + const std::array& pattern) { + ThreadingContext ctx = MakeContext(); + + // Just over 4 MiB, so the conversion has a final partial chunk. + constexpr size_t kCols = 1024; + constexpr size_t kRows = 4 * 1024 * 1024 / (kCols * sizeof(T)) + 1; + const size_t num = kRows * kCols; + std::vector source(num); + for (size_t i = 0; i < num; ++i) { + source[i] = MakeSourceValue(pattern[i % pattern.size()]); + } + + TemporaryBlob blob; + { + BlobWriter writer(blob.path(), ctx); + writer.Add("embedding", source.data(), source.size() * sizeof(T)); + writer.Finalize(); + } + BlobReader reader(blob.path()); + const BlobRange* range = reader.Find("embedding"); + ASSERT_NE(range, nullptr); + + MatPtr mat("embedding", Type::kSFP, Extents2D(kRows, kCols)); + mat.SetScale(prev_scale); + MatOwner owner; + owner.AllocateFor(mat, ctx.allocator, MatPadding::kOdd); + + const weights_internal::TensorToRead tensor{ + .mat = &mat, + .range = *range, + .padding = MatPadding::kOdd, + .to_sfp = true, + .prev_type = source_type, + .prev_packed_bytes = source.size() * sizeof(T), + }; + weights_internal::ReadAllToSFP({tensor}, reader, ctx); + + float source_maxabs = 0.0f; + for (const T value : source) { + source_maxabs = + std::max(source_maxabs, std::abs(SourceValue(value) * prev_scale)); + } + const float expected_scale = + source_maxabs <= SfpStream::kMax ? 1.0f : source_maxabs / SfpStream::kMax; + EXPECT_FLOAT_EQ(mat.Scale(), expected_scale); + + const MatPtrT sfp(mat); + const std::array samples = { + 0, + 1, + kCols - 1, + (kRows - 1) * kCols - 1, + (kRows - 1) * kCols, + (kRows - 1) * kCols + 1, + num - 2, + num - 1, + }; + for (const size_t i : samples) { + const size_t row = i / kCols; + const size_t col = i % kCols; + const float expected = SourceValue(source[i]) * prev_scale; + const float actual = HWY_NAMESPACE::CompressTraits::ToFloatSlow( + sfp.Row(row)[col]) * + mat.Scale(); + const float tolerance = std::max(1E-6f, std::abs(expected) * 0.13f); + EXPECT_NEAR(actual, expected, tolerance) << "index " << i; + } +} + +TEST(WeightsTest, ReadBF16EmbeddingToSFPWithScaleAndPartialChunk) { + TestSFPConversion(Type::kBF16, 1.25f, + {-2.5f, -1.25f, -0.25f, 0.0f, 0.25f, 1.0f, 2.5f}); +} + +TEST(WeightsTest, ReadF32EmbeddingToSFPWithoutScaleAndPartialChunk) { + TestSFPConversion(Type::kF32, 1.0f, + {-1.75f, -1.0f, -0.125f, 0.0f, 0.125f, 1.0f, 1.75f}); +} + +TEST(WeightsTest, RejectsMismatchedEmbeddingBlobSize) { + ThreadingContext ctx = MakeContext(); + TemporaryBlob blob; + const std::array source = { + hwy::BF16FromF32(-1.0f), hwy::BF16FromF32(0.0f), hwy::BF16FromF32(1.0f), + hwy::BF16FromF32(2.0f)}; + { + BlobWriter writer(blob.path(), ctx); + writer.Add("embedding", source.data(), sizeof(source)); + writer.Finalize(); + } + BlobReader reader(blob.path()); + const BlobRange* range = reader.Find("embedding"); + ASSERT_NE(range, nullptr); + + MatPtr mat("embedding", Type::kSFP, Extents2D(2, 2)); + const weights_internal::TensorToRead tensor{ + .mat = &mat, + .range = *range, + .padding = MatPadding::kOdd, + .to_sfp = true, + .prev_type = Type::kBF16, + .prev_packed_bytes = sizeof(source) + sizeof(BF16), + }; + + EXPECT_DEATH(weights_internal::ReadAllToSFP({tensor}, reader, ctx), + "tensor.range.bytes == tensor.prev_packed_bytes"); +} + +TEST(WeightsTest, ExplicitSFPDisablesOnlyAutomaticMapping) { + ThreadingContext ctx = MakeContext(); + InferenceArgs inference; + LoaderArgs loader("", ""); + loader.to_bf16 = Tristate::kFalse; + + const uint64_t file_mib = ctx.allocator.TotalMiB() / 3 + 1; + const uint64_t file_bytes = hwy::RoundUpTo( + file_mib << 20, static_cast(ctx.allocator.BasePageBytes())); + + // Establish that the normal automatic heuristic would map this file. + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kMap); + + loader.sfp_embedding = Tristate::kTrue; + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kRead); + + // An explicit mapping request still wins and is warned about by the loader. + loader.map = Tristate::kTrue; + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kMap); + + // SFP embedding conversion composes with explicit conversion of other + // tensors to BF16. + loader.map = Tristate::kDefault; + loader.to_bf16 = Tristate::kTrue; + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kReadBF16); +} + +} // namespace +} // namespace gcpp diff --git a/util/zones.cc b/util/zones.cc index 2c63f61c..9cd3475f 100644 --- a/util/zones.cc +++ b/util/zones.cc @@ -116,6 +116,8 @@ const char* ZoneName(Zones zone) { return "Ops.Softmax"; case Zones::kStartupWeightsReadAllToBF16: return "Startup.Weights.ReadAllToBF16"; + case Zones::kStartupWeightsReadAllToSFP: + return "Startup.Weights.ReadAllToSFP"; case Zones::kStartupWeightsReadBatches: return "Startup.Weights.ReadBatches"; default: @@ -207,6 +209,8 @@ const char* CallerName(Callers caller) { return "Ops.RMSNormNoScaleInplaceBatched"; case Callers::kReadAllToBF16: return "ReadAllToBF16"; + case Callers::kReadAllToSFP: + return "ReadAllToSFP"; case Callers::kReadBatches: return "ReadBatches"; case Callers::kSampleAndStream: diff --git a/util/zones.h b/util/zones.h index 01adb005..533182a9 100644 --- a/util/zones.h +++ b/util/zones.h @@ -66,6 +66,7 @@ enum class Zones { // Keep sorted kOpsRopeAndMulBy, kOpsSoftmax, kStartupWeightsReadAllToBF16, + kStartupWeightsReadAllToSFP, kStartupWeightsReadBatches, kNumZones // must be last }; @@ -116,6 +117,7 @@ enum class Callers { // Keep sorted kOpsRMSNormInplaceBatched, kOpsRMSNormNoScaleInplaceBatched, kReadAllToBF16, + kReadAllToSFP, kReadBatches, kSampleAndStream, kTensorStats,