Repository: ollama/ollama
Pull request: fix: the problem of loading models with large difference in initial and final layer sizes

This PR is related to the bug reported for Deepseek R1 deepseek-r1:671b model loading failure on MI210 3X, 4X and possibly more GPU configurations.

**Reported Bug:**
https://github.com/ollama/ollama/issues/8776

**Observations:**
I have noticed that the Ollama go-lang script uses initial layers for finding the number of layers off-loadable to GPU, while the llama.cpp code is assigning the models fnial layers on to the GPU devices as shown in the below line

_Counting number of layers_
_https://github.com/ollama/ollama/blob/main/llm/memory.go#L213_

_Assigning layers to GPU_
_https://github.com/ollama/ollama/blob/main/llama/llama.cpp/src/llama.cpp#L330_

This can cause problems if there is a big difference in the initial layer sizes and final layer sizes, as in the case of Deepseek R1.
The go script estimates more number of layers than can be fit onto the GPU. 

This PR fixes the issue by updating the go script to use the final layers and allocate buffer size equals to the maximum layer size. 


**Logs**:
Error debug logs
[VRAM allocation error.txt](https://github.com/user-attachments/files/18881223/VRAM.allocation.error.txt)

Logs after this fix
[After Changes - Success log.txt](https://github.com/user-attachments/files/18881224/After.Changes.-.Success.log.txt)




Changed files:
filename: llm/memory.go
patch: @@ -1,9 +1,12 @@
 package llm
 
 import (
+	"cmp"
 	"fmt"
 	"log/slog"
+	"maps"
 	"os"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -120,12 +123,10 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 	}
 
 	layers := f.Tensors().GroupLayers()
-	// add one layer worth of memory as a buffer
-	if blk0, ok := layers["blk.0"]; ok {
-		layerSize = blk0.Size()
-	} else {
-		slog.Warn("model missing blk.0 layer size")
-	}
+	// add one layer (chosing the max layer) worth of memory as a buffer
+	layerSize = slices.MaxFunc(slices.Collect(maps.Values(layers)), func(a, b ggml.Layer) int {
+		return cmp.Compare(a.Size(), b.Size())
+	}).Size()
 
 	var kvct string
 	if envconfig.FlashAttention() &&
@@ -219,7 +220,7 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 	}
 
 	// For all the layers, find where they can fit on the GPU(s)
-	for i := range int(f.KV().BlockCount()) {
+	for i := int(f.KV().BlockCount()) - 1; i >= 0; i-- {
 		// Some models have inconsistent layer sizes
 		if blk, ok := layers[fmt.Sprintf("blk.%d", i)]; ok {
 			layerSize = blk.Size()
@@ -229,6 +230,7 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 
 		if opts.NumGPU >= 0 && layerCount >= opts.NumGPU {
 			// Stop allocating on GPU(s) once we hit the users target NumGPU
+			overflow += layerSize
 			continue
 		}
 
@@ -245,13 +247,13 @@ func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []strin
 				gpusWithSpace = append(gpusWithSpace[:i%j], gpusWithSpace[i%j+1:]...)
 			}
 		}
+
+		if len(gpusWithSpace) == 0 {
+			overflow += layerSize
+		}
 	}
 	if layerCount >= int(f.KV().BlockCount()) {
 		fullyLoaded = true
-	} else {
-		for i := layerCount; i < int(f.KV().BlockCount()); i++ {
-			overflow += layerSize
-		}
 	}
 
 	// Determine if we need to consider output then find where it fits