From 86e37076b6c75c607825f18d5bb941b3b3d05d6a Mon Sep 17 00:00:00 2001 From: Phil Culliton Date: Fri, 11 Sep 2026 12:35:45 -0700 Subject: [PATCH] Internal change. FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/google/gemma.cpp/pull/1032 from Mikyx-1:fix/mmap-partial-page-length 7408fa2f3980d0a07fe8589ebb8c4ef9489b8e94 PiperOrigin-RevId: 979960879 --- BUILD.bazel | 5 +- compression/compress.cc | 1 + compression/python/compression_clif_aux.cc | 6 +- compression/types.h | 19 +- evals/model_health_test.cc | 3 +- gemma/activations.h | 3 +- gemma/configs.cc | 2 +- gemma/gemma4_moe.cc | 20 +- gemma/tensor_info.h | 3 + gemma/weights.cc | 29 +- gemma/weights_test.cc | 66 +++++ io/blob_store.cc | 12 +- io/blob_store.h | 5 + ops/fast_ops-inl.h | 183 ++++++++++++ python/configs.cc | 3 +- python/convert_from_safetensors.py | 326 ++++++++++++++++++++- 16 files changed, 645 insertions(+), 41 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index 4f58cfe1..0b9d52e5 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1067,7 +1067,8 @@ cc_test( ":configs", ":cross_entropy", ":gemma_lib", - "@googletest//:gtest_main", # buildcleaner: keep + ":test_util", + "//testing/base/public:gunit_for_library_testonly", # buildcleaner: keep "@highway//:hwy", "@highway//:hwy_test_util", ], @@ -1087,7 +1088,7 @@ cc_test( ":benchmark_helper", ":gemma_lib", ":test_util", - "@googletest//:gtest_main", # buildcleaner: keep + "//testing/base/public:gunit_for_library_testonly", # buildcleaner: keep "@highway//:hwy", "@highway//:hwy_test_util", "@highway//:nanobenchmark", diff --git a/compression/compress.cc b/compression/compress.cc index 6ef89900..3525d95b 100644 --- a/compression/compress.cc +++ b/compression/compress.cc @@ -15,6 +15,7 @@ #include "compression/compress.h" +#include #include #include diff --git a/compression/python/compression_clif_aux.cc b/compression/python/compression_clif_aux.cc index 2080e757..f9b0305f 100644 --- a/compression/python/compression_clif_aux.cc +++ b/compression/python/compression_clif_aux.cc @@ -85,8 +85,10 @@ class SbsWriterImpl : public ISbsWriter { } HWY_ASSERT(weights.size() == mat.Extents().Area()); - Compress(weights.data(), weights.size(), working_set_, mat.Span(), - /*packed_ofs=*/0, ctx_); + { + Compress(weights.data(), weights.size(), working_set_, mat.Span(), + /*packed_ofs=*/0, ctx_); + } writer_.Add(name, mat.Packed(), mat.PackedBytes()); } diff --git a/compression/types.h b/compression/types.h index 429865bc..0b82a70c 100644 --- a/compression/types.h +++ b/compression/types.h @@ -269,7 +269,8 @@ constexpr bool IsMxFp4Stream() { template constexpr bool IsPacked() { return IsNuqStream() || IsI8Stream() || - IsQ4_0Stream() || IsMxFp4Stream(); + IsQ4_0Stream() || IsMxFp4Stream() + ; } template @@ -294,12 +295,14 @@ enum class Type { kInt8, kQ4_0, kMXFP4, + kReserved14, }; // These are used in `ModelConfig.Specifier`, hence the strings will not // change, though new ones may be added. static constexpr const char* kTypeStrings[] = { - "unknown", "f32", "bf16", "sfp", "nuq", "f64", "u32", - "u64", "i8", "u16", "u8", "int8", "q4_0", "mxfp4"}; + "unknown", "f32", "bf16", "sfp", "nuq", "f64", "u32", "u64", + "i8", "u16", "u8", "int8", "q4_0", "mxfp4", "reserved14" +}; static constexpr size_t kNumTypes = sizeof(kTypeStrings) / sizeof(kTypeStrings[0]); static constexpr size_t kTypeBits[] = { @@ -317,6 +320,7 @@ static constexpr size_t kTypeBits[] = { 8 * sizeof(int8_t), 4 /* Q4_0Stream, actually 4.5 */, 4 /* MxFp4Stream, actually 4.25 */, + 0 /* reserved */, }; static inline bool EnumValid(Type type) { @@ -376,17 +380,20 @@ constexpr bool IsCompressed() { hwy::IsSame, NuqStream>() || hwy::IsSame, I8Stream>() || hwy::IsSame, Q4_0Stream>() || - hwy::IsSame, MxFp4Stream>(); + hwy::IsSame, MxFp4Stream>() + ; } static inline bool IsCompressed(Type type) { return type == Type::kSFP || type == Type::kNUQ || type == Type::kI8 || - type == Type::kQ4_0 || type == Type::kMXFP4; + type == Type::kQ4_0 || type == Type::kMXFP4 + ; } static inline bool IsPacked(Type type) { return type == Type::kNUQ || type == Type::kI8 || type == Type::kQ4_0 || - type == Type::kMXFP4; + type == Type::kMXFP4 + ; } static inline bool SupportsPointerArithmetic(Type type) { diff --git a/evals/model_health_test.cc b/evals/model_health_test.cc index 6475082f..ed01280d 100644 --- a/evals/model_health_test.cc +++ b/evals/model_health_test.cc @@ -44,6 +44,7 @@ #include "evals/cross_entropy.h" #include "gemma/configs.h" #include "gemma/gemma.h" +#include "util/test_util.h" #include "hwy/base.h" #include "hwy/tests/hwy_gtest.h" @@ -485,7 +486,7 @@ TEST_F(ModelHealthTest, DeterministicGeneration) { } // namespace gcpp int main(int argc, char** argv) { - testing::InitGoogleTest(&argc, argv); + gcpp::InternalInitTest(); gcpp::ModelHealthTest::InitEnv(argc, argv); int ret = RUN_ALL_TESTS(); gcpp::ModelHealthTest::DeleteEnv(); diff --git a/gemma/activations.h b/gemma/activations.h index 341875ae..9a875cca 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -801,7 +801,8 @@ struct Activations { config.model_dim, allocator)), mla_o_in(MatFactory("mla_o_in", mla_dims.o_in_dim > 0 ? batch_size : 0, - mla_dims.o_in_dim, allocator)) { + mla_dims.o_in_dim, allocator)) + { moe_C1.AllocateAndAttachRowPtrs(row_ptrs); moe_C2.AllocateAndAttachRowPtrs(row_ptrs); ffw_expert_in.AllocateAndAttachRowPtrs(row_ptrs); diff --git a/gemma/configs.cc b/gemma/configs.cc index def520b1..3c3cbac8 100644 --- a/gemma/configs.cc +++ b/gemma/configs.cc @@ -481,7 +481,7 @@ static LayerConfig LayerConfigGemma4_26B_MoE_LM(size_t model_dim) { static ModelConfig ConfigGemma4_26B_MoE() { ModelConfig config = ConfigBaseGemmaV4(); config.display_name = "Gemma4_26B_MoE"; - config.final_cap = 0.0f; + config.final_cap = 30.0f; config.att_cap = 0.0f; config.model = Model::GEMMA4_26B_MOE; config.wrapping = PromptWrapping::GEMMA_IT; diff --git a/gemma/gemma4_moe.cc b/gemma/gemma4_moe.cc index d961b580..84e53237 100644 --- a/gemma/gemma4_moe.cc +++ b/gemma/gemma4_moe.cc @@ -48,6 +48,7 @@ #include "gemma/attention.h" // includes highway.h #include "gemma/tiled_attention.h" #include "gemma/gemma-inl.h" +#include "ops/fast_ops-inl.h" #include "ops/ops-inl.h" HWY_BEFORE_NAMESPACE(); @@ -471,9 +472,12 @@ void Gemma4MoETransformerLayer(size_t num_tokens, size_t layer_idx, /*is_attention=*/true, env.ctx); // Dual-Path FFW - pre_norm( - layer.pre_ffw2_ns.HasPtr() ? layer.pre_ffw2_ns : layer.pre_ffw_norm_scale, - activations.pre_ffw_rms_out); + const MatPtr& shared_norm = layer.pre_ffw2_ns.HasPtr() ? layer.pre_ffw2_ns : layer.pre_ffw_norm_scale; + const MatPtr& moe_norm = layer.pre_ffw_norm_scale; + const MatPtr& shared_post_norm = layer.post_ffw2_ns; + const MatPtr& moe_post_norm = layer.post_ffw1_ns; + + pre_norm(shared_norm, activations.pre_ffw_rms_out); // Shared MLP Path FFWNoVit(layer, activations, env); // writes to activations.ffw_out @@ -487,17 +491,17 @@ void Gemma4MoETransformerLayer(size_t num_tokens, size_t layer_idx, } } - if (layer.post_ffw2_ns.HasPtr()) { - rms_norm_inplace(layer.post_ffw2_ns, activations.attention.att_sums); + if (shared_post_norm.HasPtr()) { + rms_norm_inplace(shared_post_norm, activations.attention.att_sums); } // MoE Path - pre_norm(layer.pre_ffw_norm_scale, activations.pre_ffw_rms_out); + pre_norm(moe_norm, activations.pre_ffw_rms_out); Gemma4MoE::MoEFFW(layer, activations, env); // writes to activations.ffw_out - if (layer.post_ffw1_ns.HasPtr()) { - rms_norm_inplace(layer.post_ffw1_ns, activations.ffw_out); + if (moe_post_norm.HasPtr()) { + rms_norm_inplace(moe_post_norm, activations.ffw_out); } // Combine & Final Norm (Fix for dual-path combination) diff --git a/gemma/tensor_info.h b/gemma/tensor_info.h index 15339f8e..d020bb93 100644 --- a/gemma/tensor_info.h +++ b/gemma/tensor_info.h @@ -68,6 +68,9 @@ struct TensorInfo { // If false, then [10, 20, 30] -> [10*20, 30] and [30] -> [1, 30]. // If true, then [10, 20, 30] -> [10, 20*30] and [30] -> [1, 30]. bool cols_take_extra_dims = false; + // Optional pre-computed scale (e.g. for kW2_UL weights from QAFT). + // If > 0.0, used directly instead of re-estimating via ScaleWeightsW2UL. + float scale = 0.0f; }; // Collapses/expands the tensor dims into 2D extents, which may be 0, 0 for diff --git a/gemma/weights.cc b/gemma/weights.cc index 86b22463..3bf668ae 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -20,6 +20,7 @@ #include #include +#include #include // NOLINT #include #include @@ -129,7 +130,10 @@ static void SplitPackedMatrix(MatPtr& parent, size_t split_row, MatPtr& w1, void LayerWeightsPtrs::SplitW1() { // Used for Gemma layers; FFWVit uses different tensors. if (layer_config.type == LayerAttentionType::kVit) return; - if (layer_config.IsMoE()) return; + if (layer_config.IsMoE() && !gating_einsum_w.HasPtr() && + !gating_einsum_w1.HasPtr()) { + return; + } // Files have both or neither of w1 and w2. HWY_ASSERT(gating_einsum_w1.HasPtr() == gating_einsum_w2.HasPtr()); @@ -543,6 +547,18 @@ void LayerWeightsPtrs::Fixup(Model model, std::vector& mat_owners, const size_t elem_bytes = qkv_einsum_w2.ElementBytes(); const size_t old_row_bytes = old_stride * elem_bytes; const size_t kv_heads = layer_config.kv_heads; + const size_t qkv_dim = layer_config.qkv_dim; + + // In Gemma 4 global layers, attention_k_eq_v is true (K0 == V0). + // If already interleaved by exporter: [K0, V0, K1, V1], slice 0 == slice 1. + // If not interleaved: [K0, K1, V0, V1], slice 0 (K0) != slice 1 (K1). + const uint8_t* slice0 = qkv_einsum_w2.RowBytes(0); + const uint8_t* slice1 = qkv_einsum_w2.RowBytes(qkv_dim); + if (std::memcmp(slice0, slice1, old_row_bytes) == 0) { + // Exporter already emitted interleaved layout; skip fixup. + return; + } + const size_t total_bytes = qkv_einsum_w2.Rows() * old_row_bytes; hwy::AlignedFreeUniquePtr tmp = hwy::AllocateAligned(total_bytes); @@ -556,7 +572,6 @@ void LayerWeightsPtrs::Fixup(Model model, std::vector& mat_owners, } const size_t new_row_bytes = qkv_einsum_w2.Cols() * elem_bytes; - const size_t qkv_dim = layer_config.qkv_dim; const uint8_t* src_ptr = tmp.get(); for (size_t i = 0; i < kv_heads; ++i) { for (size_t row = 0; row < qkv_dim; ++row) { @@ -676,12 +691,12 @@ WeightsPtrs::Mode weights_internal::ChooseMode(uint64_t file_bytes, map = Tristate::kFalse; } - // Disable mapping if not padded to the base page size. - if (file_bytes % allocator.BasePageBytes() != 0) { + // Kernels can read multiple vectors ahead, so retain blob alignment even + // though mapping itself does not require a page-aligned file length. + if (file_bytes % kBlobAlign != 0) { if (map == Tristate::kTrue) { // Only complain if explicitly requested. - HWY_WARN("Unable to map non-padded file (%zu, %zu), reading instead.", - static_cast(file_bytes >> 10), - allocator.BasePageBytes()); + HWY_WARN("File size %zu is not a multiple of %zu bytes, reading instead.", + static_cast(file_bytes), kBlobAlign); } map = Tristate::kFalse; } diff --git a/gemma/weights_test.cc b/gemma/weights_test.cc index 4cec58bd..4485b448 100644 --- a/gemma/weights_test.cc +++ b/gemma/weights_test.cc @@ -219,5 +219,71 @@ TEST(WeightsTest, ExplicitSFPDisablesOnlyAutomaticMapping) { WeightsPtrs::Mode::kReadBF16); } +TEST(WeightsTest, MapsBlobAlignedNonPageAlignedFile) { + ThreadingContext ctx = MakeContext(); + const size_t file_bytes = ctx.allocator.BasePageBytes() + kBlobAlign; + ASSERT_EQ(file_bytes % kBlobAlign, 0); + ASSERT_NE(file_bytes % ctx.allocator.BasePageBytes(), 0); + std::vector contents(file_bytes); + contents.front() = 1; + contents.back() = 2; + + TemporaryBlob blob; + auto file = OpenFileOrAbort(blob.path(), "w+"); + ASSERT_TRUE(file->Write(contents.data(), contents.size(), 0)); + MapPtr mapped = file->Map(); + ASSERT_NE(mapped, nullptr); + EXPECT_EQ(mapped[0], contents.front()); + EXPECT_EQ(mapped[file_bytes - 1], contents.back()); + + InferenceArgs inference; + LoaderArgs loader("", ""); + loader.map = Tristate::kTrue; + loader.to_bf16 = Tristate::kFalse; + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kMap); +} + +TEST(WeightsTest, MappingRequiresBlobAlignment) { + ThreadingContext ctx = MakeContext(); + InferenceArgs inference; + LoaderArgs loader("", ""); + loader.to_bf16 = Tristate::kFalse; + + // Large enough for the automatic mapping heuristic to select kMap. + const uint64_t file_mib = ctx.allocator.TotalMiB() / 3 + 1; + const uint64_t page_bytes = ctx.allocator.BasePageBytes(); + const uint64_t page_aligned_bytes = + hwy::RoundUpTo(file_mib << 20, page_bytes); + + for (const Tristate map : {Tristate::kDefault, Tristate::kTrue}) { + loader.map = map; + for (const uint64_t file_bytes : + {page_aligned_bytes, page_aligned_bytes + kBlobAlign}) { + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kMap) + << "file_bytes=" << file_bytes; + } + // Cover both an almost-full page (Jan's 4090-byte example) and a file + // whose final page has plenty of space but lacks blob alignment. + for (const uint64_t file_bytes : + {page_bytes - 6, page_bytes + 1, page_aligned_bytes - 6, + page_aligned_bytes + 1}) { + EXPECT_EQ(weights_internal::ChooseMode(file_bytes, loader, inference, + ctx.allocator), + WeightsPtrs::Mode::kRead) + << "file_bytes=" << file_bytes; + } + } + + loader.map = Tristate::kTrue; + loader.to_bf16 = Tristate::kTrue; + EXPECT_EQ(weights_internal::ChooseMode(page_aligned_bytes + 1, loader, + inference, ctx.allocator), + WeightsPtrs::Mode::kReadBF16); +} + } // namespace } // namespace gcpp diff --git a/io/blob_store.cc b/io/blob_store.cc index 4d11fcb8..8a8f05ae 100644 --- a/io/blob_store.cc +++ b/io/blob_store.cc @@ -35,15 +35,9 @@ namespace gcpp { static_assert(HWY_IS_LITTLE_ENDIAN, "Assumes little endian"); -// Each blob offset is a multiple of this, an upper bound on SVE vectors and -// usually also larger than L2 cache lines. This is useful when memory mapping -// the entire file, because offset alignment then determines the alignment of -// the blob in memory. Aligning each blob to the (largest) page size would be -// too wasteful, see `kEndAlign`. -constexpr size_t kBlobAlign = 256; // test also hard-codes this value - -// Linux mmap requires the file to be a multiple of the (base) page size, which -// can be up to 64 KiB on Arm. Apple uses 16 KiB, most others use 4 KiB. +// Pad newly written files to cover base page sizes up to 64 KiB on Arm. +// Apple uses 16 KiB, most others use 4 KiB. Mapping also accepts older files +// padded only to kBlobAlign. constexpr size_t kEndAlign = 64 * 1024; constexpr size_t kU128Bytes = sizeof(hwy::uint128_t); diff --git a/io/blob_store.h b/io/blob_store.h index 82c2357b..3b6e8e0c 100644 --- a/io/blob_store.h +++ b/io/blob_store.h @@ -34,6 +34,11 @@ namespace gcpp { +// Blob offsets and padded sizes are multiples of this, an upper bound on SVE +// vector bytes and usually also larger than L2 cache lines. This preserves blob +// alignment when mapping the entire file without requiring page-sized padding. +constexpr size_t kBlobAlign = 256; // test also hard-codes this value + // One blob's extents within the file. struct BlobRange { uint64_t End() const { return offset + bytes; } diff --git a/ops/fast_ops-inl.h b/ops/fast_ops-inl.h index 3ca5c1a0..1f49084b 100644 --- a/ops/fast_ops-inl.h +++ b/ops/fast_ops-inl.h @@ -314,9 +314,192 @@ static HWY_NOINLINE HWY_MAYBE_UNUSED void FastSigmoid(T* HWY_RESTRICT x, }); } +// 16-point in-register Fast Walsh-Hadamard Transform on a 16-lane F32 vector. +template , HWY_IF_F32_D(DF)> +HWY_INLINE VF FWHT16(DF df, VF v) { + // Stage 1: h = 1 + const auto e1 = hn::DupEven(v); + const auto o1 = hn::DupOdd(v); + v = hn::OddEven(hn::Sub(e1, o1), hn::Add(e1, o1)); + + // Stage 2: h = 2 + const hn::Repartition du64; + const auto vu64 = hn::BitCast(du64, v); + const auto e2 = hn::BitCast(df, hn::DupEven(vu64)); + const auto o2 = hn::BitCast(df, hn::DupOdd(vu64)); + v = hn::BitCast( + df, hn::OddEven(hn::BitCast(du64, hn::Sub(e2, o2)), + hn::BitCast(du64, hn::Add(e2, o2)))); + + // Stage 3: h = 4 + const hn::Half dfh; + const hn::Half dfq; + auto lo_h = hn::LowerHalf(dfh, v); + auto hi_h = hn::UpperHalf(dfh, v); + auto q0 = hn::LowerHalf(dfq, lo_h); + auto q1 = hn::UpperHalf(dfq, lo_h); + auto q2 = hn::LowerHalf(dfq, hi_h); + auto q3 = hn::UpperHalf(dfq, hi_h); + lo_h = hn::Combine(dfh, hn::Sub(q0, q1), hn::Add(q0, q1)); + hi_h = hn::Combine(dfh, hn::Sub(q2, q3), hn::Add(q2, q3)); + + // Stage 4: h = 8 + return hn::Combine(df, hn::Sub(lo_h, hi_h), hn::Add(lo_h, hi_h)); +} + +template > +HWY_INLINE void FastWalshHadamard128(D d, T* HWY_RESTRICT data, size_t length) { + HWY_DASSERT(length % 128 == 0); + constexpr float kNorm = 0.08838834764831845f; // 1.0f / sqrt(128.0f) + const hn::ScalableTag df; + const hn::Repartition dbf; + using VF = hn::Vec; + const auto vnorm = hn::Set(df, kNorm); + + if constexpr (hn::MaxLanes(df) == 16 && !HWY_HAVE_SCALABLE) { + // Optimal in-register path for AVX-512 (H128 = H8 x H16): + for (size_t block = 0; block < length; block += 128) { + VF v0, v1, v2, v3, v4, v5, v6, v7; + if constexpr (IsBF16()) { + const auto b0 = hn::Load(dbf, data + block + 0); + const auto b1 = hn::Load(dbf, data + block + 32); + const auto b2 = hn::Load(dbf, data + block + 64); + const auto b3 = hn::Load(dbf, data + block + 96); + v0 = hn::PromoteLowerTo(df, b0); + v1 = hn::PromoteUpperTo(df, b0); + v2 = hn::PromoteLowerTo(df, b1); + v3 = hn::PromoteUpperTo(df, b1); + v4 = hn::PromoteLowerTo(df, b2); + v5 = hn::PromoteUpperTo(df, b2); + v6 = hn::PromoteLowerTo(df, b3); + v7 = hn::PromoteUpperTo(df, b3); + } else { + v0 = hn::Load(df, data + block + 0); + v1 = hn::Load(df, data + block + 16); + v2 = hn::Load(df, data + block + 32); + v3 = hn::Load(df, data + block + 48); + v4 = hn::Load(df, data + block + 64); + v5 = hn::Load(df, data + block + 80); + v6 = hn::Load(df, data + block + 96); + v7 = hn::Load(df, data + block + 112); + } + + // Inter-vector stages (H8): + // h = 64 + auto t0 = v0; v0 = hn::Add(t0, v4); v4 = hn::Sub(t0, v4); + auto t1 = v1; v1 = hn::Add(t1, v5); v5 = hn::Sub(t1, v5); + auto t2 = v2; v2 = hn::Add(t2, v6); v6 = hn::Sub(t2, v6); + auto t3 = v3; v3 = hn::Add(t3, v7); v7 = hn::Sub(t3, v7); + + // h = 32 + t0 = v0; v0 = hn::Add(t0, v2); v2 = hn::Sub(t0, v2); + t1 = v1; v1 = hn::Add(t1, v3); v3 = hn::Sub(t1, v3); + auto t4 = v4; v4 = hn::Add(t4, v6); v6 = hn::Sub(t4, v6); + auto t5 = v5; v5 = hn::Add(t5, v7); v7 = hn::Sub(t5, v7); + + // h = 16 + t0 = v0; v0 = hn::Add(t0, v1); v1 = hn::Sub(t0, v1); + auto t2_ = v2; v2 = hn::Add(t2_, v3); v3 = hn::Sub(t2_, v3); + t4 = v4; v4 = hn::Add(t4, v5); v5 = hn::Sub(t4, v5); + auto t6 = v6; v6 = hn::Add(t6, v7); v7 = hn::Sub(t6, v7); + + // Intra-vector stages (H16) and normalization: + v0 = hn::Mul(FWHT16(df, v0), vnorm); + v1 = hn::Mul(FWHT16(df, v1), vnorm); + v2 = hn::Mul(FWHT16(df, v2), vnorm); + v3 = hn::Mul(FWHT16(df, v3), vnorm); + v4 = hn::Mul(FWHT16(df, v4), vnorm); + v5 = hn::Mul(FWHT16(df, v5), vnorm); + v6 = hn::Mul(FWHT16(df, v6), vnorm); + v7 = hn::Mul(FWHT16(df, v7), vnorm); + + if constexpr (IsBF16()) { + hn::Store(hn::OrderedDemote2To(dbf, v0, v1), dbf, data + block + 0); + hn::Store(hn::OrderedDemote2To(dbf, v2, v3), dbf, data + block + 32); + hn::Store(hn::OrderedDemote2To(dbf, v4, v5), dbf, data + block + 64); + hn::Store(hn::OrderedDemote2To(dbf, v6, v7), dbf, data + block + 96); + } else { + hn::Store(v0, df, data + block + 0); + hn::Store(v1, df, data + block + 16); + hn::Store(v2, df, data + block + 32); + hn::Store(v3, df, data + block + 48); + hn::Store(v4, df, data + block + 64); + hn::Store(v5, df, data + block + 80); + hn::Store(v6, df, data + block + 96); + hn::Store(v7, df, data + block + 112); + } + } + } else { + // Portable SIMD fallback for architectures with < 16 F32 lanes: + const size_t lanes = hn::Lanes(df); + HWY_ALIGN float buf[128]; + for (size_t block = 0; block < length; block += 128) { + if constexpr (IsBF16()) { + for (size_t i = 0; i < 128; i += 2 * lanes) { + const auto b = hn::Load(dbf, data + block + i); + hn::Store(hn::PromoteLowerTo(df, b), df, buf + i); + hn::Store(hn::PromoteUpperTo(df, b), df, buf + i + lanes); + } + } else { + for (size_t i = 0; i < 128; i += lanes) { + hn::Store(hn::Load(df, data + block + i), df, buf + i); + } + } + + for (size_t step = 1; step < 128; step <<= 1) { + const size_t jump = step << 1; + if (step >= lanes) { + for (size_t i = 0; i < 128; i += jump) { + for (size_t j = 0; j < step; j += lanes) { + const auto u = hn::Load(df, buf + i + j); + const auto v = hn::Load(df, buf + i + j + step); + hn::Store(hn::Add(u, v), df, buf + i + j); + hn::Store(hn::Sub(u, v), df, buf + i + j + step); + } + } + } else { + for (size_t i = 0; i < 128; i += jump) { + for (size_t j = 0; j < step; ++j) { + const float u = buf[i + j]; + const float v = buf[i + j + step]; + buf[i + j] = u + v; + buf[i + j + step] = u - v; + } + } + } + } + + for (size_t i = 0; i < 128; i += lanes) { + hn::Store(hn::Mul(hn::Load(df, buf + i), vnorm), df, buf + i); + } + + if constexpr (IsBF16()) { + for (size_t i = 0; i < 128; i += 2 * lanes) { + const auto f0 = hn::Load(df, buf + i); + const auto f1 = hn::Load(df, buf + i + lanes); + hn::Store(hn::OrderedDemote2To(dbf, f0, f1), dbf, data + block + i); + } + } else { + for (size_t i = 0; i < 128; i += lanes) { + hn::Store(hn::Load(df, buf + i), df, data + block + i); + } + } + } + } +} + +template +static HWY_NOINLINE HWY_MAYBE_UNUSED void FastWalshHadamard128(T* HWY_RESTRICT x, + size_t size) { + namespace hn = hwy::HWY_NAMESPACE; + const hn::ScalableTag d; + FastWalshHadamard128(d, x, size); +} + // NOLINTNEXTLINE(google-readability-namespace-comments) } // namespace HWY_NAMESPACE } // namespace gcpp HWY_AFTER_NAMESPACE(); + #endif // NOLINT diff --git a/python/configs.cc b/python/configs.cc index e264f451..e5b67185 100644 --- a/python/configs.cc +++ b/python/configs.cc @@ -131,7 +131,8 @@ PYBIND11_MODULE(configs, py_module) { .def_readwrite("min_size", &gcpp::TensorInfo::min_size) .def_readwrite("scaled_softplus", &gcpp::TensorInfo::scaled_softplus) .def_readwrite("cols_take_extra_dims", - &gcpp::TensorInfo::cols_take_extra_dims); + &gcpp::TensorInfo::cols_take_extra_dims) + .def_readwrite("scale", &gcpp::TensorInfo::scale); class_(py_module, "TensorInfoRegistry") .def(init()) diff --git a/python/convert_from_safetensors.py b/python/convert_from_safetensors.py index 1e2dbe3c..7a0bfed5 100644 --- a/python/convert_from_safetensors.py +++ b/python/convert_from_safetensors.py @@ -1725,6 +1725,322 @@ def add_gating_einsum(i): csv.writer(csv_handle).writerows(metadata) +def export_gemma4_moe_sbs( + model_specifier: str, + load_path: str, + tokenizer_file: str, + csv_file: str, + sbs_file: str, +) -> None: + """Exports an sbs file from a Gemma 4 MoE safetensors checkpoint.""" + if load_path.endswith(".json"): + with open(load_path, "r") as f: + j_obj = json.load(f) + files = list(set(j_obj["weight_map"].values())) + files = [os.path.join(os.path.dirname(load_path), f) for f in files] + else: + files = [load_path] + + params = {} + for file in files: + with safetensors.safe_open(file, framework="pt") as f: + for k in f.keys(): + if ( + k.startswith("vision_tower.") + or k.startswith("multi_modal_projector.") + or k.startswith("model.vision_tower.") + or k.startswith("model.embed_vision.") + ): + continue + params[k] = f.get_tensor(k) + + if "model.language_model.embed_tokens.weight" in params: + llm_prefix = "model.language_model." + elif "language_model.model.embed_tokens.weight" in params: + llm_prefix = "language_model.model." + elif "model.embed_tokens.weight" in params: + llm_prefix = "model." + elif "embed_tokens.weight" in params: + llm_prefix = "" + else: + raise ValueError( + "Could not locate embed_tokens.weight in Gemma 4 checkpoint." + ) + + embed_tokens = params[f"{llm_prefix}embed_tokens.weight"] + vocab_size, model_dim = embed_tokens.shape + head_dim = 256 + num_layers = 30 + ff_hidden_dim = 2112 + +weight_type = configs.Type.kSFP + sbs_config = configs.ModelConfig( + configs.Model.GEMMA4_26B_MOE, weight_type, configs.PromptWrapping.GEMMA_IT + ) + sbs_config.final_cap = 30.0 + + writer = compression.SbsWriter(sbs_file) + metadata = [] + scales = {} + + def add_data( + param_name, + data, + expected_shape, + sbs_name, + layer_index=None, + is_w2_ul=False, + ): + if expected_shape is not None: + if not isinstance(expected_shape, tuple): + expected_shape = (expected_shape,) + assert ( + data.shape == expected_shape + ), f"{param_name}: got {data.shape}, expected {expected_shape}" + + assert isinstance(data, torch.Tensor) + data = data.to(torch.float32).numpy() + data = np.array(data) + + if layer_index is not None: + sbs_name = sbs_name + f"_{layer_index}" + + value = flatten_f32(data) + scale = compute_scale(value) + both_names = param_name + "::" + sbs_name + metadata.append((both_names, data.dtype, data.shape, scale)) + + if _is_float_param(sbs_name): + packed = configs.Type.kF32 + elif ( + _is_bf16_param(sbs_name) + or sbs_name.startswith("router_scale") + or sbs_name.startswith("p_expert_sc") + or sbs_name.startswith("skip_scale") + ): + packed = configs.Type.kBF16 + else: + packed = configs.Type.kSFP + scales[sbs_name] = scale + + info = configs.TensorInfo() + info.name = sbs_name + info.shape = data.shape + writer.insert(sbs_name, value, packed, info) + + # Embeddings & Final Norm + add_data( + f"{llm_prefix}embed_tokens.weight", + params.pop(f"{llm_prefix}embed_tokens.weight"), + (vocab_size, model_dim), + "c_embedding", + ) + add_data( + f"{llm_prefix}norm.weight", + params.pop(f"{llm_prefix}norm.weight") - 1.0, + (model_dim,), + "c_final_norm", + ) + + for i in range(num_layers): + layer_head_dim = 512 if (i % 6 == 5) else 256 + + # Attention + if f"{llm_prefix}layers.{i}.self_attn.o_proj.weight" in params: + o = params.pop(f"{llm_prefix}layers.{i}.self_attn.o_proj.weight") + n_heads = o.shape[1] // layer_head_dim + o = o.reshape(model_dim, n_heads, layer_head_dim).permute(1, 0, 2) + add_data( + f"{llm_prefix}layers.{i}.self_attn.o_proj.weight", + o, + (n_heads, model_dim, layer_head_dim), + "att_ein", + i, + ) + + if f"{llm_prefix}layers.{i}.self_attn.q_proj.weight" in params: + q = params.pop(f"{llm_prefix}layers.{i}.self_attn.q_proj.weight") + k = params.pop(f"{llm_prefix}layers.{i}.self_attn.k_proj.weight") + if f"{llm_prefix}layers.{i}.self_attn.v_proj.weight" in params: + v = params.pop(f"{llm_prefix}layers.{i}.self_attn.v_proj.weight") + else: + # For full_attention layers where attention_k_eq_v is true + v = k.clone() + n_q = q.shape[0] // layer_head_dim + n_kv = k.shape[0] // layer_head_dim + q = q.reshape(n_q, layer_head_dim, model_dim) + k = k.reshape(n_kv, layer_head_dim, model_dim) + v = v.reshape(n_kv, layer_head_dim, model_dim) + stacked = ( + torch.stack((k, v), dim=0) + .transpose(0, 1) + .reshape(2 * n_kv, layer_head_dim, model_dim) + ) + qkv = torch.cat([q, stacked], dim=0) + add_data( + f"{llm_prefix}layers.{i}.self_attn.qkv_proj.weight", + qkv, + (n_q + 2 * n_kv, layer_head_dim, model_dim), + "qkv_ein", + i, + ) + + # Norms + for norm_name, sbs_norm in [ + ("input_layernorm.weight", "pre_att_ns"), + ("post_attention_layernorm.weight", "post_att_ns"), + ("pre_feedforward_layernorm.weight", "pre_ffw2_ns"), + ("post_feedforward_layernorm.weight", "post_ff_ns"), + ("post_feedforward_layernorm_1.weight", "post_ffw2_ns"), + ("post_feedforward_layernorm_2.weight", "post_ffw1_ns"), + ("pre_feedforward_layernorm_2.weight", "pre_ff_ns"), + ]: + key = f"{llm_prefix}layers.{i}.{norm_name}" + if key in params: + add_data(key, params.pop(key) - 1.0, (model_dim,), sbs_norm, i) + + for qk_name, sbs_qk in [ + ("self_attn.q_norm.weight", "query_norm"), + ("self_attn.k_norm.weight", "key_norm"), + ]: + key = f"{llm_prefix}layers.{i}.{qk_name}" + if key in params: + qk_t = params.pop(key) - 1.0 + add_data(key, qk_t, (qk_t.shape[0],), sbs_qk, i) + + # Shared MLP + shared_gate_key = f"{llm_prefix}layers.{i}.mlp.gate_proj.weight" + shared_up_key = f"{llm_prefix}layers.{i}.mlp.up_proj.weight" + if shared_gate_key in params and shared_up_key in params: + shared_gate = params.pop(shared_gate_key) + shared_up = params.pop(shared_up_key) + add_data( + shared_gate_key, + shared_gate, + (ff_hidden_dim, model_dim), + "gating1_w", + i, + ) + add_data( + shared_up_key, + shared_up, + (ff_hidden_dim, model_dim), + "gating2_w", + i, + ) + + shared_down_key = f"{llm_prefix}layers.{i}.mlp.down_proj.weight" + if shared_down_key in params: + shared_down = params.pop(shared_down_key) + add_data( + shared_down_key, + shared_down, + (model_dim, ff_hidden_dim), + "linear_w", + i, + ) + + scalar_key = f"{llm_prefix}layers.{i}.layer_scalar" + if scalar_key in params: + scalar = params.pop(scalar_key).reshape(1) + add_data( + scalar_key, + scalar, + (1,), + "skip_scale", + i, + ) + + # MoE Experts + gate_up_key = f"{llm_prefix}layers.{i}.experts.gate_up_proj" + if gate_up_key in params: + gate_up = params.pop(gate_up_key) + num_exp = gate_up.shape[0] + hidden_e = gate_up.shape[1] // 2 + gate_up = gate_up.reshape(num_exp, 2, hidden_e, model_dim) + for e in range(num_exp): + add_data( + f"{gate_up_key}_{e}_1", + gate_up[e, 0, :, :], + (hidden_e, model_dim), + f"gating1_w_{i}_{e}", + layer_index=None, + ) + add_data( + f"{gate_up_key}_{e}_2", + gate_up[e, 1, :, :], + (hidden_e, model_dim), + f"gating2_w_{i}_{e}", + layer_index=None, + ) + + down_key = f"{llm_prefix}layers.{i}.experts.down_proj" + if down_key in params: + down = params.pop(down_key) + num_exp = down.shape[0] + hidden_e = down.shape[2] + for e in range(num_exp): + add_data( + f"{down_key}_{e}", + down[e, :, :], + (model_dim, hidden_e), + f"linear_w_{i}_{e}", + layer_index=None, + ) + + router_key = f"{llm_prefix}layers.{i}.router.proj.weight" + if router_key in params: + router = params.pop(router_key) + add_data( + router_key, + router, + (router.shape[0], model_dim), + "moe_router", + i, + ) + + router_scale_key = f"{llm_prefix}layers.{i}.router.scale" + if router_scale_key in params: + router_scale = params.pop(router_scale_key) + add_data( + router_scale_key, + router_scale, + (model_dim,), + "router_scale", + i, + ) + + p_expert_sc_key = f"{llm_prefix}layers.{i}.router.per_expert_scale" + if p_expert_sc_key in params: + p_expert_sc = params.pop(p_expert_sc_key) + add_data( + p_expert_sc_key, + p_expert_sc, + (p_expert_sc.shape[0],), + "p_expert_sc", + i, + ) + + if params: + logging.info( + "Unconsumed parameters (e.g. vision or shared embeddings): %s", + list(params.keys())[:10], + ) + + if tokenizer_file.endswith(".json"): + sbs_config.tokenizer_kind = configs.TokenizerKind.kHfBpe + tokenizer_blob = pack_bpe_tokenizer(tokenizer_file) + else: + sbs_config.tokenizer_kind = configs.TokenizerKind.kSentencePiece + with open(tokenizer_file, "rb") as f: + tokenizer_blob = f.read() + writer.write(sbs_config, tokenizer_blob) + + with open(csv_file, "w") as csv_handle: + csv.writer(csv_handle).writerows(metadata) + logging.info("Successfully exported Gemma 4 MoE SBS to %s", sbs_file) + + def main(argv: Sequence[str]) -> None: if len(argv) > 1: raise app.UsageError("Too many command-line arguments.") @@ -1772,6 +2088,10 @@ def main(argv: Sequence[str]) -> None: export_gemma4_lm_sbs( model_specifier, load_path, tokenizer_file, metadata_file, sbs_file ) + elif model_specifier.startswith("gemma4-"): + export_gemma4_moe_sbs( + model_specifier, load_path, tokenizer_file, metadata_file, sbs_file + ) elif model_specifier.startswith("t5gemma"): export_t5gemma_sbs( model_specifier, @@ -1787,9 +2107,9 @@ def main(argv: Sequence[str]) -> None: ) else: raise app.UsageError( - f"Unsupported model_specifier {model_specifier!r}. Expected a " - "'paligemma*', 'gemma3-*-lm-*', 'gemma4-*-lm-*', 'qwen3-*', or " - "'t5gemma*' specifier." + f"Unsupported model_specifier {model_specifier!r}. Expected a" + " 'paligemma*', 'gemma3-*-lm-*', 'gemma4-*-lm-*', 'gemma4-*'," + " 'qwen3-*', or 't5gemma*' specifier." )