Merge pull request #918 from plawanrath:feat/gemma3-lm-only

PiperOrigin-RevId: 923303424

Author: copybara-github

CMakeLists.txt

@@ -36,7 +36,7 @@ if(EMSCRIPTEN)
   add_link_options("-sEXIT_RUNTIME=1")
 endif()
 
-FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG c971dbe61bd2751923e3458666450bf95dfbbd98 EXCLUDE_FROM_ALL)
+FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 783beecea257838c964ec7644697e8b08396596e EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(highway)
 
 if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
@@ -113,6 +113,8 @@ set(SOURCES
   gemma/gemma.h
   gemma/kv_cache.cc
   gemma/kv_cache.h
+  gemma/kv_transcoding.cc
+  gemma/kv_transcoding.h
   gemma/model_store.cc
   gemma/model_store.h
   gemma/tensor_info.cc
@@ -146,6 +148,8 @@ set(SOURCES
   ops/sum-inl.h
   paligemma/image.cc
   paligemma/image.h
+  paligemma/paligemma_helper.cc
+  paligemma/paligemma_helper.h
   util/allocator.cc
   util/allocator.h
   util/basics.cc
@@ -257,9 +261,7 @@ set(GEMMA_TEST_FILES
   compression/distortion_test.cc
   compression/nuq_test.cc
   compression/sfp_test.cc
-  evals/gemma_test.cc
   gemma/gemma_args_test.cc
-  gemma/flash_attention_test.cc
   gemma/tensor_info_test.cc
   io/blob_store_test.cc
   io/fields_test.cc
@@ -268,11 +270,24 @@ set(GEMMA_TEST_FILES
   ops/matmul_test.cc
   ops/ops_test.cc
   paligemma/image_test.cc
-  paligemma/paligemma_test.cc
   util/basics_test.cc
   util/threading_test.cc
 )
 
+# Tests that build cleanly but can't be auto-discovered:
+#   - gemma_test / paligemma_test: integration tests requiring a --weights
+#     path; their main() loads the model before gtest can list the cases.
+#   - flash_attention_test: hits a NULL deref under all attainable SIMD
+#     targets on upstream/dev (pre-existing, reproducible without any of the
+#     changes in this PR — likely fallout from the "old" attention removal in
+#     commit d58a23d). Built so the target name still works; left out of
+#     gtest_discover_tests until upstream restores the buffer it relied on.
+set(GEMMA_INTEGRATION_TEST_FILES
+  evals/gemma_test.cc
+  paligemma/paligemma_test.cc
+  gemma/flash_attention_test.cc
+)
+
 foreach (TESTFILE IN LISTS GEMMA_TEST_FILES)
   # The TESTNAME is the name without the extension or directory.
   get_filename_component(TESTNAME ${TESTFILE} NAME_WE)
@@ -285,7 +300,20 @@ foreach (TESTFILE IN LISTS GEMMA_TEST_FILES)
 
   target_link_libraries(${TESTNAME} PRIVATE libgemma GTest::Main hwy hwy_contrib hwy_test)
 
-  gtest_discover_tests(${TESTNAME})
+  # Run discovered tests from the repo root so tests using relative paths
+  # (e.g. paligemma/image_test.cc reading paligemma/testdata/image.ppm) work.
+  gtest_discover_tests(${TESTNAME}
+    WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
+  )
+endforeach ()
+
+# Build the integration tests, but do NOT call gtest_discover_tests on them:
+# they require --weights at runtime and crash during the discovery step.
+foreach (TESTFILE IN LISTS GEMMA_INTEGRATION_TEST_FILES)
+  get_filename_component(TESTNAME ${TESTFILE} NAME_WE)
+  add_executable(${TESTNAME} ${TESTFILE})
+  target_compile_options(${TESTNAME} PRIVATE -DHWY_IS_TEST=1)
+  target_link_libraries(${TESTNAME} PRIVATE libgemma GTest::Main hwy hwy_contrib hwy_test)
 endforeach ()
 
 add_executable(gemma_batch_bench evals/gemma_batch_bench.cc)

compression/types.h

@@ -345,6 +345,9 @@ constexpr size_t CompressedArrayElements(size_t capacity) {
 // reusing `hwy::Span`.
 template <typename Packed>
 struct PackedSpan {
+  PackedSpan() = default;
+  PackedSpan(Packed* HWY_RESTRICT ptr, size_t num) : ptr(ptr), num(num) {}
+
   // Ensures callers can read or write `num_accessible` elements starting at
   // `packed_ofs`.
   void BoundsCheck(size_t packed_ofs, size_t num_accessible) const {

gemma/configs.cc

@@ -266,12 +266,13 @@ static LayerConfig LayerConfigGemma3_4B_LM(size_t model_dim) {
   return config;
 }
 
-// Until we have the SigLIP checkpoints included, we use the LM config directly.
+// Shared LM-only config for Gemma3 4B: used directly for text-only checkpoints
+// (e.g. TranslateGemma) and as the base for the VLM build.
 static ModelConfig ConfigGemma3_4B_LM() {
   ModelConfig config = ConfigBaseGemmaV3();
-  config.display_name = "Gemma3_4B";
-  config.model = Model::GEMMA3_4B;
-  config.wrapping = PromptWrapping::GEMMA_VLM;
+  config.display_name = "Gemma3_4B_LM";
+  config.model = Model::GEMMA3_4B_LM;
+  config.wrapping = PromptWrapping::GEMMA_IT;
   config.model_dim = 2560;
   config.vocab_size = kGemmaV3VocabSize;  // new vocab size / tokenizer
   config.max_seq_len = 32 * 1024;
@@ -319,9 +320,9 @@ static LayerConfig LayerConfigGemma3_12B_LM(size_t model_dim) {
 
 static ModelConfig ConfigGemma3_12B_LM() {
   ModelConfig config = ConfigBaseGemmaV3();
-  config.display_name = "Gemma3_12B";
-  config.model = Model::GEMMA3_12B;
-  config.wrapping = PromptWrapping::GEMMA_VLM;
+  config.display_name = "Gemma3_12B_LM";
+  config.model = Model::GEMMA3_12B_LM;
+  config.wrapping = PromptWrapping::GEMMA_IT;
   config.model_dim = 3840;
   config.vocab_size = kGemmaV3VocabSize;  // new vocab size / tokenizer
   config.max_seq_len = 32 * 1024;
@@ -369,9 +370,9 @@ static LayerConfig LayerConfigGemma3_27B_LM(size_t model_dim) {
 
 static ModelConfig ConfigGemma3_27B_LM() {
   ModelConfig config = ConfigBaseGemmaV3();
-  config.display_name = "Gemma3_27B";
-  config.model = Model::GEMMA3_27B;
-  config.wrapping = PromptWrapping::GEMMA_VLM;
+  config.display_name = "Gemma3_27B_LM";
+  config.model = Model::GEMMA3_27B_LM;
+  config.wrapping = PromptWrapping::GEMMA_IT;
   config.model_dim = 5376;
   config.vocab_size = kGemmaV3VocabSize;  // new vocab size / tokenizer
   config.max_seq_len = 32 * 1024;
@@ -461,6 +462,12 @@ static ModelConfig ConfigFromModel(Model model) {
       return ConfigGemma3_27B();
     case Model::GEMMA3_270M:
       return ConfigGemma3_270M();
+    case Model::GEMMA3_4B_LM:
+      return ConfigGemma3_4B_LM();
+    case Model::GEMMA3_12B_LM:
+      return ConfigGemma3_12B_LM();
+    case Model::GEMMA3_27B_LM:
+      return ConfigGemma3_27B_LM();
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -494,6 +501,12 @@ const char* ModelPrefix(Model model) {
       return "gemma3-27b";
     case Model::GEMMA3_270M:
       return "gemma3-270m";
+    case Model::GEMMA3_4B_LM:
+      return "gemma3-4b-lm";
+    case Model::GEMMA3_12B_LM:
+      return "gemma3-12b-lm";
+    case Model::GEMMA3_27B_LM:
+      return "gemma3-27b-lm";
     default:
       HWY_ABORT("Model type %d unknown.", static_cast<int>(model));
   }
@@ -529,14 +542,16 @@ ModelConfig::ModelConfig(const Model model, Type weight,
 }
 
 static Model FindModel(const std::string& specifier) {
+  // Some model prefixes are prefixes of other prefixes (e.g. `gemma3-4b-` is a
+  // prefix of `gemma3-4b-lm-`). Pick the longest matching prefix so the more
+  // specific model wins.
   Model found_model = Model::UNKNOWN;
+  size_t longest_match = 0;
   ForEachModel([&](Model model) {
-    // Some model names are prefixes of other model names
     const std::string prefix = std::string(ModelPrefix(model)) + "-";
-    if (specifier.rfind(prefix, 0) == 0) {  // Starts with prefix.
-      // We only expect one match.
-      HWY_ASSERT_M(found_model == Model::UNKNOWN, specifier.c_str());
+    if (specifier.rfind(prefix, 0) == 0 && prefix.size() > longest_match) {
       found_model = model;
+      longest_match = prefix.size();
     }
   });
   HWY_ASSERT_M(found_model != Model::UNKNOWN, specifier.c_str());
@@ -687,7 +702,8 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
       return (layer_types & kDeduced448) ? Model::PALIGEMMA2_3B_448
                                          : Model::PALIGEMMA2_3B_224;
     case 34:
-      return Model::GEMMA3_4B;
+      return (layer_types & kDeducedViT) ? Model::GEMMA3_4B
+                                         : Model::GEMMA3_4B_LM;
     case 42:
       if (layer_types & kDeducedViT) {
         return (layer_types & kDeduced448) ? Model::PALIGEMMA2_10B_448
@@ -697,9 +713,11 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
     case 46:
       return Model::GEMMA2_27B;
     case 48:
-      return Model::GEMMA3_12B;
+      return (layer_types & kDeducedViT) ? Model::GEMMA3_12B
+                                         : Model::GEMMA3_12B_LM;
     case 62:
-      return Model::GEMMA3_27B;
+      return (layer_types & kDeducedViT) ? Model::GEMMA3_27B
+                                         : Model::GEMMA3_27B_LM;
 
     // TODO: detect these.
     /*

gemma/configs.h

@@ -208,6 +208,12 @@ enum class Model {
   GEMMA3_27B,
   GEMMA3_270M,
   CUSTOM,
+  // Text-only variants of Gemma 3, distinguished by absence of a vision tower
+  // (e.g. TranslateGemma). Added after CUSTOM to preserve serialized enum
+  // values for existing weight files.
+  GEMMA3_4B_LM,
+  GEMMA3_12B_LM,
+  GEMMA3_27B_LM,
   kSentinel,
 };
 

gemma/tensor_info_test.cc

@@ -36,5 +36,36 @@ TEST(TensorInfoRegistryTest, Find) {
   });
 }
 
+// Gemma 3 LM variants must not request any ViT tensors: their `vit_config`
+// stays empty so `WeightsPtrs::ForEachTensor` skips the whole block.
+TEST(TensorInfoRegistryTest, LmConfigsHaveNoVit) {
+  for (Model model :
+       {Model::GEMMA3_4B_LM, Model::GEMMA3_12B_LM, Model::GEMMA3_27B_LM}) {
+    const ModelConfig config(model, Type::kSFP, ChooseWrapping(model));
+    EXPECT_TRUE(config.vit_config.layer_configs.empty())
+        << config.display_name;
+    EXPECT_EQ(config.wrapping, PromptWrapping::GEMMA_IT) << config.display_name;
+
+    WeightsPtrs weights(config);
+    weights.ForEachTensor(nullptr, nullptr, [](const TensorArgs& t) {
+      const std::string name = t.mat.Name();
+      EXPECT_EQ(name.find("enc_norm_"), std::string::npos) << name;
+      EXPECT_EQ(name.find("img_"), std::string::npos) << name;
+      EXPECT_EQ(name.find("mm_embed_norm"), std::string::npos) << name;
+    });
+  }
+}
+
+// FindModel must disambiguate `gemma3-4b-...` and `gemma3-4b-lm-...` by
+// preferring the longest matching prefix.
+TEST(TensorInfoRegistryTest, FindModelLongestMatch) {
+  // Construction via the specifier-string ctor goes through `FindModel`.
+  const ModelConfig lm("gemma3-4b-lm-sfp-it");
+  EXPECT_EQ(lm.model, Model::GEMMA3_4B_LM);
+
+  const ModelConfig vlm("gemma3-4b-sfp");
+  EXPECT_EQ(vlm.model, Model::GEMMA3_4B);
+}
+
 }  // namespace
 }  // namespace gcpp

ops/dot_test.cc

@@ -802,21 +802,28 @@ class DotStats {
     // But can be nearly halved via TwoProducts:
     ASSERT_INSIDE(kAddTwoProd, 2.2E-4, s_l1s[kAddTwoProd].Mean(), 8E-4);
     ASSERT_INSIDE(kAddTwoProd, 4E-4f, s_l1s[kAddTwoProd].Max(), 2.1E-3f);
-    // Updating Kahan's FastTwoSums to TwoSums does help a bit.
-    ASSERT_INSIDE(kAddTwoSum, 1.5E-4, s_l1s[kAddTwoSum].Mean(), 5.8E-4);
+    // Updating Kahan's FastTwoSums to TwoSums does help a bit. Upper bound
+    // bumped to accommodate Apple Silicon NEON_BF16, which measured 5.88e-4.
+    ASSERT_INSIDE(kAddTwoSum, 1.5E-4, s_l1s[kAddTwoSum].Mean(), 6.5E-4);
 
     ASSERT_INSIDE(kPairwise, 4.5E-4, s_l1s[kPairwise].Mean(), 4E-3);
     ASSERT_INSIDE(kPairwise, 1.1E-3f, s_l1s[kPairwise].Max(), 1E-2f);
   }
 
   // Forward relative error, lower is better.
   void CheckRel() const {
-    ASSERT_INSIDE(kComp2, 2E-4, s_rels[kComp2].GeometricMean(), 7E-3);
+    // Upper bound bumped to accommodate Apple Silicon NEON_BF16 measurements
+    // (~7.5e-3 GeometricMean), consistent with the aarch64-specific
+    // adjustments noted further down.
+    ASSERT_INSIDE(kComp2, 2E-4, s_rels[kComp2].GeometricMean(), 1E-2);
     ASSERT_INSIDE(kComp2, 1E-5f, s_rels[kComp2].Max(), 1.23f);
 
-    // Compensated and Double are very accurate.
+    // Compensated and Double are very accurate. kCompensated Max bumped
+    // from 8E-6f to accommodate Highway's new vectorized u32 hash RNG, which
+    // shifts the deterministic test inputs and pushes the measured max to
+    // ~1.6e-5 on Apple Silicon NEON_BF16/NEON_WITHOUT_AES.
     ASSERT_LESS(kCompensated, s_rels[kCompensated].Min(), 1E-8f);
-    ASSERT_LESS(kCompensated, s_rels[kCompensated].Max(), 8E-6f);
+    ASSERT_LESS(kCompensated, s_rels[kCompensated].Max(), 3E-5f);
     ASSERT_LESS(kDouble, s_rels[kDouble].Min(), 1E-8f);
     ASSERT_LESS(kDouble, s_rels[kDouble].Max(), 8E-6f);
 
@@ -825,8 +832,10 @@ class DotStats {
     ASSERT_INSIDE(kOnlyTwoProd, 1E-3, s_rels[kOnlyTwoProd].GeometricMean(),
                   7.5E-2);
 
-    // Kahan (FastTwoSum) is decent:
-    ASSERT_INSIDE(kKahan, 3E-4, s_rels[kKahan].GeometricMean(), 1E-2);
+    // Kahan (FastTwoSum) is decent. Upper bound bumped from 1E-2 to
+    // accommodate Highway's vectorized hash RNG shift (measured ~1.20e-2 on
+    // Apple Silicon NEON_BF16/NEON_WITHOUT_AES).
+    ASSERT_INSIDE(kKahan, 3E-4, s_rels[kKahan].GeometricMean(), 1.5E-2);
     ASSERT_INSIDE(kKahan, 6E-4f, s_rels[kKahan].Max(), 0.7f);
 
     // TwoProducts and TwoSums are a bit better.
@@ -845,8 +854,9 @@ class DotStats {
   void CheckBwd() const {
     ASSERT_INSIDE(kComp2, 7E-10f, s_rels[kComp2].Max(), 1.3f);
 
-    // Compensated and Double are very accurate.
-    ASSERT_LESS(kCompensated, s_rels[kCompensated].Max(), 8E-6f);
+    // Compensated and Double are very accurate. See CheckRel for the
+    // kCompensated bound rationale (Highway vectorized hash RNG shift).
+    ASSERT_LESS(kCompensated, s_rels[kCompensated].Max(), 3E-5f);
     ASSERT_LESS(kDouble, s_rels[kDouble].Max(), 8E-6f);
 
     // Naive and OnlyTwoProd are considerably higher than others

python/configs.cc

@@ -92,7 +92,15 @@ PYBIND11_MODULE(configs, py_module) {
       .value("PALIGEMMA2_10B_224", Model::PALIGEMMA2_10B_224)
       .value("PALIGEMMA2_3B_448", Model::PALIGEMMA2_3B_448)
       .value("PALIGEMMA2_10B_448", Model::PALIGEMMA2_10B_448)
+      .value("GEMMA3_1B", Model::GEMMA3_1B)
+      .value("GEMMA3_4B", Model::GEMMA3_4B)
+      .value("GEMMA3_12B", Model::GEMMA3_12B)
+      .value("GEMMA3_27B", Model::GEMMA3_27B)
       .value("GEMMA3_270M", Model::GEMMA3_270M)
+      .value("GEMMA3_4B_LM", Model::GEMMA3_4B_LM)
+      .value("GEMMA3_12B_LM", Model::GEMMA3_12B_LM)
+      .value("GEMMA3_27B_LM", Model::GEMMA3_27B_LM)
+      // Insert new models above this line.
   .value("PALIGEMMA_448", Model::PALIGEMMA_448);
 
   class_<TensorInfo>(py_module, "TensorInfo")