llm.go 4.52 KB
Newer Older
1
2
3
package llm

import (
4
	"context"
5
	"fmt"
6
	"log/slog"
7
	"os"
8
	"runtime"
9
	"slices"
10

11
12
	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/gpu"
13
14
15
)

type LLM interface {
Bruce MacDonald's avatar
Bruce MacDonald committed
16
	Predict(context.Context, PredictOpts, func(PredictResult)) error
17
18
19
	Embedding(context.Context, string) ([]float64, error)
	Encode(context.Context, string) ([]int, error)
	Decode(context.Context, []int) (string, error)
20
21
22
	Close()
}

23
24
25
26
var cpuOnlyFamilies = []string{
	"mamba",
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
27
func New(model string, adapters, projectors []string, opts api.Options) (LLM, error) {
28
29
30
31
32
33
34
35
	if _, err := os.Stat(model); err != nil {
		return nil, err
	}

	f, err := os.Open(model)
	if err != nil {
		return nil, err
	}
Michael Yang's avatar
Michael Yang committed
36
	defer f.Close()
37

Bruce MacDonald's avatar
Bruce MacDonald committed
38
	ggml, err := DecodeGGML(f)
39
40
41
42
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
43
	if opts.NumCtx > int(ggml.NumCtx()) {
44
		slog.Warn(fmt.Sprintf("requested context length is greater than model's max context length (%d > %d), using %d instead", opts.NumCtx, ggml.NumCtx(), ggml.NumCtx()))
Michael Yang's avatar
Michael Yang committed
45
46
47
		opts.NumCtx = int(ggml.NumCtx())
	}

48
49
50
51
	if opts.NumCtx < 4 {
		opts.NumCtx = 4
	}

52
53
	vram, _ := gpu.CheckVRAM()
	size := ggml.Size
54
55

	// fp16 k,v matrices require = n_ctx * n_layer * n_embd / n_head * n_head_kv * 2 bytes each * 2 key and value
56
	kv := 2 * 2 * int64(opts.NumCtx) * int64(ggml.NumLayers()) * int64(ggml.NumEmbed()) * int64(ggml.NumHeadKv()) / int64(max(ggml.NumHead(), 1))
57

58
	// this amount is the overhead + tensors in memory
Michael Yang's avatar
typo  
Michael Yang committed
59
	// TODO: get this from the llama.cpp's graph calculations instead of
60
	// estimating it's 1/6 * kv_cache_size * num_gqa
61
	graph := int64(ggml.NumGQA()) * kv / 6
62

63
64
65
66
67
	// certain model architectures don't support gpu inference yet
	if slices.Contains(cpuOnlyFamilies, ggml.ModelFamily()) {
		opts.NumGPU = 0
	}

68
	info := gpu.GetGPUInfo()
69
70
71
72
73
	switch runtime.GOOS {
	case "darwin":
		if opts.NumGPU == 0 {
			break
		}
74

75
		if size+kv+graph > vram {
76
			slog.Info("not enough vram available, setting num_gpu=0")
77
78
79
80
			opts.NumGPU = 0
			break
		}

81
82
		// TODO: implement layer splitting on macOS
		opts.NumGPU = 999
83
	default:
84
		if info.Library == "cpu" {
85
			slog.Info("GPU not available, falling back to CPU")
86
87
88
89
90
91
			opts.NumGPU = 0
			break
		}

		// don't use GPU at all if no layers are loaded
		if opts.NumGPU == 0 {
92
93
			info.Library = "cpu"
			info.Variant = gpu.GetCPUVariant()
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
			break
		}

		// user-defined GPU count
		if opts.NumGPU != -1 {
			break
		}

		// the "main" GPU needs the most memory and determines the limit
		// of how many layers can be loaded. It needs to fit:
		// 1. the full compute graph allocation for all devices (graph)
		// 2. the proportional kv cache for all devices (kv * % layers)
		// 3. the proportional model (size * % layers / # devices)
		// This estimates the number of layers
		maxlayers := int64(ggml.NumLayers()) + 1
		devices := int64(info.DeviceCount)
		avg := vram / devices
		layers := maxlayers * (avg - graph) / (kv + size/devices)
		if layers > maxlayers {
			layers = maxlayers
		}
115

116
117
118
		// 1 + 2 must fit on the main gpu
		min := graph + kv*layers/maxlayers
		if layers <= 0 || min > avg {
119
			slog.Info("not enough vram available, falling back to CPU only")
120
121
			info.Library = "cpu"
			info.Variant = gpu.GetCPUVariant()
122
123
			opts.NumGPU = 0
			break
124
		}
125
126

		opts.NumGPU = int(layers)
127
128
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
129
130
	opts.RopeFrequencyBase = 0.0
	opts.RopeFrequencyScale = 0.0
Daniel Hiltgen's avatar
Daniel Hiltgen committed
131
	return newLlmServer(info, model, adapters, projectors, opts)
132
133
134
}

// Give any native cgo implementations an opportunity to initialize
Daniel Hiltgen's avatar
Daniel Hiltgen committed
135
136
func Init() error {
	return nativeInit()
137
}
138

Daniel Hiltgen's avatar
Daniel Hiltgen committed
139
func newLlmServer(gpuInfo gpu.GpuInfo, model string, adapters, projectors []string, opts api.Options) (LLM, error) {
140
	dynLibs := getDynLibs(gpuInfo)
141
142
143
144

	// Check to see if the user has requested a specific library instead of auto-detecting
	demandLib := os.Getenv("OLLAMA_LLM_LIBRARY")
	if demandLib != "" {
145
		libPath := availableDynLibs[demandLib]
146
		if libPath == "" {
147
			slog.Info(fmt.Sprintf("Invalid OLLAMA_LLM_LIBRARY %s - not found", demandLib))
148
		} else {
149
			slog.Info(fmt.Sprintf("Loading OLLAMA_LLM_LIBRARY=%s", demandLib))
150
			dynLibs = []string{libPath}
151
152
153
		}
	}

154
155
156
157
	// We stage into a temp directory, and if we've been idle for a while, it may have been reaped
	_, err := os.Stat(dynLibs[0])
	if err != nil {
		slog.Info(fmt.Sprintf("%s has disappeared, reloading libraries", dynLibs[0]))
Daniel Hiltgen's avatar
Daniel Hiltgen committed
158
		err = nativeInit()
159
160
161
162
163
		if err != nil {
			return nil, err
		}
	}

164
165
166
	err2 := fmt.Errorf("unable to locate suitable llm library")
	for _, dynLib := range dynLibs {
		srv, err := newDynExtServer(dynLib, model, adapters, projectors, opts)
167
168
169
		if err == nil {
			return srv, nil
		}
170
		slog.Warn(fmt.Sprintf("Failed to load dynamic library %s  %s", dynLib, err))
171
		err2 = err
172
173
	}

174
	return nil, err2
175
}