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

import (
4
	"cmp"
5
	"fmt"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
6
	"log/slog"
7
	"maps"
8
	"os"
9
	"slices"
10
11
	"strconv"
	"strings"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
12
13

	"github.com/ollama/ollama/api"
14
	"github.com/ollama/ollama/discover"
15
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
16
	"github.com/ollama/ollama/format"
Michael Yang's avatar
Michael Yang committed
17
	"github.com/ollama/ollama/fs/ggml"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
18
19
20
)

// This algorithm looks for a complete fit to determine if we need to unload other models
21
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
22
	// Split up the GPUs by type and try them
23
	var estimatedVRAM uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
24
25
	for _, gpus := range allGpus.ByLibrary() {
		var layerCount int
26
		estimate := EstimateGPULayers(gpus, f, projectors, opts, numParallel)
27
		layerCount, estimatedVRAM = estimate.Layers, estimate.VRAMSize
Daniel Hiltgen's avatar
Daniel Hiltgen committed
28
		if opts.NumGPU < 0 {
Michael Yang's avatar
Michael Yang committed
29
			if layerCount > 0 && layerCount >= int(f.KV().BlockCount()+1) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
30
31
32
33
34
35
36
37
38
39
40
				return true, estimatedVRAM
			}
		} else {
			if layerCount > 0 && layerCount >= opts.NumGPU {
				return true, estimatedVRAM
			}
		}
	}
	return false, estimatedVRAM
}

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
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
	TensorSplit string

	// For multi-GPU scenarios, this is the size in bytes per GPU
	GPUSizes []uint64
59
60
61
62
63
64
65
66
67
68
69
70

	// 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
71
72

	projectorWeights, projectorGraph uint64
73
74
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
75
// 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
76
// The GPUs provided must all be the same Library
77
func EstimateGPULayers(gpus []discover.GpuInfo, f *ggml.GGML, projectors []string, opts api.Options, numParallel int) MemoryEstimate {
78
79
80
81
82
83
84
85
86
87
	// 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
88
89
	var projectorWeights uint64
	var projectorGraph uint64
90
91
92
93

	// Conditional output size on GPU 0
	var memoryLayerOutput uint64

Daniel Hiltgen's avatar
Daniel Hiltgen committed
94
95
	// The sizes of a layer
	var layerSize uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
96

97
98
99
100
101
102
103
104
105
	// 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

106
	overhead := envconfig.GpuOverhead()
107
108
109
110
111
	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
112
113

	for _, projector := range projectors {
114
115
116
		weight, graph := projectorMemoryRequirements(projector)
		projectorWeights += weight
		projectorGraph += graph
Daniel Hiltgen's avatar
Daniel Hiltgen committed
117
118
119
120

		// multimodal models require at least 2048 context
		opts.NumCtx = max(opts.NumCtx, 2048)
	}
121
122
123
	if projectorWeights == 0 && projectorGraph == 0 {
		projectorWeights, projectorGraph = f.VisionGraphSize()
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
124

Michael Yang's avatar
Michael Yang committed
125
	layers := f.Tensors().GroupLayers()
126
127
128
129
	// add one layer (chosing the max layer) worth of memory as a buffer
	layerSize = slices.MaxFunc(slices.Collect(maps.Values(layers)), func(a, b ggml.Layer) int {
		return cmp.Compare(a.Size(), b.Size())
	}).Size()
Michael Yang's avatar
Michael Yang committed
130

131
	var kvct string
Michael Yang's avatar
Michael Yang committed
132
133
134
	if envconfig.FlashAttention() &&
		discover.GetGPUInfo().FlashAttentionSupported() &&
		f.SupportsFlashAttention() {
135
		requested := strings.ToLower(envconfig.KvCacheType())
Michael Yang's avatar
Michael Yang committed
136
		if requested != "" && f.SupportsKVCacheType(requested) {
137
138
139
140
			kvct = requested
		}
	}

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

143
144
145
146
147
148
149
150
	if len(kv) > 0 {
		layerSize += kv[0]
	}

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
152
	if graphPartialOffload == 0 {
153
		graphPartialOffload = f.KV().GQA() * kvTotal / 6
Daniel Hiltgen's avatar
Daniel Hiltgen committed
154
155
156
157
158
	}
	if graphFullOffload == 0 {
		graphFullOffload = graphPartialOffload
	}

159
160
161
	// on metal there's no partial offload overhead
	if gpus[0].Library == "metal" {
		graphPartialOffload = graphFullOffload
Daniel Hiltgen's avatar
Daniel Hiltgen committed
162
163
164
	} else if len(gpus) > 1 {
		// multigpu should always use the partial graph size
		graphFullOffload = graphPartialOffload
165
166
	}

167
	if layer, ok := layers["output_norm"]; ok {
Michael Yang's avatar
Michael Yang committed
168
		memoryLayerOutput += layer.Size()
169
170
	}
	if layer, ok := layers["output"]; ok {
Michael Yang's avatar
Michael Yang committed
171
		memoryLayerOutput += layer.Size()
172
	} else if layer, ok := layers["token_embd"]; ok {
Michael Yang's avatar
Michael Yang committed
173
		memoryLayerOutput += layer.Size()
Michael Yang's avatar
Michael Yang committed
174
175
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
176
	// Output layer handled at the end if we have space
177
	gpuZeroOverhead := projectorWeights + projectorGraph
178
179

	// 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
180
	var layerCount int
181
182
183
184
	layerCounts := make([]int, len(gpus))
	gpuAllocations := make([]uint64, len(gpus))
	type gs struct {
		i int
185
		g *discover.GpuInfo
186
187
188
189
190
191
192
193
	}
	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
194
		if gpus[i].FreeMemory < overhead+gzo+max(graphPartialOffload, graphFullOffload)+gpus[i].MinimumMemory+2*layerSize {
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
			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),
			)
210
211
212
			continue
		}
		gpusWithSpace = append(gpusWithSpace, gs{i, &gpus[i]})
Daniel Hiltgen's avatar
Daniel Hiltgen committed
213
		gpuAllocations[i] += gpus[i].MinimumMemory + layerSize // We hold off on graph until we know partial vs. full
214
215
216
217
218
219
220
221
	}

	var gpuZeroID int
	if len(gpusWithSpace) > 0 {
		gpuZeroID = gpusWithSpace[0].i
		gpuAllocations[gpuZeroID] += gpuZeroOverhead
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
222
	// For all the layers, find where they can fit on the GPU(s)
223
	for i := int(f.KV().BlockCount()) - 1; i >= 0; i-- {
224
225
		// Some models have inconsistent layer sizes
		if blk, ok := layers[fmt.Sprintf("blk.%d", i)]; ok {
Michael Yang's avatar
Michael Yang committed
226
			layerSize = blk.Size()
227
			layerSize += kv[i]
Michael Yang's avatar
Michael Yang committed
228
			memoryWeights += blk.Size()
229
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
230

231
232
		if opts.NumGPU >= 0 && layerCount >= opts.NumGPU {
			// Stop allocating on GPU(s) once we hit the users target NumGPU
233
			overflow += layerSize
234
235
236
237
238
239
240
			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)
241
			if g.g.FreeMemory > overhead+used+layerSize {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
242
				gpuAllocations[g.i] += layerSize
243
				layerCounts[g.i]++
Michael Yang's avatar
typo  
Michael Yang committed
244
				layerCount++
245
246
247
				break
			} else {
				gpusWithSpace = append(gpusWithSpace[:i%j], gpusWithSpace[i%j+1:]...)
Michael Yang's avatar
typo  
Michael Yang committed
248
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
249
		}
250
251
252
253

		if len(gpusWithSpace) == 0 {
			overflow += layerSize
		}
254
	}
Michael Yang's avatar
Michael Yang committed
255
	if layerCount >= int(f.KV().BlockCount()) {
256
257
		fullyLoaded = true
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
258
259

	// Determine if we need to consider output then find where it fits
260
	if memoryLayerOutput > 0 && (opts.NumGPU < 0 || layerCount < opts.NumGPU) {
261
262
263
		for j := len(gpusWithSpace); j > 0; j-- {
			g := gpusWithSpace[layerCount%j]
			used := gpuAllocations[g.i] + max(graphPartialOffload, graphFullOffload)
264
			if g.g.FreeMemory > overhead+used+memoryLayerOutput {
265
266
267
268
269
270
				gpuAllocations[g.i] += memoryLayerOutput
				layerCounts[g.i]++
				layerCount++
				break
			}
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
271

Michael Yang's avatar
Michael Yang committed
272
		if layerCount < int(f.KV().BlockCount())+1 {
273
274
275
			fullyLoaded = false
			overflow += memoryLayerOutput
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
276
277
	}

278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
	// Add the applicable (full or partial) graph allocations
	for i := range gpus {
		if layerCounts[i] <= 0 {
			continue
		}
		if fullyLoaded {
			gpuAllocations[i] += graphFullOffload
		} else {
			gpuAllocations[i] += graphPartialOffload
		}
	}
	if fullyLoaded {
		graphOffload = graphFullOffload
	} else {
		graphOffload = graphPartialOffload
Daniel Hiltgen's avatar
Daniel Hiltgen committed
293
294
	}

295
296
297
298
	// Summaries for the log
	var memoryRequiredPartial, memoryRequiredTotal uint64
	for i := range gpuAllocations {
		memoryRequiredPartial += gpuAllocations[i]
Daniel Hiltgen's avatar
Daniel Hiltgen committed
299
	}
300
	memoryRequiredTotal = memoryRequiredPartial + overflow
Daniel Hiltgen's avatar
Daniel Hiltgen committed
301

302
303
304
305
306
307
308
309
310
311
312
313
	tensorSplit := ""
	if len(gpus) > 1 {
		splits := make([]string, len(gpus))
		for i, count := range layerCounts {
			splits[i] = strconv.Itoa(count)
		}
		tensorSplit = strings.Join(splits, ",")
	}
	allocationsList := []string{}
	for _, a := range gpuAllocations {
		allocationsList = append(allocationsList, format.HumanBytes2(a))
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
314

315
316
317
318
319
320
321
322
323
	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
324
		layersModel:         int(f.KV().BlockCount()) + 1,
325
		availableList:       availableList,
326
		kv:                  kvTotal,
327
328
329
330
331
		allocationsList:     allocationsList,
		memoryWeights:       memoryWeights,
		memoryLayerOutput:   memoryLayerOutput,
		graphFullOffload:    graphFullOffload,
		graphPartialOffload: graphPartialOffload,
332
333
		projectorWeights:    projectorWeights,
		projectorGraph:      projectorGraph,
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
	}

	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
352
353
354
func (m MemoryEstimate) LogValue() slog.Value {
	attrs := []slog.Attr{
		slog.String("library", m.inferenceLibrary),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
355
356
		slog.Group(
			"layers",
Michael Yang's avatar
Michael Yang committed
357
			// requested number of layers to offload
358
			"requested", m.layersRequested,
359
			// The number of layers the model has (including output)
360
			"model", m.layersModel,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
361
			// estimated number of layers that can be offloaded
362
363
364
			"offload", m.Layers,
			// multi-gpu split for tensors
			"split", m.TensorSplit,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
365
366
367
		),
		slog.Group(
			"memory",
368
			// memory available by GPU for offloading
369
			"available", m.availableList,
Michael Yang's avatar
Michael Yang committed
370
			"gpu_overhead", format.HumanBytes2(envconfig.GpuOverhead()),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
371
372
373
			slog.Group(
				"required",
				// memory required for full offloading
374
				"full", format.HumanBytes2(m.TotalSize),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
375
				// memory required to offload layers.estimate layers
376
				"partial", format.HumanBytes2(m.VRAMSize),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
377
				// memory of KV cache
378
				"kv", format.HumanBytes2(m.kv),
379
				// Allocations across the GPUs
380
				"allocations", m.allocationsList,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
381
382
383
384
			),
			slog.Group(
				"weights",
				// memory of the weights
385
				"total", format.HumanBytes2(m.memoryWeights+m.memoryLayerOutput),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
386
				// memory of repeating layers
Michael Yang's avatar
Michael Yang committed
387
				"repeating", format.HumanBytes2(m.memoryWeights),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
388
				// memory of non-repeating layers
389
				"nonrepeating", format.HumanBytes2(m.memoryLayerOutput),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
390
391
392
393
			),
			slog.Group(
				"graph",
				// memory of graph when fully offloaded
394
				"full", format.HumanBytes2(m.graphFullOffload),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
395
				// memory of graph when not fully offloaded
396
				"partial", format.HumanBytes2(m.graphPartialOffload),
Daniel Hiltgen's avatar
Daniel Hiltgen committed
397
398
			),
		),
Michael Yang's avatar
Michael Yang committed
399
400
401
402
403
404
405
406
407
408
409
	}

	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
410
}
411
412
413
414
415
416
417
418

func projectorMemoryRequirements(filename string) (weights, graphSize uint64) {
	file, err := os.Open(filename)
	if err != nil {
		return 0, 0
	}
	defer file.Close()

419
	ggml, _, err := ggml.Decode(file, 1024)
420
421
422
423
	if err != nil {
		return 0, 0
	}

Michael Yang's avatar
Michael Yang committed
424
425
	for _, layer := range ggml.Tensors().GroupLayers() {
		weights += layer.Size()
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
	}

	switch arch := ggml.KV().Architecture(); arch {
	case "mllama":
		kv := func(n string) uint64 {
			if v, ok := ggml.KV()[arch+".vision."+n].(uint32); ok {
				return uint64(v)
			}

			return 0
		}

		imageSize := kv("image_size")

		maxNumTiles := kv("max_num_tiles")
		embeddingLength := kv("embedding_length")
		headCount := kv("attention.head_count")

		numPatches := (imageSize / kv("patch_size")) * (imageSize / kv("patch_size"))
Michael Yang's avatar
Michael Yang committed
445
		if _, ok := ggml.Tensors().GroupLayers()["v"]["class_embd"]; ok {
446
447
448
449
450
451
452
453
454
455
456
457
458
459
			numPatches++
		}

		numPaddedPatches := numPatches + 8 - (numPatches%8)%8

		graphSize = 4 * (8 +
			imageSize*imageSize*kv("num_channels")*maxNumTiles +
			embeddingLength*numPatches*maxNumTiles +
			9*embeddingLength*numPaddedPatches*maxNumTiles +
			numPaddedPatches*maxNumTiles*numPaddedPatches*maxNumTiles*headCount)
	}

	return weights, graphSize
}