fix gemma, command-r layer weights

f81f3081 · Michael Yang · 7fea1ecd · f81f3081
Commit f81f3081 authored Apr 26, 2024 by Michael Yang
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 4 deletions

llm/memory.go llm/memory.go +8 -4

No files found.
--- 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 {