memory.go 14.8 KB
Newer Older
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1
2
3
package llm

import (
4
	"fmt"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
5
	"log/slog"
6
	"os"
Jesse Gross's avatar
Jesse Gross committed
7
	"sort"
8
	"strings"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
9
10

	"github.com/ollama/ollama/api"
11
	"github.com/ollama/ollama/discover"
12
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
13
	"github.com/ollama/ollama/format"
Michael Yang's avatar
Michael Yang committed
14
	"github.com/ollama/ollama/fs/ggml"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
15
16
)

Jesse Gross's avatar
Jesse Gross committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
// pickBestFullFitByLibrary will try to find the optimal placement of the model in the available GPUs where the model fully fits
// The list of GPUs returned will always be the same brand (library)
// If the model can not be fit fully within the available GPU(s) nil is returned
func pickBestFullFitByLibrary(f *ggml.GGML, modelPath string, projectors []string, adapters []string, opts api.Options, gpus discover.GpuInfoList, numParallel int) discover.GpuInfoList {
	for _, gl := range gpus.ByLibrary() {
		sgl := append(make(discover.GpuInfoList, 0, len(gl)), gl...)

		// TODO - potentially sort by performance capability, existing models loaded, etc.
		// TODO - Eliminate any GPUs that already have envconfig.MaxRunners loaded on them
		// Note: at present, this will favor most current available VRAM descending and ignoring faster GPU speed in mixed setups
		sort.Sort(sort.Reverse(discover.ByFreeMemory(sgl)))

		if !envconfig.SchedSpread() {
			// Try to pack into as few GPUs as possible, starting from 1 GPU
			for numGPUs := 1; numGPUs <= len(sgl); numGPUs++ {
				gpuSubset := sgl[:numGPUs]
				ok, estimatedVRAM := PredictServerFit(gpuSubset, f, adapters, projectors, opts, numParallel)

				if ok {
					slog.Info("new model will fit in available VRAM across minimum required GPUs, loading",
						"model", modelPath,
						"library", sgl[0].Library,
						"parallel", numParallel,
						"required", format.HumanBytes2(estimatedVRAM),
						"gpus", numGPUs)
					return gpuSubset
				}
			}
		} else {
			// TODO future refinements
			// - if multiple Libraries, see if any single GPU in any Library will fit
			// - try subsets of GPUs instead of just falling back to 1 or all in a family

			// Now try all the GPUS (OLLAMA_SCHED_SPREAD is set)
			if ok, estimatedVRAM := PredictServerFit(sgl, f, adapters, projectors, opts, numParallel); ok {
				slog.Info("new model will fit in available VRAM, loading",
					"model", modelPath,
					"library", sgl[0].Library,
					"parallel", numParallel,
					"required", format.HumanBytes2(estimatedVRAM),
					"gpus", len(sgl))
				return sgl
			}
		}
	}
	return nil
}

// If multiple Libraries are detected, pick the Library which loads the most layers for the model
func pickBestPartialFitByLibrary(f *ggml.GGML, projectors []string, adapters []string, opts api.Options, gpus discover.GpuInfoList, numParallel int) discover.GpuInfoList {
	byLibrary := gpus.ByLibrary()
	if len(byLibrary) <= 1 {
		return gpus
	}
	var bestEstimate uint64
	var bestFit int
	for i, gl := range byLibrary {
		_, estimatedVRAM := PredictServerFit(gl, f, adapters, projectors, opts, numParallel)
		if estimatedVRAM > bestEstimate {
			bestEstimate = estimatedVRAM
			bestFit = i
		}
	}
	return byLibrary[bestFit]
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
83
// This algorithm looks for a complete fit to determine if we need to unload other models
84
func PredictServerFit(allGpus discover.GpuInfoList, f *ggml.GGML, adapters, projectors []string, opts api.Options, numParallel int) (bool, uint64) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
85
	// Split up the GPUs by type and try them
86
	var estimatedVRAM uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
87
88
	for _, gpus := range allGpus.ByLibrary() {
		var layerCount int
Jesse Gross's avatar
Jesse Gross committed
89
		estimate := estimateGPULayers(gpus, f, projectors, opts, numParallel)
90
		layerCount, estimatedVRAM = estimate.Layers, estimate.VRAMSize
Daniel Hiltgen's avatar
Daniel Hiltgen committed
91
		if opts.NumGPU < 0 {
Michael Yang's avatar
Michael Yang committed
92
			if layerCount > 0 && layerCount >= int(f.KV().BlockCount()+1) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
93
94
95
96
97
98
99
100
101
102
103
				return true, estimatedVRAM
			}
		} else {
			if layerCount > 0 && layerCount >= opts.NumGPU {
				return true, estimatedVRAM
			}
		}
	}
	return false, estimatedVRAM
}

104
105
106
107
108
109
110
111
112
113
114
115
116
117
type MemoryEstimate struct {
	// How many layers we predict we can load
	Layers int

	// The size of the graph which occupies the main GPU
	Graph uint64

	// How much VRAM will be allocated given the number of layers we predict
	VRAMSize uint64

	// The total size of the model if loaded into VRAM.  If all layers are loaded, VRAMSize == TotalSize
	TotalSize uint64

	// For multi-GPU scenarios, this provides the tensor split parameter
Jesse Gross's avatar
Jesse Gross committed
118
	TensorSplit []int
119
120
121

	// For multi-GPU scenarios, this is the size in bytes per GPU
	GPUSizes []uint64
122
123
124
125
126
127
128
129
130
131
132
133

	// internal fields for logging purposes
	inferenceLibrary    string
	layersRequested     int
	layersModel         int
	availableList       []string
	kv                  uint64
	allocationsList     []string
	memoryWeights       uint64
	memoryLayerOutput   uint64
	graphFullOffload    uint64
	graphPartialOffload uint64
134
135

	projectorWeights, projectorGraph uint64
136
137
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
138
// Given a model and one or more GPU targets, predict how many layers and bytes we can load, and the total size
Daniel Hiltgen's avatar
Daniel Hiltgen committed
139
// The GPUs provided must all be the same Library
Jesse Gross's avatar
Jesse Gross committed
140
func estimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []string, opts api.Options, numParallel int) MemoryEstimate {
141
142
143
144
145
146
147
148
149
150
	// Graph size for a partial offload, applies to all GPUs
	var graphPartialOffload uint64

	// Graph size when all layers are offloaded, applies to all GPUs
	var graphFullOffload uint64

	// Final graph offload once we know full or partial
	var graphOffload uint64

	// Projectors loaded into GPU0 only
151
152
153
154
155
	var llamaEngineProjectorWeights uint64

	// Projectors loaded with output layer
	var ollamaEngineProjectorWeights uint64
	var ollamaEngineProjectorGraph uint64
156
157
158
159

	// Conditional output size on GPU 0
	var memoryLayerOutput uint64

Daniel Hiltgen's avatar
Daniel Hiltgen committed
160
161
	// The sizes of a layer
	var layerSize uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
162

163
164
165
166
167
168
169
170
171
	// The sum of all the layer sizes (just for logging)
	var memoryWeights uint64

	// True if all the layers are loaded
	var fullyLoaded bool

	// Overflow that didn't fit into the GPU
	var overflow uint64

172
	overhead := envconfig.GpuOverhead()
173
174
175
176
177
	availableList := make([]string, len(gpus))
	for i, gpu := range gpus {
		availableList[i] = format.HumanBytes2(gpu.FreeMemory)
	}
	slog.Debug("evaluating", "library", gpus[0].Library, "gpu_count", len(gpus), "available", availableList)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
178
179

	for _, projector := range projectors {
180
		llamaEngineProjectorWeights += projectorMemoryRequirements(projector)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
181
	}
182
183
	if llamaEngineProjectorWeights == 0 {
		ollamaEngineProjectorWeights, ollamaEngineProjectorGraph = f.VisionGraphSize()
184
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
185

Michael Yang's avatar
Michael Yang committed
186
	layers := f.Tensors().GroupLayers()
Michael Yang's avatar
typo  
Michael Yang committed
187
188
	// add one layer worth of memory as a buffer
	if blk0, ok := layers["blk.0"]; ok {
Michael Yang's avatar
Michael Yang committed
189
		layerSize = blk0.Size()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
190
191
	} else {
		slog.Warn("model missing blk.0 layer size")
Michael Yang's avatar
typo  
Michael Yang committed
192
	}
Michael Yang's avatar
Michael Yang committed
193

194
	var kvct string
Michael Yang's avatar
Michael Yang committed
195
196
197
	if envconfig.FlashAttention() &&
		discover.GetGPUInfo().FlashAttentionSupported() &&
		f.SupportsFlashAttention() {
198
		requested := strings.ToLower(envconfig.KvCacheType())
Michael Yang's avatar
Michael Yang committed
199
		if requested != "" && f.SupportsKVCacheType(requested) {
200
201
202
203
			kvct = requested
		}
	}

204
	kv, graphPartialOffload, graphFullOffload := f.GraphSize(uint64(opts.NumCtx), uint64(min(opts.NumCtx, opts.NumBatch)), numParallel, kvct)
205

206
207
208
209
210
211
212
213
	if len(kv) > 0 {
		layerSize += kv[0]
	}

	var kvTotal uint64
	for _, kvLayer := range kv {
		kvTotal += kvLayer
	}
214

Daniel Hiltgen's avatar
Daniel Hiltgen committed
215
	if graphPartialOffload == 0 {
216
217
218
219
220
221
		headsKV := f.KV().HeadCountKVMin()
		if headsKV == 0 {
			headsKV = 1
		}
		gqa := f.KV().HeadCountMax() / headsKV
		graphPartialOffload = gqa * kvTotal / 6
Daniel Hiltgen's avatar
Daniel Hiltgen committed
222
223
224
225
226
	}
	if graphFullOffload == 0 {
		graphFullOffload = graphPartialOffload
	}

227
228
229
	// on metal there's no partial offload overhead
	if gpus[0].Library == "metal" {
		graphPartialOffload = graphFullOffload
Daniel Hiltgen's avatar
Daniel Hiltgen committed
230
231
232
	} else if len(gpus) > 1 {
		// multigpu should always use the partial graph size
		graphFullOffload = graphPartialOffload
233
234
	}

235
	// Output layer handled at the end if we have space
236
	if layer, ok := layers["output_norm"]; ok {
Michael Yang's avatar
Michael Yang committed
237
		memoryLayerOutput += layer.Size()
238
239
	}
	if layer, ok := layers["output"]; ok {
Michael Yang's avatar
Michael Yang committed
240
		memoryLayerOutput += layer.Size()
241
	} else if layer, ok := layers["token_embd"]; ok {
Michael Yang's avatar
Michael Yang committed
242
		memoryLayerOutput += layer.Size()
Michael Yang's avatar
Michael Yang committed
243
244
	}

245
	gpuZeroOverhead := llamaEngineProjectorWeights
246
247

	// Reduce set of GPUs to only those that have sufficient space to fit overhead and at least one layer
Michael Yang's avatar
Michael Yang committed
248
	var layerCount int
Jesse Gross's avatar
Jesse Gross committed
249
	tensorSplit := make([]int, len(gpus))
250
251
252
	gpuAllocations := make([]uint64, len(gpus))
	type gs struct {
		i int
253
		g *discover.GpuInfo
254
255
256
257
258
259
260
261
	}
	gpusWithSpace := []gs{}
	for i := range gpus {
		var gzo uint64
		if len(gpusWithSpace) == 0 {
			gzo = gpuZeroOverhead
		}
		// Only include GPUs that can fit the graph, gpu minimum, the layer buffer and at least more layer
262
		if gpus[i].FreeMemory < overhead+gzo+max(graphPartialOffload, graphFullOffload)+gpus[i].MinimumMemory+2*layerSize {
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
			slog.Debug("gpu has too little memory to allocate any layers",
				"id", gpus[i].ID,
				"library", gpus[i].Library,
				"variant", gpus[i].Variant,
				"compute", gpus[i].Compute,
				"driver", fmt.Sprintf("%d.%d", gpus[i].DriverMajor, gpus[i].DriverMinor),
				"name", gpus[i].Name,
				"total", format.HumanBytes2(gpus[i].TotalMemory),
				"available", format.HumanBytes2(gpus[i].FreeMemory),
				"minimum_memory", gpus[i].MinimumMemory,
				"layer_size", format.HumanBytes2(layerSize),
				"gpu_zer_overhead", format.HumanBytes2(gzo),
				"partial_offload", format.HumanBytes2(graphPartialOffload),
				"full_offload", format.HumanBytes2(graphFullOffload),
			)
278
279
280
			continue
		}
		gpusWithSpace = append(gpusWithSpace, gs{i, &gpus[i]})
Daniel Hiltgen's avatar
Daniel Hiltgen committed
281
		gpuAllocations[i] += gpus[i].MinimumMemory + layerSize // We hold off on graph until we know partial vs. full
282
283
284
285
286
287
	}

	var gpuZeroID int
	if len(gpusWithSpace) > 0 {
		gpuZeroID = gpusWithSpace[0].i
		gpuAllocations[gpuZeroID] += gpuZeroOverhead
288
289
	} else {
		overflow += gpuZeroOverhead
290
291
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
292
	// For all the layers, find where they can fit on the GPU(s)
293
	for i := int(f.KV().BlockCount()) - 1; i >= 0; i-- {
294
295
		// Some models have inconsistent layer sizes
		if blk, ok := layers[fmt.Sprintf("blk.%d", i)]; ok {
Michael Yang's avatar
Michael Yang committed
296
			layerSize = blk.Size()
297
			layerSize += kv[i]
Michael Yang's avatar
Michael Yang committed
298
			memoryWeights += blk.Size()
299
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
300

301
302
		if opts.NumGPU >= 0 && layerCount >= opts.NumGPU {
			// Stop allocating on GPU(s) once we hit the users target NumGPU
303
			overflow += layerSize
304
305
306
307
308
309
310
			continue
		}

		// distribute the layers across the GPU(s) that have space
		for j := len(gpusWithSpace); j > 0; j-- {
			g := gpusWithSpace[i%j]
			used := gpuAllocations[g.i] + max(graphPartialOffload, graphFullOffload)
311
			if g.g.FreeMemory > overhead+used+layerSize {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
312
				gpuAllocations[g.i] += layerSize
Jesse Gross's avatar
Jesse Gross committed
313
				tensorSplit[g.i]++
Michael Yang's avatar
typo  
Michael Yang committed
314
				layerCount++
315
316
317
				break
			} else {
				gpusWithSpace = append(gpusWithSpace[:i%j], gpusWithSpace[i%j+1:]...)
Michael Yang's avatar
typo  
Michael Yang committed
318
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
319
		}
320
321
322
323

		if len(gpusWithSpace) == 0 {
			overflow += layerSize
		}
324
	}
Michael Yang's avatar
Michael Yang committed
325
	if layerCount >= int(f.KV().BlockCount()) {
326
327
		fullyLoaded = true
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
328
329

	// Determine if we need to consider output then find where it fits
330
331
	memoryLastLayer := memoryLayerOutput + ollamaEngineProjectorWeights + ollamaEngineProjectorGraph
	if memoryLastLayer > 0 {
332
333
334
335
		if opts.NumGPU < 0 || layerCount < opts.NumGPU {
			for j := len(gpusWithSpace); j > 0; j-- {
				g := gpusWithSpace[layerCount%j]
				used := gpuAllocations[g.i] + max(graphPartialOffload, graphFullOffload)
336
337
				if g.g.FreeMemory > overhead+used+memoryLastLayer {
					gpuAllocations[g.i] += memoryLastLayer
Jesse Gross's avatar
Jesse Gross committed
338
					tensorSplit[g.i]++
339
340
341
					layerCount++
					break
				}
342
343
			}
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
344

Michael Yang's avatar
Michael Yang committed
345
		if layerCount < int(f.KV().BlockCount())+1 {
346
			fullyLoaded = false
347
			overflow += memoryLastLayer
348
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
349
350
	}

351
352
	// Add the applicable (full or partial) graph allocations
	for i := range gpus {
Jesse Gross's avatar
Jesse Gross committed
353
		if tensorSplit[i] <= 0 {
354
355
356
357
358
359
360
361
362
363
364
365
			continue
		}
		if fullyLoaded {
			gpuAllocations[i] += graphFullOffload
		} else {
			gpuAllocations[i] += graphPartialOffload
		}
	}
	if fullyLoaded {
		graphOffload = graphFullOffload
	} else {
		graphOffload = graphPartialOffload
Daniel Hiltgen's avatar
Daniel Hiltgen committed
366
367
	}

368
369
370
371
	// Summaries for the log
	var memoryRequiredPartial, memoryRequiredTotal uint64
	for i := range gpuAllocations {
		memoryRequiredPartial += gpuAllocations[i]
Daniel Hiltgen's avatar
Daniel Hiltgen committed
372
	}
373
	memoryRequiredTotal = memoryRequiredPartial + overflow
Daniel Hiltgen's avatar
Daniel Hiltgen committed
374

375
376
377
378
	allocationsList := []string{}
	for _, a := range gpuAllocations {
		allocationsList = append(allocationsList, format.HumanBytes2(a))
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
379

380
381
382
383
384
385
386
387
388
	estimate := MemoryEstimate{
		TotalSize: memoryRequiredTotal,
		Layers:    0,
		Graph:     0,
		VRAMSize:  0,
		GPUSizes:  []uint64{},

		inferenceLibrary:    gpus[0].Library,
		layersRequested:     opts.NumGPU,
Michael Yang's avatar
Michael Yang committed
389
		layersModel:         int(f.KV().BlockCount()) + 1,
390
		availableList:       availableList,
391
		kv:                  kvTotal,
392
393
394
395
396
		allocationsList:     allocationsList,
		memoryWeights:       memoryWeights,
		memoryLayerOutput:   memoryLayerOutput,
		graphFullOffload:    graphFullOffload,
		graphPartialOffload: graphPartialOffload,
397
398
		projectorWeights:    llamaEngineProjectorWeights + ollamaEngineProjectorWeights,
		projectorGraph:      ollamaEngineProjectorGraph,
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
	}

	if gpus[0].Library == "cpu" {
		return estimate
	}
	if layerCount == 0 {
		slog.Debug("insufficient VRAM to load any model layers")
		return estimate
	}
	estimate.Layers = layerCount
	estimate.Graph = graphOffload
	estimate.VRAMSize = memoryRequiredPartial
	estimate.TotalSize = memoryRequiredTotal
	estimate.TensorSplit = tensorSplit
	estimate.GPUSizes = gpuAllocations
	return estimate
}

Michael Yang's avatar
Michael Yang committed
417
418
419
func (m MemoryEstimate) LogValue() slog.Value {
	attrs := []slog.Attr{
		slog.String("library", m.inferenceLibrary),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
420
421
		slog.Group(
			"layers",
Michael Yang's avatar
Michael Yang committed
422
			// requested number of layers to offload
423
			"requested", m.layersRequested,
424
			// The number of layers the model has (including output)
425
			"model", m.layersModel,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
426
			// estimated number of layers that can be offloaded
427
428
429
			"offload", m.Layers,
			// multi-gpu split for tensors
			"split", m.TensorSplit,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
430
431
432
		),
		slog.Group(
			"memory",
433
			// memory available by GPU for offloading
434
			"available", m.availableList,
Michael Yang's avatar
Michael Yang committed
435
			"gpu_overhead", format.HumanBytes2(envconfig.GpuOverhead()),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
436
437
438
			slog.Group(
				"required",
				// memory required for full offloading
439
				"full", format.HumanBytes2(m.TotalSize),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
440
				// memory required to offload layers.estimate layers
441
				"partial", format.HumanBytes2(m.VRAMSize),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
442
				// memory of KV cache
443
				"kv", format.HumanBytes2(m.kv),
444
				// Allocations across the GPUs
445
				"allocations", m.allocationsList,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
446
447
448
449
			),
			slog.Group(
				"weights",
				// memory of the weights
450
				"total", format.HumanBytes2(m.memoryWeights+m.memoryLayerOutput),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
451
				// memory of repeating layers
Michael Yang's avatar
Michael Yang committed
452
				"repeating", format.HumanBytes2(m.memoryWeights),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
453
				// memory of non-repeating layers
454
				"nonrepeating", format.HumanBytes2(m.memoryLayerOutput),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
455
456
457
458
			),
			slog.Group(
				"graph",
				// memory of graph when fully offloaded
459
				"full", format.HumanBytes2(m.graphFullOffload),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
460
				// memory of graph when not fully offloaded
461
				"partial", format.HumanBytes2(m.graphPartialOffload),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
462
463
			),
		),
Michael Yang's avatar
Michael Yang committed
464
465
466
467
468
469
470
471
472
473
474
	}

	if m.projectorWeights > 0 {
		attrs = append(attrs, slog.Group(
			"projector",
			"weights", format.HumanBytes2(m.projectorWeights),
			"graph", format.HumanBytes2(m.projectorGraph),
		))
	}

	return slog.GroupValue(attrs...)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
475
}
476

477
func projectorMemoryRequirements(filename string) (weights uint64) {
478
479
	file, err := os.Open(filename)
	if err != nil {
480
		return 0
481
482
483
	}
	defer file.Close()

484
	ggml, err := ggml.Decode(file, 1024)
485
	if err != nil {
486
		return 0
487
488
	}

Michael Yang's avatar
Michael Yang committed
489
490
	for _, layer := range ggml.Tensors().GroupLayers() {
		weights += layer.Size()
491
492
	}

493
	return weights
494
}