ggml.go 8.5 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
)

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

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

21
22
type KV map[string]any

Michael Yang's avatar
Michael Yang committed
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
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
48
func (kv KV) FileType() fileType {
Michael Yang's avatar
Michael Yang committed
49
	if u64 := kv.u64("general.file_type"); u64 > 0 {
Michael Yang's avatar
Michael Yang committed
50
		return fileType(uint32(u64))
Michael Yang's avatar
Michael Yang committed
51
52
	}

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

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
65
66
67
68
69
	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
70
71
72
}

func (kv KV) GQA() uint64 {
Michael Yang's avatar
Michael Yang committed
73
	return kv.HeadCount() / kv.HeadCountKV()
Michael Yang's avatar
Michael Yang committed
74
75
76
77
78
79
80
81
82
83
}

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
84
85
86
87
88
89
90
type Tensors []*Tensor

func (ts Tensors) Layers() map[string]Layer {
	layers := make(map[string]Layer)
	for _, t := range ts {
		parts := strings.Split(t.Name, ".")
		if parts[0] == "blk" {
91
92
			// 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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
		}

		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
109
		size += t.Size()
Michael Yang's avatar
Michael Yang committed
110
111
112
113
114
	}

	return size
}

115
type Tensor struct {
Michael Yang's avatar
Michael Yang committed
116
117
118
	Name   string `json:"name"`
	Kind   uint32 `json:"kind"`
	Offset uint64 `json:"-"`
119
120

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

Michael Yang's avatar
Michael Yang committed
123
	io.WriterTo `json:"-"`
124
125
126
}

func (t Tensor) blockSize() uint64 {
127
128
	switch t.Kind {
	case 0, 1, 24, 25, 26, 27, 28, 31: // F32, F16, I8, I16, I32, I64, F64, BF16
129
		return 1
130
	case 2, 3, 8, 9, 20: // Q4_0, Q4_1, Q8_0, Q8_1, IQ4_NL
131
		return 32
132
	default: // All others
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
		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
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
		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
197
198
199
200
201
202
203
204
205
206
207
208
209
	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
210
func (t Tensor) Size() uint64 {
211
212
213
	return t.parameters() * t.typeSize() / t.blockSize()
}

214
215
type container interface {
	Name() string
Michael Yang's avatar
Michael Yang committed
216
	Decode(io.ReadSeeker) (model, error)
217
218
219
}

const (
Bruce MacDonald's avatar
Bruce MacDonald committed
220
	// Magic constant for `ggml` files (unversioned).
221
	FILE_MAGIC_GGML = 0x67676d6c
Bruce MacDonald's avatar
Bruce MacDonald committed
222
	// Magic constant for `ggml` files (versioned, ggmf).
223
	FILE_MAGIC_GGMF = 0x67676d66
Bruce MacDonald's avatar
Bruce MacDonald committed
224
	// Magic constant for `ggml` files (versioned, ggjt).
225
	FILE_MAGIC_GGJT = 0x67676a74
Bruce MacDonald's avatar
Bruce MacDonald committed
226
	// Magic constant for `ggla` files (LoRA adapter).
227
	FILE_MAGIC_GGLA = 0x67676C61
Bruce MacDonald's avatar
Bruce MacDonald committed
228
	// Magic constant for `gguf` files (versioned, gguf)
Michael Yang's avatar
ggufv3  
Michael Yang committed
229
230
	FILE_MAGIC_GGUF_LE = 0x46554747
	FILE_MAGIC_GGUF_BE = 0x47475546
231
232
)

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

Michael Yang's avatar
Michael Yang committed
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
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"
	case FILE_MAGIC_GGUF_LE, FILE_MAGIC_GGUF_BE:
		return "gguf"
	default:
		return ""
	}
}

Michael Yang's avatar
Michael Yang committed
252
func DecodeGGML(rs io.ReadSeeker) (*GGML, int64, error) {
253
	var magic uint32
Michael Yang's avatar
Michael Yang committed
254
	if err := binary.Read(rs, binary.LittleEndian, &magic); err != nil {
Michael Yang's avatar
Michael Yang committed
255
		return nil, 0, err
256
257
258
259
	}

	var c container
	switch magic {
Bruce MacDonald's avatar
Bruce MacDonald committed
260
	case FILE_MAGIC_GGML, FILE_MAGIC_GGMF, FILE_MAGIC_GGJT:
Michael Yang's avatar
Michael Yang committed
261
		return nil, 0, ErrUnsupportedFormat
262
	case FILE_MAGIC_GGLA:
263
		c = &containerGGLA{}
Michael Yang's avatar
ggufv3  
Michael Yang committed
264
	case FILE_MAGIC_GGUF_LE:
265
		c = &containerGGUF{ByteOrder: binary.LittleEndian}
Michael Yang's avatar
ggufv3  
Michael Yang committed
266
	case FILE_MAGIC_GGUF_BE:
267
		c = &containerGGUF{ByteOrder: binary.BigEndian}
268
	default:
Michael Yang's avatar
Michael Yang committed
269
		return nil, 0, errors.New("invalid file magic")
270
271
	}

Michael Yang's avatar
Michael Yang committed
272
	model, err := c.Decode(rs)
Michael Yang's avatar
Michael Yang committed
273
274
275
	if errors.Is(err, io.EOF) {
		// noop
	} else if err != nil {
Michael Yang's avatar
Michael Yang committed
276
		return nil, 0, err
277
278
	}

Michael Yang's avatar
Michael Yang committed
279
280
	offset, err := rs.Seek(0, io.SeekCurrent)
	if err != nil {
Michael Yang's avatar
Michael Yang committed
281
		return nil, 0, err
Michael Yang's avatar
Michael Yang committed
282
283
	}

284
	// final model type
285
286
287
	return &GGML{
		container: c,
		model:     model,
Michael Yang's avatar
Michael Yang committed
288
	}, offset, nil
289
}
Michael Yang's avatar
Michael Yang committed
290

Michael Yang's avatar
Michael Yang committed
291
292
293
294
295
func (llm GGML) GraphSize(context, batch uint64) (partialOffload, fullOffload uint64) {
	embedding := llm.KV().EmbeddingLength()
	heads := llm.KV().HeadCount()
	headsKV := llm.KV().HeadCountKV()
	vocab := uint64(len(llm.KV()["tokenizer.ggml.tokens"].([]any)))
Michael Yang's avatar
Michael Yang committed
296

Michael Yang's avatar
Michael Yang committed
297
298
	layers := llm.Tensors().Layers()

Michael Yang's avatar
Michael Yang committed
299
300
	switch llm.KV().Architecture() {
	case "llama":
Michael Yang's avatar
Michael Yang committed
301
302
303
304
305
306
307
		fullOffload = 4 * batch * (1 + 4*embedding + context*(1+heads))

		partialOffload = 4 * batch * embedding
		partialOffload += max(
			4*batch*(1+embedding+max(context, embedding))+embedding*embedding*9/16+4*context*(batch*heads+embedding/heads*headsKV),
			4*batch*(embedding+vocab)+embedding*vocab*105/128,
		)
Michael Yang's avatar
Michael Yang committed
308

Michael Yang's avatar
Michael Yang committed
309
310
311
312
		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
313
				3*ffnGateExpsWeight.Size()+4*batch*(2*ff+headsKV+embedding+context+embedding/heads*headsKV),
Michael Yang's avatar
Michael Yang committed
314
315
316
317
				4*(context*batch*heads+context*embedding/heads*headsKV+batch*1024+embedding/heads*headsKV*batch),
			)
		} else if ffnGateWeight, ok := layers["blk.0"]["ffn_gate.0.weight"]; ok {
			// mixtral 8x7b
Michael Yang's avatar
Michael Yang committed
318
319
320
321
322
323
324
			ffnGateWeight1 := ffnGateWeight.Shape[1]
			fullOffload = 4 * batch * (2 + 3*embedding + context*(1+heads) + 2*headsKV + ffnGateWeight1)
			partialOffload = max(
				4*batch*(3+embedding/heads*headsKV+embedding+context*(1+heads)+ffnGateWeight1)+(embedding*embedding+3*embedding*headsKV*ffnGateWeight1)*9/16,
				4*batch*(1+2*embedding+context*(1+heads))+embedding*(6*context*headsKV/heads+embedding*9/16),
			)
		}
Michael Yang's avatar
Michael Yang committed
325
326
327
328
329
330
331
332
333
334
335
	case "gemma":
		fullOffload = 4 * batch * (embedding + vocab)
		partialOffload = 4*batch*(2*embedding+vocab+1) + embedding*vocab*105/128
	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
336
			4*batch*(1+2*embedding+context*(1+heads))+4*embedding*context+embedding*embedding*9/16,
Michael Yang's avatar
Michael Yang committed
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
		)
	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
353

Michael Yang's avatar
Michael Yang committed
354
355
356
357
		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
358
359
360
361
362
363
	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
364
365
	}

Michael Yang's avatar
Michael Yang committed
366
	return
Michael Yang's avatar
Michael Yang committed
367
}