gered/ollama
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

clean up num_gpu calculation code (#673)

Browse source
This commit is contained in:
Bruce MacDonald 2023-10-02 14:53:42 -04:00 committed by GitHub
parent 1fbf3585d6
commit b1f7123301
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
llm/llama.go
View file

@ -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