ollama/llm/ggml.go

243 lines
4.5 KiB
Go
Raw Normal View History

2023-07-21 20:33:56 +00:00
package llm
import (
"encoding/binary"
"errors"
"io"
)
type GGML struct {
container
model
2023-11-24 19:57:20 +00:00
Size int64
}
2023-07-21 20:33:56 +00:00
const (
fileTypeF32 uint32 = iota
fileTypeF16
fileTypeQ4_0
fileTypeQ4_1
fileTypeQ4_1_F16
fileTypeQ8_0 uint32 = iota + 2
fileTypeQ5_0
fileTypeQ5_1
fileTypeQ2_K
fileTypeQ3_K_S
fileTypeQ3_K_M
fileTypeQ3_K_L
fileTypeQ4_K_S
fileTypeQ4_K_M
fileTypeQ5_K_S
fileTypeQ5_K_M
fileTypeQ6_K
2024-02-21 00:06:29 +00:00
fileTypeIQ2_XXS
fileTypeIQ2_XS
fileTypeQ2_K_S
fileTypeQ3_K_XS
fileTypeIQ3_XXS
2023-07-21 20:33:56 +00:00
)
func fileType(fileType uint32) string {
switch fileType {
case fileTypeF32:
return "F32"
case fileTypeF16:
return "F16"
case fileTypeQ4_0:
return "Q4_0"
case fileTypeQ4_1:
return "Q4_1"
case fileTypeQ4_1_F16:
return "Q4_1_F16"
case fileTypeQ8_0:
return "Q8_0"
case fileTypeQ5_0:
return "Q5_0"
case fileTypeQ5_1:
return "Q5_1"
case fileTypeQ2_K:
return "Q2_K"
case fileTypeQ3_K_S:
return "Q3_K_S"
case fileTypeQ3_K_M:
return "Q3_K_M"
case fileTypeQ3_K_L:
return "Q3_K_L"
case fileTypeQ4_K_S:
return "Q4_K_S"
case fileTypeQ4_K_M:
return "Q4_K_M"
case fileTypeQ5_K_S:
return "Q5_K_S"
case fileTypeQ5_K_M:
return "Q5_K_M"
case fileTypeQ6_K:
return "Q6_K"
2024-02-21 00:06:29 +00:00
case fileTypeIQ2_XXS:
return "IQ2_XXS"
case fileTypeIQ2_XS:
return "IQ2_XS"
case fileTypeQ2_K_S:
return "Q2_K_S"
case fileTypeQ3_K_XS:
return "Q3_K_XS"
case fileTypeIQ3_XXS:
return "IQ3_XXS"
2023-08-17 18:37:27 +00:00
default:
2023-10-03 02:52:25 +00:00
return "unknown"
2023-08-17 18:37:27 +00:00
}
}
2023-07-21 20:33:56 +00:00
2023-08-17 18:37:27 +00:00
type model interface {
ModelFamily() string
ModelType() string
FileType() string
NumLayers() uint32
NumGQA() uint32
NumEmbed() uint32
NumHead() uint32
NumHeadKv() uint32
2024-01-12 22:54:01 +00:00
NumCtx() uint32
2023-07-21 20:33:56 +00:00
}
type KV map[string]any
type Tensor struct {
Name string
Kind uint32
Offset uint64
// Shape is the number of elements in each dimension
Shape []uint64
io.WriterTo
}
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
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
return 2 + 3*blockSize/8
default:
return 0
}
}
func (t Tensor) parameters() uint64 {
var count uint64 = 1
for _, n := range t.Shape {
count *= n
}
return count
}
func (t Tensor) size() uint64 {
return t.parameters() * t.typeSize() / t.blockSize()
}
2023-07-21 20:33:56 +00:00
type container interface {
Name() string
2024-03-09 20:28:36 +00:00
Decode(io.ReadSeeker) (model, error)
2023-07-21 20:33:56 +00:00
}
const (
2023-09-07 17:55:37 +00:00
// Magic constant for `ggml` files (unversioned).
2023-07-21 20:33:56 +00:00
FILE_MAGIC_GGML = 0x67676d6c
2023-09-07 17:55:37 +00:00
// Magic constant for `ggml` files (versioned, ggmf).
2023-07-21 20:33:56 +00:00
FILE_MAGIC_GGMF = 0x67676d66
2023-09-07 17:55:37 +00:00
// Magic constant for `ggml` files (versioned, ggjt).
2023-07-21 20:33:56 +00:00
FILE_MAGIC_GGJT = 0x67676a74
2023-09-07 17:55:37 +00:00
// Magic constant for `ggla` files (LoRA adapter).
2023-07-21 20:33:56 +00:00
FILE_MAGIC_GGLA = 0x67676C61
2023-09-07 17:55:37 +00:00
// Magic constant for `gguf` files (versioned, gguf)
FILE_MAGIC_GGUF_LE = 0x46554747
FILE_MAGIC_GGUF_BE = 0x47475546
2023-07-21 20:33:56 +00:00
)
var ErrUnsupportedFormat = errors.New("unsupported model format")
2024-03-09 20:28:36 +00:00
func DecodeGGML(rs io.ReadSeeker) (*GGML, error) {
2023-11-24 19:57:20 +00:00
var magic uint32
2024-03-09 20:28:36 +00:00
if err := binary.Read(rs, binary.LittleEndian, &magic); err != nil {
2023-11-24 19:57:20 +00:00
return nil, err
}
var c container
switch magic {
case FILE_MAGIC_GGML, FILE_MAGIC_GGMF, FILE_MAGIC_GGJT:
return nil, ErrUnsupportedFormat
2023-07-21 20:33:56 +00:00
case FILE_MAGIC_GGLA:
c = &containerGGLA{}
case FILE_MAGIC_GGUF_LE:
c = &containerGGUF{ByteOrder: binary.LittleEndian}
case FILE_MAGIC_GGUF_BE:
c = &containerGGUF{ByteOrder: binary.BigEndian}
2023-07-21 20:33:56 +00:00
default:
return nil, errors.New("invalid file magic")
}
2024-03-09 20:28:36 +00:00
model, err := c.Decode(rs)
2024-03-08 23:38:53 +00:00
if errors.Is(err, io.EOF) {
// noop
} else if err != nil {
2023-07-21 20:33:56 +00:00
return nil, err
}
2024-03-09 20:28:36 +00:00
offset, err := rs.Seek(0, io.SeekCurrent)
if err != nil {
return nil, err
}
2023-07-21 20:33:56 +00:00
// final model type
2023-11-24 19:57:20 +00:00
return &GGML{
container: c,
model: model,
2024-03-09 20:28:36 +00:00
Size: offset,
2023-11-24 19:57:20 +00:00
}, nil
}