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

Merge pull request #768 from jmorganca/mxyng/bytes

Browse source 
fix memory check
This commit is contained in:
Michael Yang 2023-10-16 12:42:41 -07:00 committed by GitHub
commit 9ef2fce33a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 50 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
api/client.go
View file

@ -14,6 +14,7 @@ import (
"runtime"
"strings"
"github.com/jmorganca/ollama/format"
"github.com/jmorganca/ollama/version"
)
@ -127,7 +128,7 @@ func (c *Client) do(ctx context.Context, method, path string, reqData, respData
return nil
}
const maxBufferSize = 512 * 1000 // 512KB
const maxBufferSize = 512 * format.KiloByte
func (c *Client) stream(ctx context.Context, method, path string, data any, fn func([]byte) error) error {
var buf *bytes.Buffer

19
format/bytes.go
View file

@ -2,14 +2,21 @@ package format
import "fmt"
const (
Byte = 1
KiloByte = Byte * 1000
MegaByte = KiloByte * 1000
GigaByte = MegaByte * 1000
)
func HumanBytes(b int64) string {
switch {
case b > 1000*1000*1000:
return fmt.Sprintf("%d GB", b/1000/1000/1000)
case b > 1000*1000:
return fmt.Sprintf("%d MB", b/1000/1000)
case b > 1000:
return fmt.Sprintf("%d KB", b/1000)
case b > GigaByte:
return fmt.Sprintf("%d GB", b/GigaByte)
case b > MegaByte:
return fmt.Sprintf("%d MB", b/MegaByte)
case b > KiloByte:
return fmt.Sprintf("%d KB", b/KiloByte)
default:
return fmt.Sprintf("%d B", b)
}

24
llm/llama.go
View file

@ -24,6 +24,7 @@ import (
"time"
"github.com/jmorganca/ollama/api"
"github.com/jmorganca/ollama/format"
)
//go:embed llama.cpp/*/build/*/bin/*
@ -197,7 +198,7 @@ type llama struct {
var errNoGPU = errors.New("nvidia-smi command failed")
// CheckVRAM returns the available VRAM in MiB on Linux machines with NVIDIA GPUs
// CheckVRAM returns the free VRAM in bytes on Linux machines with NVIDIA GPUs
func CheckVRAM() (int64, error) {
cmd := exec.Command("nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits")
var stdout bytes.Buffer
@ -207,7 +208,7 @@ func CheckVRAM() (int64, error) {
return 0, errNoGPU
}
var free int64
var freeMiB int64
scanner := bufio.NewScanner(&stdout)
for scanner.Scan() {
line := scanner.Text()
@ -216,15 +217,16 @@ func CheckVRAM() (int64, error) {
return 0, fmt.Errorf("failed to parse available VRAM: %v", err)
}
free += vram
freeMiB += vram
}
if free*1024*1024 < 2*1000*1000*1000 {
freeBytes := freeMiB * 1024 * 1024
if freeBytes < 2*format.GigaByte {
log.Printf("less than 2 GB VRAM available, falling back to CPU only")
free = 0
freeMiB = 0
}
return free, nil
return freeBytes, nil
}
func NumGPU(numLayer, fileSizeBytes int64, opts api.Options) int {
@ -232,7 +234,7 @@ func NumGPU(numLayer, fileSizeBytes int64, opts api.Options) int {
return opts.NumGPU
}
if runtime.GOOS == "linux" {
vramMib, err := CheckVRAM()
freeBytes, err := CheckVRAM()
if err != nil {
if err.Error() != "nvidia-smi command failed" {
log.Print(err.Error())
@ -241,15 +243,13 @@ func NumGPU(numLayer, fileSizeBytes int64, opts api.Options) int {
return 0
}
freeVramBytes := int64(vramMib) * 1024 * 1024 // 1 MiB = 1024^2 bytes
// Calculate bytes per layer
// TODO: this is a rough heuristic, better would be to calculate this based on number of layers and context size
bytesPerLayer := fileSizeBytes / numLayer
// max number of layers we can fit in VRAM, subtract 8% to prevent consuming all available VRAM and running out of memory
layers := int(freeVramBytes/bytesPerLayer) * 92 / 100
log.Printf("%d MiB VRAM available, loading up to %d GPU layers", vramMib, layers)
layers := int(freeBytes/bytesPerLayer) * 92 / 100
log.Printf("%d MiB VRAM available, loading up to %d GPU layers", freeBytes, layers)
return layers
}
@ -509,7 +509,7 @@ type PredictRequest struct {
Stop []string `json:"stop,omitempty"`
}
const maxBufferSize = 512 * 1000 // 512KB
const maxBufferSize = 512 * format.KiloByte
func (llm *llama) Predict(ctx context.Context, prevContext []int, prompt string, fn func(api.GenerateResponse)) error {
prevConvo, err := llm.Decode(ctx, prevContext)

44
llm/llm.go
View file

@ -10,6 +10,7 @@ import (
"github.com/pbnjay/memory"
"github.com/jmorganca/ollama/api"
"github.com/jmorganca/ollama/format"
)
type LLM interface {
@ -55,39 +56,30 @@ func New(workDir, model string, adapters []string, opts api.Options) (LLM, error
opts.NumGPU = 0
}
}
}
totalResidentMemory := memory.TotalMemory()
var requiredMemory int64
var f16Multiplier int64 = 2
switch ggml.ModelType() {
case "3B", "7B":
if ggml.FileType() == "F16" && totalResidentMemory < 16*1000*1000 {
return nil, fmt.Errorf("F16 model requires at least 16 GB of memory")
} else if totalResidentMemory < 8*1000*1000 {
return nil, fmt.Errorf("model requires at least 8 GB of memory")
}
requiredMemory = 8 * format.GigaByte
case "13B":
if ggml.FileType() == "F16" && totalResidentMemory < 32*1000*1000 {
return nil, fmt.Errorf("F16 model requires at least 32 GB of memory")
} else if totalResidentMemory < 16*1000*1000 {
return nil, fmt.Errorf("model requires at least 16 GB of memory")
}
requiredMemory = 16 * format.GigaByte
case "30B", "34B", "40B":
if ggml.FileType() == "F16" && totalResidentMemory < 64*1000*1000 {
return nil, fmt.Errorf("F16 model requires at least 64 GB of memory")
} else if totalResidentMemory < 32*1000*1000 {
return nil, fmt.Errorf("model requires at least 32 GB of memory")
}
requiredMemory = 32 * format.GigaByte
case "65B", "70B":
if ggml.FileType() == "F16" && totalResidentMemory < 128*1000*1000 {
return nil, fmt.Errorf("F16 model requires at least 128 GB of memory")
} else if totalResidentMemory < 64*1000*1000 {
return nil, fmt.Errorf("model requires at least 64 GB of memory")
}
requiredMemory = 64 * format.GigaByte
case "180B":
if ggml.FileType() == "F16" && totalResidentMemory < 512*1000*1000 {
return nil, fmt.Errorf("F16 model requires at least 512GB of memory")
} else if totalResidentMemory < 128*1000*1000 {
return nil, fmt.Errorf("model requires at least 128GB of memory")
requiredMemory = 128 * format.GigaByte
f16Multiplier = 4
}
systemMemory := int64(memory.TotalMemory())
if ggml.FileType() == "F16" && requiredMemory*f16Multiplier > systemMemory {
return nil, fmt.Errorf("F16 model requires at least %s of total memory", format.HumanBytes(requiredMemory))
} else if requiredMemory > systemMemory {
return nil, fmt.Errorf("model requires at least %s of total memory", format.HumanBytes(requiredMemory))
}
}