193 lines
5.9 KiB
Go
193 lines
5.9 KiB
Go
package llm
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
"os"
|
|
"strconv"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
"github.com/ollama/ollama/format"
|
|
"github.com/ollama/ollama/gpu"
|
|
)
|
|
|
|
// This algorithm looks for a complete fit to determine if we need to unload other models
|
|
func PredictServerFit(allGpus gpu.GpuInfoList, ggml *GGML, adapters, projectors []string, opts api.Options) (bool, uint64) {
|
|
var estimatedVRAM uint64
|
|
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
|
|
}
|
|
|
|
// Split up the GPUs by type and try them
|
|
for _, gpus := range allGpus.ByLibrary() {
|
|
var layerCount int
|
|
layerCount, estimatedVRAM = EstimateGPULayers(gpus, ggml, projectors, opts)
|
|
if opts.NumGPU < 0 {
|
|
if layerCount > 0 && layerCount >= int(ggml.KV().BlockCount()+1) {
|
|
return true, estimatedVRAM
|
|
}
|
|
} else {
|
|
if layerCount > 0 && layerCount >= opts.NumGPU {
|
|
return true, estimatedVRAM
|
|
}
|
|
}
|
|
}
|
|
return false, estimatedVRAM
|
|
}
|
|
|
|
// Given a model and one or more GPU targets, predict how many layers and bytes we can load
|
|
// The GPUs provided must all be the same Library
|
|
func EstimateGPULayers(gpus []gpu.GpuInfo, ggml *GGML, projectors []string, opts api.Options) (int, uint64) {
|
|
if gpus[0].Library == "cpu" {
|
|
return 0, 0
|
|
}
|
|
var memoryAvailable uint64
|
|
for _, info := range gpus {
|
|
memoryAvailable += info.FreeMemory
|
|
}
|
|
userLimit := os.Getenv("OLLAMA_MAX_VRAM")
|
|
if userLimit != "" {
|
|
avail, err := strconv.ParseUint(userLimit, 10, 64)
|
|
if err != nil {
|
|
slog.Error("invalid setting, ignoring", "OLLAMA_MAX_VRAM", userLimit, "error", err)
|
|
} else {
|
|
slog.Info("user override memory limit", "OLLAMA_MAX_VRAM", avail, "actual", memoryAvailable)
|
|
memoryAvailable = avail
|
|
}
|
|
}
|
|
|
|
slog.Debug("evaluating", "library", gpus[0].Library, "gpu_count", len(gpus), "available", format.HumanBytes2(memoryAvailable))
|
|
|
|
// TODO - this is probably wrong, first GPU vs secondaries will have different overheads
|
|
memoryMinimum := gpus[0].MinimumMemory
|
|
|
|
for _, projector := range projectors {
|
|
memoryMinimum += 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
|
|
var kv uint64 = 2 * 2 * uint64(opts.NumCtx) * ggml.KV().BlockCount() * ggml.KV().EmbeddingLength() / ggml.KV().HeadCount() * ggml.KV().HeadCountKV()
|
|
|
|
graphPartialOffload, graphFullOffload := ggml.GraphSize(uint64(opts.NumCtx), uint64(min(opts.NumCtx, opts.NumBatch)))
|
|
if graphPartialOffload == 0 {
|
|
graphPartialOffload = ggml.KV().GQA() * kv / 6
|
|
}
|
|
|
|
if graphFullOffload == 0 {
|
|
graphFullOffload = graphPartialOffload
|
|
}
|
|
|
|
graphFullOffload *= uint64(len(gpus))
|
|
graphPartialOffload *= uint64(len(gpus))
|
|
|
|
// on metal there's no partial offload overhead
|
|
if gpus[0].Library == "metal" {
|
|
graphPartialOffload = graphFullOffload
|
|
}
|
|
|
|
// memoryRequiredTotal represents the memory required for full GPU offloading (all layers)
|
|
memoryRequiredTotal := memoryMinimum + graphFullOffload
|
|
|
|
// memoryRequiredPartial represents the memory required for partial GPU offloading (n > 0, n < layers)
|
|
memoryRequiredPartial := memoryMinimum + graphPartialOffload
|
|
|
|
if memoryRequiredPartial > memoryAvailable {
|
|
slog.Debug("insufficient VRAM to load any model layers")
|
|
return 0, 0
|
|
}
|
|
|
|
layers := ggml.Tensors().Layers()
|
|
|
|
var memoryLayerOutput uint64
|
|
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 {
|
|
// memory is preallocated for output tensors
|
|
memoryRequiredTotal += memoryLayerOutput
|
|
memoryRequiredPartial += memoryLayerOutput
|
|
}
|
|
|
|
var layerCount int
|
|
for i := 0; i < int(ggml.KV().BlockCount()); i++ {
|
|
memoryLayer := layers[fmt.Sprintf("blk.%d", i)].size()
|
|
|
|
// KV is proportional to the number of layers
|
|
memoryLayer += kv / ggml.KV().BlockCount()
|
|
|
|
memoryRequiredTotal += memoryLayer
|
|
if memoryAvailable > memoryRequiredPartial+memoryLayer {
|
|
memoryRequiredPartial += memoryLayer
|
|
layerCount++
|
|
}
|
|
}
|
|
|
|
if gpus[0].Library != "metal" || !opts.UseMMap {
|
|
// memory was not preallocated for output tensors
|
|
memoryRequiredTotal += memoryLayerOutput
|
|
}
|
|
|
|
if memoryAvailable > memoryRequiredTotal {
|
|
layerCount = int(ggml.KV().BlockCount()) + 1
|
|
memoryRequiredPartial = memoryRequiredTotal
|
|
}
|
|
|
|
memoryWeights := memoryRequiredTotal - memoryMinimum - graphFullOffload - kv
|
|
|
|
slog.Info(
|
|
"offload to gpu",
|
|
slog.Group(
|
|
"layers",
|
|
// actual number of layers offloaded
|
|
"real", opts.NumGPU,
|
|
// estimated number of layers that can be offloaded
|
|
"estimate", layerCount,
|
|
),
|
|
slog.Group(
|
|
"memory",
|
|
// memory available for offloading
|
|
"available", format.HumanBytes2(memoryAvailable),
|
|
slog.Group(
|
|
"required",
|
|
// memory required for full offloading
|
|
"full", format.HumanBytes2(memoryRequiredTotal),
|
|
// memory required to offload layers.estimate layers
|
|
"partial", format.HumanBytes2(memoryRequiredPartial),
|
|
// memory of KV cache
|
|
"kv", format.HumanBytes2(kv),
|
|
),
|
|
slog.Group(
|
|
"weights",
|
|
// memory of the weights
|
|
"total", format.HumanBytes2(memoryWeights),
|
|
// memory of repeating layers
|
|
"repeating", format.HumanBytes2(memoryWeights-memoryLayerOutput),
|
|
// memory of non-repeating layers
|
|
"nonrepeating", format.HumanBytes2(memoryLayerOutput),
|
|
),
|
|
slog.Group(
|
|
"graph",
|
|
// memory of graph when fully offloaded
|
|
"full", format.HumanBytes2(graphFullOffload),
|
|
// memory of graph when not fully offloaded
|
|
"partial", format.HumanBytes2(graphPartialOffload),
|
|
),
|
|
),
|
|
)
|
|
return layerCount, uint64(memoryRequiredPartial)
|
|
}
|