Repository: ollama/ollama
File: llm/llm.go (Go)

Code before:
		slog.Warn("requested context length is greater than model max context length", "requested", opts.NumCtx, "model", ggml.KV().ContextLength())
		opts.NumCtx = int(ggml.KV().ContextLength())
	}

	if opts.NumCtx < 4 {
		opts.NumCtx = 4
	}

	vram, _ := gpu.CheckVRAM()
	info := gpu.GetGPUInfo()

	mem := int64(396 * 1024 * 1024)
	for _, projector := range projectors {
		mem += projectorMemoryRequirements(projector)
	}

	// fp16 k,v = (1 (k) + 1 (v)) * sizeof(float16) * n_ctx * n_layer * n_embd / n_head * n_head_kv
	kv := 2 * 2 * int64(opts.NumCtx) * int64(ggml.KV().BlockCount()) * int64(ggml.KV().EmbeddingLength()) / int64(ggml.KV().HeadCount()) * int64(ggml.KV().HeadCountKV())

	// this amount is the overhead + tensors in memory
	// TODO: get this from the llama.cpp's graph calculations instead of
	// estimating it's 1/6 * kv_cache_size * num_gqa
	graph := int64(ggml.KV().GQA()) * kv / 6
	mem += graph

	if mem > vram {
		opts.NumGPU = 0
	}

	if opts.NumGPU < 0 && info.Library != "cpu" {
		for opts.NumGPU = 0; opts.NumGPU < int(ggml.KV().BlockCount()) && vram > mem; opts.NumGPU++ {
			layer := ggml.LayerSize(fmt.Sprintf("blk.%d.", opts.NumGPU))
			layerKV := kv / int64(ggml.KV().BlockCount())
			if vram <= mem+layer+layerKV {
				break
			}

			mem += layer + layerKV
		}

		// only offload output tensors if the repeating tensors have been offloaded
		if layer := ggml.LayerSize("output."); opts.NumGPU >= int(ggml.KV().BlockCount()) && vram > mem+layer {
			opts.NumGPU++
			mem += layer
		}

		opts.VRAMUsed = mem
		opts.VRAMTotal = vram
	}

	if opts.NumGPU == 0 || slices.Contains(cpuOnlyFamilies, ggml.KV().Architecture()) {

Review comment (suggestion):
This confused me at first.  This might warrant a comment explaining that this is the minimum threshold to even consider loading into the GPU before we calculate the layers.  I'd also suggest logging in this scenarios to help users understand why we completely bypassed their GPU for the given model.

Code after:
	if opts.NumCtx > int(ggml.KV().ContextLength()) {
		slog.Warn("requested context length is greater than model max context length", "requested", opts.NumCtx, "model", ggml.KV().ContextLength())
		opts.NumCtx = int(ggml.KV().ContextLength())
	}

	if opts.NumCtx < 4 {
		opts.NumCtx = 4
	}

	availableMemory, _ := gpu.CheckVRAM()
	info := gpu.GetGPUInfo()

	usedMemory := info.MinimumMemory
	for _, projector := range projectors {
		usedMemory += projectorMemoryRequirements(projector)

		// multimodal models require at least 2048 context
		opts.NumCtx = max(opts.NumCtx, 2048)
	}

	// fp16 k,v = (1 (k) + 1 (v)) * sizeof(float16) * n_ctx * n_layer * n_embd / n_head * n_head_kv
	kv := 2 * 2 * int64(opts.NumCtx) * int64(ggml.KV().BlockCount()) * int64(ggml.KV().EmbeddingLength()) / int64(ggml.KV().HeadCount()) * int64(ggml.KV().HeadCountKV())

	// this amount is the overhead + tensors in memory
	// TODO: get this from the llama.cpp's graph calculations instead of
	// estimating it's 1/6 * kv_cache_size * num_gqa
	graph := int64(ggml.KV().GQA()) * kv / 6
	usedMemory += graph

	if usedMemory > availableMemory || slices.Contains(cpuOnlyFamilies, ggml.KV().Architecture()) {
		info.Library = "cpu"
	}

	requiredMemory := usedMemory

	var layers int
	for i := 0; i < int(ggml.KV().BlockCount()); i++ {
		layerMemory := ggml.LayerSize(fmt.Sprintf("blk.%d.", i)) + kv/int64(ggml.KV().BlockCount())
		requiredMemory += layerMemory

		if availableMemory > usedMemory+layerMemory && (opts.NumGPU < 0 || layers < opts.NumGPU) {
			usedMemory += layerMemory
			layers++
		}
	}

	memOutputLayer := ggml.LayerSize("output.")
	requiredMemory += memOutputLayer

	// only offload output layer if all repeating layers are offloaded
	if layers >= int(ggml.KV().BlockCount()) && availableMemory > usedMemory+memOutputLayer {