matt/ollama37
1
0
Fork 0
mirror of https://github.com/dogkeeper886/ollama37.git synced 2025-12-10 15:57:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add gemma safetensors conversion (#3250)

Browse Source 
Co-authored-by: Michael Yang <mxyng@pm.me>
This commit is contained in:
Patrick Devine
2024-03-28 18:54:01 -07:00
committed by GitHub
parent 97ae517fbf
commit 5a5efee46b
11 changed files with 949 additions and 833 deletions
Download Patch File Download Diff File Expand all files Collapse all files

392
convert/convert.go
View File

@@ -12,8 +12,13 @@ import (
"path/filepath"
"regexp"
"slices"
"strings"
"github.com/d4l3k/go-bfloat16"
"github.com/mitchellh/mapstructure"
"github.com/pdevine/tensor"
"github.com/pdevine/tensor/native"
"github.com/x448/float16"
"google.golang.org/protobuf/proto"
"github.com/ollama/ollama/convert/sentencepiece"
@@ -33,6 +38,15 @@ type Params struct {
RopeFreqBase float64 `json:"rope_theta"`
BoSTokenID int `json:"bos_token_id"`
EoSTokenID int `json:"eos_token_id"`
HeadDimension int `json:"head_dim"`
PaddingTokenID int `json:"pad_token_id"`
ByteOrder
}
type ByteOrder interface {
binary.ByteOrder
binary.AppendByteOrder
}
type MetaData struct {
@@ -41,27 +55,29 @@ type MetaData struct {
Offsets []int `mapstructure:"data_offsets"`
}
func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
func ReadSafeTensors(fn string, offset uint64, params *Params) ([]llm.Tensor, uint64, error) {
f, err := os.Open(fn)
if err != nil {
return []llm.Tensor{}, 0, err
return nil, 0, err
}
defer f.Close()
var jsonSize uint64
binary.Read(f, binary.LittleEndian, &jsonSize)
if err := binary.Read(f, binary.LittleEndian, &jsonSize); err != nil {
return nil, 0, err
}
buf := make([]byte, jsonSize)
_, err = io.ReadFull(f, buf)
if err != nil {
return []llm.Tensor{}, 0, err
return nil, 0, err
}
d := json.NewDecoder(bytes.NewBuffer(buf))
d.UseNumber()
var parsed map[string]interface{}
if err = d.Decode(&parsed); err != nil {
return []llm.Tensor{}, 0, err
return nil, 0, err
}
var keys []string
@@ -78,7 +94,7 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
vals := parsed[k].(map[string]interface{})
var data MetaData
if err = mapstructure.Decode(vals, &data); err != nil {
return []llm.Tensor{}, 0, err
return nil, 0, err
}
var size uint64
@@ -100,7 +116,7 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
ggufName, err := GetTensorName(k)
if err != nil {
slog.Error("%v", err)
return []llm.Tensor{}, 0, err
return nil, 0, err
}
shape := []uint64{0, 0, 0, 0}
@@ -109,14 +125,24 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
}
t := llm.Tensor{
Name: ggufName,
Kind: kind,
Offset: offset,
Shape: shape[:],
FileName: fn,
OffsetPadding: 8 + jsonSize,
FileOffsets: []uint64{uint64(data.Offsets[0]), uint64(data.Offsets[1])},
Name: ggufName,
Kind: kind,
Offset: offset,
Shape: shape[:],
}
t.WriterTo = safetensorWriterTo{
t: &t,
params: params,
bo: params.ByteOrder,
headCount: uint32(params.AttentionHeads),
headCountKV: uint32(params.KeyValHeads),
filename: fn,
start: uint64(data.Offsets[0]),
end: uint64(data.Offsets[1]),
padding: 8 + jsonSize,
}
slog.Debug(fmt.Sprintf("%v", t))
tensors = append(tensors, t)
offset += size
@@ -124,21 +150,21 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
return tensors, offset, nil
}
func GetSafeTensors(dirpath string) ([]llm.Tensor, error) {
func GetSafeTensors(dirpath string, params *Params) ([]llm.Tensor, error) {
var tensors []llm.Tensor
files, err := filepath.Glob(filepath.Join(dirpath, "/model-*.safetensors"))
if err != nil {
return []llm.Tensor{}, err
return nil, err
}
var offset uint64
for _, f := range files {
var t []llm.Tensor
var err error
t, offset, err = ReadSafeTensors(f, offset)
t, offset, err = ReadSafeTensors(f, offset, params)
if err != nil {
slog.Error("%v", err)
return []llm.Tensor{}, err
return nil, err
}
tensors = append(tensors, t...)
}
@@ -160,6 +186,7 @@ func GetParams(dirpath string) (*Params, error) {
return nil, err
}
params.ByteOrder = binary.LittleEndian
return &params, nil
}
@@ -171,7 +198,7 @@ type Vocab struct {
Types []int32
}
func LoadTokens(dirpath string) (*Vocab, error) {
func LoadTokens(dirpath string, params *Params) (*Vocab, error) {
slog.Info(fmt.Sprintf("reading vocab from %s", filepath.Join(dirpath, "tokenizer.model")))
in, err := os.ReadFile(filepath.Join(dirpath, "tokenizer.model"))
if err != nil {
@@ -196,6 +223,14 @@ func LoadTokens(dirpath string) (*Vocab, error) {
v.Tokens = append(v.Tokens, p.GetPiece())
v.Scores = append(v.Scores, p.GetScore())
t := p.GetType()
switch t {
case sentencepiece.ModelProto_SentencePiece_UNKNOWN:
case sentencepiece.ModelProto_SentencePiece_CONTROL:
case sentencepiece.ModelProto_SentencePiece_UNUSED:
case sentencepiece.ModelProto_SentencePiece_BYTE:
default:
t = sentencepiece.ModelProto_SentencePiece_NORMAL
}
v.Types = append(v.Types, int32(t))
}
@@ -243,6 +278,16 @@ func LoadTokens(dirpath string) (*Vocab, error) {
}
slog.Info(fmt.Sprintf("vocab size w/ extra tokens: %d", len(v.Tokens)))
if params.VocabSize > len(v.Tokens) {
missingTokens := params.VocabSize - len(v.Tokens)
slog.Warn(fmt.Sprintf("vocab is missing %d tokens", missingTokens))
for cnt := 0; cnt < missingTokens; cnt++ {
v.Tokens = append(v.Tokens, fmt.Sprintf("<dummy%05d>", cnt+1))
v.Scores = append(v.Scores, -1)
v.Types = append(v.Types, int32(llm.GGUFTokenUserDefined))
}
}
return v, nil
}
@@ -279,42 +324,287 @@ func GetTensorName(n string) (string, error) {
return "", fmt.Errorf("couldn't find a layer name for '%s'", n)
}
func WriteGGUF(name string, tensors []llm.Tensor, params *Params, vocab *Vocab) (string, error) {
c := llm.ContainerGGUF{
ByteOrder: binary.LittleEndian,
type safetensorWriterTo struct {
t *llm.Tensor
params *Params
bo ByteOrder
headCount uint32
headCountKV uint32
filename string
start, end, padding uint64
}
func (r safetensorWriterTo) addOnes(data []float32) ([]float32, error) {
n := tensor.New(tensor.WithShape(int(r.t.Shape[0])), tensor.WithBacking(data))
ones := tensor.Ones(tensor.Float32, int(r.t.Shape[0]))
var err error
n, err = n.Add(ones)
if err != nil {
return []float32{}, err
}
m := llm.NewGGUFModel(&c)
m.Tensors = tensors
m.KV["general.architecture"] = "llama"
m.KV["general.name"] = name
m.KV["llama.context_length"] = uint32(params.ContextSize)
m.KV["llama.embedding_length"] = uint32(params.HiddenSize)
m.KV["llama.block_count"] = uint32(params.HiddenLayers)
m.KV["llama.feed_forward_length"] = uint32(params.IntermediateSize)
m.KV["llama.rope.dimension_count"] = uint32(128)
m.KV["llama.attention.head_count"] = uint32(params.AttentionHeads)
m.KV["llama.attention.head_count_kv"] = uint32(params.KeyValHeads)
m.KV["llama.attention.layer_norm_rms_epsilon"] = float32(params.NormEPS)
m.KV["llama.rope.freq_base"] = float32(params.RopeFreqBase)
m.KV["general.file_type"] = uint32(1)
m.KV["tokenizer.ggml.model"] = "llama"
newN, err := native.SelectF32(n, 0)
if err != nil {
return []float32{}, err
}
m.KV["tokenizer.ggml.tokens"] = vocab.Tokens
m.KV["tokenizer.ggml.scores"] = vocab.Scores
m.KV["tokenizer.ggml.token_type"] = vocab.Types
var fullTensor []float32
for _, v := range newN {
fullTensor = append(fullTensor, v...)
}
m.KV["tokenizer.ggml.bos_token_id"] = uint32(params.BoSTokenID)
m.KV["tokenizer.ggml.eos_token_id"] = uint32(params.EoSTokenID)
m.KV["tokenizer.ggml.unknown_token_id"] = uint32(0)
m.KV["tokenizer.ggml.add_bos_token"] = true
m.KV["tokenizer.ggml.add_eos_token"] = false
return fullTensor, nil
}
// llamacpp sets the chat template, however we don't need to set it since we pass it in through a layer
// m.KV["tokenizer.chat_template"] = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}" // XXX removeme
func (r safetensorWriterTo) repack(data []uint16, heads int) ([]uint16, error) {
n := tensor.New(tensor.WithShape(int(r.t.Shape[0]), int(r.t.Shape[1])), tensor.WithBacking(data))
origShape := n.Shape().Clone()
c.V3.NumTensor = uint64(len(tensors))
c.V3.NumKV = uint64(len(m.KV))
// reshape the tensor and swap axes 1 and 2 to unpack the layer for gguf
if err := n.Reshape(heads, 2, origShape[0]/heads/2, origShape[1]); err != nil {
return nil, err
}
if err := n.T(0, 2, 1, 3); err != nil {
return nil, err
}
if err := n.Reshape(origShape...); err != nil {
return nil, err
}
if err := n.Transpose(); err != nil {
return nil, err
}
newN, err := native.SelectU16(n, 1)
if err != nil {
return nil, err
}
var fullTensor []uint16
for _, v := range newN {
fullTensor = append(fullTensor, v...)
}
return fullTensor, nil
}
func (r safetensorWriterTo) WriteTo(w io.Writer) (n int64, err error) {
arch, err := getArchFromParams(r.params)
if err != nil {
return 0, err
}
f, err := os.Open(r.filename)
if err != nil {
return 0, err
}
defer f.Close()
if _, err = f.Seek(int64(r.padding+r.start), 0); err != nil {
return 0, err
}
switch arch {
case "llama":
pattern := `^blk\.[0-9]+\.attn_(?P<layer>q|k)\.weight$`
re, err := regexp.Compile(pattern)
if err != nil {
return 0, err
}
matches := re.FindAllStringSubmatch(r.t.Name, -1)
if len(matches) > 0 {
layerSize := r.end - r.start
var err error
tData := make([]uint16, layerSize/2)
if err = binary.Read(f, r.bo, tData); err != nil {
return 0, err
}
layerType := matches[0][re.SubexpIndex("layer")]
var heads uint32
switch layerType {
case "q":
heads = r.headCount
case "k":
heads = r.headCountKV
if heads == 0 {
heads = r.headCount
}
}
tData, err = r.repack(tData, int(heads))
if err != nil {
return 0, err
}
var buf []byte
for _, n := range tData {
buf = r.bo.AppendUint16(buf, n)
}
tempBuf := make([]uint16, len(tData))
tDataF32 := bfloat16.DecodeFloat32(buf)
for cnt, v := range tDataF32 {
tDataF16 := float16.Fromfloat32(v)
tempBuf[cnt] = uint16(tDataF16)
}
if err = binary.Write(w, r.bo, tempBuf); err != nil {
return 0, err
}
return 0, nil
}
case "gemma":
if strings.HasSuffix(r.t.Name, "norm.weight") {
slog.Debug(fmt.Sprintf("converting '%s'", r.t.Name))
data := make([]byte, r.end-r.start)
if err = binary.Read(f, r.bo, data); err != nil {
return 0, err
}
tDataF32 := bfloat16.DecodeFloat32(data)
var err error
tDataF32, err = r.addOnes(tDataF32)
if err != nil {
return 0, err
}
if err := binary.Write(w, r.bo, tDataF32); err != nil {
return 0, err
}
return 0, nil
}
}
remaining := r.end - r.start
bufSize := uint64(10240)
var finished bool
for {
data := make([]byte, min(bufSize, remaining))
b, err := io.ReadFull(f, data)
remaining -= uint64(b)
if err == io.EOF || remaining <= 0 {
finished = true
} else if err != nil {
return 0, err
}
// convert bfloat16 -> ieee float32
tDataF32 := bfloat16.DecodeFloat32(data)
switch r.t.Kind {
case 0:
if err := binary.Write(w, r.bo, tDataF32); err != nil {
return 0, err
}
case 1:
// convert float32 -> float16
tempBuf := make([]uint16, len(data)/2)
for cnt, v := range tDataF32 {
tDataF16 := float16.Fromfloat32(v)
tempBuf[cnt] = uint16(tDataF16)
}
if err := binary.Write(w, binary.LittleEndian, tempBuf); err != nil {
return 0, err
}
}
if finished {
break
}
}
return 0, nil
}
func getArchFromParams(params *Params) (string, error) {
var arch string
switch len(params.Architectures) {
case 0:
return "", fmt.Errorf("No architecture specified to convert")
case 1:
switch params.Architectures[0] {
case "MistralForCausalLM":
arch = "llama"
case "GemmaForCausalLM":
arch = "gemma"
default:
return "", fmt.Errorf("Models based on '%s' are not yet supported", params.Architectures[0])
}
default:
return "", fmt.Errorf("Multimodal models are not yet supported")
}
return arch, nil
}
func WriteGGUF(name string, tensors []llm.Tensor, params *Params, vocab *Vocab) (string, error) {
arch, err := getArchFromParams(params)
if err != nil {
return "", err
}
kv := llm.KV{
"general.architecture": arch,
"general.name": name,
}
switch arch {
case "llama":
kv["llama.context_length"] = uint32(params.ContextSize)
kv["llama.embedding_length"] = uint32(params.HiddenSize)
kv["llama.block_count"] = uint32(params.HiddenLayers)
kv["llama.feed_forward_length"] = uint32(params.IntermediateSize)
kv["llama.rope.dimension_count"] = uint32(params.HiddenSize / params.AttentionHeads)
slog.Debug(fmt.Sprintf("rope dim count = %d", kv["llama.rope.dimension_count"]))
kv["llama.attention.head_count"] = uint32(params.AttentionHeads)
kv["llama.attention.head_count_kv"] = uint32(params.KeyValHeads)
kv["llama.attention.layer_norm_rms_epsilon"] = float32(params.NormEPS)
kv["llama.rope.freq_base"] = float32(params.RopeFreqBase)
case "gemma":
kv["gemma.context_length"] = uint32(params.ContextSize)
kv["gemma.embedding_length"] = uint32(params.HiddenSize)
kv["gemma.block_count"] = uint32(params.HiddenLayers)
kv["gemma.feed_forward_length"] = uint32(params.IntermediateSize)
kv["gemma.attention.head_count"] = uint32(params.AttentionHeads)
kv["gemma.attention.head_count_kv"] = uint32(params.KeyValHeads)
kv["gemma.attention.layer_norm_rms_epsilon"] = float32(params.NormEPS)
kv["gemma.attention.key_length"] = uint32(params.HeadDimension)
kv["gemma.attention.value_length"] = uint32(params.HeadDimension)
}
kv["general.file_type"] = uint32(1)
kv["tokenizer.ggml.model"] = "llama"
kv["tokenizer.ggml.tokens"] = vocab.Tokens
kv["tokenizer.ggml.scores"] = vocab.Scores
kv["tokenizer.ggml.token_type"] = vocab.Types
kv["tokenizer.ggml.bos_token_id"] = uint32(params.BoSTokenID)
kv["tokenizer.ggml.eos_token_id"] = uint32(params.EoSTokenID)
switch arch {
case "llama":
kv["tokenizer.ggml.unknown_token_id"] = uint32(0)
case "gemma":
kv["tokenizer.ggml.padding_token_id"] = uint32(params.PaddingTokenID)
kv["tokenizer.ggml.unknown_token_id"] = uint32(3)
}
kv["tokenizer.ggml.add_bos_token"] = true
kv["tokenizer.ggml.add_eos_token"] = false
f, err := os.CreateTemp("", "ollama-gguf")
if err != nil {
@@ -322,8 +612,8 @@ func WriteGGUF(name string, tensors []llm.Tensor, params *Params, vocab *Vocab)
}
defer f.Close()
err = m.Encode(f)
if err != nil {
m := llm.NewGGUFV3(params.ByteOrder)
if err := m.Encode(f, kv, tensors); err != nil {
return "", err
}