Skip to content

GitLab

Unverified Commit b1f71233 authored by Bruce MacDonald's avatar Bruce MacDonald Committed by GitHub
Browse files

clean up num_gpu calculation code (#673)

parent 1fbf3585
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 4 deletions
llm/llama.go
View file @ b1f71233
......@@ -218,7 +218,6 @@ func NumGPU(numLayer, fileSizeBytes int64, opts api.Options) int {
if opts.NumGPU != -1 {
return opts.NumGPU
}
n := 1 // default to enable metal on macOS
if runtime.GOOS == "linux" {
vramMib, err := CheckVRAM()
if err != nil {
......@@ -235,10 +234,11 @@ func NumGPU(numLayer, fileSizeBytes int64, opts api.Options) int {
// TODO: this is a rough heuristic, better would be to calculate this based on number of layers and context size
bytesPerLayer := fileSizeBytes / numLayer
// set n to the max number of layers we can fit in VRAM
return int(totalVramBytes / bytesPerLayer)
// max number of layers we can fit in VRAM
layers := int(totalVramBytes / bytesPerLayer)
log.Printf("%d MiB VRAM available, loading up to %d GPU layers", vramMib, layers)
log.Printf("%d MiB VRAM available, loading up to %d GPU layers", vramMib, n)
return layers
}
// default to enable metal on macOS
return 1
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment