From 00ba208a5c036eee72d4a631b4f57c126095cb03 Mon Sep 17 00:00:00 2001
From: Kawrakow <iwankawrakow@gmail.com>
Date: Sun, 19 Apr 2026 14:25:05 +0200
Subject: [PATCH] Fix Gemma4 partial offload (#1657)

* Fix Gemma4 partial offload

* Also here
---
 src/llama-load-tensors.cpp |  2 +-
 src/llama-model.cpp        |  6 +++---
 src/llama.cpp              | 10 +++++++---
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp
index 3e099a37..7d7b385b 100644
--- a/src/llama-load-tensors.cpp
+++ b/src/llama-load-tensors.cpp
@@ -235,7 +235,7 @@ create_tensors_helper::create_tensors_helper(llama_model_loader & _ml, llama_mod
         if (model.split_mode == LLAMA_SPLIT_MODE_ATTN || model.split_mode == LLAMA_SPLIT_MODE_GRAPH || ml.ncmoe >= n_layer || model.devices.size() < 2) {
             int nmax = std::min(ml.ncmoe, n_layer);
             for (int i = 0; i < nmax; ++i) {
-                std::string pattern = "blk\\." + std::to_string(i) + "\\.(ffn_(up|down|gate|gate_up)_exps\\.weight)";
+                std::string pattern = "blk\\." + std::to_string(i) + "\\.ffn_(up|down|gate|gate_up)_exps\\.(weight|scale)";
                 this->overrides.emplace_back(std::make_pair(std::regex(pattern), buft));
             }
         }
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 94f266c2..d5088eea 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -1890,14 +1890,14 @@ llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int i
     for (auto & entry : it->second) {
         auto base_name = ::format(entry.second.c_str(), il);
         auto this_name = base_name + ".weight";
-        if (tensor_name.find(this_name) == 0) {
+        if (tensor_name == this_name) {
             return entry.first;
         }
         this_name = base_name + ".bias";
-        if (tensor_name.find(this_name) == 0) {
+        if (tensor_name == this_name) {
             return entry.first;
         }
-        if (tensor_name.find(base_name) == 0) {
+        if (tensor_name == base_name) {
             return entry.first;
         }
     }
diff --git a/src/llama.cpp b/src/llama.cpp
index afab028e..47836f3d 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -2119,7 +2119,11 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
             continue;
         }
         result[il] += size;
-        if (auto pos = name.rfind(".bias"); pos < name.size() && name.size() - pos == 4) {
+        if (auto pos = name.rfind(".bias"); pos < name.size() && name.size() - pos == 5) {
+            // bias, we don't need to account for those
+            continue;
+        }
+        if (auto pos = name.rfind(".scale"); pos < name.size() && name.size() - pos == 6) {
             // bias, we don't need to account for those
             continue;
         }
@@ -2475,7 +2479,7 @@ static bool llm_load_tensors(
                     }
                     if (has_experts) {
                         LLAMA_LOG_INFO("Adding experts CPU overrides for layer %d\n", il);
-                        std::string pattern = "blk\\." + std::to_string(il) + "\\.(ffn_(up|down|gate|gate_up)_exps\\.weight)";
+                        std::string pattern = "blk\\." + std::to_string(il) + "\\.(ffn_(up|down|gate|gate_up)_exps\\.(weight|scale))";
                         auto & o = overrides.emplace_back();
                         o.pattern = strdup(pattern.c_str());
                         o.buft = buft;
@@ -2550,7 +2554,7 @@ static bool llm_load_tensors(
                             }
                             if (has_experts) {
                                 LLAMA_LOG_INFO("Adding experts CPU overrides for layer %d in device %d\n", il, id);
-                                std::string pattern = "blk\\." + std::to_string(il) + "\\.(ffn_(up|down|gate|gate_up)_exps\\.weight)";
+                                std::string pattern = "blk\\." + std::to_string(il) + "\\.(ffn_(up|down|gate|gate_up)_exps\\.(weight|scale))";
                                 auto & o = overrides.emplace_back();
                                 o.pattern = strdup(pattern.c_str());
                                 o.buft = buft;