Add note on weight update and improve error message

PiperOrigin-RevId: 621849989

Author: jan-wassenberg

README.md

@@ -8,6 +8,11 @@ For additional information about Gemma, see
 specific artifacts, are [available on
 kaggle](https://www.kaggle.com/models/google/gemma).
 
+NOTE: 2024-04-04: if using 2B models, please re-download weights from Kaggle and
+ensure you have the latest version (-mqa or version 3). We are changing the code
+to match the new weights. If you wish to use old weights, change `ConfigGemma2B`
+in `configs.h` back to `kVocabSize = 256128` and `kKVHeads = 8`.
+
 ## Who is this project for?
 
 Modern LLM inference engines are sophisticated systems, often with bespoke

compression/blob_store.cc

@@ -378,7 +378,13 @@ BlobError BlobReader::Enqueue(hwy::uint128_t key, void* data, size_t size) {
   uint64_t offset;
   size_t actual_size;
   if (!blob_store_->FindKey(key, offset, actual_size)) return __LINE__;
-  if (actual_size != size) return __LINE__;
+  if (actual_size != size) {
+    fprintf(stderr,
+            "Mismatch between expected %d and actual %d KiB size. Please see "
+            "README.md on how to update the weights.\n",
+            static_cast<int>(size >> 10), static_cast<int>(actual_size >> 10));
+    return __LINE__;
+  }
 
   EnqueueChunkRequests(offset, actual_size, reinterpret_cast<uint8_t*>(data),
                        requests_);

gemma.cc

@@ -109,6 +109,12 @@ hwy::AlignedUniquePtr<Weights<TConfig>> LoadWeights(const Path& checkpoint) {
   weights->layers =
       hwy::MakeUniqueAlignedArray<Layer<TConfig>>(TConfig::kLayers);
 
+  if (checkpoint.path.empty()) {
+    HWY_ABORT(
+        "Loading --compressed_weights failed; we require a --weights argument. "
+        "Please see issue #11 on how to create this file.\n");
+  }
+
   FILE* fptr;
   fptr = fopen(checkpoint.path.c_str(), "rb");
   if (fptr == nullptr) {