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

skip gpu if less than 2GB VRAM are available (#1059)

Browse source
This commit is contained in:
Bruce MacDonald 2023-11-09 13:16:16 -08:00 committed by GitHub
parent db8bf336fc
commit 1ae84bc2a2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
llm/llama.go
View file

@ -196,7 +196,10 @@ type llama struct {
Running Running
} }
var errNoGPU = errors.New("nvidia-smi command failed") var (
errNvidiaSMI = errors.New("nvidia-smi command failed")
errAvailableVRAM = errors.New("not enough VRAM available, falling back to CPU only")
)
// CheckVRAM returns the free VRAM in bytes on Linux machines with NVIDIA GPUs // CheckVRAM returns the free VRAM in bytes on Linux machines with NVIDIA GPUs
func CheckVRAM() (int64, error) { func CheckVRAM() (int64, error) {
@ -205,7 +208,7 @@ func CheckVRAM() (int64, error) {
cmd.Stdout = &stdout cmd.Stdout = &stdout
err := cmd.Run() err := cmd.Run()
if err != nil { if err != nil {
return 0, errNoGPU return 0, errNvidiaSMI
} }
var freeMiB int64 var freeMiB int64
@ -226,8 +229,8 @@ func CheckVRAM() (int64, error) {
freeBytes := freeMiB * 1024 * 1024 freeBytes := freeMiB * 1024 * 1024
if freeBytes < 2*format.GigaByte { if freeBytes < 2*format.GigaByte {
log.Printf("less than 2 GB VRAM available, falling back to CPU only") log.Printf("less than 2 GB VRAM available")
freeMiB = 0 return 0, errAvailableVRAM
} }
return freeBytes, nil return freeBytes, nil
@ -240,7 +243,7 @@ func NumGPU(numLayer, fileSizeBytes int64, opts api.Options) int {
if runtime.GOOS == "linux" { if runtime.GOOS == "linux" {
freeBytes, err := CheckVRAM() freeBytes, err := CheckVRAM()
if err != nil { if err != nil {
if err.Error() != "nvidia-smi command failed" { if !errors.Is(err, errNvidiaSMI) {
log.Print(err.Error()) log.Print(err.Error())
} }
// nvidia driver not installed or no nvidia GPU found // nvidia driver not installed or no nvidia GPU found