From 344846b4fc6fff5988f215e7cf839f880d04decf Mon Sep 17 00:00:00 2001 From: Lorenzo Bona Date: Fri, 30 Oct 2026 05:36:27 -0300 Subject: [PATCH] gemma3n, gemma4 : create token_embd before the tied output (#31394) The duplicated output looks up token_embd in its buffer context, so token_embd must exist first. With the old order a load where input and output share a buffer type (CPU-only, -ngl 0) held the table twice: gemma-4-E2B-it Q4_K_M, no mmap, CPU model buffer 3163.73 -> 2947.73 MiB. --- src/models/gemma3n.cpp | 4 ++-- src/models/gemma4.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/models/gemma3n.cpp b/src/models/gemma3n.cpp index f2e5ae6c8..7848f9f59 100644 --- a/src/models/gemma3n.cpp +++ b/src/models/gemma3n.cpp @@ -25,14 +25,14 @@ void llama_model_gemma3n::load_arch_tensors(llama_model_loader &) { const int64_t laurel_rank = hparams.laurel_rank; const int64_t n_embd_altup = hparams.n_embd_altup; + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); // if output is NULL, init from the input tok embed if (output == NULL) { output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); } - tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); - altup_proj = create_tensor(tn(LLM_TENSOR_ALTUP_PROJ, "weight"), {n_embd, n_embd, n_altup - 1}, 0); altup_unembd_proj = create_tensor(tn(LLM_TENSOR_ALTUP_UNEMBD_PROJ, "weight"), {n_embd, n_embd, n_altup - 1}, 0); diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index f4f80d396..42d6a7d47 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -47,14 +47,14 @@ void llama_model_gemma4::load_arch_tensors(llama_model_loader &) { throw std::runtime_error("Gemma 4 requires n_embd_head_k_swa == n_embd_head_v_swa"); } + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); // if output is NULL, init from the input tok embed if (output == NULL) { output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); } - tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); - if (n_embd_per_layer > 0) { per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, TENSOR_READ_LAZY); per_layer_model_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight", 0), {n_embd, n_embd_per_layer * n_layer}, 0);