From f81f30811878ee3b59deaa2319eec489fafb39ef Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Fri, 26 Apr 2024 15:00:54 -0700 Subject: [PATCH] fix gemma, command-r layer weights --- llm/memory.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/llm/memory.go b/llm/memory.go index 7ac7d8e7..d1e79e26 100644 --- a/llm/memory.go +++ b/llm/memory.go @@ -102,10 +102,14 @@ func EstimateGPULayers(gpus []gpu.GpuInfo, ggml *GGML, projectors []string, opts layers := ggml.Tensors().Layers() var memoryLayerOutput uint64 - for k, v := range layers { - if k == "output" || k == "output_norm" { - memoryLayerOutput += v.size() - } + if layer, ok := layers["output_norm"]; ok { + memoryLayerOutput += layer.size() + } + + if layer, ok := layers["output"]; ok { + memoryLayerOutput += layer.size() + } else if layer, ok := layers["token_embd"]; ok { + memoryLayerOutput += layer.size() } if gpus[0].Library == "metal" && opts.UseMMap {