convert: capture head_dim for mistral (ollama#5818)

jay74jung · Jul 22, 2024 · d835368 · d835368
1 parent 5784c05
commit d835368
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/convert/mistral.go b/convert/mistral.go
@@ -71,6 +71,11 @@ func (m *MistralModel) WriteGGUF(ws io.WriteSeeker) error {
 		"tokenizer.ggml.unknown_token_id": uint32(0),
 	}
 
+	if m.Params.HeadDimension > 0 {
+		kv["llama.attention.key_length"] = uint32(m.Params.HeadDimension)
+		kv["llama.attention.value_length"] = uint32(m.Params.HeadDimension)
+	}
+
 	return llm.NewGGUFV3(m.Params.ByteOrder).Encode(ws, kv, m.Tensors)
 }