ggml.go 11 KB
Newer Older
1
2
3
4
5
package llm

import (
	"encoding/binary"
	"errors"
Michael Yang's avatar
Michael Yang committed
6
	"fmt"
7
	"io"
Michael Yang's avatar
Michael Yang committed
8
	"strings"
9
10

	"github.com/ollama/ollama/util/bufioutil"
11
12
)

Michael Yang's avatar
Michael Yang committed
13
14
15
16
type GGML struct {
	container
	model
}
17

Michael Yang's avatar
Michael Yang committed
18
type model interface {
Michael Yang's avatar
Michael Yang committed
19
	KV() KV
Michael Yang's avatar
Michael Yang committed
20
	Tensors() Tensors
21
22
}

23
24
type KV map[string]any

Michael Yang's avatar
Michael Yang committed
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
func (kv KV) u64(key string) uint64 {
	switch v := kv[key].(type) {
	case uint64:
		return v
	case uint32:
		return uint64(v)
	case float64:
		return uint64(v)
	default:
		return 0
	}
}

func (kv KV) Architecture() string {
	if s, ok := kv["general.architecture"].(string); ok {
		return s
	}

	return "unknown"
}

func (kv KV) ParameterCount() uint64 {
	return kv.u64("general.parameter_count")
}

Michael Yang's avatar
Michael Yang committed
50
func (kv KV) FileType() fileType {
Michael Yang's avatar
Michael Yang committed
51
	if u64 := kv.u64("general.file_type"); u64 > 0 {
Michael Yang's avatar
Michael Yang committed
52
		return fileType(uint32(u64))
Michael Yang's avatar
Michael Yang committed
53
54
	}

Michael Yang's avatar
Michael Yang committed
55
	return fileTypeUnknown
Michael Yang's avatar
Michael Yang committed
56
57
58
59
60
61
62
63
64
65
66
}

func (kv KV) BlockCount() uint64 {
	return kv.u64(fmt.Sprintf("%s.block_count", kv.Architecture()))
}

func (kv KV) HeadCount() uint64 {
	return kv.u64(fmt.Sprintf("%s.attention.head_count", kv.Architecture()))
}

func (kv KV) HeadCountKV() uint64 {
Michael Yang's avatar
Michael Yang committed
67
68
69
70
71
	if headCountKV := kv.u64(fmt.Sprintf("%s.attention.head_count_kv", kv.Architecture())); headCountKV > 0 {
		return headCountKV
	}

	return 1
Michael Yang's avatar
Michael Yang committed
72
73
}

Michael Yang's avatar
Michael Yang committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
func (kv KV) EmbeddingHeadCount() uint64 {
	if heads := kv.HeadCount(); heads > 0 {
		return kv.EmbeddingLength() / kv.HeadCount()
	}

	return 0
}

func (kv KV) EmbeddingHeadCountK() uint64 {
	if k := kv.u64(fmt.Sprintf("%s.attention.key_length", kv.Architecture())); k > 0 {
		return k
	}

	return kv.EmbeddingHeadCount()
}

func (kv KV) EmbeddingHeadCountV() uint64 {
	if v := kv.u64(fmt.Sprintf("%s.attention.value_length", kv.Architecture())); v > 0 {
		return v
	}

	return kv.EmbeddingHeadCount()
}

Michael Yang's avatar
Michael Yang committed
98
func (kv KV) GQA() uint64 {
Michael Yang's avatar
Michael Yang committed
99
	return kv.HeadCount() / kv.HeadCountKV()
Michael Yang's avatar
Michael Yang committed
100
101
102
103
104
105
106
107
108
109
}

func (kv KV) EmbeddingLength() uint64 {
	return kv.u64(fmt.Sprintf("%s.embedding_length", kv.Architecture()))
}

func (kv KV) ContextLength() uint64 {
	return kv.u64(fmt.Sprintf("%s.context_length", kv.Architecture()))
}

Michael Yang's avatar
Michael Yang committed
110
111
112
113
114
func (kv KV) ChatTemplate() string {
	s, _ := kv["tokenizer.chat_template"].(string)
	return s
}

115
116
117
118
type Tensors struct {
	Items  []*Tensor
	Offset uint64
}
Michael Yang's avatar
Michael Yang committed
119
120
121

func (ts Tensors) Layers() map[string]Layer {
	layers := make(map[string]Layer)
122
	for _, t := range ts.Items {
Michael Yang's avatar
Michael Yang committed
123
124
		parts := strings.Split(t.Name, ".")
		if parts[0] == "blk" {
125
126
			// join first and second part, e.g. blk.%d
			parts = append([]string{fmt.Sprintf("%s.%s", parts[0], parts[1])}, parts[2:]...)
Michael Yang's avatar
Michael Yang committed
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
		}

		if _, ok := layers[parts[0]]; !ok {
			layers[parts[0]] = make(Layer)
		}

		layers[parts[0]][strings.Join(parts[1:], ".")] = t
	}

	return layers
}

type Layer map[string]*Tensor

func (l Layer) size() (size uint64) {
	for _, t := range l {
Michael Yang's avatar
Michael Yang committed
143
		size += t.Size()
Michael Yang's avatar
Michael Yang committed
144
145
146
147
148
	}

	return size
}

149
type Tensor struct {
Michael Yang's avatar
Michael Yang committed
150
151
152
	Name   string `json:"name"`
	Kind   uint32 `json:"kind"`
	Offset uint64 `json:"-"`
153
154

	// Shape is the number of elements in each dimension
Michael Yang's avatar
Michael Yang committed
155
	Shape []uint64 `json:"shape"`
156

Michael Yang's avatar
Michael Yang committed
157
	io.WriterTo `json:"-"`
158
159
160
}

func (t Tensor) blockSize() uint64 {
161
	switch t.Kind {
Michael Yang's avatar
Michael Yang committed
162
	case 0, 1, 24, 25, 26, 27, 28, 30: // F32, F16, I8, I16, I32, I64, F64, BF16
163
		return 1
Michael Yang's avatar
Michael Yang committed
164
	case 2, 3, 4, 5, 6, 7, 8, 9, 20: // Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q8_1, IQ4_NL
165
		return 32
166
	default: // All others
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
		return 256
	}
}

func (t Tensor) typeSize() uint64 {
	blockSize := t.blockSize()

	switch t.Kind {
	case 0: // FP32
		return 4
	case 1: // FP16
		return 2
	case 2: // Q4_0
		return 2 + blockSize/2
	case 3: // Q4_1
		return 2 + 2 + blockSize/2
	case 6: // Q5_0
		return 2 + 4 + blockSize/2
	case 7: // Q5_1
		return 2 + 2 + 4 + blockSize/2
	case 8: // Q8_0
		return 2 + blockSize
	case 9: // Q8_1
		return 4 + 4 + blockSize
	case 10: // Q2_K
		return blockSize/16 + blockSize/4 + 2 + 2
	case 11: // Q3_K
		return blockSize/8 + blockSize/4 + 12 + 2
	case 12: // Q4_K
		return 2 + 2 + 12 + blockSize/2
	case 13: // Q5_K
		return 2 + 2 + 12 + blockSize/8 + blockSize/2
	case 14: // Q6_K
		return blockSize/2 + blockSize/4 + blockSize/16 + 2
	case 15: // Q8_K
		return 2 + blockSize + 2*blockSize/16
	case 16: // IQ2_XXS
		return 2 + 2*blockSize/8
	case 17: // IQ2_XS
		return 2 + 2*blockSize/8 + blockSize/32
	case 18: // IQ3_XXS
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
		return 2 + blockSize/4 + blockSize/8
	case 19: // IQ1_S
		return 2 + blockSize/8 + blockSize/16
	case 20: // IQ4_NL
		return 2 + blockSize/2
	case 21: // IQ3_S
		return 2 + blockSize/4 + blockSize/8 + blockSize/32 + 4
	case 22: // IQ2_S
		return 2 + blockSize/4 + blockSize/16
	case 23: // IQ4_XS
		return 2 + 2 + blockSize/2 + blockSize/64
	case 24: // I8
		return 1
	case 25: // I16
		return 2
	case 26: // I32
		return 4
	case 27: // I64
		return 8
	case 28: // F64
		return 8
	case 29: // IQ1_M
		return blockSize/8 + blockSize/16 + blockSize/32
231
232
233
234
235
236
237
238
239
240
241
242
243
	default:
		return 0
	}
}

func (t Tensor) parameters() uint64 {
	var count uint64 = 1
	for _, n := range t.Shape {
		count *= n
	}
	return count
}

Michael Yang's avatar
Michael Yang committed
244
func (t Tensor) Size() uint64 {
245
246
247
	return t.parameters() * t.typeSize() / t.blockSize()
}

248
249
type container interface {
	Name() string
Michael Yang's avatar
Michael Yang committed
250
	Decode(io.ReadSeeker) (model, error)
251
252
253
}

const (
Bruce MacDonald's avatar
Bruce MacDonald committed
254
	// Magic constant for `ggml` files (unversioned).
255
	FILE_MAGIC_GGML = 0x67676d6c
Bruce MacDonald's avatar
Bruce MacDonald committed
256
	// Magic constant for `ggml` files (versioned, ggmf).
257
	FILE_MAGIC_GGMF = 0x67676d66
Bruce MacDonald's avatar
Bruce MacDonald committed
258
	// Magic constant for `ggml` files (versioned, ggjt).
259
	FILE_MAGIC_GGJT = 0x67676a74
Bruce MacDonald's avatar
Bruce MacDonald committed
260
	// Magic constant for `ggla` files (LoRA adapter).
261
	FILE_MAGIC_GGLA = 0x67676C61
Bruce MacDonald's avatar
Bruce MacDonald committed
262
	// Magic constant for `gguf` files (versioned, gguf)
263
264
	FILE_MAGIC_GGUF_LE = 0x46554747
	FILE_MAGIC_GGUF_BE = 0x47475546
265
266
)

Bruce MacDonald's avatar
Bruce MacDonald committed
267
268
var ErrUnsupportedFormat = errors.New("unsupported model format")

Michael Yang's avatar
Michael Yang committed
269
270
271
272
273
274
275
276
277
278
func DetectGGMLType(b []byte) string {
	switch binary.LittleEndian.Uint32(b[:4]) {
	case FILE_MAGIC_GGML:
		return "ggml"
	case FILE_MAGIC_GGMF:
		return "ggmf"
	case FILE_MAGIC_GGJT:
		return "ggjt"
	case FILE_MAGIC_GGLA:
		return "ggla"
279
	case FILE_MAGIC_GGUF_LE, FILE_MAGIC_GGUF_BE:
Michael Yang's avatar
Michael Yang committed
280
281
282
283
284
285
		return "gguf"
	default:
		return ""
	}
}

286
287
288
289
290
291
292
293
294
295
296
297
// DecodeGGML decodes a GGML model from the given reader.
//
// It collects array values for arrays with a size less than or equal to
// maxArraySize. If maxArraySize is 0, the default value of 1024 is used. If
// the maxArraySize is negative, all arrays are collected.
func DecodeGGML(rs io.ReadSeeker, maxArraySize int) (*GGML, int64, error) {
	if maxArraySize == 0 {
		maxArraySize = 1024
	}

	rs = bufioutil.NewBufferedSeeker(rs, 32<<10)

298
	var magic uint32
Michael Yang's avatar
Michael Yang committed
299
	if err := binary.Read(rs, binary.LittleEndian, &magic); err != nil {
Michael Yang's avatar
Michael Yang committed
300
		return nil, 0, err
301
302
303
	}

	var c container
304
	switch magic {
Bruce MacDonald's avatar
Bruce MacDonald committed
305
	case FILE_MAGIC_GGML, FILE_MAGIC_GGMF, FILE_MAGIC_GGJT:
Michael Yang's avatar
Michael Yang committed
306
		return nil, 0, ErrUnsupportedFormat
307
	case FILE_MAGIC_GGLA:
308
		c = &containerGGLA{}
309
	case FILE_MAGIC_GGUF_LE:
310
		c = &containerGGUF{ByteOrder: binary.LittleEndian, maxArraySize: maxArraySize}
311
	case FILE_MAGIC_GGUF_BE:
312
		c = &containerGGUF{ByteOrder: binary.BigEndian, maxArraySize: maxArraySize}
313
	default:
Michael Yang's avatar
Michael Yang committed
314
		return nil, 0, errors.New("invalid file magic")
315
316
	}

Michael Yang's avatar
Michael Yang committed
317
	model, err := c.Decode(rs)
318
	if err != nil {
Michael Yang's avatar
Michael Yang committed
319
		return nil, 0, err
320
321
	}

Michael Yang's avatar
Michael Yang committed
322
323
	offset, err := rs.Seek(0, io.SeekCurrent)
	if err != nil {
Michael Yang's avatar
Michael Yang committed
324
		return nil, 0, err
Michael Yang's avatar
Michael Yang committed
325
326
	}

327
	// final model type
328
329
330
	return &GGML{
		container: c,
		model:     model,
Michael Yang's avatar
Michael Yang committed
331
	}, offset, nil
332
}
Michael Yang's avatar
Michael Yang committed
333

Michael Yang's avatar
Michael Yang committed
334
335
336
337
func (llm GGML) GraphSize(context, batch uint64) (partialOffload, fullOffload uint64) {
	embedding := llm.KV().EmbeddingLength()
	heads := llm.KV().HeadCount()
	headsKV := llm.KV().HeadCountKV()
338
	vocab := uint64(llm.KV()["tokenizer.ggml.tokens"].(*array).size)
Michael Yang's avatar
Michael Yang committed
339

Michael Yang's avatar
Michael Yang committed
340
341
342
	embeddingHeads := llm.KV().EmbeddingHeadCount()
	embeddingHeadsK := llm.KV().EmbeddingHeadCountK()

Michael Yang's avatar
Michael Yang committed
343
344
	layers := llm.Tensors().Layers()

Michael Yang's avatar
Michael Yang committed
345
346
	switch llm.KV().Architecture() {
	case "llama":
Michael Yang's avatar
Michael Yang committed
347
348
349
350
		fullOffload = 4 * batch * (1 + 4*embedding + context*(1+heads))

		partialOffload = 4 * batch * embedding
		partialOffload += max(
351
			// 4*batch*(4+6*embedding+context*(2*heads)+llm.KV().GQA()),
Michael Yang's avatar
Michael Yang committed
352
			4*batch*(1+embedding+max(context, embedding))+embedding*embedding*9/16+4*context*(batch*heads+embeddingHeads*headsKV),
Michael Yang's avatar
Michael Yang committed
353
354
			4*batch*(embedding+vocab)+embedding*vocab*105/128,
		)
Michael Yang's avatar
Michael Yang committed
355

Michael Yang's avatar
Michael Yang committed
356
357
358
359
		if ffnGateExpsWeight, ok := layers["blk.0"]["ffn_gate_exps.weight"]; ok {
			// mixtral 8x22b
			ff := uint64(llm.KV()["llama.feed_forward_length"].(uint32))
			partialOffload = max(
Michael Yang's avatar
Michael Yang committed
360
361
				3*ffnGateExpsWeight.Size()+4*batch*(2*ff+headsKV+embedding+context+embeddingHeads*headsKV),
				4*(context*batch*heads+context*embeddingHeads*headsKV+batch*1024+embeddingHeads*headsKV*batch),
Michael Yang's avatar
Michael Yang committed
362
363
364
			)
		} else if ffnGateWeight, ok := layers["blk.0"]["ffn_gate.0.weight"]; ok {
			// mixtral 8x7b
Michael Yang's avatar
Michael Yang committed
365
366
367
			ffnGateWeight1 := ffnGateWeight.Shape[1]
			fullOffload = 4 * batch * (2 + 3*embedding + context*(1+heads) + 2*headsKV + ffnGateWeight1)
			partialOffload = max(
Michael Yang's avatar
Michael Yang committed
368
				4*batch*(3+embeddingHeads*headsKV+embedding+context*(1+heads)+ffnGateWeight1)+(embedding*embedding+3*embedding*headsKV*ffnGateWeight1)*9/16,
Michael Yang's avatar
Michael Yang committed
369
370
371
				4*batch*(1+2*embedding+context*(1+heads))+embedding*(6*context*headsKV/heads+embedding*9/16),
			)
		}
Michael Yang's avatar
Michael Yang committed
372
373
374
375
376
377
378
379
380
381
382
383
	case "gemma", "gemma2":
		fullOffload = max(
			4*batch*(embedding+vocab),
			4*batch*(2+context+context*heads+2*embedding+2*embeddingHeadsK*heads),
		)

		partialOffload = max(
			4*embedding*batch+embedding*vocab*105/128+4*vocab*batch,
			4*batch*(2*embedding+1+2*embeddingHeadsK*heads+context+context*heads)+
				4*embeddingHeadsK*context*8+
				embedding*embeddingHeadsK*heads*9/16,
		)
Michael Yang's avatar
Michael Yang committed
384
385
386
387
388
389
390
391
	case "command-r":
		fullOffload = max(
			4*batch*(embedding+vocab),
			4*batch*(2+4*embedding+context*(1+heads)),
		)

		partialOffload = max(
			4*batch*(embedding+vocab)+embedding*vocab*105/128,
Michael Yang's avatar
Michael Yang committed
392
			4*batch*(1+2*embedding+context*(1+heads))+4*embedding*context+embedding*embedding*9/16,
Michael Yang's avatar
Michael Yang committed
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
		)
	case "qwen2":
		fullOffload = max(
			4*batch*(embedding+vocab),
			4*batch*(1+2*embedding+context+context*heads),
		)

		partialOffload = max(
			4*batch*(embedding+vocab)+embedding*vocab*105/128,
			4*(batch*(1+2*embedding+context*(1+heads))+embedding*(1+context)),
		)
	case "phi2":
		fullOffload = max(
			4*batch*(embedding+vocab),
			4*batch*(1+4*embedding+context+context*heads),
		)
Michael Yang's avatar
Michael Yang committed
409

Michael Yang's avatar
Michael Yang committed
410
411
412
413
		partialOffload = max(
			4*batch*(2*embedding+vocab)+embedding*vocab*105/128,
			4*batch*(2+3*embedding+context+context*heads),
		)
Michael Yang's avatar
Michael Yang committed
414
415
416
417
418
419
	case "stablelm":
		fullOffload = 4 * batch * (context*(1+heads) + 3*embedding + 2)
		partialOffload = max(
			4*batch*(vocab+2*embedding),
			fullOffload,
		)
Michael Yang's avatar
Michael Yang committed
420
421
422
	case "deepseek2":
		fullOffload = max(
			4*batch*(3*embedding+vocab),
Michael Yang's avatar
Michael Yang committed
423
			4*batch*(3*embedding+2+context*(1+headsKV)+2*embeddingHeadsK*headsKV),
Michael Yang's avatar
Michael Yang committed
424
425
426
427
		)

		partialOffload = max(
			4*batch*(3*embedding+vocab)+embedding*vocab*105/128,
Michael Yang's avatar
Michael Yang committed
428
			4*batch*(2*embedding+1+2*embeddingHeadsK*headsKV+context+context*headsKV)+4*embeddingHeadsK*context*headsKV+embedding*embeddingHeadsK*headsKV*9/16,
Michael Yang's avatar
Michael Yang committed
429
		)
Michael Yang's avatar
Michael Yang committed
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
	case "chatglm":
		fullOffload = 4 * batch * (embedding + vocab)
		partialOffload = 4*batch*(embedding+vocab) + embedding*vocab*105/128
		if qkvBias, ok := layers["blk.0"]["attn_qkv.bias"]; ok {
			fullOffload = max(
				fullOffload,
				4*batch*(2+
					2*embedding+
					context+
					context*heads+
					embeddingHeadsK*heads+
					qkvBias.Shape[0]),
			)

			partialOffload = max(
				partialOffload,
				4*batch*(1+
					2*embedding+
					embeddingHeadsK*heads+
					context+
					context*heads)+
					4*embeddingHeadsK*context+
					4*context*embeddingHeadsK+
					4*qkvBias.Shape[0],
			)
		}
Michael Yang's avatar
Michael Yang committed
456
457
	}

Michael Yang's avatar
Michael Yang committed
458
	return
Michael Yang's avatar
Michael Yang committed
459
}