diff --git a/README.md b/README.md index 86faaeac..48b77323 100644 --- a/README.md +++ b/README.md @@ -298,6 +298,7 @@ See the [API documentation](./docs/api.md) for all endpoints. ### Web & Desktop - [Open WebUI](https://github.com/open-webui/open-webui) +- [SwiftChat (macOS with ReactNative)](https://github.com/aws-samples/swift-chat) - [Enchanted (macOS native)](https://github.com/AugustDev/enchanted) - [Hollama](https://github.com/fmaclen/hollama) - [Lollms-Webui](https://github.com/ParisNeo/lollms-webui) @@ -455,6 +456,7 @@ See the [API documentation](./docs/api.md) for all endpoints. ### Apple Vision Pro +- [SwiftChat](https://github.com/aws-samples/swift-chat) (Cross-platform AI chat app supporting Apple Vision Pro via "Designed for iPad") - [Enchanted](https://github.com/AugustDev/enchanted) ### Database @@ -532,6 +534,7 @@ See the [API documentation](./docs/api.md) for all endpoints. ### Mobile +- [SwiftChat](https://github.com/aws-samples/swift-chat) (Lightning-fast Cross-platform AI chat app with native UI for Android, iOS and iPad) - [Enchanted](https://github.com/AugustDev/enchanted) - [Maid](https://github.com/Mobile-Artificial-Intelligence/maid) - [Ollama App](https://github.com/JHubi1/ollama-app) (Modern and easy-to-use multi-platform client for Ollama) @@ -583,12 +586,14 @@ See the [API documentation](./docs/api.md) for all endpoints. - [TextLLaMA](https://github.com/adarshM84/TextLLaMA) A Chrome Extension that helps you write emails, correct grammar, and translate into any language - [Simple-Discord-AI](https://github.com/zyphixor/simple-discord-ai) - [LLM Telegram Bot](https://github.com/innightwolfsleep/llm_telegram_bot) (telegram bot, primary for RP. Oobabooga-like buttons, [A1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui) API integration e.t.c) +- [mcp-llm](https://github.com/sammcj/mcp-llm) (MCP Server to allow LLMs to call other LLMs) ### Supported backends - [llama.cpp](https://github.com/ggerganov/llama.cpp) project founded by Georgi Gerganov. ### Observability +- [Opik](https://www.comet.com/docs/opik/cookbook/ollama) is an open-source platform to debug, evaluate, and monitor your LLM applications, RAG systems, and agentic workflows with comprehensive tracing, automated evaluations, and production-ready dashboards. Opik supports native intergration to Ollama. - [Lunary](https://lunary.ai/docs/integrations/ollama) is the leading open-source LLM observability platform. It provides a variety of enterprise-grade features such as real-time analytics, prompt templates management, PII masking, and comprehensive agent tracing. - [OpenLIT](https://github.com/openlit/openlit) is an OpenTelemetry-native tool for monitoring Ollama Applications & GPUs using traces and metrics. - [HoneyHive](https://docs.honeyhive.ai/integrations/ollama) is an AI observability and evaluation platform for AI agents. Use HoneyHive to evaluate agent performance, interrogate failures, and monitor quality in production. diff --git a/convert/convert.go b/convert/convert.go index 015303e7..7b9fe31f 100644 --- a/convert/convert.go +++ b/convert/convert.go @@ -13,8 +13,13 @@ import ( ) type ModelParameters struct { - Architectures []string `json:"architectures"` - VocabSize uint32 `json:"vocab_size"` + Architectures []string `json:"architectures"` + VocabSize uint32 `json:"vocab_size"` + TextModel TextParameters `json:"text_config"` +} + +type TextParameters struct { + VocabSize uint32 `json:"vocab_size"` } type AdapterParameters struct { @@ -185,6 +190,8 @@ func ConvertModel(fsys fs.FS, ws io.WriteSeeker) error { conv = &gemmaModel{} case "Gemma2ForCausalLM": conv = &gemma2Model{} + case "Gemma3ForCausalLM", "Gemma3ForConditionalGeneration": + conv = &gemma3Model{Architecture: p.Architectures[0]} case "Phi3ForCausalLM": conv = &phi3Model{} case "Qwen2ForCausalLM": @@ -213,7 +220,14 @@ func ConvertModel(fsys fs.FS, ws io.WriteSeeker) error { } vocabSize := int(p.VocabSize) + if vocabSize == 0 { + tVocabSize := int(p.TextModel.VocabSize) + vocabSize = tVocabSize + } + switch { + case vocabSize == 0: + slog.Warn("vocabulary size was not explicitly set by the model", "default size", len(t.Vocabulary.Tokens)) case vocabSize > len(t.Vocabulary.Tokens): slog.Warn("vocabulary is smaller than expected, padding with dummy tokens", "expect", vocabSize, "actual", len(t.Vocabulary.Tokens)) for i := range vocabSize - len(t.Vocabulary.Tokens) { diff --git a/convert/convert_gemma.go b/convert/convert_gemma.go index 6c04145f..2f329943 100644 --- a/convert/convert_gemma.go +++ b/convert/convert_gemma.go @@ -45,7 +45,7 @@ func (p *gemmaModel) KV(t *Tokenizer) ggml.KV { func (p *gemmaModel) Tensors(ts []Tensor) []ggml.Tensor { var out []ggml.Tensor for _, t := range ts { - if strings.HasSuffix(t.Name(), "_norm.weight") { + if !strings.HasPrefix(t.Name(), "v.") && strings.HasSuffix(t.Name(), "_norm.weight") { t.SetRepacker(p.addOne) } diff --git a/convert/convert_gemma3.go b/convert/convert_gemma3.go new file mode 100644 index 00000000..c82800c5 --- /dev/null +++ b/convert/convert_gemma3.go @@ -0,0 +1,142 @@ +package convert + +import ( + "cmp" + + "github.com/ollama/ollama/fs/ggml" +) + +type gemma3Model struct { + gemmaModel + Architecture string + TextModel struct { + HeadDim uint32 `json:"head_dim"` + HiddenSize uint32 `json:"hidden_size"` + HiddenLayers uint32 `json:"num_hidden_layers"` + IntermediateSize uint32 `json:"intermediate_size"` + SlidingWindow uint32 `json:"sliding_window"` + } `json:"text_config"` + VisionModel struct { + NumAttentionHeads uint32 `json:"num_attention_heads"` // attention.head_count 16 + LayerNormEpsilon float32 `json:"layer_norm_eps"` // attention.layer_norm_epsilon 1e-05 + NumHiddenLayers uint32 `json:"num_hidden_layers"` // block_count 32 + HiddenSize uint32 `json:"hidden_size"` // embedding_length 1280 + IntermediateSize uint32 `json:"intermediate_size"` // feed_forward_length 5120 + ImageSize uint32 `json:"image_size"` // image_size 560 + NumChannels uint32 `json:"num_channels"` // num_channels 3 + PatchSize uint32 `json:"patch_size"` // patch_size 14 + } `json:"vision_config"` + MaxPositionEmbeddings uint32 `json:"max_position_embeddings"` + NumAttentionHeads uint32 `json:"num_attention_heads"` + NumKeyValueHeads uint32 `json:"num_key_value_heads"` + RMSNormEPS float32 `json:"rms_norm_eps"` + HeadDim uint32 `json:"head_dim"` + FinalLogitSoftcap float32 `json:"final_logit_softcapping"` + RopeLocalTheta float32 `json:"rope_local_base_freq"` + RopeGlobalTheta float32 `json:"rope_global_base_freq"` + SlidingWindow uint32 `json:"sliding_window"` + MultiModalTokensPerImage uint32 `json:"mm_tokens_per_image"` +} + +const ( + gemma4BLayerCount = 34 + gemma12BLayerCount = 48 + gemma27BLayerCount = 62 +) + +func (p *gemma3Model) KV(t *Tokenizer) ggml.KV { + kv := p.ModelParameters.KV(t) + kv["general.architecture"] = "gemma3" + + numBlocks := cmp.Or(p.HiddenLayers, p.TextModel.HiddenLayers) + kv["gemma3.block_count"] = numBlocks + + var ( + numHeads uint32 + numKVHeads uint32 + ) + + switch numBlocks { + case gemma4BLayerCount: + numHeads = 8 + numKVHeads = 4 + case gemma12BLayerCount: + numHeads = 16 + numKVHeads = 8 + case gemma27BLayerCount: + numHeads = 32 + numKVHeads = 16 + default: + numHeads = p.NumAttentionHeads + numKVHeads = p.NumKeyValueHeads + } + + kv["gemma3.attention.head_count"] = numHeads + kv["gemma3.attention.head_count_kv"] = numKVHeads + + switch p.Architecture { + case "Gemma3ForCausalLM": + kv["gemma3.context_length"] = p.MaxPositionEmbeddings + kv["gemma3.attention.layer_norm_rms_epsilon"] = p.RMSNormEPS + kv["gemma3.attention.key_length"] = p.HeadDim + kv["gemma3.attention.value_length"] = p.HeadDim + kv["gemma3.attention.sliding_window"] = p.SlidingWindow + kv["gemma3.final_logit_softcapping"] = cmp.Or(p.FinalLogitSoftcap, 30) + kv["gemma3.rope.local.freq_base"] = cmp.Or(p.RopeLocalTheta, 10000.0) + kv["gemma3.rope.global.freq_base"] = cmp.Or(p.RopeGlobalTheta, 1000000.0) + kv["gemma3.embedding_length"] = p.HiddenSize + kv["gemma3.feed_forward_length"] = p.IntermediateSize + default: + kv["gemma3.context_length"] = cmp.Or(p.MaxPositionEmbeddings, 8192) + kv["gemma3.embedding_length"] = p.TextModel.HiddenSize + kv["gemma3.feed_forward_length"] = p.TextModel.IntermediateSize + kv["gemma3.attention.sliding_window"] = p.TextModel.SlidingWindow + kv["gemma3.vision.block_count"] = p.VisionModel.NumHiddenLayers + kv["gemma3.vision.embedding_length"] = p.VisionModel.HiddenSize + kv["gemma3.vision.feed_forward_length"] = p.VisionModel.IntermediateSize + kv["gemma3.vision.image_size"] = p.VisionModel.ImageSize + kv["gemma3.vision.patch_size"] = p.VisionModel.PatchSize + kv["gemma3.vision.num_channels"] = cmp.Or(p.VisionModel.NumChannels, 3) + kv["gemma3.vision.attention.head_count"] = p.VisionModel.NumAttentionHeads + kv["gemma3.vision.attention.layer_norm_epsilon"] = cmp.Or(p.VisionModel.LayerNormEpsilon, 1e-6) + kv["gemma3.attention.key_length"] = cmp.Or(p.TextModel.HeadDim, 256) + kv["gemma3.attention.value_length"] = cmp.Or(p.TextModel.HeadDim, 256) + } + + if p.MultiModalTokensPerImage > 0 { + kv["gemma3.mm.tokens_per_image"] = p.MultiModalTokensPerImage + } + + return kv +} + +func (p *gemma3Model) Replacements() []string { + return []string{ + "lm_head", "output", + "model.embed_tokens", "token_embd", + "model.norm", "output_norm", + "vision_tower.vision_model.embeddings", "v", + "vision_tower.vision_model", "v", + "vision_model.vision_model.embeddings", "v", + "vision_model.vision_model", "v", + "language_model.", "", + "model.layers", "blk", + "encoder.layers", "blk", + "input_layernorm", "attn_norm", + "self_attn.q_proj", "attn_q", + "self_attn.q_norm", "attn_q_norm", + "self_attn.k_proj", "attn_k", + "self_attn.k_norm", "attn_k_norm", + "self_attn.v_proj", "attn_v", + "self_attn.o_proj", "attn_output", + "self_attn.out_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", + "input_projection_weight", "input_projection.weight", + "multi_modal_projector", "mm", + } +} diff --git a/convert/tokenizer_spm.go b/convert/tokenizer_spm.go index 5e506087..340c3d58 100644 --- a/convert/tokenizer_spm.go +++ b/convert/tokenizer_spm.go @@ -6,7 +6,9 @@ import ( "errors" "fmt" "io/fs" + "log/slog" "os" + "reflect" "slices" "google.golang.org/protobuf/proto" @@ -15,6 +17,8 @@ import ( ) func parseSentencePiece(fsys fs.FS) (*Vocabulary, error) { + slog.Debug("using spm vocabulary") + ast, err := parseAdditionalSpecialTokens(fsys) if err != nil { return nil, err @@ -43,10 +47,19 @@ func parseSentencePiece(fsys fs.FS) (*Vocabulary, error) { v.Types = append(v.Types, int32(t)) default: tt := int32(sentencepiece.ModelProto_SentencePiece_NORMAL) - if slices.Contains(ast, piece.GetPiece()) { + + // temporary fix to handle gemma3 broken configs + if slices.Contains([]string{"", ""}, piece.GetPiece()) { tt = int32(sentencepiece.ModelProto_SentencePiece_CONTROL) } + for _, t := range ast { + if t.Content == piece.GetPiece() { + tt = int32(sentencepiece.ModelProto_SentencePiece_CONTROL) + break + } + } + v.Types = append(v.Types, tt) } } @@ -78,10 +91,16 @@ func parseSentencePiece(fsys fs.FS) (*Vocabulary, error) { return cmp.Compare(i.id, j.id) }) - n := len(v.Tokens) - for i, t := range ts { - if t.id != i+n { - return nil, fmt.Errorf("invalid token id: %d", t.id) + for _, t := range ts { + if t.id < len(v.Tokens) { + if v.Tokens[t.id] == t.content { + slog.Warn("tokenizer", "duplicate token", t.content, "id", t.id) + continue + } + return nil, fmt.Errorf("token mismatch: %s != %s at pos [%d]", t.content, v.Tokens[t.id], t.id) + } + if t.id != len(v.Tokens) { + return nil, fmt.Errorf("invalid token id: [%d] as pos [%d]", t.id, len(v.Tokens)) } v.Tokens = append(v.Tokens, t.content) @@ -92,7 +111,15 @@ func parseSentencePiece(fsys fs.FS) (*Vocabulary, error) { return &v, nil } -func parseAdditionalSpecialTokens(fsys fs.FS) ([]string, error) { +type specialToken struct { + Content string `json:"content"` + Lstrip bool `json:"lstrip"` + Normalized bool `json:"normalized"` + Rstrip bool `json:"rstrip"` + SingleWord bool `json:"single_word"` +} + +func parseAdditionalSpecialTokens(fsys fs.FS) ([]specialToken, error) { f, err := fsys.Open("special_tokens_map.json") if errors.Is(err, os.ErrNotExist) { return nil, nil @@ -102,12 +129,43 @@ func parseAdditionalSpecialTokens(fsys fs.FS) ([]string, error) { defer f.Close() var m struct { - AdditionalSpecialTokens []string `json:"additional_special_tokens"` + AdditionalSpecialTokens any `json:"additional_special_tokens"` } if err := json.NewDecoder(f).Decode(&m); err != nil { return nil, err } - return m.AdditionalSpecialTokens, nil + var ast []specialToken + + switch st := m.AdditionalSpecialTokens.(type) { + case []string: + for _, s := range st { + ast = append(ast, specialToken{Content: s}) + } + case []any: + for _, s := range st { + // marshal and unmarshal the object to get the special token + tMap := s.(map[string]any) + data, err := json.Marshal(tMap) + if err != nil { + return nil, err + } + + var token specialToken + err = json.Unmarshal(data, &token) + if err != nil { + return nil, err + } + + ast = append(ast, token) + } + + default: + slog.Warn("special token", "unknown token", reflect.TypeOf(st)) + } + + slog.Debug("spm tokenizer", "additional tokens", ast) + + return ast, nil } diff --git a/docs/faq.md b/docs/faq.md index 04e8433d..4aaccc2e 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -20,7 +20,7 @@ Please refer to the [GPU docs](./gpu.md). ## How can I specify the context window size? -By default, Ollama uses a context window size of 2048 tokens. +By default, Ollama uses a context window size of 2048 tokens. This can be overridden with the `OLLAMA_CONTEXT_LENGTH` environment variable. For example, to set the default context length to 8K, use: `OLLAMA_CONTEXT_LENGTH=8192 ollama serve`. To change this when using `ollama run`, use `/set parameter`: diff --git a/docs/linux.md b/docs/linux.md index 12581bdd..2dda87f3 100644 --- a/docs/linux.md +++ b/docs/linux.md @@ -75,7 +75,7 @@ RestartSec=3 Environment="PATH=$PATH" [Install] -WantedBy=default.target +WantedBy=multi-user.target ``` Then start the service: diff --git a/fs/ggml/ggml.go b/fs/ggml/ggml.go index 8662c3b0..d32296d9 100644 --- a/fs/ggml/ggml.go +++ b/fs/ggml/ggml.go @@ -124,6 +124,19 @@ func (kv KV) Uints(key string, defaultValue ...[]uint32) []uint32 { return s } +func (kv KV) Floats(key string, defaultValue ...[]float32) []float32 { + r := keyValue(kv, key, &array{}) + s := make([]float32, r.size) + for i := range r.size { + s[i] = float32(r.values[i].(float32)) + } + return s +} + +func (kv KV) OllamaEngineRequired() bool { + return kv.Architecture() == "gemma3" +} + func keyValue[T string | uint32 | uint64 | float32 | *array | bool](kv KV, key string, defaultValue ...T) T { if !strings.HasPrefix(key, "tokenizer.") && !strings.HasPrefix(key, "general.") { key = kv.Architecture() + "." + key @@ -476,7 +489,7 @@ func (f GGML) GraphSize(context, batch uint64, kvCacheType string) (kv, partialO // vocab graph 4*batch*(embedding+vocab)+embedding*vocab*105/128, ) - case "gemma", "gemma2": + case "gemma", "gemma2", "gemma3": fullOffload = max( 4*batch*(embedding+vocab), 4*batch*(2+context+context*heads+2*embedding+2*embeddingHeadsK*heads), diff --git a/kvcache/cache.go b/kvcache/cache.go index 2541f7c1..d3548905 100644 --- a/kvcache/cache.go +++ b/kvcache/cache.go @@ -4,6 +4,7 @@ import ( "errors" "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/model/input" ) var ( @@ -51,7 +52,7 @@ type Cache interface { // StartForward is called before the start of the model's forward pass. // For each token in the coming batch, there must be a corresponding // entry in positions and seqs. - StartForward(ctx ml.Context, positions []int32, seqs []int) error + StartForward(ctx ml.Context, opts input.Options) error // CopyPrefix copies tokens in the range [0, len) from srcSeq to dstSeq CopyPrefix(srcSeq, dstSeq int, len int32) diff --git a/kvcache/causal.go b/kvcache/causal.go index 9a79fa57..edf6666d 100644 --- a/kvcache/causal.go +++ b/kvcache/causal.go @@ -8,6 +8,7 @@ import ( "slices" "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/model/input" ) type shiftFn func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) @@ -20,9 +21,10 @@ type shiftFn func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, e type Causal struct { DType ml.DType Capacity int32 - causal bool windowSize int32 + opts CausalOptions + // config controls mostly backend-specific optimizations config *ml.CacheConfig @@ -78,7 +80,6 @@ type cellRange struct { func NewCausalCache(shift shiftFn) *Causal { return &Causal{ - causal: true, windowSize: math.MaxInt32, shiftFn: shift, ctxs: make(map[int]ml.Context), @@ -89,7 +90,6 @@ func NewCausalCache(shift shiftFn) *Causal { func NewSWACache(windowSize int32, shift shiftFn) *Causal { return &Causal{ - causal: true, windowSize: windowSize, shiftFn: shift, ctxs: make(map[int]ml.Context), @@ -140,10 +140,11 @@ func (c *Causal) Close() { } } -func (c *Causal) StartForward(ctx ml.Context, positions []int32, seqs []int) error { - c.curBatchSize = len(positions) - c.curSequences = seqs - c.curPositions = positions +func (c *Causal) StartForward(ctx ml.Context, opts input.Options) error { + c.curBatchSize = len(opts.Positions) + c.curSequences = opts.Sequences + c.curPositions = opts.Positions + c.opts.Except = nil var err error c.curLoc, err = c.findStartLoc() @@ -156,8 +157,8 @@ func (c *Causal) StartForward(ctx ml.Context, positions []int32, seqs []int) err } c.curCellRange = newRange() - for i, pos := range positions { - seq := seqs[i] + for i, pos := range opts.Positions { + seq := opts.Sequences[i] c.cells[c.curLoc+i] = cacheCell{pos: pos, sequences: []int{seq}} @@ -234,9 +235,10 @@ func (c *Causal) buildMask(ctx ml.Context) (ml.Tensor, error) { mask := make([]float32, batchSize*length) for i := range c.curBatchSize { + enabled := !slices.Contains(c.opts.Except, i) for j := c.curCellRange.min; j <= c.curCellRange.max; j++ { if !slices.Contains(c.cells[j].sequences, c.curSequences[i]) || - (c.causal && c.cells[j].pos > c.curPositions[i]) || + (enabled && c.cells[j].pos > c.curPositions[i]) || c.cells[j].pos < c.curPositions[i]-c.windowSize { mask[i*length+(j-c.curCellRange.min)] = float32(math.Inf(-1)) } @@ -403,15 +405,16 @@ func (c *Causal) SetLayer(layer int) { c.curLayer = layer } -// SetCausal enables or disables causal mask generation for subsequent calls to Get. -// This state carries over to future forward passes. The default value is true. -// -// ctx may be set to nil if this is called from outside of a forward pass, for -// example, when initializing the cache. -func (c *Causal) SetCausal(ctx ml.Context, causal bool) { - if c.causal != causal { - c.causal = causal +type CausalOptions struct { + // Enabled controls whether the causal mask is generated for a particular index in a batch + Except []int +} +// SetCausal disables causal mask generation for a particular range of indicies in +// the current batch for subsequent calls to Get. The state resets for the next forward pass. +func (c *Causal) SetCausal(ctx ml.Context, opts CausalOptions) { + if !slices.Equal(c.opts.Except, opts.Except) { + c.opts = opts if ctx != nil { var err error c.curMask, err = c.buildMask(ctx) diff --git a/kvcache/causal_test.go b/kvcache/causal_test.go index 412f33e3..56d85ceb 100644 --- a/kvcache/causal_test.go +++ b/kvcache/causal_test.go @@ -6,6 +6,7 @@ import ( "testing" "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/model/input" ) type testCase struct { @@ -269,7 +270,7 @@ func testCache(t *testing.T, backend ml.Backend, cache Cache, tests []testCase) context := backend.NewContext() defer context.Close() - err := cache.StartForward(context, test.pos, test.seqs) + err := cache.StartForward(context, input.Options{Positions: test.pos, Sequences: test.seqs}) if err != nil { panic(err) } @@ -440,11 +441,19 @@ func (t *testTensor) Scale(ctx ml.Context, s float64) ml.Tensor { panic("not implemented") } +func (t *testTensor) AvgPool1D(ctx ml.Context, k, s, p int) ml.Tensor { + panic("not implemented") +} + +func (t *testTensor) AvgPool2D(ctx ml.Context, k, s int, p float32) ml.Tensor { + panic("not implemented") +} + func (t *testTensor) Conv2D(ctx ml.Context, weight ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor { panic("not implemented") } -func (t *testTensor) RoPE(ctx ml.Context, positionIDs, ropeFactors ml.Tensor, dim uint32, base, scale float32) ml.Tensor { +func (t *testTensor) RoPE(ctx ml.Context, positionIDs, ropeFactors ml.Tensor, dim, ropeType uint32, base, scale float32) ml.Tensor { panic("not implemented") } @@ -494,6 +503,10 @@ func (t *testTensor) Contiguous(ctx ml.Context) ml.Tensor { panic("not implemented") } +func (t *testTensor) Set(ctx ml.Context, t2 ml.Tensor, offset int, strides ...int) ml.Tensor { + panic("not implemented") +} + func (t *testTensor) Pad(ctx ml.Context, shape ...int) ml.Tensor { panic("not implemented") } diff --git a/kvcache/encoder.go b/kvcache/encoder.go index 867ee37a..6a9df2ab 100644 --- a/kvcache/encoder.go +++ b/kvcache/encoder.go @@ -4,6 +4,7 @@ import ( "fmt" "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/model/input" ) // Encoder cache stores K and V tensors that are position independent @@ -78,9 +79,11 @@ func (c *EncoderCache) Close() { } } -func (c *EncoderCache) StartForward(ctx ml.Context, positions []int32, seqs []int) error { - // The image is always in the first position - c.curPos = positions[0] +func (c *EncoderCache) StartForward(ctx ml.Context, opts input.Options) error { + // We work with the most recent image + if len(opts.Multimodal) > 0 { + c.curPos = opts.Positions[opts.Multimodal[len(opts.Multimodal)-1].Index] + } return nil } diff --git a/kvcache/wrapper.go b/kvcache/wrapper.go index 76956a88..aaccd166 100644 --- a/kvcache/wrapper.go +++ b/kvcache/wrapper.go @@ -4,6 +4,7 @@ import ( "math" "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/model/input" ) // Wrapper cache is a container for multiple types of caches, @@ -40,14 +41,14 @@ func (c *WrapperCache) Close() { } } -func (c *WrapperCache) StartForward(ctx ml.Context, positions []int32, seqs []int) error { +func (c *WrapperCache) StartForward(ctx ml.Context, opts input.Options) error { for i, cache := range c.caches { - err := cache.StartForward(ctx, positions, seqs) + err := cache.StartForward(ctx, opts) if err != nil { // unwind on error - Remove with endIndex set to math.MaxInt32 does not fail for j := i - 1; j >= 0; j-- { - for k := range positions { - _ = c.caches[j].Remove(seqs[k], positions[k], math.MaxInt32) + for k := range opts.Positions { + _ = c.caches[j].Remove(opts.Sequences[k], opts.Positions[k], math.MaxInt32) } } return err diff --git a/llama/llama.go b/llama/llama.go index bb5028bd..a026bee2 100644 --- a/llama/llama.go +++ b/llama/llama.go @@ -245,6 +245,20 @@ func LoadModelFromFile(modelPath string, params ModelParams) (*Model, error) { return &m, nil } +func LoadVocabFromFile(path string) (*Vocab, error) { + mp := C.CString(path) + defer C.free(unsafe.Pointer(mp)) + v := Vocab{c: C.llama_load_vocab_from_file(mp)} + if v.c == nil { + return nil, fmt.Errorf("unable to load vocab: %s", path) + } + return &v, nil +} + +func FreeVocab(vocab *Vocab) { + C.llama_free_vocab(vocab.c) +} + func FreeModel(model *Model) { C.llama_model_free(model.c) } @@ -293,6 +307,10 @@ func (m *Model) ApplyLoraFromFile(context *Context, loraPath string, scale float return nil } +type Vocab struct { + c *C.struct_llama_vocab +} + func (m *Model) Vocab() *C.struct_llama_vocab { return C.llama_model_get_vocab(m.c) } @@ -669,3 +687,53 @@ func SchemaToGrammar(schema []byte) []byte { } return buf[:n] } + +type Sampler struct { + c *C.struct_llama_sampler +} + +func NewGrammarSampler(vocab *Vocab, grammar string) *Sampler { + cGrammar := C.CString(grammar) + cRoot := C.CString("root") + defer C.free(unsafe.Pointer(cGrammar)) + defer C.free(unsafe.Pointer(cRoot)) + + sampler := &Sampler{c: C.llama_sampler_init_grammar(vocab.c, cGrammar, cRoot)} + + return sampler +} + +func (s *Sampler) Accept(token int32) { + C.llama_sampler_accept(s.c, C.llama_token(token)) +} + +type TokenData struct { + Id int32 + Logit float32 +} + +func (s *Sampler) Apply(tokens []TokenData) { + tds := make([]C.struct_llama_token_data, len(tokens)) + for i, token := range tokens { + tds[i] = C.struct_llama_token_data{ + id: C.int32_t(token.Id), + logit: C.float(token.Logit), + p: C.float(0.0), + } + } + tda := &C.llama_token_data_array{ + data: (*C.struct_llama_token_data)(unsafe.Pointer(&tds[0])), + size: C.size_t(len(tokens)), + selected: C.int64_t(-1), + sorted: C.bool(false), + } + + var pinner runtime.Pinner + pinner.Pin(&tds[0]) + defer pinner.Unpin() + + C.llama_sampler_apply(s.c, tda) + for i := range tokens { + tokens[i].Logit = float32(tds[i].logit) + } +} diff --git a/llama/patches/0020-ollama-debug-tensor.patch b/llama/patches/0020-ollama-debug-tensor.patch new file mode 100644 index 00000000..b9f2e4ab --- /dev/null +++ b/llama/patches/0020-ollama-debug-tensor.patch @@ -0,0 +1,33 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: Michael Yang +Date: Sun, 9 Mar 2025 14:44:16 -0700 +Subject: [PATCH] ollama debug tensor + +--- + ggml/src/ggml-cpu/ggml-cpu.c | 6 ++++++ + 1 file changed, 6 insertions(+) + +diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c +index 2f606d82..ec60e8fc 100644 +--- a/ggml/src/ggml-cpu/ggml-cpu.c ++++ b/ggml/src/ggml-cpu/ggml-cpu.c +@@ -11,6 +11,8 @@ + #include "ggml-threading.h" + #include "ggml.h" + ++#include "ollama-debug.h" ++ + #if defined(_MSC_VER) || defined(__MINGW32__) + #include // using malloc.h with MSC/MINGW + #elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) +@@ -14103,6 +14105,10 @@ static thread_ret_t ggml_graph_compute_thread(void * data) { + + ggml_compute_forward(¶ms, node); + ++#ifdef OLLAMA_DEBUG ++ ollama_debug(node, true); ++#endif ++ + if (state->ith == 0 && cplan->abort_callback && + cplan->abort_callback(cplan->abort_callback_data)) { + atomic_store_explicit(&tp->abort, node_n + 1, memory_order_relaxed); diff --git a/llama/sampling_ext.cpp b/llama/sampling_ext.cpp index 0f137dc8..b816cedd 100644 --- a/llama/sampling_ext.cpp +++ b/llama/sampling_ext.cpp @@ -2,6 +2,9 @@ #include "sampling.h" #include "sampling_ext.h" #include "json-schema-to-grammar.h" +#include "llama.h" +#include "llama-model.h" +#include "llama-model-loader.h" struct common_sampler *common_sampler_cinit(const struct llama_model *model, struct common_sampler_cparams *params) { try { @@ -64,3 +67,22 @@ int schema_to_grammar(const char *json_schema, char *grammar, size_t max_len) return 0; } } + +struct llama_vocab * llama_load_vocab_from_file(const char * fname) { + llama_vocab * vocab = new llama_vocab(); + try { + const auto kv = LLM_KV(LLM_ARCH_UNKNOWN); + std::vector splits = {}; + llama_model_loader ml(std::string(fname), splits, false, false, nullptr); + vocab->load(ml, kv); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error loading model: %s\n", __func__, err.what()); + return nullptr; + } + + return vocab; +} + +void llama_free_vocab(struct llama_vocab * vocab) { + delete vocab; +} diff --git a/llama/sampling_ext.h b/llama/sampling_ext.h index 39f499f1..9be7c100 100644 --- a/llama/sampling_ext.h +++ b/llama/sampling_ext.h @@ -35,6 +35,9 @@ extern "C" int schema_to_grammar(const char *json_schema, char *grammar, size_t max_len); + struct llama_vocab * llama_load_vocab_from_file(const char * fname); + void llama_free_vocab(struct llama_vocab * vocab); + #ifdef __cplusplus } #endif diff --git a/llm/server.go b/llm/server.go index 9553ba8f..c6f11712 100644 --- a/llm/server.go +++ b/llm/server.go @@ -271,7 +271,7 @@ func NewLlamaServer(gpus discover.GpuInfoList, modelPath string, f *ggml.GGML, a var llamaModel *llama.Model var textProcessor model.TextProcessor - if envconfig.NewEngine() { + if envconfig.NewEngine() || f.KV().OllamaEngineRequired() { textProcessor, err = model.NewTextProcessor(modelPath) if err != nil { // To prepare for opt-out mode, instead of treating this as an error, we fallback to the old runner @@ -729,29 +729,24 @@ func (s *llmServer) Completion(ctx context.Context, req CompletionRequest, fn fu } if len(req.Format) > 0 { - format := string(req.Format) - if format != `null` && format != `""` { - if s.textProcessor != nil { - // New engine handles this on the backend - request["format"] = req.Format - } else { - // old engine - switch format { - case `"json"`: - request["grammar"] = grammarJSON - default: - if req.Format[0] != '{' { - return fmt.Errorf("invalid format: %q; expected \"json\" or a valid JSON Schema object", req.Format) - } - - // User provided a JSON schema - g := llama.SchemaToGrammar(req.Format) - if g == nil { - return fmt.Errorf("invalid JSON schema in format") - } - request["grammar"] = string(g) - } + switch string(req.Format) { + case `null`, `""`: + // Field was set, but "missing" a value. We accept + // these as "not set". + break + case `"json"`: + request["grammar"] = grammarJSON + default: + if req.Format[0] != '{' { + return fmt.Errorf("invalid format: %q; expected \"json\" or a valid JSON Schema object", req.Format) } + + // User provided a JSON schema + g := llama.SchemaToGrammar(req.Format) + if g == nil { + return fmt.Errorf("invalid JSON schema in format") + } + request["grammar"] = string(g) } } diff --git a/ml/backend.go b/ml/backend.go index 3abacbf1..c63c73d4 100644 --- a/ml/backend.go +++ b/ml/backend.go @@ -5,6 +5,7 @@ import ( "encoding/binary" "fmt" "os" + "slices" "strconv" "strings" ) @@ -18,6 +19,7 @@ type Config interface { Strings(string, ...[]string) []string Uints(string, ...[]uint32) []uint32 + Floats(string, ...[]float32) []float32 } type Backend interface { @@ -133,8 +135,10 @@ type Tensor interface { RMSNorm(ctx Context, weight Tensor, eps float32) Tensor Scale(ctx Context, s float64) Tensor + AvgPool2D(ctx Context, k, s int, p float32) Tensor Conv2D(ctx Context, weight Tensor, s0, s1, p0, p1, d0, d1 int) Tensor - RoPE(ctx Context, positionIDs, ropeFactors Tensor, dim uint32, base, scale float32) Tensor + + RoPE(ctx Context, positionIDs, ropeFactors Tensor, dim, ropeType uint32, base, scale float32) Tensor Tanh(ctx Context) Tensor GELU(ctx Context) Tensor @@ -144,6 +148,7 @@ type Tensor interface { View(ctx Context, offset int, shape ...int) Tensor Permute(ctx Context, shape ...int) Tensor Contiguous(ctx Context) Tensor + Set(ctx Context, t2 Tensor, offset int, strides ...int) Tensor Pad(ctx Context, shape ...int) Tensor Unpad(ctx Context, shape ...int) Tensor @@ -241,16 +246,17 @@ func dump[S ~[]E, E number](ctx Context, t Tensor, items int, fn func(E) string) } shape := t.Shape() + slices.Reverse(shape) var sb strings.Builder var f func([]int, int) f = func(dims []int, stride int) { prefix := strings.Repeat(" ", len(shape)-len(dims)+1) - fmt.Fprint(&sb, "[") - defer func() { fmt.Fprint(&sb, "]") }() + sb.WriteString("[") + defer func() { sb.WriteString("]") }() for i := 0; i < dims[0]; i++ { if i >= items && i < dims[0]-items { - fmt.Fprint(&sb, "..., ") + sb.WriteString("..., ") // skip to next printable element skip := dims[0] - 2*items if len(dims) > 1 { @@ -265,9 +271,14 @@ func dump[S ~[]E, E number](ctx Context, t Tensor, items int, fn func(E) string) fmt.Fprint(&sb, ",", strings.Repeat("\n", len(dims)-1), prefix) } } else { - fmt.Fprint(&sb, fn(s[stride+i])) + text := fn(s[stride+i]) + if len(text) > 0 && text[0] != '-' { + sb.WriteString(" ") + } + + sb.WriteString(text) if i < dims[0]-1 { - fmt.Fprint(&sb, ", ") + sb.WriteString(", ") } } } diff --git a/ml/backend/ggml/ggml.go b/ml/backend/ggml/ggml.go index 74512f33..03b9acb3 100644 --- a/ml/backend/ggml/ggml.go +++ b/ml/backend/ggml/ggml.go @@ -240,11 +240,22 @@ func New(r *os.File, params ml.BackendParams) (ml.Backend, error) { switch { case contains(t.Name, "position_embd", "token_embd", "token_norm_embd", "token_types"): createTensor(tensor{source: t}, input.bts) + if _, ok := meta.Tensors().GroupLayers()["output"]; !ok && t.Name == "token_embd.weight" { + createTensor(tensor{source: t, target: "output.weight"}, output.bts) + } case contains(t.Name, "cls", "output", "output_norm"): createTensor(tensor{source: t}, output.bts) case strings.HasPrefix(t.Name, "v.") || strings.HasPrefix(t.Name, "mm."): // TODO: assign vision tensors to the gpu if possible - createTensor(tensor{source: t}, input.bts) + createTensor(tensor{source: t}, output.bts) + case contains(t.Name, "rope_freqs", "rope_factors_long", "rope_factors_short"): + // these tensors should be repeated per layer + for i, layer := range layers { + createTensor(tensor{ + source: t, + target: "blk." + strconv.Itoa(i) + "." + t.Name, + }, layer.bts) + } default: layerIndex := -1 if fields := strings.FieldsFunc(t.Name, func(r rune) bool { return !unicode.IsNumber(r) }); len(fields) > 0 { @@ -256,14 +267,8 @@ func New(r *os.File, params ml.BackendParams) (ml.Backend, error) { if layerIndex >= 0 { createTensor(tensor{source: t}, layers[layerIndex].bts) } else { - // this is a repeating tensor that doesn't explicitly associated with a layer so - // duplicate it for each layer - for i, layer := range layers { - createTensor(tensor{ - source: t, - target: "blk." + strconv.Itoa(i) + "." + t.Name, - }, layer.bts) - } + // load all other tensors on the cpu + createTensor(tensor{source: t}, input.bts) } } } @@ -352,7 +357,7 @@ func New(r *os.File, params ml.BackendParams) (ml.Backend, error) { if C.ggml_backend_is_cpu(b) { // set number of threads for cpu backend - C.ggml_backend_cpu_set_n_threads(b, C.int(params.NumThreads)) + C.ggml_backend_cpu_set_n_threads(b, C.int(Threads(params.NumThreads))) } } @@ -893,10 +898,13 @@ func (t *Tensor) View(ctx ml.Context, offset int, shape ...int) ml.Tensor { } const ( - ropeTypeNorm C.int = iota + ropeTypeNorm C.int = 0 + ropeTypeNeox C.int = 2 + ropeTypeMrope C.int = 8 + ropeTypeVision C.int = 24 ) -func (t *Tensor) RoPE(ctx ml.Context, positionIDs, ropeFactors ml.Tensor, ropeDim uint32, ropeBase, ropeScale float32) ml.Tensor { +func (t *Tensor) RoPE(ctx ml.Context, positionIDs, ropeFactors ml.Tensor, ropeDim, ropeType uint32, ropeBase, ropeScale float32) ml.Tensor { if ropeFactors == nil { ropeFactors = &Tensor{b: t.b} } @@ -911,8 +919,8 @@ func (t *Tensor) RoPE(ctx ml.Context, positionIDs, ropeFactors ml.Tensor, ropeDi t: C.ggml_rope_ext( ctx.(*Context).ctx, dequant, positionIDs.(*Tensor).t, ropeFactors.(*Tensor).t, C.int(ropeDim), - 131072, // YaRN n_ctx_train - ropeTypeNorm, // ROPE_TYPE_NORM + C.int(ropeType), + 131072, // YaRN n_ctx_train C.float(ropeBase), C.float(ropeScale), 0., // YaRN ext_factor @@ -944,6 +952,27 @@ func (t *Tensor) Conv2D(ctx ml.Context, t2 ml.Tensor, s0, s1, p0, p1, d0, d1 int } } +func (t *Tensor) AvgPool2D(ctx ml.Context, k, s int, p float32) ml.Tensor { + return &Tensor{ + b: t.b, + t: C.ggml_pool_2d(ctx.(*Context).ctx, t.t, C.GGML_OP_POOL_AVG, C.int(k), C.int(k), C.int(s), C.int(s), C.float(p), C.float(p)), + } +} + +func (t *Tensor) Set(ctx ml.Context, t2 ml.Tensor, offset int, strides ...int) ml.Tensor { + var tt *C.struct_ggml_tensor + switch len(strides) { + case 0: + tt = C.ggml_set_1d(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.size_t(offset)) + case 1: + tt = C.ggml_set_2d(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.size_t(offset), C.size_t(strides[0])) + default: + panic("unsupported number of dimensions") + } + + return &Tensor{b: t.b, t: tt} +} + func (t *Tensor) ScaledDotProductAttention(ctx ml.Context, key, value, mask ml.Tensor, scale float64) ml.Tensor { var kqMask *C.struct_ggml_tensor if mask != nil { diff --git a/ml/backend/ggml/ggml/include/ollama-debug.h b/ml/backend/ggml/ggml/include/ollama-debug.h new file mode 100644 index 00000000..36a2e241 --- /dev/null +++ b/ml/backend/ggml/ggml/include/ollama-debug.h @@ -0,0 +1,11 @@ +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +void ollama_debug(const struct ggml_tensor *tensor, bool verbose); + +#ifdef __cplusplus +} +#endif diff --git a/ml/backend/ggml/ggml/src/ggml-cpu/cpu_debug.go b/ml/backend/ggml/ggml/src/ggml-cpu/cpu_debug.go new file mode 100644 index 00000000..7eab9813 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ggml-cpu/cpu_debug.go @@ -0,0 +1,6 @@ +//go:build debug + +package cpu + +// #cgo CPPFLAGS: -DOLLAMA_DEBUG +import "C" diff --git a/ml/backend/ggml/ggml/src/ggml-cpu/ggml-cpu.c b/ml/backend/ggml/ggml/src/ggml-cpu/ggml-cpu.c index 2f606d82..ec60e8fc 100644 --- a/ml/backend/ggml/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ml/backend/ggml/ggml/src/ggml-cpu/ggml-cpu.c @@ -11,6 +11,8 @@ #include "ggml-threading.h" #include "ggml.h" +#include "ollama-debug.h" + #if defined(_MSC_VER) || defined(__MINGW32__) #include // using malloc.h with MSC/MINGW #elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) @@ -14103,6 +14105,10 @@ static thread_ret_t ggml_graph_compute_thread(void * data) { ggml_compute_forward(¶ms, node); +#ifdef OLLAMA_DEBUG + ollama_debug(node, true); +#endif + if (state->ith == 0 && cplan->abort_callback && cplan->abort_callback(cplan->abort_callback_data)) { atomic_store_explicit(&tp->abort, node_n + 1, memory_order_relaxed); diff --git a/ml/backend/ggml/ggml/src/ollama-debug.c b/ml/backend/ggml/ggml/src/ollama-debug.c new file mode 100644 index 00000000..b0e9d7f0 --- /dev/null +++ b/ml/backend/ggml/ggml/src/ollama-debug.c @@ -0,0 +1,115 @@ +#include + +#include "ollama-debug.h" + +static int mul(int64_t *dims, int ndims) { + int result = 1; + for (int i = 0; i < ndims; i++) { + result *= dims[i]; + } + + return result; +} + +static void repeat(char c, int n) { + for (int i = 0; i < n; i++) { + fprintf(stderr, "%c", c); + } +} + +static void print_tensor(const void *tensor, void (*cb)(const void *, int), + int shape, + int64_t *dims, int ndims, int stride, + int nitems, int pad) { + fprintf(stderr, "["); + for (int i = 0; i < dims[0]; i++) { + if (i >= nitems && i < dims[0] - nitems) { + fprintf(stderr, "... (%lld more), ", dims[0] - 2 * nitems); + int skip = dims[0] - 2 * nitems; + if (ndims > 1) { + stride += mul(dims + 1, ndims - 1) * skip; + repeat('\n', ndims - 1); + repeat(' ', shape - ndims + 1 + pad); + } + i += skip - 1; + } else if (ndims > 1) { + print_tensor(tensor, cb, shape, dims + 1, ndims - 1, stride, + nitems, pad); + stride += mul(dims + 1, ndims - 1); + if (i < dims[0] - 1) { + fprintf(stderr, ", "); + repeat('\n', ndims - 1); + repeat(' ', shape - ndims + 1 + pad); + } + } else { + cb(tensor, stride + i); + if (i < dims[0] - 1) { + fprintf(stderr, ", "); + } + } + } + fprintf(stderr, "]"); +} + +static void print_tensor_f16(const void *tensor, int i) { + float value = ggml_fp16_to_fp32(((const ggml_fp16_t *)tensor)[i]); + fprintf(stderr, "%s%f", value < 0 ? "" : " ", value); +} + +static void print_tensor_f32(const void *tensor, int i) { + float value = ((const float *)tensor)[i]; + fprintf(stderr, "%s%f", value < 0 ? "" : " ", value); +} + +static void print_tensor_i32(const void *tensor, int i) { + int32_t value = ((const int32_t *)tensor)[i]; + fprintf(stderr, "%s%d", value < 0 ? "" : " ", value); +} + +static void ollama_debug_tensor(const struct ggml_tensor *tensor, bool verbose, const char *prefix, int indent) { + fprintf(stderr, "%s%s %s (%s): [%lld %lld %lld %lld]\n", prefix, tensor->name, + ggml_op_name(tensor->op), ggml_type_name(tensor->type), tensor->ne[0], + tensor->ne[1], tensor->ne[2], tensor->ne[3]); + + if (!verbose) { + return; + } + + for (int i = 0; i < indent; i++) { + fprintf(stderr, " "); + } + + switch (tensor->type) { + case GGML_TYPE_F16: + print_tensor(ggml_get_data(tensor), print_tensor_f16, ggml_n_dims(tensor), + (int64_t *)tensor->ne, ggml_n_dims(tensor), 0, 3, indent); + break; + case GGML_TYPE_F32: + print_tensor(ggml_get_data(tensor), print_tensor_f32, ggml_n_dims(tensor), + (int64_t *)tensor->ne, ggml_n_dims(tensor), 0, 3, indent); + break; + case GGML_TYPE_I32: + print_tensor(ggml_get_data(tensor), print_tensor_i32, ggml_n_dims(tensor), + (int64_t *)tensor->ne, ggml_n_dims(tensor), 0, 3, indent); + break; + default: + fprintf(stderr, "\n"); + return; + } + + fprintf(stderr, "\n"); +} + +void ollama_debug(const struct ggml_tensor *tensor, bool verbose) { + ollama_debug_tensor(tensor, verbose, ">>> ", 4); + + for (int i = 0; i < GGML_MAX_SRC && tensor->src[i] != NULL; ++i) { + char src[8]; + const int n = snprintf(src, sizeof(src), " src%d ", i); + if (n >= sizeof(src)) { + src[sizeof(src) - 1] = '\0'; + } + + ollama_debug_tensor(tensor->src[i], verbose, src, 4); + } +} diff --git a/ml/backend/ggml/threads.go b/ml/backend/ggml/threads.go new file mode 100644 index 00000000..cbc524f7 --- /dev/null +++ b/ml/backend/ggml/threads.go @@ -0,0 +1,7 @@ +//go:build !debug + +package ggml + +func Threads(n int) int { + return n +} diff --git a/ml/backend/ggml/threads_debug.go b/ml/backend/ggml/threads_debug.go new file mode 100644 index 00000000..cfd334bd --- /dev/null +++ b/ml/backend/ggml/threads_debug.go @@ -0,0 +1,7 @@ +//go:build debug + +package ggml + +func Threads(_ int) int { + return 1 +} diff --git a/model/input/input.go b/model/input/input.go new file mode 100644 index 00000000..0cb3f3f4 --- /dev/null +++ b/model/input/input.go @@ -0,0 +1,37 @@ +package input + +// Input represents one token in the input stream +type Input struct { + // Token is a single element of text. + Token int32 + + // Multimodal is opaque data representing a non-text + // element such as an image (or part of one if the image + // can be processed in pieces). It may be either together + // with Token or on its own. + Multimodal any + + // MultimodalHash is a unique representation of the data + // stored in Multimodal, used for caching and comparing + // equality. + MultimodalHash uint64 +} + +// MultimodalIndex is a multimodal element (such as an image) +// together with an index into the slice of Inputs with the +// corresponding token. Note that the index is not the same +// as the position - to find that use the index with the +// Positions slice. +type MultimodalIndex struct { + Index int + Multimodal any +} + +// Options contains the inputs for a model forward pass +type Options struct { + Inputs []int32 + Multimodal []MultimodalIndex + Positions []int32 + Sequences []int + Outputs []int32 +} diff --git a/model/model.go b/model/model.go index 75b7f639..89b6c803 100644 --- a/model/model.go +++ b/model/model.go @@ -19,66 +19,12 @@ import ( "github.com/ollama/ollama/kvcache" "github.com/ollama/ollama/ml" _ "github.com/ollama/ollama/ml/backend" + "github.com/ollama/ollama/model/input" ) -// Input represents one token in the input stream -type Input struct { - // Token is a single element of text. - Token int32 - - // Multimodal is opaque data representing a non-text - // element such as an image (or part of one if the image - // can be processed in pieces). It may be either together - // with Token or on its own. - Multimodal any - - // MultimodalHash is a unique representation of the data - // stored in Multimodal, used for caching and comparing - // equality. - MultimodalHash uint64 -} - -// MultimodalIndex is a multimodal element (such as an image) -// together with an index into the slice of Inputs with the -// corresponding token. Note that the index is not the same -// as the position - to find that use the index with the -// Positions slice. -type MultimodalIndex struct { - Index int - Multimodal any -} - -// Options contains the inputs for a model forward pass -type Options struct { - Inputs []int32 - Multimodal []MultimodalIndex - Positions []int32 - Sequences []int - Outputs []int32 -} - -type config struct { - Cache kvcache.Cache -} - -// Base implements the common fields and methods for all models -type Base struct { - b ml.Backend - config -} - -// Backend returns the underlying backend that will run the model -func (m *Base) Backend() ml.Backend { - return m.b -} - -func (m *Base) Config() config { - return m.config -} - // Model implements a specific model architecture, defining the forward pass and any model-specific configuration type Model interface { - Forward(ml.Context, Options) (ml.Tensor, error) + Forward(ml.Context, input.Options) (ml.Tensor, error) Backend() ml.Backend Config() config @@ -112,7 +58,26 @@ type MultimodalProcessor interface { // This function is also responsible for updating MultimodalHash for any Multimodal // that is modified to ensure that there is a unique hash value that accurately // represents the contents. - PostTokenize(ml.Context, []Input) ([]Input, error) + PostTokenize(ml.Context, []input.Input) ([]input.Input, error) +} + +// Base implements the common fields and methods for all models +type Base struct { + b ml.Backend + config +} + +type config struct { + Cache kvcache.Cache +} + +// Backend returns the underlying backend that will run the model +func (m *Base) Backend() ml.Backend { + return m.b +} + +func (m *Base) Config() config { + return m.config } var models = make(map[string]func(ml.Config) (Model, error)) @@ -313,7 +278,7 @@ func canNil(t reflect.Type) bool { t.Kind() == reflect.Slice } -func Forward(ctx ml.Context, m Model, opts Options) (ml.Tensor, error) { +func Forward(ctx ml.Context, m Model, opts input.Options) (ml.Tensor, error) { if len(opts.Positions) != len(opts.Sequences) { return nil, fmt.Errorf("length of positions (%v) must match length of seqs (%v)", len(opts.Positions), len(opts.Sequences)) } @@ -324,7 +289,7 @@ func Forward(ctx ml.Context, m Model, opts Options) (ml.Tensor, error) { cache := m.Config().Cache if cache != nil { - err := cache.StartForward(ctx, opts.Positions, opts.Sequences) + err := cache.StartForward(ctx, opts) if err != nil { return nil, err } diff --git a/model/model_test.go b/model/model_test.go index 8761817e..354dd1d8 100644 --- a/model/model_test.go +++ b/model/model_test.go @@ -11,6 +11,7 @@ import ( "github.com/ollama/ollama/ml" "github.com/ollama/ollama/ml/backend/ggml" "github.com/ollama/ollama/ml/nn" + "github.com/ollama/ollama/model/input" ) func TestParseTags(t *testing.T) { @@ -162,7 +163,7 @@ func TestGetTextProcessor(t *testing.T) { type notTextProcessorModel struct{} -func (notTextProcessorModel) Forward(ml.Context, Options) (ml.Tensor, error) { +func (notTextProcessorModel) Forward(ml.Context, input.Options) (ml.Tensor, error) { panic("unimplemented") } diff --git a/model/models/gemma2/model.go b/model/models/gemma2/model.go new file mode 100644 index 00000000..2b8597c4 --- /dev/null +++ b/model/models/gemma2/model.go @@ -0,0 +1,220 @@ +package gemma2 + +import ( + "math" + + "github.com/ollama/ollama/kvcache" + "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/ml/nn" + "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" +) + +type Options struct { + hiddenSize, numHeads, numKVHeads int + attnKeyLen, attnValLen int + eps, ropeBase, ropeScale float32 + attnLogitSoftcap float32 + finalLogitSoftcap float32 + largeModelScaling bool +} + +type Model struct { + model.Base + model.SentencePieceModel + + TokenEmbedding *nn.Embedding `gguf:"token_embd"` + Layers []Layer `gguf:"blk"` + OutputNorm *nn.RMSNorm `gguf:"output_norm"` + Output *nn.Linear `gguf:"output,alt:token_embd"` // just set to token_embd? + + *Options +} + +const ( + gemma27BLayerCount = 46 +) + +func New(c ml.Config) (model.Model, error) { + m := Model{ + SentencePieceModel: model.NewSentencePieceModel( + c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`), + &model.Vocabulary{ + Values: c.Strings("tokenizer.ggml.tokens"), + Scores: c.Floats("tokenizer.ggml.scores"), + Types: c.Uints("tokenizer.ggml.token_type"), + BOS: int32(c.Uint("tokenizer.ggml.bos_token_id")), + EOS: int32(c.Uint("tokenizer.ggml.eos_token_id")), + }, + ), + Layers: make([]Layer, c.Uint("block_count")), + Options: &Options{ + hiddenSize: int(c.Uint("embedding_length")), + numHeads: int(c.Uint("attention.head_count")), + numKVHeads: int(c.Uint("attention.head_count_kv")), + attnKeyLen: int(c.Uint("attention.key_length")), + attnValLen: int(c.Uint("attention.value_length")), + eps: c.Float("attention.layer_norm_rms_epsilon"), + ropeBase: c.Float("rope.freq_base", 10000.0), + ropeScale: c.Float("rope.freq_scale", 1.0), + attnLogitSoftcap: c.Float("attn_logit_softcapping"), + finalLogitSoftcap: c.Float("final_logit_softcapping"), + }, + } + + slidingWindowLen := int32(c.Uint("attention.sliding_window")) + m.Cache = kvcache.NewWrapperCache(kvcache.NewSWACache(slidingWindowLen, m.Shift), kvcache.NewCausalCache(m.Shift)) + m.Cache.SetConfig(ml.CacheConfig{}) + + return &m, nil +} + +type SelfAttention struct { + Query *nn.Linear `gguf:"attn_q"` + Key *nn.Linear `gguf:"attn_k"` + Value *nn.Linear `gguf:"attn_v"` + Output *nn.Linear `gguf:"attn_output"` +} + +func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor { + batchSize := hiddenState.Dim(1) + ropeType := uint32(2) + + q := sa.Query.Forward(ctx, hiddenState) + q = q.Reshape(ctx, opts.attnKeyLen, opts.numHeads, batchSize) + q = q.RoPE(ctx, positionIDs, nil, uint32(opts.attnKeyLen), ropeType, opts.ropeBase, opts.ropeScale) + + if opts.largeModelScaling { + q = q.Scale(ctx, 1.0/math.Sqrt(float64(opts.hiddenSize/opts.numHeads))) + } else { + q = q.Scale(ctx, 1.0/math.Sqrt(float64(opts.attnKeyLen))) + } + + k := sa.Key.Forward(ctx, hiddenState) + k = k.Reshape(ctx, opts.attnKeyLen, opts.numKVHeads, batchSize) + k = k.RoPE(ctx, positionIDs, nil, uint32(opts.attnKeyLen), ropeType, opts.ropeBase, opts.ropeScale) + + v := sa.Value.Forward(ctx, hiddenState) + v = v.Reshape(ctx, opts.attnValLen, opts.numKVHeads, batchSize) + + cache.Put(ctx, k, v) + k, v, mask := cache.Get(ctx) + + q = q.Permute(ctx, 0, 2, 1, 3) + k = k.Permute(ctx, 0, 2, 1, 3) + v = v.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx) + + kq := k.Mulmat(ctx, q) + + // logit softcap + kq = kq.Scale(ctx, 1.0/float64(opts.attnLogitSoftcap)) + kq = kq.Tanh(ctx) + kq = kq.Scale(ctx, float64(opts.attnLogitSoftcap)) + + kq = kq.Add(ctx, mask) + kq = kq.Softmax(ctx) + + kqv := v.Mulmat(ctx, kq) + kqv = kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) + kqv = kqv.Reshape(ctx, opts.attnValLen*opts.numHeads, batchSize) + + return sa.Output.Forward(ctx, kqv) +} + +func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { + return key.RoPE(ctx, shift, nil, uint32(m.Options.attnKeyLen), uint32(2), m.Options.ropeBase, m.Options.ropeScale), nil +} + +type MLP struct { + Up *nn.Linear `gguf:"ffn_up"` + Down *nn.Linear `gguf:"ffn_down"` + Gate *nn.Linear `gguf:"ffn_gate"` +} + +func (mlp *MLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *Options) ml.Tensor { + hiddenState = mlp.Gate.Forward(ctx, hiddenState).GELU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenState)) + return mlp.Down.Forward(ctx, hiddenState) +} + +type Layer struct { + AttentionNorm *nn.RMSNorm `gguf:"attn_norm"` + SelfAttention *SelfAttention + PostAttentionNorm *nn.RMSNorm `gguf:"post_attention_norm"` + MLPNorm *nn.RMSNorm `gguf:"ffn_norm"` + MLP *MLP + PostMLPNorm *nn.RMSNorm `gguf:"post_ffw_norm"` +} + +func (l *Layer) Forward(ctx ml.Context, hiddenState, positionIDs, outputs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor { + residual := hiddenState + + hiddenState = l.AttentionNorm.Forward(ctx, hiddenState, opts.eps) + hiddenState = l.SelfAttention.Forward(ctx, hiddenState, positionIDs, cache, opts) + hiddenState = l.PostAttentionNorm.Forward(ctx, hiddenState, opts.eps) + + // In the final layer (outputs != nil), optimize by pruning to just the token positions + // we need logits for. + if outputs != nil { + hiddenState = hiddenState.Rows(ctx, outputs) + residual = residual.Rows(ctx, outputs) + } + + hiddenState = hiddenState.Add(ctx, residual) + residual = hiddenState + + hiddenState = l.MLPNorm.Forward(ctx, hiddenState, opts.eps) + hiddenState = l.MLP.Forward(ctx, hiddenState, opts) + hiddenState = l.PostMLPNorm.Forward(ctx, hiddenState, opts.eps) + return hiddenState.Add(ctx, residual) +} + +func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) { + inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs)) + if err != nil { + return nil, err + } + + positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions)) + if err != nil { + return nil, err + } + + outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs)) + if err != nil { + return nil, err + } + + hiddenState := m.TokenEmbedding.Forward(ctx, inputs) + hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(m.Options.hiddenSize))) + + if len(m.Layers) == gemma27BLayerCount { + m.Options.largeModelScaling = true + } + + for i, layer := range m.Layers { + cacheType := i % 2 + m.Cache.SetLayer(i) + wc := m.Cache.(*kvcache.WrapperCache) + wc.SetLayerType(cacheType) + + var lastLayerOutputs ml.Tensor + if i == len(m.Layers)-1 { + lastLayerOutputs = outputs + } + + hiddenState = layer.Forward(ctx, hiddenState, positions, lastLayerOutputs, m.Cache, m.Options) + } + + hiddenState = m.OutputNorm.Forward(ctx, hiddenState, m.eps) + hiddenState = m.Output.Forward(ctx, hiddenState) + + // final logit softcap + hiddenState = hiddenState.Scale(ctx, 1.0/float64(m.Options.finalLogitSoftcap)) + hiddenState = hiddenState.Tanh(ctx) + hiddenState = hiddenState.Scale(ctx, float64(m.Options.finalLogitSoftcap)) + return hiddenState.Rows(ctx, outputs), nil +} + +func init() { + model.Register("gemma2", New) +} diff --git a/model/models/gemma3/model.go b/model/models/gemma3/model.go new file mode 100644 index 00000000..b5311f18 --- /dev/null +++ b/model/models/gemma3/model.go @@ -0,0 +1,173 @@ +package gemma3 + +import ( + "bytes" + "encoding/binary" + "hash/fnv" + "image" + "math" + + "github.com/ollama/ollama/kvcache" + "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/ml/nn" + "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" +) + +type Model struct { + model.Base + model.SentencePieceModel + + *VisionModel `gguf:"v,vision"` + *TextModel + + *MultiModalProjector `gguf:"mm"` + + ImageProcessor +} + +var _ model.MultimodalProcessor = (*Model)(nil) + +type MultiModalProjector struct { + SoftEmbNorm *nn.RMSNorm `gguf:"mm_soft_emb_norm"` + InputProjection *nn.Linear `gguf:"mm_input_projection"` + + tokensPerImage int +} + +func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, imageSize, patchSize int, eps float32) ml.Tensor { + l := visionOutputs.Dim(0) + + visionOutputs = visionOutputs.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) + patchesPerImage := imageSize / patchSize + visionOutputs = visionOutputs.Reshape(ctx, patchesPerImage, patchesPerImage, l) + + kernelSize := patchesPerImage / int(math.Sqrt(float64(p.tokensPerImage))) + visionOutputs = visionOutputs.AvgPool2D(ctx, kernelSize, kernelSize, 0) + visionOutputs = visionOutputs.Reshape(ctx, visionOutputs.Dim(0)*visionOutputs.Dim(1), l) + visionOutputs = visionOutputs.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) + visionOutputs = p.SoftEmbNorm.Forward(ctx, visionOutputs, eps) + + // TODO: inputProjection must be transposed since they're incompatible with visionOutputs + visionOutputs = p.InputProjection.Weight.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx).Mulmat(ctx, visionOutputs) + return visionOutputs +} + +func New(c ml.Config) (model.Model, error) { + m := Model{ + SentencePieceModel: model.NewSentencePieceModel( + c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`), + &model.Vocabulary{ + Values: c.Strings("tokenizer.ggml.tokens"), + Scores: c.Floats("tokenizer.ggml.scores"), + Types: c.Uints("tokenizer.ggml.token_type"), + BOS: int32(c.Uint("tokenizer.ggml.bos_token_id")), + AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true), + EOS: int32(1), + AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false), + EOT: int32(106), + AddEOT: c.Bool("tokenizer.ggml.add_eot_token", false), + }, + ), + ImageProcessor: newImageProcessor(c), + VisionModel: newVisionModel(c), + TextModel: newTextModel(c), + MultiModalProjector: &MultiModalProjector{ + tokensPerImage: int(c.Uint("mm_tokens_per_image", 256)), + }, + } + + slidingWindowLen := int32(c.Uint("attention.sliding_window")) + m.Cache = kvcache.NewWrapperCache(kvcache.NewSWACache(slidingWindowLen, m.Shift), kvcache.NewCausalCache(m.Shift)) + + return &m, nil +} + +func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, error) { + image, _, err := image.Decode(bytes.NewReader(multimodalData)) + if err != nil { + return nil, err + } + + f32s, err := m.ImageProcessor.ProcessImage(image) + if err != nil { + return nil, err + } + + pixelValues, err := ctx.Input().FromFloatSlice(f32s, + m.ImageProcessor.imageSize, + m.ImageProcessor.imageSize, + m.ImageProcessor.numChannels, + ) + if err != nil { + return nil, err + } + + visionOutputs := m.VisionModel.Forward(ctx, pixelValues) + visionOutputs = m.MultiModalProjector.Forward(ctx, visionOutputs, m.imageSize, m.patchSize, m.VisionModel.eps) + return visionOutputs, nil +} + +type imageToken struct { + embedding ml.Tensor + index int +} + +func (m *Model) PostTokenize(ctx ml.Context, inputs []input.Input) ([]input.Input, error) { + var result []input.Input + fnvHash := fnv.New64a() + + for _, inp := range inputs { + if inp.Multimodal == nil { + result = append(result, inp) + } else { + imageInputs := []input.Input{ + {Token: 108}, // "\n\n" + {Token: 255999}, // """ + } + result = append(result, imageInputs...) + + // add image embeddings + inputMultimodal := inp.Multimodal.(ml.Tensor) + + for i := range inputMultimodal.Dim(1) { + fnvHash.Reset() + binary.Write(fnvHash, binary.NativeEndian, inp.MultimodalHash) + fnvHash.Write([]byte{byte(i)}) + + imageToken := imageToken{embedding: inputMultimodal, index: i} + result = append(result, input.Input{Multimodal: imageToken, MultimodalHash: fnvHash.Sum64()}) + } + + result = append(result, + input.Input{Token: 256000}, // + input.Input{Token: 108}, // "\n\n" + ) + } + } + + return result, nil +} + +func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) { + inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs)) + if err != nil { + return nil, err + } + + positions, err := ctx.Input().FromIntSlice(opts.Positions, len(opts.Positions)) + if err != nil { + return nil, err + } + + outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs)) + if err != nil { + return nil, err + } + + return m.TextModel.Forward(ctx, inputs, positions, outputs, opts, m.Cache), nil +} + +func init() { + model.Register("gemma3", New) +} diff --git a/model/models/gemma3/model_text.go b/model/models/gemma3/model_text.go new file mode 100644 index 00000000..5b5e2d6e --- /dev/null +++ b/model/models/gemma3/model_text.go @@ -0,0 +1,254 @@ +package gemma3 + +import ( + "math" + + "github.com/ollama/ollama/kvcache" + "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/ml/nn" + "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" +) + +type TextOptions struct { + hiddenSize, numHeads, numKVHeads int + attnKeyLen, attnValLen int + eps, ropeScale float32 + ropeLocalBase, ropeGlobalBase float32 + finalLogitSoftcap float32 + largeModelScaling bool +} + +type TextModel struct { + model.Base + model.SentencePieceModel + + TokenEmbedding *nn.Embedding `gguf:"token_embd"` + Layers []TextLayer `gguf:"blk"` + OutputNorm *nn.RMSNorm `gguf:"output_norm"` + Output *nn.Linear `gguf:"output,alt:token_embd"` + + *TextOptions +} + +const ( + gemmaGlobalCacheCount = 6 + gemma27BLayerCount = 62 +) + +const ( + cacheTypeSWA = iota + cacheTypeCausal +) + +func newTextModel(c ml.Config) *TextModel { + numBlocks := int(c.Uint("block_count")) + + m := TextModel{ + SentencePieceModel: model.NewSentencePieceModel( + c.String("tokenizer.ggml.pretokenizer", `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`), + &model.Vocabulary{ + Values: c.Strings("tokenizer.ggml.tokens"), + Scores: c.Floats("tokenizer.ggml.scores"), + Types: c.Uints("tokenizer.ggml.token_type"), + BOS: int32(c.Uint("tokenizer.ggml.bos_token_id")), + EOS: int32(c.Uint("tokenizer.ggml.eos_token_id")), + }, + ), + Layers: make([]TextLayer, numBlocks), + TextOptions: &TextOptions{ + hiddenSize: int(c.Uint("embedding_length")), + numHeads: int(c.Uint("attention.head_count")), + numKVHeads: int(c.Uint("attention.head_count_kv")), + attnKeyLen: int(c.Uint("attention.key_length", 256)), + attnValLen: int(c.Uint("attention.value_length", 256)), + eps: c.Float("attention.layer_norm_rms_epsilon", 1e-06), + ropeLocalBase: c.Float("rope.local.freq_base", 10000.0), + ropeGlobalBase: c.Float("rope.global.freq_base", 1000000.0), + ropeScale: c.Float("rope.freq_scale", 1.0), + finalLogitSoftcap: c.Float("final_logit_softcapping", 30.0), + }, + } + + if numBlocks == gemma27BLayerCount { + m.largeModelScaling = true + } + + return &m +} + +type TextSelfAttention struct { + Query *nn.Linear `gguf:"attn_q"` + QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"` + Key *nn.Linear `gguf:"attn_k"` + KeyNorm *nn.RMSNorm `gguf:"attn_k_norm"` + Value *nn.Linear `gguf:"attn_v"` + Output *nn.Linear `gguf:"attn_output"` +} + +func (sa *TextSelfAttention) Forward(ctx ml.Context, layer int, hiddenState, positionIDs ml.Tensor, cache kvcache.Cache, opts *TextOptions) ml.Tensor { + batchSize := hiddenState.Dim(1) + ropeType := uint32(2) + + ropeBase := opts.ropeLocalBase + if (layer+1)%gemmaGlobalCacheCount == 0 { + ropeBase = opts.ropeGlobalBase + } + + q := sa.Query.Forward(ctx, hiddenState) + q = q.Reshape(ctx, opts.attnKeyLen, opts.numHeads, batchSize) + q = sa.QueryNorm.Forward(ctx, q, opts.eps) + q = q.RoPE(ctx, positionIDs, nil, uint32(opts.attnKeyLen), ropeType, ropeBase, opts.ropeScale) + + if opts.largeModelScaling { + q = q.Scale(ctx, 1.0/math.Sqrt(float64(opts.hiddenSize/opts.numHeads))) + } else { + q = q.Scale(ctx, 1.0/math.Sqrt(float64(opts.attnKeyLen))) + } + + k := sa.Key.Forward(ctx, hiddenState) + k = k.Reshape(ctx, opts.attnKeyLen, opts.numKVHeads, batchSize) + k = sa.KeyNorm.Forward(ctx, k, opts.eps) + k = k.RoPE(ctx, positionIDs, nil, uint32(opts.attnKeyLen), ropeType, ropeBase, opts.ropeScale) + + v := sa.Value.Forward(ctx, hiddenState) + v = v.Reshape(ctx, opts.attnValLen, opts.numKVHeads, batchSize) + + scaleFactor := 1.0 + kqv := nn.Attention(ctx, q, k, v, scaleFactor, cache) + kqv = kqv.Reshape(ctx, opts.attnValLen*opts.numHeads, batchSize) + + return sa.Output.Forward(ctx, kqv) +} + +func (m *TextModel) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { + ropeBase := m.TextOptions.ropeLocalBase + if (layer+1)%gemmaGlobalCacheCount == 0 { + ropeBase = m.TextOptions.ropeGlobalBase + } + + return key.RoPE(ctx, shift, nil, uint32(m.TextOptions.attnKeyLen), uint32(2), ropeBase, m.TextOptions.ropeScale), nil +} + +type TextMLP struct { + Up *nn.Linear `gguf:"ffn_up"` + Down *nn.Linear `gguf:"ffn_down"` + Gate *nn.Linear `gguf:"ffn_gate"` +} + +func (mlp *TextMLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *TextOptions) ml.Tensor { + hiddenState = mlp.Gate.Forward(ctx, hiddenState).GELU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenState)) + return mlp.Down.Forward(ctx, hiddenState) +} + +type TextLayer struct { + AttentionNorm *nn.RMSNorm `gguf:"attn_norm"` + SelfAttention *TextSelfAttention + PostAttentionNorm *nn.RMSNorm `gguf:"post_attention_norm"` + MLPNorm *nn.RMSNorm `gguf:"ffn_norm"` + MLP *TextMLP + PostMLPNorm *nn.RMSNorm `gguf:"post_ffw_norm"` +} + +func (l *TextLayer) Forward(ctx ml.Context, layer int, hiddenState, positionIDs, outputs ml.Tensor, cache kvcache.Cache, opts *TextOptions) ml.Tensor { + residual := hiddenState + + hiddenState = l.AttentionNorm.Forward(ctx, hiddenState, opts.eps) + hiddenState = l.SelfAttention.Forward(ctx, layer, hiddenState, positionIDs, cache, opts) + hiddenState = l.PostAttentionNorm.Forward(ctx, hiddenState, opts.eps) + + // In the final layer (outputs != nil), optimize by pruning to just the token positions + // we need logits for. + if outputs != nil { + hiddenState = hiddenState.Rows(ctx, outputs) + residual = residual.Rows(ctx, outputs) + } + + hiddenState = hiddenState.Add(ctx, residual) + residual = hiddenState + + hiddenState = l.MLPNorm.Forward(ctx, hiddenState, opts.eps) + hiddenState = l.MLP.Forward(ctx, hiddenState, opts) + hiddenState = l.PostMLPNorm.Forward(ctx, hiddenState, opts.eps) + return hiddenState.Add(ctx, residual) +} + +func setImageEmbeddings(ctx ml.Context, hiddenState ml.Tensor, multimodal []input.MultimodalIndex) []int { + var embedding ml.Tensor + var src, dst, length int + var except []int + + for _, image := range multimodal { + imageToken := image.Multimodal.(imageToken) + imageSrc := imageToken.index + imageDst := image.Index + + if embedding == nil { + embedding = imageToken.embedding + src = imageSrc + dst = imageDst + length = 1 + } else if embedding == imageToken.embedding && imageSrc+1 == src && imageDst+1 == dst { + src = imageSrc + dst = imageDst + length++ + } else if embedding == imageToken.embedding && src+length == imageSrc && dst+length == imageDst { + length++ + } else { + visionOutputs := embedding.View(ctx, src*embedding.Stride(1), length*embedding.Dim(0)) + ctx.Forward(visionOutputs.Copy(ctx, hiddenState.View(ctx, dst*hiddenState.Stride(1), length*hiddenState.Dim(0)))) + + embedding = imageToken.embedding + src = imageSrc + dst = imageDst + length = 1 + } + + except = append(except, imageDst) + } + + if embedding != nil { + visionOutputs := embedding.View(ctx, src*embedding.Stride(1), length*embedding.Dim(0)) + ctx.Forward(visionOutputs.Copy(ctx, hiddenState.View(ctx, dst*hiddenState.Stride(1), length*hiddenState.Dim(0)))) + } + + return except +} + +func (m *TextModel) Forward(ctx ml.Context, inputs, positions, outputs ml.Tensor, opts input.Options, cache kvcache.Cache) ml.Tensor { + hiddenState := m.TokenEmbedding.Forward(ctx, inputs) + hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(m.TextOptions.hiddenSize))) + + except := setImageEmbeddings(ctx, hiddenState, opts.Multimodal) + + for i, layer := range m.Layers { + // gemma alternates between the sliding window (local) and causal (global) + // kv cache every 6 layers + cacheType := cacheTypeSWA + if (i+1)%gemmaGlobalCacheCount == 0 { + cacheType = cacheTypeCausal + } + cache.SetLayer(i) + wc := cache.(*kvcache.WrapperCache) + wc.SetLayerType(cacheType) + + if causal, ok := wc.UnderlyingCache().(*kvcache.Causal); ok { + causal.SetCausal(ctx, kvcache.CausalOptions{Except: except}) + } + + var lastLayerOutputs ml.Tensor + if i == len(m.Layers)-1 { + lastLayerOutputs = outputs + } + + hiddenState = layer.Forward(ctx, i, hiddenState, positions, lastLayerOutputs, cache, m.TextOptions) + } + + hiddenState = m.OutputNorm.Forward(ctx, hiddenState, m.eps) + hiddenState = m.Output.Forward(ctx, hiddenState) + + // final logit softcap + hiddenState = hiddenState.Scale(ctx, 1.0/float64(m.TextOptions.finalLogitSoftcap)) + hiddenState = hiddenState.Tanh(ctx) + return hiddenState.Scale(ctx, float64(m.TextOptions.finalLogitSoftcap)) +} diff --git a/model/models/gemma3/model_vision.go b/model/models/gemma3/model_vision.go new file mode 100644 index 00000000..94aa27bd --- /dev/null +++ b/model/models/gemma3/model_vision.go @@ -0,0 +1,127 @@ +package gemma3 + +import ( + "math" + + "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/ml/nn" +) + +var batchSize int = 1 + +type VisionSelfAttention struct { + Query *nn.Linear `gguf:"attn_q"` + Key *nn.Linear `gguf:"attn_k"` + Value *nn.Linear `gguf:"attn_v"` + Output *nn.Linear `gguf:"attn_output"` +} + +func (sa *VisionSelfAttention) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionModelOptions) ml.Tensor { + headDim := opts.hiddenSize / opts.numHeads + + query := sa.Query.Forward(ctx, hiddenState) + key := sa.Key.Forward(ctx, hiddenState) + value := sa.Value.Forward(ctx, hiddenState) + + query = query.Reshape(ctx, headDim, opts.numHeads, query.Dim(1), batchSize) + key = key.Reshape(ctx, headDim, opts.numHeads, key.Dim(1), batchSize) + value = value.Reshape(ctx, headDim, opts.numHeads, value.Dim(1), batchSize) + + attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(headDim)), nil) + attention = attention.Reshape(ctx, opts.hiddenSize, attention.Dim(2), batchSize) + + hiddenState = sa.Output.Forward(ctx, attention) + return hiddenState +} + +type VisionMLP struct { + FC1 *nn.Linear `gguf:"fc1"` + FC2 *nn.Linear `gguf:"fc2"` +} + +func (mlp *VisionMLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionModelOptions) ml.Tensor { + hiddenState = mlp.FC1.Forward(ctx, hiddenState).GELU(ctx) + hiddenState = mlp.FC2.Forward(ctx, hiddenState) + return hiddenState +} + +type VisionEncoderLayer struct { + LayerNorm1 *nn.LayerNorm `gguf:"layer_norm1"` + SelfAttention *VisionSelfAttention + + LayerNorm2 *nn.LayerNorm `gguf:"layer_norm2"` + MLP *VisionMLP `gguf:"mlp"` +} + +func (e *VisionEncoderLayer) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionModelOptions) ml.Tensor { + residual := hiddenState + + // self attention + hiddenState = e.LayerNorm1.Forward(ctx, hiddenState, opts.eps) + hiddenState = e.SelfAttention.Forward(ctx, hiddenState, opts) + hiddenState = hiddenState.Add(ctx, residual) + residual = hiddenState + + // feed forward + hiddenState = e.LayerNorm2.Forward(ctx, hiddenState, opts.eps) + hiddenState = e.MLP.Forward(ctx, hiddenState, opts) + return hiddenState.Add(ctx, residual) +} + +type VisionModelOptions struct { + hiddenSize, numHeads int + imageSize, patchSize int + eps float32 +} + +type VisionModel struct { + PatchEmbedding *nn.Conv2D `gguf:"patch_embedding"` + PositionEmbedding *nn.Embedding `gguf:"position_embedding"` + PostLayerNorm *nn.LayerNorm `gguf:"post_layernorm"` + + Layers []VisionEncoderLayer `gguf:"blk"` + + *VisionModelOptions +} + +func (m *VisionModel) Forward(ctx ml.Context, pixelValues ml.Tensor) ml.Tensor { + numPatches := (m.imageSize / m.patchSize) * (m.imageSize / m.patchSize) + + hiddenState := m.PatchEmbedding.Forward(ctx, pixelValues, m.patchSize, m.patchSize, 0, 0, 1, 1) + hiddenState = hiddenState.Reshape(ctx, numPatches, m.hiddenSize) + hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) + + positions := make([]int32, numPatches) + for i := range positions { + positions[i] = int32(i) + } + + positionIDs, err := ctx.Input().FromIntSlice(positions, len(positions)) + if err != nil { + panic(err) + } + + hiddenState = hiddenState.Add(ctx, m.PositionEmbedding.Forward(ctx, positionIDs)) + + for _, layer := range m.Layers { + hiddenState = layer.Forward(ctx, hiddenState, m.VisionModelOptions) + } + + hiddenState = m.PostLayerNorm.Forward(ctx, hiddenState, m.eps) + return hiddenState +} + +func newVisionModel(c ml.Config) *VisionModel { + return &VisionModel{ + Layers: make([]VisionEncoderLayer, c.Uint("vision.block_count")), + VisionModelOptions: &VisionModelOptions{ + hiddenSize: int(c.Uint("vision.embedding_length")), + numHeads: int(c.Uint("vision.attention.head_count")), + + imageSize: int(c.Uint("vision.image_size")), + patchSize: int(c.Uint("vision.patch_size")), + + eps: c.Float("vision.attention.layer_norm_epsilon"), + }, + } +} diff --git a/model/models/gemma3/process_image.go b/model/models/gemma3/process_image.go new file mode 100644 index 00000000..fe8269a3 --- /dev/null +++ b/model/models/gemma3/process_image.go @@ -0,0 +1,58 @@ +package gemma3 + +import ( + "image" + + "github.com/ollama/ollama/ml" + "github.com/ollama/ollama/model/imageproc" +) + +type ImageProcessor struct { + imageSize, patchSize, numChannels int +} + +func newImageProcessor(c ml.Config) ImageProcessor { + return ImageProcessor{ + imageSize: int(c.Uint("vision.image_size")), + patchSize: int(c.Uint("vision.patch_size")), + numChannels: int(c.Uint("vision.num_channels")), + } +} + +func (p *ImageProcessor) pack(img image.Image, mean, std [3]float32) []float32 { + var pixelVals, rVals, gVals, bVals []float32 + + bounds := img.Bounds() + for y := bounds.Min.Y; y < bounds.Max.Y; y++ { + for x := bounds.Min.X; x < bounds.Max.X; x++ { + c := img.At(x, y) + r, g, b, _ := c.RGBA() + rVal := float32(r>>8) / 255.0 + gVal := float32(g>>8) / 255.0 + bVal := float32(b>>8) / 255.0 + + rVal = (rVal - mean[0]) / std[0] + gVal = (gVal - mean[1]) / std[1] + bVal = (bVal - mean[2]) / std[2] + + rVals = append(rVals, rVal) + gVals = append(gVals, gVal) + bVals = append(bVals, bVal) + } + } + + pixelVals = append(pixelVals, rVals...) + pixelVals = append(pixelVals, gVals...) + pixelVals = append(pixelVals, bVals...) + + return pixelVals +} + +func (p ImageProcessor) ProcessImage(img image.Image) ([]float32, error) { + outputSize := image.Point{p.imageSize, p.imageSize} + newImage := imageproc.Composite(img) + newImage = imageproc.Resize(newImage, outputSize, imageproc.ResizeBilinear) + + data := p.pack(newImage, imageproc.ImageNetStandardMean, imageproc.ImageNetStandardSTD) + return data, nil +} diff --git a/model/models/llama/model.go b/model/models/llama/model.go index 9ccfff61..19a2ab8c 100644 --- a/model/models/llama/model.go +++ b/model/models/llama/model.go @@ -9,6 +9,7 @@ import ( "github.com/ollama/ollama/ml" "github.com/ollama/ollama/ml/nn" "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" ) type Options struct { @@ -75,14 +76,15 @@ type SelfAttention struct { func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor { batchSize := hiddenState.Dim(1) headDim := opts.hiddenSize / opts.numHeads + ropeType := uint32(0) q := sa.Query.Forward(ctx, hiddenState) q = q.Reshape(ctx, headDim, opts.numHeads, batchSize) - q = q.RoPE(ctx, positionIDs, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale) + q = q.RoPE(ctx, positionIDs, sa.RopeFactors, opts.ropeDim, ropeType, opts.ropeBase, opts.ropeScale) k := sa.Key.Forward(ctx, hiddenState) k = k.Reshape(ctx, headDim, opts.numKVHeads, batchSize) - k = k.RoPE(ctx, positionIDs, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale) + k = k.RoPE(ctx, positionIDs, sa.RopeFactors, opts.ropeDim, ropeType, opts.ropeBase, opts.ropeScale) v := sa.Value.Forward(ctx, hiddenState) v = v.Reshape(ctx, headDim, opts.numKVHeads, batchSize) @@ -95,7 +97,7 @@ func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positionIDs ml.Ten } func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { - return key.RoPE(ctx, shift, m.Layers[layer].SelfAttention.RopeFactors, m.ropeDim, m.ropeBase, m.ropeScale), nil + return key.RoPE(ctx, shift, m.Layers[layer].SelfAttention.RopeFactors, uint32(0), m.ropeDim, m.ropeBase, m.ropeScale), nil } type MLP struct { @@ -137,7 +139,7 @@ func (l *Layer) Forward(ctx ml.Context, hiddenState, positionIDs, outputs ml.Ten return hiddenState.Add(ctx, residual) } -func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) { +func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) { inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs)) if err != nil { return nil, err diff --git a/model/models/mllama/model.go b/model/models/mllama/model.go index 54c63296..31ba15df 100644 --- a/model/models/mllama/model.go +++ b/model/models/mllama/model.go @@ -12,6 +12,7 @@ import ( "github.com/ollama/ollama/ml" "github.com/ollama/ollama/ml/nn" "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" ) type Model struct { @@ -101,8 +102,8 @@ func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) (any, er return m.Projector.Forward(ctx, crossAttentionStates), nil } -func (m *Model) PostTokenize(ctx ml.Context, inputs []model.Input) ([]model.Input, error) { - var images []model.Input +func (m *Model) PostTokenize(ctx ml.Context, inputs []input.Input) ([]input.Input, error) { + var images []input.Input fnvHash := fnv.New64a() for i := range inputs { @@ -125,15 +126,15 @@ func (m *Model) PostTokenize(ctx ml.Context, inputs []model.Input) ([]model.Inpu } } - inputs = slices.DeleteFunc(inputs, func(input model.Input) bool { return input.Token == -1 }) + inputs = slices.DeleteFunc(inputs, func(input input.Input) bool { return input.Token == -1 }) return inputs, nil } -func (m *Model) Forward(ctx ml.Context, opts model.Options) (ml.Tensor, error) { +func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) { var crossAttentionStates ml.Tensor - if opts.Multimodal != nil { - crossAttentionStates = opts.Multimodal[0].Multimodal.(ml.Tensor) + if len(opts.Multimodal) > 0 { + crossAttentionStates = opts.Multimodal[len(opts.Multimodal)-1].Multimodal.(ml.Tensor) } inputs, err := ctx.Input().FromIntSlice(opts.Inputs, len(opts.Inputs)) diff --git a/model/models/mllama/model_text.go b/model/models/mllama/model_text.go index 373589f9..1cf30d89 100644 --- a/model/models/mllama/model_text.go +++ b/model/models/mllama/model_text.go @@ -20,14 +20,15 @@ type TextSelfAttention struct { func (sa *TextSelfAttention) Forward(ctx ml.Context, hiddenState, positions, _ ml.Tensor, cache *kvcache.WrapperCache, opts *TextModelOptions) ml.Tensor { batchSize := hiddenState.Dim(1) headDim := opts.hiddenSize / opts.numHeads + ropeType := uint32(0) query := sa.Query.Forward(ctx, hiddenState) query = query.Reshape(ctx, headDim, opts.numHeads, batchSize) - query = query.RoPE(ctx, positions, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale) + query = query.RoPE(ctx, positions, sa.RopeFactors, opts.ropeDim, ropeType, opts.ropeBase, opts.ropeScale) key := sa.Key.Forward(ctx, hiddenState) key = key.Reshape(ctx, headDim, opts.numKVHeads, batchSize) - key = key.RoPE(ctx, positions, sa.RopeFactors, opts.ropeDim, opts.ropeBase, opts.ropeScale) + key = key.RoPE(ctx, positions, sa.RopeFactors, opts.ropeDim, ropeType, opts.ropeBase, opts.ropeScale) value := sa.Value.Forward(ctx, hiddenState) value = value.Reshape(ctx, headDim, opts.numKVHeads, batchSize) @@ -40,8 +41,9 @@ func (sa *TextSelfAttention) Forward(ctx ml.Context, hiddenState, positions, _ m } func (m *TextModel) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { + // This will only get called for layers in the cache, which are just the self attention layers if sa, ok := m.Transformer.Layers[layer].(*TextSelfAttentionDecoderLayer); ok { - return key.RoPE(ctx, shift, sa.SelfAttention.RopeFactors, m.ropeDim, m.ropeBase, m.ropeScale), nil + return key.RoPE(ctx, shift, sa.SelfAttention.RopeFactors, m.ropeDim, uint32(0), m.ropeBase, m.ropeScale), nil } return key, nil diff --git a/model/models/mllama/process_image.go b/model/models/mllama/process_image.go index c298d6ae..c94d14a6 100644 --- a/model/models/mllama/process_image.go +++ b/model/models/mllama/process_image.go @@ -144,8 +144,6 @@ func (p *ImageProcessor) splitToTiles(img image.Image, numTilesSize image.Point) return images } -// remove the "alpha" channel by drawing over a prefilled image -// // remove the "alpha" channel by drawing over a prefilled image // //nolint:unused diff --git a/model/models/models.go b/model/models/models.go index d0b68b32..ce1d2ce0 100644 --- a/model/models/models.go +++ b/model/models/models.go @@ -1,6 +1,8 @@ package models import ( + _ "github.com/ollama/ollama/model/models/gemma2" + _ "github.com/ollama/ollama/model/models/gemma3" _ "github.com/ollama/ollama/model/models/llama" _ "github.com/ollama/ollama/model/models/mllama" ) diff --git a/model/process_text.go b/model/process_text.go index 0d75a0ed..01af65b6 100644 --- a/model/process_text.go +++ b/model/process_text.go @@ -4,6 +4,7 @@ import ( "cmp" "iter" "log/slog" + "slices" "strings" "sync" @@ -18,6 +19,15 @@ const ( SpecialEOS ) +const ( + TOKEN_TYPE_NORMAL = iota + 1 + TOKEN_TYPE_UNKNOWN + TOKEN_TYPE_CONTROL + TOKEN_TYPE_USER_DEFINED + TOKEN_TYPE_UNUSED + TOKEN_TYPE_BYTE +) + type TextProcessor interface { Encode(s string, addSpecial bool) ([]int32, error) Decode([]int32) (string, error) @@ -27,11 +37,11 @@ type TextProcessor interface { type Vocabulary struct { Values []string Types []uint32 - Scores []uint32 + Scores []float32 Merges []string - BOS, EOS int32 - AddBOS, AddEOS bool + BOS, EOS, EOT int32 + AddBOS, AddEOS, AddEOT bool specialOnce sync.Once special []string @@ -48,7 +58,7 @@ func (v *Vocabulary) Is(id int32, special Special) bool { case SpecialBOS: return id == v.BOS case SpecialEOS: - return id == v.EOS + return id == v.EOS || id == v.EOT default: return false } @@ -76,7 +86,9 @@ func (v *Vocabulary) Decode(id int32) string { func (v *Vocabulary) SpecialVocabulary() []string { v.specialOnce.Do(func() { for i := range v.Values { - if v.Types[i] == 3 { + if slices.Contains([]int{105, 106}, i) { + v.special = append(v.special, v.Values[i]) + } else if v.Types[i] == TOKEN_TYPE_CONTROL { v.special = append(v.special, v.Values[i]) } } diff --git a/model/process_text_spm.go b/model/process_text_spm.go new file mode 100644 index 00000000..68e3ed01 --- /dev/null +++ b/model/process_text_spm.go @@ -0,0 +1,246 @@ +package model + +import ( + "iter" + "log/slog" + "strings" + + "github.com/dlclark/regexp2" + queue "github.com/emirpasic/gods/v2/queues/priorityqueue" +) + +const spmWhitespaceSep = "▁" + +func replaceWhitespaceBySeperator(s string) string { + return strings.ReplaceAll(s, " ", spmWhitespaceSep) +} + +type SentencePieceModel struct { + maxTokenLen int + pre *regexp2.Regexp + vocab *Vocabulary +} + +var _ TextProcessor = (*SentencePieceModel)(nil) + +func NewSentencePieceModel(pre string, vocab *Vocabulary) SentencePieceModel { + slog.Debug("Tokens", "num tokens", len(vocab.Values), "vals", vocab.Values[:5], "scores", vocab.Scores[:5], "types", vocab.Types[:5]) + + counter := map[int]int{} + var maxTokenLen int + for cnt := range vocab.Types { + switch vocab.Types[cnt] { + case TOKEN_TYPE_NORMAL, TOKEN_TYPE_USER_DEFINED, TOKEN_TYPE_UNUSED: + maxTokenLen = max(maxTokenLen, len(vocab.Values[cnt])) + fallthrough + default: + counter[int(vocab.Types[cnt])] += 1 + } + } + + slog.Debug("Token counts", "normal", counter[TOKEN_TYPE_NORMAL], "unknown", counter[TOKEN_TYPE_UNKNOWN], "control", counter[TOKEN_TYPE_CONTROL], + "user defined", counter[TOKEN_TYPE_USER_DEFINED], "unused", counter[TOKEN_TYPE_UNUSED], "byte", counter[TOKEN_TYPE_BYTE], + "max token len", maxTokenLen) + + return SentencePieceModel{ + maxTokenLen: maxTokenLen, + pre: regexp2.MustCompile(pre, regexp2.Unicode|regexp2.RE2), + vocab: vocab, + } +} + +func (spm SentencePieceModel) Is(id int32, special Special) bool { + return spm.vocab.Is(id, special) +} + +func (spm *SentencePieceModel) split(s string) iter.Seq[string] { + return func(yield func(string) bool) { + for m, _ := spm.pre.FindStringMatch(s); m != nil; m, _ = spm.pre.FindNextMatch(m) { + if !yield(m.String()) { + break + } + } + } +} + +func (spm SentencePieceModel) Encode(s string, addSpecial bool) ([]int32, error) { + fragments := []fragment{{value: s}} + for _, special := range spm.vocab.SpecialVocabulary() { + // TODO: process special tokens concurrently + id := spm.vocab.Encode(special) + for i := 0; i < len(fragments); i++ { + frag := fragments[i] + if len(frag.ids) > 0 { + continue + } + + var middle []fragment + switch i := strings.Index(frag.value, special); { + case i < 0: + middle = append(middle, frag) + case i > 0: + middle = append(middle, fragment{value: frag.value[:i]}) + fallthrough + default: + middle = append(middle, fragment{value: special, ids: []int32{id}}) + if rest := frag.value[i+len(special):]; rest != "" { + middle = append(middle, fragment{value: rest}) + } + } + + fragments = append(fragments[:i], append(middle, fragments[i+1:]...)...) + } + } + slog.Debug("fragments", "frags", fragments) + + var ids []int32 + for _, frag := range fragments { + if len(frag.ids) > 0 { + ids = append(ids, frag.ids...) + continue + } + + for split := range spm.split(frag.value) { + split = replaceWhitespaceBySeperator(split) + + var sb strings.Builder + sb.Write([]byte(split)) + if id := spm.vocab.Encode(sb.String()); id >= 0 { + ids = append(ids, id) + continue + } + + runes := []rune(sb.String()) + pq := queue.NewWith(func(a, b any) int { + priA := a.(*candidate) + priB := b.(*candidate) + if priA.score > priB.score || (priA.score == priB.score && priA.a < priB.a) { + return -1 + } + return 1 + }) + + merges := make([]merge, len(runes)) + for r := range runes { + merges[r] = merge{ + p: r - 1, + n: r + 1, + runes: []rune{runes[r]}, + } + } + + slog.Debug("tokenizer", "merges", merges) + + pairwise := func(a, b int) *candidate { + if a < 0 || b >= len(runes) { + return nil + } + + left, right := string(merges[a].runes), string(merges[b].runes) + if id := spm.vocab.Encode(left + right); id >= 0 { + return &candidate{ + a: a, + b: b, + score: spm.vocab.Scores[id], + } + } + return nil + } + + for i := range len(runes) - 1 { + if pair := pairwise(i, i+1); pair != nil { + pq.Enqueue(pair) + } + } + + pqv := pq.Values() + for _, v := range pqv { + e := v.(*candidate) + slog.Debug("candidate", "candidate", e) + } + + for !pq.Empty() { + v, _ := pq.Dequeue() + pair := v.(*candidate) + left, right := merges[pair.a], merges[pair.b] + + slog.Debug("pair", "left", left, "right", right) + if len(left.runes) == 0 || len(right.runes) == 0 { + continue + } + + if id := spm.vocab.Encode(string(left.runes) + string(right.runes)); id < 0 { + continue + } + + merges[pair.a].runes = append(left.runes, right.runes...) + merges[pair.b].runes = nil + merges[pair.a].n = right.n + if right.n < len(merges) { + merges[right.n].p = pair.a + } + + if pair := pairwise(merges[pair.a].p, pair.a); pair != nil { + pq.Enqueue(pair) + } + + if pair := pairwise(pair.a, merges[pair.a].n); pair != nil { + pq.Enqueue(pair) + } + } + + slog.Debug("merges", "merges", merges) + + for _, merge := range merges { + if len(merge.runes) > 0 { + if id := spm.vocab.Encode(string(merge.runes)); id >= 0 { + ids = append(ids, id) + } else { + slog.Debug("missing token", "token", string(merge.runes)) + } + } + } + } + } + + if addSpecial && len(ids) > 0 { + if spm.vocab.AddBOS { + if ids[0] == spm.vocab.BOS { + slog.Warn("adding bos token to prompt which already has it", "id", spm.vocab.BOS) + } + + slog.Debug("adding bos token to prompt", "id", spm.vocab.BOS) + ids = append([]int32{spm.vocab.BOS}, ids...) + } + + if spm.vocab.AddEOS { + if ids[len(ids)-1] == spm.vocab.EOS { + slog.Warn("adding eos token to prompt which already has it", "id", spm.vocab.EOS) + } + + slog.Debug("adding eos token to prompt", "id", spm.vocab.EOS) + ids = append(ids, spm.vocab.EOS) + } + } + + return ids, nil +} + +type candidate struct { + a, b int + score float32 +} + +func (spm SentencePieceModel) Decode(ids []int32) (string, error) { + var sb strings.Builder + for _, id := range ids { + data := spm.vocab.Decode(id) + data = strings.ReplaceAll(data, spmWhitespaceSep, " ") + if _, err := sb.WriteString(data); err != nil { + return "", err + } + } + + slog.Debug("decoded", "ids", ids, "text", sb.String()) + return sb.String(), nil +} diff --git a/model/process_text_spm_test.go b/model/process_text_spm_test.go new file mode 100644 index 00000000..a43004db --- /dev/null +++ b/model/process_text_spm_test.go @@ -0,0 +1,118 @@ +package model + +import ( + "log/slog" + "os" + "path/filepath" + "slices" + "testing" + + "google.golang.org/protobuf/proto" + + "github.com/ollama/ollama/convert/sentencepiece" +) + +func loadSentencePieceVocab(t *testing.T) SentencePieceModel { + t.Helper() + + bts, err := os.ReadFile(filepath.Join("testdata", "gemma2", "tokenizer.model")) + if err != nil { + t.Fatal(err) + } + + var spm sentencepiece.ModelProto + if err := proto.Unmarshal(bts, &spm); err != nil { + t.Fatal(err) + } + + preTokenizer := `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+` + + var v Vocabulary + + for _, piece := range spm.GetPieces() { + v.Values = append(v.Values, piece.GetPiece()) + v.Scores = append(v.Scores, piece.GetScore()) + switch t := piece.GetType(); t { + case sentencepiece.ModelProto_SentencePiece_UNKNOWN, + sentencepiece.ModelProto_SentencePiece_CONTROL, + sentencepiece.ModelProto_SentencePiece_UNUSED, + sentencepiece.ModelProto_SentencePiece_BYTE: + v.Types = append(v.Types, uint32(t)) + default: + tt := uint32(sentencepiece.ModelProto_SentencePiece_NORMAL) + // todo parse the special tokens file + // - this will roundtrip correctly but the and + // tokens aren't processed + v.Types = append(v.Types, tt) + } + } + + return NewSentencePieceModel(preTokenizer, &v) +} + +func TestSentencePieceEncode(t *testing.T) { + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug})) + slog.SetDefault(logger) + + tokenizer := loadSentencePieceVocab(t) + + t.Run("basic roundtrip", func(t *testing.T) { + t.Parallel() + + cases := []string{ + "hello", + "hello ", + "hello ", + " hello", + " hello ", + " hello ", + "hello world", + "请考试我的软件！12345", + "你好", + "Hello 你好 world!", + "Special characters: !@#$%^&*()_+-=[]{}|;':\",./<>?", + "Multilingual: 你好 こんにちは Привет Hola مرحبا", + "Numbers and symbols: 123456789 +- */", + "Special tokens: text ", + "Code snippets: func main() { fmt.Println(\"Hello World\") }", + "Long text: " + "Lorem ipsum dolor sit amet, consectetur adipiscing elit. " + + "Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. " + + "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris.", + } + + for _, want := range cases { + ids, err := tokenizer.Encode(want, true) + if err != nil { + t.Fatal(err) + } + + if got, err := tokenizer.Decode(ids); err != nil { + t.Fatal(err) + } else if got != want { + t.Errorf("got %q, want %q [%#v]", got, want, ids) + } + } + }) + + t.Run("special tokens", func(t *testing.T) { + type candidate struct { + token string + ids []int32 + } + + cases := []candidate{ + {"", []int32{2}}, + {"", []int32{1}}, + } + + for _, want := range cases { + ids, err := tokenizer.Encode(want.token, true) + if err != nil { + t.Fatal(err) + } + if !slices.Equal(ids, want.ids) { + t.Errorf("got %#v, want %#v", ids, want.ids) + } + } + }) +} diff --git a/model/testdata/gemma2/tokenizer.model b/model/testdata/gemma2/tokenizer.model new file mode 100644 index 00000000..14a24226 Binary files /dev/null and b/model/testdata/gemma2/tokenizer.model differ diff --git a/runner/ollamarunner/cache.go b/runner/ollamarunner/cache.go index 3244c0b8..a411fddb 100644 --- a/runner/ollamarunner/cache.go +++ b/runner/ollamarunner/cache.go @@ -10,6 +10,7 @@ import ( "github.com/ollama/ollama/kvcache" "github.com/ollama/ollama/ml" "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" ) type InputCache struct { @@ -79,7 +80,7 @@ type InputCacheSlot struct { Id int // Inputs that are stored in the KV cache - Inputs []model.Input + Inputs []input.Input // is this cache actively being processed as part of a sequence? InUse bool @@ -88,7 +89,7 @@ type InputCacheSlot struct { lastUsed time.Time } -func (c *InputCache) LoadCacheSlot(prompt []model.Input, cachePrompt bool) (*InputCacheSlot, []model.Input, error) { +func (c *InputCache) LoadCacheSlot(prompt []input.Input, cachePrompt bool) (*InputCacheSlot, []input.Input, error) { var slot *InputCacheSlot var numPast int32 var err error @@ -139,7 +140,7 @@ func (c *InputCache) LoadCacheSlot(prompt []model.Input, cachePrompt bool) (*Inp return slot, prompt, nil } -func (c *InputCache) findLongestCacheSlot(prompt []model.Input) (*InputCacheSlot, int32, error) { +func (c *InputCache) findLongestCacheSlot(prompt []input.Input) (*InputCacheSlot, int32, error) { longest := int32(-1) var longestSlot *InputCacheSlot @@ -162,7 +163,7 @@ func (c *InputCache) findLongestCacheSlot(prompt []model.Input) (*InputCacheSlot return longestSlot, longest, nil } -func (c *InputCache) findBestCacheSlot(prompt []model.Input) (*InputCacheSlot, int32, error) { +func (c *InputCache) findBestCacheSlot(prompt []input.Input) (*InputCacheSlot, int32, error) { oldest := time.Now() var oldestSlot *InputCacheSlot @@ -198,7 +199,7 @@ func (c *InputCache) findBestCacheSlot(prompt []model.Input) (*InputCacheSlot, i if longest > 0 && longestSlot != oldestSlot { slog.Debug("forking cache slot", "src", longestSlot.Id, "dst", oldestSlot.Id, "inputs", longest, "total", len(longestSlot.Inputs)) - oldestSlot.Inputs = make([]model.Input, longest) + oldestSlot.Inputs = make([]input.Input, longest) copy(oldestSlot.Inputs, longestSlot.Inputs[:longest]) if c.cache != nil { c.cache.CopyPrefix(longestSlot.Id, oldestSlot.Id, longest) @@ -208,7 +209,7 @@ func (c *InputCache) findBestCacheSlot(prompt []model.Input) (*InputCacheSlot, i return oldestSlot, longest, nil } -func countCommonPrefix(a []model.Input, b []model.Input) int32 { +func countCommonPrefix(a []input.Input, b []input.Input) int32 { var count int32 for i := range a { diff --git a/runner/ollamarunner/cache_test.go b/runner/ollamarunner/cache_test.go index 9ce03b73..0a1b73f5 100644 --- a/runner/ollamarunner/cache_test.go +++ b/runner/ollamarunner/cache_test.go @@ -5,7 +5,7 @@ import ( "testing" "time" - "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" ) func TestCountCommon(t *testing.T) { @@ -15,50 +15,50 @@ func TestCountCommon(t *testing.T) { tests := []struct { name string - t1 []model.Input - t2 []model.Input + t1 []input.Input + t2 []input.Input expected int32 }{ { name: "Equal", - t1: []model.Input{{Token: 1}, {Token: 2}, {Token: 3}}, - t2: []model.Input{{Token: 1}, {Token: 2}, {Token: 3}}, + t1: []input.Input{{Token: 1}, {Token: 2}, {Token: 3}}, + t2: []input.Input{{Token: 1}, {Token: 2}, {Token: 3}}, expected: 3, }, { name: "Prefix", - t1: []model.Input{{Token: 1}}, - t2: []model.Input{{Token: 1}, {Token: 2}, {Token: 3}}, + t1: []input.Input{{Token: 1}}, + t2: []input.Input{{Token: 1}, {Token: 2}, {Token: 3}}, expected: 1, }, { name: "Image Prefix", - t1: []model.Input{{Multimodal: imgA, MultimodalHash: 1}}, - t2: []model.Input{{Multimodal: imgA, MultimodalHash: 1}, {Multimodal: imgB, MultimodalHash: 2}, {Multimodal: imgC, MultimodalHash: 3}}, + t1: []input.Input{{Multimodal: imgA, MultimodalHash: 1}}, + t2: []input.Input{{Multimodal: imgA, MultimodalHash: 1}, {Multimodal: imgB, MultimodalHash: 2}, {Multimodal: imgC, MultimodalHash: 3}}, expected: 1, }, { name: "Mixed", - t1: []model.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}}, - t2: []model.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}, {Token: 5}}, + t1: []input.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}}, + t2: []input.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}, {Token: 5}}, expected: 2, }, { name: "Mixed, Same Length", - t1: []model.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}}, - t2: []model.Input{{Token: 1}, {Multimodal: imgB, MultimodalHash: 2}}, + t1: []input.Input{{Token: 1}, {Multimodal: imgA, MultimodalHash: 1}}, + t2: []input.Input{{Token: 1}, {Multimodal: imgB, MultimodalHash: 2}}, expected: 1, }, { name: "Empty", - t1: []model.Input{}, - t2: []model.Input{{Token: 1}, {Token: 2}, {Token: 3}}, + t1: []input.Input{}, + t2: []input.Input{{Token: 1}, {Token: 2}, {Token: 3}}, expected: 0, }, { name: "Both Empty", - t1: []model.Input{}, - t2: []model.Input{}, + t1: []input.Input{}, + t2: []input.Input{}, expected: 0, }, } @@ -82,7 +82,7 @@ func TestFindCacheSlot(t *testing.T) { tests := []struct { name string cache InputCache - prompt []model.Input + prompt []input.Input longest expected best expected }{ @@ -91,18 +91,18 @@ func TestFindCacheSlot(t *testing.T) { cache: InputCache{slots: []InputCacheSlot{ { Id: 0, - Inputs: []model.Input{}, + Inputs: []input.Input{}, InUse: false, lastUsed: time.Time{}, }, { Id: 1, - Inputs: []model.Input{}, + Inputs: []input.Input{}, InUse: false, lastUsed: time.Time{}, }, }}, - prompt: []model.Input{{Token: 1}}, + prompt: []input.Input{{Token: 1}}, longest: expected{result: 0, len: 0}, best: expected{result: 0, len: 0}, }, @@ -111,18 +111,18 @@ func TestFindCacheSlot(t *testing.T) { cache: InputCache{slots: []InputCacheSlot{ { Id: 0, - Inputs: []model.Input{{Token: 1}}, + Inputs: []input.Input{{Token: 1}}, InUse: false, lastUsed: time.Now().Add(-time.Second), }, { Id: 1, - Inputs: []model.Input{{Token: 1}, {Token: 2}}, + Inputs: []input.Input{{Token: 1}, {Token: 2}}, InUse: false, lastUsed: time.Now().Add(-2 * time.Second), }, }}, - prompt: []model.Input{{Token: 1}, {Token: 2}}, + prompt: []input.Input{{Token: 1}, {Token: 2}}, longest: expected{result: 1, len: 2}, best: expected{result: 1, len: 2}, }, @@ -131,18 +131,18 @@ func TestFindCacheSlot(t *testing.T) { cache: InputCache{slots: []InputCacheSlot{ { Id: 0, - Inputs: []model.Input{{Token: 1}, {Token: 2}}, + Inputs: []input.Input{{Token: 1}, {Token: 2}}, InUse: false, lastUsed: time.Now().Add(-time.Second), }, { Id: 1, - Inputs: []model.Input{}, + Inputs: []input.Input{}, InUse: false, lastUsed: time.Time{}, }, }}, - prompt: []model.Input{{Token: 2}}, + prompt: []input.Input{{Token: 2}}, longest: expected{result: 0, len: 0}, best: expected{result: 1, len: 0}, }, @@ -152,19 +152,19 @@ func TestFindCacheSlot(t *testing.T) { slots: []InputCacheSlot{ { Id: 0, - Inputs: []model.Input{{Token: 1}, {Token: 2}}, + Inputs: []input.Input{{Token: 1}, {Token: 2}}, InUse: false, lastUsed: time.Now().Add(-time.Second), }, { Id: 1, - Inputs: []model.Input{}, + Inputs: []input.Input{}, InUse: false, lastUsed: time.Time{}, }, }, }, - prompt: []model.Input{{Token: 1}}, + prompt: []input.Input{{Token: 1}}, longest: expected{result: 0, len: 1}, best: expected{result: 1, len: 1}, }, @@ -173,18 +173,18 @@ func TestFindCacheSlot(t *testing.T) { cache: InputCache{slots: []InputCacheSlot{ { Id: 0, - Inputs: []model.Input{{Token: 1}}, + Inputs: []input.Input{{Token: 1}}, InUse: false, lastUsed: time.Now().Add(-time.Second), }, { Id: 1, - Inputs: []model.Input{{Token: 1}, {Token: 2}}, + Inputs: []input.Input{{Token: 1}, {Token: 2}}, InUse: false, lastUsed: time.Now().Add(-2 * time.Second), }, }}, - prompt: []model.Input{{Token: 2}, {Token: 3}}, + prompt: []input.Input{{Token: 2}, {Token: 3}}, longest: expected{result: 0, len: 0}, best: expected{result: 1, len: 0}, }, @@ -193,18 +193,18 @@ func TestFindCacheSlot(t *testing.T) { cache: InputCache{slots: []InputCacheSlot{ { Id: 0, - Inputs: []model.Input{{Token: 1}, {Token: 2}}, + Inputs: []input.Input{{Token: 1}, {Token: 2}}, InUse: true, lastUsed: time.Now().Add(-time.Second), }, { Id: 1, - Inputs: []model.Input{{Token: 1}}, + Inputs: []input.Input{{Token: 1}}, InUse: false, lastUsed: time.Now().Add(-2 * time.Second), }, }}, - prompt: []model.Input{{Token: 1}, {Token: 2}}, + prompt: []input.Input{{Token: 1}, {Token: 2}}, longest: expected{result: 1, len: 1}, best: expected{result: 1, len: 2}, }, diff --git a/runner/ollamarunner/runner.go b/runner/ollamarunner/runner.go index a51b1459..c1475cbb 100644 --- a/runner/ollamarunner/runner.go +++ b/runner/ollamarunner/runner.go @@ -26,6 +26,7 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/ml" "github.com/ollama/ollama/model" + "github.com/ollama/ollama/model/input" "github.com/ollama/ollama/runner/common" "github.com/ollama/ollama/sample" @@ -41,10 +42,10 @@ type Sequence struct { iBatch int // prompt inputs left to evaluate - inputs []model.Input + inputs []input.Input // inputs that have been added to a batch but not yet submitted to Forward - pendingInputs []model.Input + pendingInputs []input.Input // tokens that have been generated but not returned yet (e.g. for stop sequences) pendingResponses []string @@ -144,8 +145,8 @@ func (s *Server) NewSequence(prompt string, images []ImageData, params NewSequen // inputs processes the prompt and images into a list of inputs // by splitting the prompt on [img-] tags, tokenizing text and // decoding images -func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]model.Input, error) { - var inputs []model.Input +func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]input.Input, error) { + var inputs []input.Input var parts []string var matches [][]string @@ -168,7 +169,7 @@ func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]mo } for _, t := range tokens { - inputs = append(inputs, model.Input{Token: t}) + inputs = append(inputs, input.Input{Token: t}) } // image - decode and store @@ -196,7 +197,7 @@ func (s *Server) inputs(ctx ml.Context, prompt string, images []ImageData) ([]mo _, _ = s.multimodalHash.Write(images[imageIndex].Data) imageHash := s.multimodalHash.Sum64() - inputs = append(inputs, model.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash}) + inputs = append(inputs, input.Input{Multimodal: imageEmbeddings, MultimodalHash: imageHash}) postTokenize = true } } @@ -250,12 +251,15 @@ type Server struct { // KV cache cache *InputCache - // next sequence for prompt processing to avoid starvation - nextSeq int - // multimodalHash generates hashes for comparing equality // of non-text data multimodalHash maphash.Hash + + // vocab is a llama.cpp vocab required for gammar-based + // constrained generation (json mode, structured outputs) + // TODO: this is temporary until Ollama sampling supports + // constrained generation + vocab *sample.Vocab } func (s *Server) allNil() bool { @@ -329,29 +333,25 @@ func (s *Server) processBatch() error { } defer s.mu.Unlock() - var options model.Options - - seqIdx := s.nextSeq - 1 - for range s.seqs { - seqIdx = (seqIdx + 1) % len(s.seqs) - seq := s.seqs[seqIdx] + var options input.Options + for i, seq := range s.seqs { if seq == nil { continue } // if past the num predict limit if seq.numPredict > 0 && seq.numPredicted >= seq.numPredict { - s.removeSequence(seqIdx, "limit") + s.removeSequence(i, "limit") continue } if !s.cache.enabled { seq.inputs = append(seq.cache.Inputs, seq.inputs...) - seq.cache.Inputs = []model.Input{} + seq.cache.Inputs = []input.Input{} } - for i, input := range seq.inputs { + for j, inp := range seq.inputs { if int32(len(seq.cache.Inputs)+len(seq.pendingInputs)+1) > s.cache.numCtx { if len(seq.pendingInputs) == 0 { err := s.cache.ShiftCacheSlot(seq.cache, seq.numKeep) @@ -363,33 +363,23 @@ func (s *Server) processBatch() error { } } - if i >= s.batchSize { + if j >= s.batchSize { break } - // TODO(jessegross): This is a workaround for generating an attention mask and also providing a hint - // to the encoder cache. - // - // Break the batch when switching from text to images so that images are always at the beginning. - if input.Multimodal != nil && !(len(seq.pendingInputs) == 0 || - (len(options.Multimodal) > 0 && options.Multimodal[len(options.Multimodal)-1].Index == len(options.Inputs)-1)) { - s.nextSeq = seqIdx - break - } - - options.Inputs = append(options.Inputs, input.Token) - if input.Multimodal != nil { - options.Multimodal = append(options.Multimodal, model.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: input.Multimodal}) + options.Inputs = append(options.Inputs, inp.Token) + if inp.Multimodal != nil { + options.Multimodal = append(options.Multimodal, input.MultimodalIndex{Index: len(options.Inputs) - 1, Multimodal: inp.Multimodal}) } options.Positions = append(options.Positions, int32(len(seq.cache.Inputs)+len(seq.pendingInputs))) options.Sequences = append(options.Sequences, seq.cache.Id) seq.iBatch = len(options.Outputs) - if i+1 == len(seq.inputs) { + if j+1 == len(seq.inputs) { options.Outputs = append(options.Outputs, int32(len(options.Inputs)-1)) } - seq.pendingInputs = append(seq.pendingInputs, input) + seq.pendingInputs = append(seq.pendingInputs, inp) } seq.inputs = seq.inputs[len(seq.pendingInputs):] @@ -417,7 +407,7 @@ func (s *Server) processBatch() error { // After calling Forward, pending inputs are now in the cache if len(seq.pendingInputs) > 0 { seq.cache.Inputs = append(seq.cache.Inputs, seq.pendingInputs...) - seq.pendingInputs = []model.Input{} + seq.pendingInputs = []input.Input{} } // don't sample prompt processing @@ -464,7 +454,7 @@ func (s *Server) processBatch() error { return err } - seq.inputs = []model.Input{{Token: token}} + seq.inputs = []input.Input{{Token: token}} seq.pendingResponses = append(seq.pendingResponses, piece) sequence := strings.Join(seq.pendingResponses, "") @@ -590,18 +580,25 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) { return } + var grammar *sample.Grammar + var err error + if req.Grammar != "" { + grammar, err = sample.NewGrammar(s.vocab, req.Grammar) + if err != nil { + http.Error(w, "failed to load model vocabulary required for format", http.StatusInternalServerError) + return + } + } + sampler := sample.NewSampler( req.Temperature, req.TopK, req.TopP, req.MinP, req.Seed, + grammar, ) - if req.Grammar != "" { - panic("grammars are not yet supported") - } - seq, err := s.NewSequence(req.Prompt, req.Images, NewSequenceParams{ numPredict: req.NumPredict, stop: req.Stop, @@ -813,6 +810,8 @@ func (s *Server) loadModel( panic(err) } + s.vocab = sample.NewVocab(mpath) + // TODO(jessegross): LoRA loading if lpath.String() != "" { panic("loras are not yet implemented") diff --git a/sample/samplers.go b/sample/samplers.go index a5a0507c..aea99b3f 100644 --- a/sample/samplers.go +++ b/sample/samplers.go @@ -2,57 +2,105 @@ package sample import ( "errors" + "math" "math/rand/v2" "slices" + "sync" + + "github.com/ollama/ollama/llama" ) -// Sampler is not thread-safe. Each goroutine should have its own instance -type Sampler interface { - Sample([]float32) (int32, error) -} - -// logit represents information about a single token during sampling -type logit struct { +// token represents information about a single token during sampling +type token struct { id int32 // The token's unique identifier value float32 // The raw logit or probability from the model } -type weighted struct { +type Sampler struct { rng *rand.Rand - tokens []logit topK int topP float32 minP float32 temperature float32 + grammar *Grammar } -func (s *weighted) Sample(logits []float32) (int32, error) { - if len(s.tokens) < len(logits) { - s.tokens = make([]logit, len(logits)) - } - - tokens := s.tokens[:len(logits)] - - for i, v := range logits { +func (s *Sampler) Sample(logits []float32) (int32, error) { + tokens := make([]token, len(logits)) + for i := range logits { tokens[i].id = int32(i) - tokens[i].value = v + tokens[i].value = logits[i] + } + + t, err := s.sample(tokens) + if err != nil { + return -1, err + } + + if s.grammar != nil { + // optimization: first check if the max logit is accepted by the grammar + // if the max logit is rejected, apply the grammar to all logits (slower) + top := []token{t} + s.grammar.Apply(top) + if !math.IsInf(float64(top[0].value), -1) { + s.grammar.Accept(top[0].id) + return top[0].id, nil + } + + // since .sample has side effects of modifying the tokens + // we need to reset them before applying the grammar and + // sampling again + for i := range logits { + tokens[i].id = int32(i) + tokens[i].value = logits[i] + } + s.grammar.Apply(tokens) + t, err = s.sample(tokens) + if err != nil { + return -1, err + } + s.grammar.Accept(t.id) + } + + return t.id, nil +} + +// greedy returns the highest probability token from the tokens +func greedy(tokens []token) token { + max := tokens[0] + for i := 1; i < len(tokens); i++ { + if tokens[i].value > max.value { + max = tokens[i] + } + } + + return max +} + +// sample returns the highest probability token from the tokens +// given sampler parameters. It also has side effects of modifying the tokens +func (s *Sampler) sample(tokens []token) (token, error) { + if s.temperature == 0 { + return greedy(tokens), nil } - // Tokens are sorted by logits in TopK or SortTokens if s.topK > 0 { tokens = topK(tokens, s.topK) } else { sortLogits(tokens) } + // token logit values are updated to probabilities tokens = temperature(tokens, s.temperature) - tokens = softmax(tokens) tokens = topP(tokens, s.topP) tokens = minP(tokens, s.minP) + // TODO: this should fall back to greedy sampling + // or topP, topK values etc should be such that + // there are always tokens to sample from if len(tokens) == 0 { - return -1, errors.New("no valid logits found for weighted sampling") + return token{}, errors.New("no tokens to sample from") } var r float32 @@ -70,48 +118,18 @@ func (s *weighted) Sample(logits []float32) (int32, error) { } r *= tokens[len(tokens)-1].value - idx, _ := slices.BinarySearchFunc(tokens, r, func(token logit, target float32) int { - // Compare cumulative probabilities + idx, _ := slices.BinarySearchFunc(tokens, r, func(token token, target float32) int { if token.value < target { return -1 } - // First token that exceeds target return 1 }) - if idx >= len(tokens) { - idx = len(tokens) - 1 - } - - return tokens[idx].id, nil -} - -type greedy struct{} - -// Greedy sample returns the index of the maximum value in logits. -func (s greedy) Sample(logits []float32) (int32, error) { - if len(logits) == 0 { - return -1, errors.New("no logits provided for greedy sampling") - } - - maxIdx := 0 - maxVal := logits[0] - for i := 1; i < len(logits); i++ { - if logits[i] > maxVal { - maxVal = logits[i] - maxIdx = i - } - } - - return int32(maxIdx), nil + return tokens[idx], nil } // TODO(parthsareen): update sampler interface to use json unmarshal https://github.com/ollama/ollama/issues/9278 -func NewSampler(temperature float32, topK int, topP float32, minP float32, seed int) Sampler { - if temperature == 0 { - return &greedy{} - } - +func NewSampler(temperature float32, topK int, topP float32, minP float32, seed int, grammar *Grammar) Sampler { var rng *rand.Rand if seed != -1 { // PCG requires two parameters: sequence and stream @@ -120,7 +138,9 @@ func NewSampler(temperature float32, topK int, topP float32, minP float32, seed // Use golden ratio hash to generate statistically independent seeds rng = rand.New(rand.NewPCG(sequence, sequence^0x9E3779B9)) } - temperature = max(temperature, 1) + if temperature < 0.0 { + temperature = 0.0 + } if topP < 0.0 { topP = 0.0 @@ -136,11 +156,73 @@ func NewSampler(temperature float32, topK int, topP float32, minP float32, seed minP = 1.0 } - return &weighted{ + return Sampler{ rng: rng, topK: topK, topP: topP, minP: minP, temperature: temperature, + grammar: grammar, } } + +type Grammar struct { + vocab *Vocab + grammar string + sampler *llama.Sampler +} + +func NewGrammar(vocab *Vocab, grammar string) (*Grammar, error) { + v, err := vocab.Load() + if err != nil { + return nil, err + } + + return &Grammar{ + vocab: vocab, + grammar: grammar, + sampler: llama.NewGrammarSampler(v, grammar), + }, nil +} + +func (g *Grammar) Apply(tokens []token) { + tds := make([]llama.TokenData, len(tokens)) + for i, token := range tokens { + tds[i].Id = token.id + tds[i].Logit = token.value + } + + g.sampler.Apply(tds) + + for i := range tokens { + tokens[i].value = tds[i].Logit + } +} + +func (g *Grammar) Accept(token int32) { + g.sampler.Accept(token) +} + +type Vocab struct { + once sync.Once + vocab *llama.Vocab + err error + path string +} + +func NewVocab(path string) *Vocab { + return &Vocab{path: path} +} + +// Load returns the lazily-loaded vocabulary +func (v *Vocab) Load() (*llama.Vocab, error) { + v.once.Do(func() { + vocab, err := llama.LoadVocabFromFile(v.path) + if err != nil { + v.err = err + return + } + v.vocab = vocab + }) + return v.vocab, v.err +} diff --git a/sample/samplers_benchmark_test.go b/sample/samplers_benchmark_test.go index 41c0b487..cd138014 100644 --- a/sample/samplers_benchmark_test.go +++ b/sample/samplers_benchmark_test.go @@ -16,13 +16,10 @@ func BenchmarkWeightedSampler(b *testing.B) { logits[i] = float32(rand.Float64()*10 - 5) } - sampler := NewSampler(0.8, 0, 0, 0, 42) + sampler := NewSampler(0.8, 0, 0, 0, 42, nil) b.ResetTimer() for b.Loop() { - _, err := sampler.Sample(logits) - if err != nil { - b.Fatalf("Sampling failed: %v", err) - } + sampler.Sample(logits) } }) } @@ -52,30 +49,24 @@ func BenchmarkWeightedSampler(b *testing.B) { for _, tc := range configs { b.Run("Config"+tc.name, func(b *testing.B) { - sampler := NewSampler(tc.temperature, tc.topK, tc.topP, tc.minP, tc.seed) + sampler := NewSampler(tc.temperature, tc.topK, tc.topP, tc.minP, tc.seed, nil) sampler.Sample(logits) b.ResetTimer() for b.Loop() { - _, err := sampler.Sample(logits) - if err != nil { - b.Fatalf("Sampling failed: %v", err) - } + sampler.Sample(logits) } }) } // Test with combined transforms separately - topK influences performance greatly b.Run("TransformCombined", func(b *testing.B) { - sampler := NewSampler(0.8, 50, 0.9, 0.05, 42) + sampler := NewSampler(0.8, 50, 0.9, 0.05, 42, nil) b.ResetTimer() for b.Loop() { - _, err := sampler.Sample(logits) - if err != nil { - b.Fatalf("Sampling failed: %v", err) - } + sampler.Sample(logits) } }) } @@ -90,14 +81,11 @@ func BenchmarkGreedySampler(b *testing.B) { logits[i] = float32(rand.Float64()*10 - 5) } - sampler := NewSampler(0, -1, 0, 0, -1) + sampler := NewSampler(0, -1, 0, 0, -1, nil) b.ResetTimer() for b.Loop() { - _, err := sampler.Sample(logits) - if err != nil { - b.Fatalf("Sampling failed: %v", err) - } + sampler.Sample(logits) } }) } diff --git a/sample/samplers_test.go b/sample/samplers_test.go index dbbee17b..38b9b352 100644 --- a/sample/samplers_test.go +++ b/sample/samplers_test.go @@ -7,7 +7,7 @@ import ( func TestWeighted(t *testing.T) { logits := []float32{-10, 3, -10, -10} - sampler := NewSampler(0, 0, 0, 0, 0) + sampler := NewSampler(0, 0, 0, 0, 0, nil) got, err := sampler.Sample(logits) if err != nil { t.Error(err) @@ -19,7 +19,7 @@ func TestWeighted(t *testing.T) { } logits = []float32{-100, -10, 0, 10} - sampler = NewSampler(0, 0, 0, 0, 0) + sampler = NewSampler(0, 0, 0, 0, 0, nil) got, err = sampler.Sample(logits) if err != nil { t.Error(err) @@ -31,94 +31,10 @@ func TestWeighted(t *testing.T) { } } -func TestNewSampler(t *testing.T) { - tests := []struct { - name string - temperature float32 - topK int - topP float32 - minP float32 - seed int - wantGreedy bool // Instead of wantErr, check if we get greedy sampler - }{ - { - name: "temperature", - temperature: 0.5, - wantGreedy: false, - }, - { - name: "zero temperature - greedy", - temperature: 0, - wantGreedy: true, - }, - { - name: "top k", - temperature: 0.1, - topK: 10, - wantGreedy: false, - }, - { - name: "top p", - temperature: 0.1, - topP: 0.9, - wantGreedy: false, - }, - { - name: "min p", - temperature: 0.1, - minP: 0.2, - wantGreedy: false, - }, - { - name: "seed - weighted", - temperature: 0.1, - seed: 42, - wantGreedy: false, - }, - { - name: "default values", - temperature: 0.8, - topK: 40, - topP: 0.9, - minP: 0.0, - seed: 0, - wantGreedy: false, - }, - { - name: "all zeroes - greedy", - temperature: 0.0, - topK: 0, - topP: 0.0, - minP: 0.0, - seed: 0, - wantGreedy: true, - }, - { - name: "all transforms", - temperature: 0.8, - topK: 50, - topP: 0.95, - minP: 0.1, - seed: 42, - wantGreedy: false, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - sampler := NewSampler(tt.temperature, tt.topK, tt.topP, tt.minP, tt.seed) - _, isGreedy := sampler.(*greedy) - if isGreedy != tt.wantGreedy { - t.Errorf("NewSampler() got greedy = %v, want %v", isGreedy, tt.wantGreedy) - } - }) - } -} - func BenchmarkSample(b *testing.B) { - weighted := NewSampler(0.5, 10, 0.9, 0.2, -1) samplers := map[string]Sampler{ - "Greedy": NewSampler(0, 0, 0, 0, 0), // Use NewSampler with temp=0 for greedy - "Weighted": weighted, + "Greedy": NewSampler(0, 0, 0, 0, 0, nil), // Use NewSampler with temp=0 for greedy + "Weighted": NewSampler(0.5, 10, 0.9, 0.2, -1, nil), } // Generate random logits for benchmarking @@ -132,7 +48,7 @@ func BenchmarkSample(b *testing.B) { b.ResetTimer() for b.Loop() { if _, err := s.Sample(logits); err != nil { - b.Error(err) + b.Fatalf("error sampling: %v", err) } } }) diff --git a/sample/transforms.go b/sample/transforms.go index f1f4f3b1..ab62455f 100644 --- a/sample/transforms.go +++ b/sample/transforms.go @@ -5,13 +5,25 @@ import ( "slices" ) -func softmax(ts []logit) []logit { +// temperature applies scaling and softmax to the logits +func temperature(ts []token, temp float32) []token { + // Find max logit for numerical stability + maxLogit := float32(math.Inf(-1)) + for _, t := range ts { + if t.value > maxLogit { + maxLogit = t.value + } + } + + // Apply temperature and compute exp(x - max) + temp = max(temp, 1e-7) var sum float32 for i, v := range ts { - ts[i].value = float32(math.Exp(float64(v.value))) + ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp))) sum += ts[i].value } + // Normalize for i := range ts { ts[i].value /= sum } @@ -19,27 +31,6 @@ func softmax(ts []logit) []logit { return ts } -func temperature(ti []logit, t float32) []logit { - if t == 1 { - return ti - } - - temp := max(t, 1e-7) - maxLogit := float32(math.Inf(-1)) - for _, token := range ti { - if token.value > maxLogit { - maxLogit = token.value - } - } - - // subtracting max logit to avoid under/overflow - for i := range ti { - ti[i].value = (ti[i].value - maxLogit) / temp - } - - return ti -} - // siftDown maintains a min-heap property by recursively moving larger elements down the heap. // // The heap is represented as an array where for any node at index i: @@ -51,7 +42,7 @@ func temperature(ti []logit, t float32) []logit { // 1. Finds the smallest value between the node and its children // 2. If the node is not the smallest, swaps it with its smallest child // 3. Continues this process down the affected path until the min-heap property is restored -func siftDown(data []logit, start, end int) { +func siftDown(data []token, start, end int) { root := start for { child := 2*root + 1 @@ -73,7 +64,7 @@ func siftDown(data []logit, start, end int) { } // topK limits the number of tokens considered to the k highest logits -func topK(ts []logit, k int) []logit { +func topK(ts []token, k int) []token { if k >= len(ts) { return ts } @@ -99,7 +90,7 @@ func topK(ts []logit, k int) []logit { } // topP limits tokens to those with cumulative probability p -func topP(ts []logit, p float32) []logit { +func topP(ts []token, p float32) []token { if p == 1.0 { return ts } @@ -118,7 +109,7 @@ func topP(ts []logit, p float32) []logit { } // minP limits tokens to those with cumulative probability p -func minP(ts []logit, p float32) []logit { +func minP(ts []token, p float32) []token { if p == 1.0 { return ts } @@ -145,8 +136,9 @@ func minP(ts []logit, p float32) []logit { } // TODO(parthsareen): possibly replace with simpler implementation https://github.com/ollama/ollama/issues/9584 -// Conting sort implementation to sort tokens by logits -func sortLogits(tokens []logit) { +// sortLogits sorts implementation to sort tokens by logits using counting sort +// counting sort is faster than built-in sort for this use case +func sortLogits(tokens []token) { if len(tokens) <= 1 { return } @@ -187,7 +179,7 @@ func sortLogits(tokens []logit) { } // Second pass: place elements in correct position - output := make([]logit, len(tokens)) + output := make([]token, len(tokens)) // Track current positions countsCopy := counts diff --git a/sample/transforms_test.go b/sample/transforms_test.go index 950d79b3..81e8849b 100644 --- a/sample/transforms_test.go +++ b/sample/transforms_test.go @@ -7,10 +7,10 @@ import ( ) // Helper to convert float64 slice to logit slice -func toLogits(values []float64) []logit { - tokens := make([]logit, len(values)) +func toTokens(values []float64) []token { + tokens := make([]token, len(values)) for i, v := range values { - tokens[i] = logit{ + tokens[i] = token{ id: int32(i), value: float32(v), } @@ -19,7 +19,7 @@ func toLogits(values []float64) []logit { } // Helper to compare logit slices -func compareLogits(t *testing.T, name string, want []float64, got []logit) { +func compareLogits(t *testing.T, name string, want []float64, got []token) { t.Helper() if len(want) != len(got) { t.Errorf("%s: length mismatch: want %d, got %d", name, len(want), len(got)) @@ -32,17 +32,9 @@ func compareLogits(t *testing.T, name string, want []float64, got []logit) { } } -func TestTemperature(t *testing.T) { - input := []float64{2, -1, 4, -3, 1, -2, 0} - want := []float64{-4, -10, 0, -14, -6, -12, -8} // (logit - max logit) / temp - - got := temperature(toLogits(input), 0.5) - compareLogits(t, "Temperature", want, got) -} - -func TestSoftmax(t *testing.T) { - input := []float64{-3, -2, -1, 0, 1, 2, 4} - got := softmax(toLogits(input)) +func TestTemperatureAndSoftmax(t *testing.T) { + input := []float64{1, 4, -2, 0} + got := temperature(toTokens(input), 0.5) // Check probabilities sum to 1 var sum float32 @@ -53,11 +45,14 @@ func TestSoftmax(t *testing.T) { t.Errorf("probabilities don't sum to 1: got %f", sum) } - // Check relative ordering is preserved - for i := 1; i < len(got); i++ { - if got[i].value < got[i-1].value { - t.Errorf("probability ordering not preserved at index %d", i) - } + got = temperature(toTokens(input), 1) + // Check probabilities sum to 1 + sum = 0.0 + for _, token := range got { + sum += token.value + } + if math.Abs(float64(sum)-1.0) > 1e-6 { + t.Errorf("probabilities don't sum to 1: got %f", sum) } } @@ -65,7 +60,7 @@ func TestTopK(t *testing.T) { input := []float64{-3, -2, -1, 0, 1, 2, 4} // Test k=3 - got := topK(toLogits(input), 3) + got := topK(toTokens(input), 3) if len(got) != 3 { t.Errorf("topK(3): wrong length: want 3, got %d", len(got)) } @@ -74,17 +69,16 @@ func TestTopK(t *testing.T) { compareLogits(t, "topK(3)", want, got) // Test k > len - got = topK(toLogits(input), 10) + got = topK(toTokens(input), 10) compareLogits(t, "topK(10)", input, got) } func TestTopP(t *testing.T) { input := []float64{-3, -2, -1, 0, 1, 2, 4} - tokens := toLogits(input) + tokens := toTokens(input) // First apply temperature and softmax to get probabilities tokens = temperature(tokens, 1) - tokens = softmax(tokens) sortLogits(tokens) // Then apply topP @@ -99,11 +93,10 @@ func TestTopP(t *testing.T) { func TestMinP(t *testing.T) { input := []float64{-3, -2, -1, 0, 1, 2, 4, 3} - tokens := toLogits(input) + tokens := toTokens(input) // First apply temperature and softmax tokens = temperature(tokens, 1) - tokens = softmax(tokens) // Then apply minP got := minP(tokens, 0.2) @@ -116,7 +109,7 @@ func TestMinP(t *testing.T) { func TestSortLogits(t *testing.T) { input := []float64{3, 1, 4, 2, -1, 0, -2} - tokens := toLogits(input) + tokens := toTokens(input) sortLogits(tokens) @@ -133,15 +126,15 @@ func TestSortLogits(t *testing.T) { func BenchmarkTransforms(b *testing.B) { // Generate random logits - tokens := make([]logit, 1<<16) + tokens := make([]token, 1<<16) for i := range tokens { - tokens[i] = logit{ + tokens[i] = token{ id: int32(i), value: rand.Float32(), } } - tokensCopy := make([]logit, len(tokens)) + tokensCopy := make([]token, len(tokens)) b.Run("Temperature", func(b *testing.B) { b.ResetTimer() diff --git a/scripts/build_windows.ps1 b/scripts/build_windows.ps1 index 62930d7f..60485df8 100644 --- a/scripts/build_windows.ps1 +++ b/scripts/build_windows.ps1 @@ -80,13 +80,14 @@ function checkEnv() { function buildOllama() { + mkdir -Force -path "${script:DIST_DIR}\" if ($script:ARCH -ne "arm64") { Remove-Item -ea 0 -recurse -force -path "${script:SRC_DIR}\dist\windows-${script:ARCH}" New-Item "${script:SRC_DIR}\dist\windows-${script:ARCH}\lib\ollama\" -ItemType Directory -ea 0 & cmake --fresh --preset CPU --install-prefix $script:DIST_DIR if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} - & cmake --build --preset CPU --parallel $script:JOBS + & cmake --build --preset CPU --config Release --parallel $script:JOBS if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} & cmake --install build --component CPU --strip if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} @@ -101,7 +102,7 @@ function buildOllama() { # to avoid 2022 (or newer) from being used as the default & cmake --fresh --preset "CUDA 11" -G "Visual Studio 16 2019" --install-prefix $script:DIST_DIR if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} - & cmake --build --preset "CUDA 11" --parallel $script:JOBS + & cmake --build --preset "CUDA 11" --config Release --parallel $script:JOBS if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} & cmake --install build --component "CUDA" --strip if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} @@ -112,7 +113,7 @@ function buildOllama() { write-host "Building CUDA v12 backend libraries" & cmake --fresh --preset "CUDA 12" --install-prefix $script:DIST_DIR if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} - & cmake --build --preset "CUDA 12" --parallel $script:JOBS + & cmake --build --preset "CUDA 12" --config Release --parallel $script:JOBS if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} & cmake --install build --component "CUDA" --strip if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} @@ -131,7 +132,7 @@ function buildOllama() { $env:HIPCXX="" $env:HIP_PLATFORM="" $env:CMAKE_PREFIX_PATH="" - & cmake --build --preset "ROCm" --parallel $script:JOBS + & cmake --build --preset "ROCm" --config Release --parallel $script:JOBS if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} & cmake --install build --component "HIP" --strip if ($LASTEXITCODE -ne 0) { exit($LASTEXITCODE)} diff --git a/server/prompt.go b/server/prompt.go index 5b5b958f..d053f2a8 100644 --- a/server/prompt.go +++ b/server/prompt.go @@ -26,6 +26,7 @@ func chatPrompt(ctx context.Context, m *Model, tokenize tokenizeFunc, opts *api. var system []api.Message isMllama := checkMllamaModelFamily(m) + isGemma3 := checkGemma3ModelFamily(m) var imageNumTokens int // TODO: Ideally we would compute this from the projector metadata but some pieces are implementation dependent @@ -40,7 +41,7 @@ func chatPrompt(ctx context.Context, m *Model, tokenize tokenizeFunc, opts *api. n := len(msgs) - 1 // in reverse, find all messages that fit into context window for i := n; i >= 0; i-- { - if isMllama && len(msgs[i].Images) > 1 { + if (isMllama || isGemma3) && len(msgs[i].Images) > 1 { return "", nil, errTooManyImages } @@ -157,3 +158,12 @@ func checkMllamaModelFamily(m *Model) bool { } return false } + +func checkGemma3ModelFamily(m *Model) bool { + for _, arch := range m.Config.ModelFamilies { + if arch == "gemma3" { + return true + } + } + return false +}