gguf.go 9.53 KB
Newer Older
Bruce MacDonald's avatar
Bruce MacDonald committed
1
2
3
4
5
6
7
package llm

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
8
9

	"github.com/jmorganca/ollama/format"
Bruce MacDonald's avatar
Bruce MacDonald committed
10
11
12
)

type containerGGUF struct {
Michael Yang's avatar
ggufv3  
Michael Yang committed
13
14
	bo binary.ByteOrder

Bruce MacDonald's avatar
Bruce MacDonald committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
	Version uint32

	V1 struct {
		NumTensor uint32
		NumKV     uint32
	}

	V2 struct {
		NumTensor uint64
		NumKV     uint64
	}
}

func (c *containerGGUF) Name() string {
	return "gguf"
}

Michael Yang's avatar
Michael Yang committed
32
33
func (c *containerGGUF) Decode(rso *readSeekOffset) (model, error) {
	binary.Read(rso, c.bo, &c.Version)
Bruce MacDonald's avatar
Bruce MacDonald committed
34
35
36

	switch c.Version {
	case 1:
Michael Yang's avatar
Michael Yang committed
37
		binary.Read(rso, c.bo, &c.V1)
Bruce MacDonald's avatar
Bruce MacDonald committed
38
	default:
Michael Yang's avatar
Michael Yang committed
39
		binary.Read(rso, c.bo, &c.V2)
Bruce MacDonald's avatar
Bruce MacDonald committed
40
41
42
	}

	model := newGGUFModel(c)
Michael Yang's avatar
Michael Yang committed
43
	if err := model.Decode(rso); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
		return nil, err
	}

	return model, nil
}

const (
	ggufTypeUint8 uint32 = iota
	ggufTypeInt8
	ggufTypeUint16
	ggufTypeInt16
	ggufTypeUint32
	ggufTypeInt32
	ggufTypeFloat32
	ggufTypeBool
	ggufTypeString
	ggufTypeArray
	ggufTypeUint64
	ggufTypeInt64
	ggufTypeFloat64
)

type kv map[string]any

68
69
70
71
72
73
74
75
76
type tensor struct {
	name   string
	kind   uint32
	offset uint64

	// shape is the number of elements in each dimension
	shape [4]uint64
}

Michael Yang's avatar
Michael Yang committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
func (t tensor) blockSize() uint64 {
	switch {
	case t.kind < 2:
		return 1
	case t.kind < 10:
		return 32
	default:
		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
	default:
		return 0
	}
}

func (t tensor) parameters() uint64 {
	return t.shape[0] * t.shape[1] * t.shape[2] * t.shape[3]
}

func (t tensor) size() uint64 {
	return t.parameters() * t.typeSize() / t.blockSize()
}

Bruce MacDonald's avatar
Bruce MacDonald committed
131
132
type ggufModel struct {
	*containerGGUF
133

Bruce MacDonald's avatar
Bruce MacDonald committed
134
	kv
135
136
137
	tensors []tensor

	parameters uint64
Bruce MacDonald's avatar
Bruce MacDonald committed
138
139
140
141
142
143
144
145
146
}

func newGGUFModel(container *containerGGUF) *ggufModel {
	return &ggufModel{
		containerGGUF: container,
		kv:            make(kv),
	}
}

147
148
149
150
151
152
153
154
func (llm *ggufModel) NumTensor() uint64 {
	if llm.Version == 1 {
		return uint64(llm.V1.NumTensor)
	}

	return llm.V2.NumTensor
}

Bruce MacDonald's avatar
Bruce MacDonald committed
155
156
157
158
159
160
161
162
func (llm *ggufModel) NumKV() uint64 {
	if llm.Version == 1 {
		return uint64(llm.V1.NumKV)
	}

	return llm.V2.NumKV
}

Michael Yang's avatar
Michael Yang committed
163
func (llm *ggufModel) ModelFamily() string {
Michael Yang's avatar
Michael Yang committed
164
	if t, ok := llm.kv["general.architecture"].(string); ok {
Michael Yang's avatar
Michael Yang committed
165
		return t
Bruce MacDonald's avatar
Bruce MacDonald committed
166
167
	}

Michael Yang's avatar
Michael Yang committed
168
	return "unknown"
Bruce MacDonald's avatar
Bruce MacDonald committed
169
170
}

Michael Yang's avatar
Michael Yang committed
171
func (llm *ggufModel) ModelType() string {
172
173
174
175
	if llm.parameters > 0 {
		return format.HumanNumber(llm.parameters)
	}

Michael Yang's avatar
Michael Yang committed
176
	return "unknown"
Bruce MacDonald's avatar
Bruce MacDonald committed
177
178
}

Michael Yang's avatar
Michael Yang committed
179
func (llm *ggufModel) FileType() string {
Michael Yang's avatar
Michael Yang committed
180
	if t, ok := llm.kv["general.file_type"].(uint32); ok {
Michael Yang's avatar
Michael Yang committed
181
		return fileType(t)
Bruce MacDonald's avatar
Bruce MacDonald committed
182
183
	}

Michael Yang's avatar
Michael Yang committed
184
	return "unknown"
Bruce MacDonald's avatar
Bruce MacDonald committed
185
186
}

Michael Yang's avatar
Michael Yang committed
187
func (llm *ggufModel) Decode(rso *readSeekOffset) error {
188
	// decode key-values
Bruce MacDonald's avatar
Bruce MacDonald committed
189
	for i := 0; uint64(i) < llm.NumKV(); i++ {
Michael Yang's avatar
Michael Yang committed
190
		k, err := llm.readString(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
191
192
193
194
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
195
		vtype := llm.readU32(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
196
197
198
199

		var v any
		switch vtype {
		case ggufTypeUint8:
Michael Yang's avatar
Michael Yang committed
200
			v = llm.readU8(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
201
		case ggufTypeInt8:
Michael Yang's avatar
Michael Yang committed
202
			v = llm.readI8(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
203
		case ggufTypeUint16:
Michael Yang's avatar
Michael Yang committed
204
			v = llm.readU16(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
205
		case ggufTypeInt16:
Michael Yang's avatar
Michael Yang committed
206
			v = llm.readI16(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
207
		case ggufTypeUint32:
Michael Yang's avatar
Michael Yang committed
208
			v = llm.readU32(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
209
		case ggufTypeInt32:
Michael Yang's avatar
Michael Yang committed
210
			v = llm.readI32(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
211
		case ggufTypeUint64:
Michael Yang's avatar
Michael Yang committed
212
			v = llm.readU64(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
213
		case ggufTypeInt64:
Michael Yang's avatar
Michael Yang committed
214
			v = llm.readI64(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
215
		case ggufTypeFloat32:
Michael Yang's avatar
Michael Yang committed
216
			v = llm.readF32(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
217
		case ggufTypeFloat64:
Michael Yang's avatar
Michael Yang committed
218
			v = llm.readF64(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
219
		case ggufTypeBool:
Michael Yang's avatar
Michael Yang committed
220
			v = llm.readBool(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
221
		case ggufTypeString:
Michael Yang's avatar
Michael Yang committed
222
			s, err := llm.readString(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
223
224
225
226
227
228
			if err != nil {
				return err
			}

			v = s
		case ggufTypeArray:
Michael Yang's avatar
Michael Yang committed
229
			a, err := llm.readArray(rso)
Bruce MacDonald's avatar
Bruce MacDonald committed
230
231
232
233
234
235
236
237
238
239
240
241
			if err != nil {
				return err
			}

			v = a
		default:
			return fmt.Errorf("invalid type: %d", vtype)
		}

		llm.kv[k] = v
	}

242
243
	// decode tensors
	for i := 0; uint64(i) < llm.NumTensor(); i++ {
Michael Yang's avatar
Michael Yang committed
244
		name, err := llm.readString(rso)
245
		if err != nil {
246
247
248
			return err
		}

Michael Yang's avatar
Michael Yang committed
249
		// dims is the number of dimensions in the tensor
Michael Yang's avatar
Michael Yang committed
250
		dims := llm.readU32(rso)
251
252
253

		shape := [4]uint64{1, 1, 1, 1}
		for i := 0; uint32(i) < dims; i++ {
Michael Yang's avatar
Michael Yang committed
254
			shape[i] = llm.readU64(rso)
255
256
		}

Michael Yang's avatar
Michael Yang committed
257
		tensor := tensor{
258
			name:   name,
Michael Yang's avatar
Michael Yang committed
259
260
			kind:   llm.readU32(rso),
			offset: llm.readU64(rso),
261
			shape:  shape,
Michael Yang's avatar
Michael Yang committed
262
		}
263

Michael Yang's avatar
Michael Yang committed
264
265
		llm.tensors = append(llm.tensors, tensor)
		llm.parameters += tensor.parameters()
266
267
268
269
270
271
	}

	alignment, ok := llm.kv["general.alignment"].(uint32)
	if !ok {
		alignment = 32
	}
272

Michael Yang's avatar
Michael Yang committed
273
	rso.Seek(int64(alignment)-rso.offset%int64(alignment), io.SeekCurrent)
274
	for _, tensor := range llm.tensors {
Michael Yang's avatar
Michael Yang committed
275
		padded := (int64(tensor.size()) + int64(alignment) - 1) & ^(int64(alignment) - 1)
Michael Yang's avatar
Michael Yang committed
276
		rso.Seek(padded, io.SeekCurrent)
277
278
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
279
280
281
	return nil
}

282
func (llm *ggufModel) NumLayers() uint32 {
283
284
285
286
287
	value, exists := llm.kv[fmt.Sprintf("%s.block_count", llm.ModelFamily())]
	if !exists {
		return 0
	}

288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
	return value.(uint32)
}

func (llm *ggufModel) NumHead() uint32 {
	value, exists := llm.kv[fmt.Sprintf("%s.attention.head_count", llm.ModelFamily())]
	if !exists {
		return 0
	}

	return value.(uint32)
}

func (llm *ggufModel) NumEmbed() uint32 {
	value, exists := llm.kv[fmt.Sprintf("%s.embedding_length", llm.ModelFamily())]
	if !exists {
		return 0
	}

	return value.(uint32)
}

func (llm *ggufModel) NumHeadKv() uint32 {
	value, exists := llm.kv[fmt.Sprintf("%s.attention.head_count_kv", llm.ModelFamily())]
	if !exists {
		return 0
	}

	return value.(uint32)
}

Michael Yang's avatar
Michael Yang committed
318
319
320
321
322
323
324
325
326
func (llm *ggufModel) NumCtx() uint32 {
	value, exists := llm.kv[fmt.Sprintf("%s.context_length", llm.ModelFamily())]
	if !exists {
		return 0
	}

	return value.(uint32)
}

327
328
329
330
331
332
333
func (llm *ggufModel) NumGQA() uint32 {
	numHeadKv := llm.NumHeadKv()
	if numHeadKv == 0 {
		return 0
	}

	return llm.NumHead() / numHeadKv
334
335
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
336
func (llm ggufModel) readU8(r io.Reader) uint8 {
Bruce MacDonald's avatar
Bruce MacDonald committed
337
	var u8 uint8
Michael Yang's avatar
ggufv3  
Michael Yang committed
338
	binary.Read(r, llm.bo, &u8)
Bruce MacDonald's avatar
Bruce MacDonald committed
339
340
341
	return u8
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
342
func (llm ggufModel) readI8(r io.Reader) int8 {
Bruce MacDonald's avatar
Bruce MacDonald committed
343
	var i8 int8
Michael Yang's avatar
ggufv3  
Michael Yang committed
344
	binary.Read(r, llm.bo, &i8)
Bruce MacDonald's avatar
Bruce MacDonald committed
345
346
347
	return i8
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
348
func (llm ggufModel) readU16(r io.Reader) uint16 {
Bruce MacDonald's avatar
Bruce MacDonald committed
349
	var u16 uint16
Michael Yang's avatar
ggufv3  
Michael Yang committed
350
	binary.Read(r, llm.bo, &u16)
Bruce MacDonald's avatar
Bruce MacDonald committed
351
352
353
	return u16
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
354
func (llm ggufModel) readI16(r io.Reader) int16 {
Bruce MacDonald's avatar
Bruce MacDonald committed
355
	var i16 int16
Michael Yang's avatar
ggufv3  
Michael Yang committed
356
	binary.Read(r, llm.bo, &i16)
Bruce MacDonald's avatar
Bruce MacDonald committed
357
358
359
	return i16
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
360
func (llm ggufModel) readU32(r io.Reader) uint32 {
Bruce MacDonald's avatar
Bruce MacDonald committed
361
	var u32 uint32
Michael Yang's avatar
ggufv3  
Michael Yang committed
362
	binary.Read(r, llm.bo, &u32)
Bruce MacDonald's avatar
Bruce MacDonald committed
363
364
365
	return u32
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
366
func (llm ggufModel) readI32(r io.Reader) int32 {
Bruce MacDonald's avatar
Bruce MacDonald committed
367
	var i32 int32
Michael Yang's avatar
ggufv3  
Michael Yang committed
368
	binary.Read(r, llm.bo, &i32)
Bruce MacDonald's avatar
Bruce MacDonald committed
369
370
371
	return i32
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
372
func (llm ggufModel) readU64(r io.Reader) uint64 {
Bruce MacDonald's avatar
Bruce MacDonald committed
373
	var u64 uint64
Michael Yang's avatar
ggufv3  
Michael Yang committed
374
	binary.Read(r, llm.bo, &u64)
Bruce MacDonald's avatar
Bruce MacDonald committed
375
376
377
	return u64
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
378
func (llm ggufModel) readI64(r io.Reader) int64 {
Bruce MacDonald's avatar
Bruce MacDonald committed
379
	var i64 int64
Michael Yang's avatar
ggufv3  
Michael Yang committed
380
	binary.Read(r, llm.bo, &i64)
Bruce MacDonald's avatar
Bruce MacDonald committed
381
382
383
	return i64
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
384
func (llm ggufModel) readF32(r io.Reader) float32 {
Bruce MacDonald's avatar
Bruce MacDonald committed
385
	var f32 float32
Michael Yang's avatar
ggufv3  
Michael Yang committed
386
	binary.Read(r, llm.bo, &f32)
Bruce MacDonald's avatar
Bruce MacDonald committed
387
388
389
	return f32
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
390
func (llm ggufModel) readF64(r io.Reader) float64 {
Bruce MacDonald's avatar
Bruce MacDonald committed
391
	var f64 float64
Michael Yang's avatar
ggufv3  
Michael Yang committed
392
	binary.Read(r, llm.bo, &f64)
Bruce MacDonald's avatar
Bruce MacDonald committed
393
394
395
	return f64
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
396
func (llm ggufModel) readBool(r io.Reader) bool {
Bruce MacDonald's avatar
Bruce MacDonald committed
397
	var b bool
Michael Yang's avatar
ggufv3  
Michael Yang committed
398
	binary.Read(r, llm.bo, &b)
Bruce MacDonald's avatar
Bruce MacDonald committed
399
400
401
	return b
}

Michael Yang's avatar
ggufv3  
Michael Yang committed
402
func (llm ggufModel) readStringV1(r io.Reader) (string, error) {
Bruce MacDonald's avatar
Bruce MacDonald committed
403
	var nameLength uint32
Michael Yang's avatar
ggufv3  
Michael Yang committed
404
	binary.Read(r, llm.bo, &nameLength)
Bruce MacDonald's avatar
Bruce MacDonald committed
405
406
407
408
409
410
411
412
413
414
415
416
417

	var b bytes.Buffer
	if _, err := io.CopyN(&b, r, int64(nameLength)); err != nil {
		return "", err
	}

	// gguf v1 strings are null-terminated
	b.Truncate(b.Len() - 1)

	return b.String(), nil
}

func (llm ggufModel) readString(r io.Reader) (string, error) {
418
419
420
421
	if llm.Version == 1 {
		return llm.readStringV1(r)
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
422
	var nameLength uint64
Michael Yang's avatar
ggufv3  
Michael Yang committed
423
	binary.Read(r, llm.bo, &nameLength)
Bruce MacDonald's avatar
Bruce MacDonald committed
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441

	var b bytes.Buffer
	if _, err := io.CopyN(&b, r, int64(nameLength)); err != nil {
		return "", err
	}

	return b.String(), nil
}

func (llm *ggufModel) readArrayV1(r io.Reader) (arr []any, err error) {
	atype := llm.readU32(r)
	n := llm.readU32(r)

	for i := 0; uint32(i) < n; i++ {
		switch atype {
		case ggufTypeUint8:
			arr = append(arr, llm.readU8(r))
		case ggufTypeInt8:
Michael Yang's avatar
Michael Yang committed
442
			arr = append(arr, llm.readI8(r))
Bruce MacDonald's avatar
Bruce MacDonald committed
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
		case ggufTypeUint16:
			arr = append(arr, llm.readU16(r))
		case ggufTypeInt16:
			arr = append(arr, llm.readI16(r))
		case ggufTypeUint32:
			arr = append(arr, llm.readU32(r))
		case ggufTypeInt32:
			arr = append(arr, llm.readI32(r))
		case ggufTypeFloat32:
			arr = append(arr, llm.readF32(r))
		case ggufTypeBool:
			arr = append(arr, llm.readBool(r))
		case ggufTypeString:
			s, err := llm.readStringV1(r)
			if err != nil {
				return nil, err
			}

			arr = append(arr, s)
		default:
			return nil, fmt.Errorf("invalid array type: %d", atype)
		}
	}

	return
}

func (llm *ggufModel) readArray(r io.Reader) (arr []any, err error) {
471
472
473
474
	if llm.Version == 1 {
		return llm.readArrayV1(r)
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
475
476
477
478
479
480
481
482
	atype := llm.readU32(r)
	n := llm.readU64(r)

	for i := 0; uint64(i) < n; i++ {
		switch atype {
		case ggufTypeUint8:
			arr = append(arr, llm.readU8(r))
		case ggufTypeInt8:
Michael Yang's avatar
Michael Yang committed
483
			arr = append(arr, llm.readI8(r))
Bruce MacDonald's avatar
Bruce MacDonald committed
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
		case ggufTypeUint16:
			arr = append(arr, llm.readU16(r))
		case ggufTypeInt16:
			arr = append(arr, llm.readI16(r))
		case ggufTypeUint32:
			arr = append(arr, llm.readU32(r))
		case ggufTypeInt32:
			arr = append(arr, llm.readI32(r))
		case ggufTypeUint64:
			arr = append(arr, llm.readU64(r))
		case ggufTypeInt64:
			arr = append(arr, llm.readI64(r))
		case ggufTypeFloat32:
			arr = append(arr, llm.readF32(r))
		case ggufTypeFloat64:
			arr = append(arr, llm.readF64(r))
		case ggufTypeBool:
			arr = append(arr, llm.readBool(r))
		case ggufTypeString:
			s, err := llm.readString(r)
			if err != nil {
				return nil, err
			}

			arr = append(arr, s)
		default:
			return nil, fmt.Errorf("invalid array type: %d", atype)
		}
	}

	return
}