Fix gemma2 2b conversion (ollama#6645)

anuraaga · Sep 6, 2024 · 608e87b · 608e87b
1 parent 48685c6
commit 608e87b
Show file tree

Hide file tree

Showing 4 changed files with 376 additions and 19 deletions.
diff --git a/convert/convert_gemma2.go b/convert/convert_gemma2.go
@@ -34,10 +34,20 @@ func (p *gemma2Model) KV(t *Tokenizer) llm.KV {
 }
 
 func (p *gemma2Model) Replacements() []string {
-	return append(
-		p.gemmaModel.Replacements(),
+	return []string{
+		"model.embed_tokens", "token_embd",
+		"model.norm", "output_norm",
+		"model.layers", "blk",
+		"input_layernorm", "attn_norm",
+		"self_attn.q_proj", "attn_q",
+		"self_attn.k_proj", "attn_k",
+		"self_attn.v_proj", "attn_v",
+		"self_attn.o_proj", "attn_output",
+		"mlp.gate_proj", "ffn_gate",
+		"mlp.down_proj", "ffn_down",
+		"mlp.up_proj", "ffn_up",
 		"post_attention_layernorm", "post_attention_norm",
 		"pre_feedforward_layernorm", "ffn_norm",
 		"post_feedforward_layernorm", "post_ffw_norm",
-	)
+	}
 }
diff --git a/convert/convert_test.go b/convert/convert_test.go
@@ -15,13 +15,20 @@ import (
 	"os"
 	"path/filepath"
 	"slices"
+	"strings"
 	"testing"
 
 	"golang.org/x/exp/maps"
 
 	"github.com/ollama/ollama/llm"
 )
 
+type tensorData struct {
+	Offsets []int  `json:"data_offsets"`
+	Type    string `json:"dtype"`
+	Shape   []int  `json:"shape"`
+}
+
 func convertFull(t *testing.T, fsys fs.FS) (*os.File, llm.KV, llm.Tensors) {
 	t.Helper()
 
@@ -96,6 +103,7 @@ func TestConvertModel(t *testing.T) {
 		"Mistral-7B-Instruct-v0.2",
 		"Mixtral-8x7B-Instruct-v0.1",
 		"gemma-2b-it",
+		"gemma-2-2b-it",
 		// microsoft/Phi-3-mini-128-instruct@d548c233192db00165d842bf8edff054bb3212f8
 		"Phi-3-mini-128k-instruct",
 		"all-MiniLM-L6-v2",
@@ -140,31 +148,48 @@ func TestConvertModel(t *testing.T) {
 	}
 }
 
-func TestConvertInvalidDatatype(t *testing.T) {
+func TestConvertInvalidTensorNames(t *testing.T) {
 	f, err := os.CreateTemp(t.TempDir(), "testmodel")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer f.Close()
 
 	tempDir := t.TempDir()
-	generateSafetensorTestData(t, tempDir)
+
+	td := map[string]*tensorData{}
+	offset := 4096
+
+	td["model.layers.0.self_attn.q_proj.weight"] = &tensorData{
+		Offsets: []int{0, offset},
+		Type:    "F32",
+		Shape:   []int{4096, 4096},
+	}
+	td["blk.0.attn_q.weight"] = &tensorData{
+		Offsets: []int{offset, offset * 2},
+		Type:    "F32",
+		Shape:   []int{4096, 4096},
+	}
+	generateSafetensorTestData(t, tempDir, td)
 
 	err = ConvertModel(os.DirFS(tempDir), f)
-	if err == nil || err.Error() != "unsupported safetensors model" {
+	if err == nil || !strings.HasPrefix(err.Error(), "duplicate tensor name") {
 		t.Errorf("expected error but didn't get one")
 	}
 }
 
-func generateSafetensorTestData(t *testing.T, tempDir string) {
-	type tensorData struct {
-		Offsets []int  `json:"data_offsets"`
-		Type    string `json:"dtype"`
-		Shape   []int  `json:"shape"`
+func TestConvertInvalidDatatype(t *testing.T) {
+	f, err := os.CreateTemp(t.TempDir(), "testmodel")
+	if err != nil {
+		t.Fatal(err)
 	}
-	offset := 4096 * 14336
+	defer f.Close()
+
+	tempDir := t.TempDir()
 
 	td := map[string]*tensorData{}
+	offset := 4096 * 14336
+
 	td["model.layers.0.mlp.down_proj.weight"] = &tensorData{
 		Offsets: []int{0, offset},
 		Type:    "I8",
@@ -175,8 +200,16 @@ func generateSafetensorTestData(t *testing.T, tempDir string) {
 		Type:    "U8",
 		Shape:   []int{},
 	}
+	generateSafetensorTestData(t, tempDir, td)
 
-	data, err := json.Marshal(td)
+	err = ConvertModel(os.DirFS(tempDir), f)
+	if err == nil || err.Error() != "unsupported safetensors model" {
+		t.Errorf("expected error but didn't get one")
+	}
+}
+
+func generateSafetensorTestData(t *testing.T, tempDir string, tensorData map[string]*tensorData) {
+	data, err := json.Marshal(tensorData)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -322,11 +355,6 @@ func TestConvertAdapter(t *testing.T) {
 }
 
 func generateLoraTestData(t *testing.T, tempDir string) {
-	type tensorData struct {
-		Offsets []int  `json:"data_offsets"`
-		Type    string `json:"dtype"`
-		Shape   []int  `json:"shape"`
-	}
 	offset := 4096 * 8 * 4
 
 	td := map[string]*tensorData{"__metadata__": nil}

diff --git a/convert/reader_safetensors.go b/convert/reader_safetensors.go
@@ -49,20 +49,27 @@ func parseSafetensors(fsys fs.FS, replacer *strings.Replacer, ps ...string) ([]T
 		keys := maps.Keys(headers)
 		slices.Sort(keys)
 
+		names := make(map[string]struct{}, len(keys))
+
 		for _, key := range keys {
 			if value := headers[key]; value.Type != "" {
 				// bitsandbytes quantized models are unsupported
 				if len(value.Shape) == 0 {
 					return nil, errors.New("unsupported safetensors model")
 				}
+				ggufName := replacer.Replace(key)
+				if _, ok := names[ggufName]; ok {
+					return nil, fmt.Errorf("duplicate tensor name '%s' was found for this model", ggufName)
+				}
+				names[ggufName] = struct{}{}
 				ts = append(ts, safetensor{
 					fs:     fsys,
 					path:   p,
 					dtype:  value.Type,
 					offset: safetensorsPad(n, value.Offsets[0]),
 					size:   safetensorsPad(n, value.Offsets[1]) - safetensorsPad(n, value.Offsets[0]),
 					tensorBase: &tensorBase{
-						name:  replacer.Replace(key),
+						name:  ggufName,
 						shape: value.Shape,
 					},
 				})