diff --git a/server/gemma4_test.go b/server/gemma4_test.go index 5999f9e99..9ad3b0ce4 100644 --- a/server/gemma4_test.go +++ b/server/gemma4_test.go @@ -50,6 +50,15 @@ func TestResolveGemma4Renderer(t *testing.T) { }, want: gemma4RendererLarge, }, + { + name: "legacy 12b tag resolves large", + model: &Model{ + Name: "gemma-4-12b-it", + ShortName: "gemma-4-12b-it", + Config: testConfigWithRenderer(gemma4RendererLegacy), + }, + want: gemma4RendererLarge, + }, { name: "legacy model type resolves small", model: &Model{ @@ -64,6 +73,13 @@ func TestResolveGemma4Renderer(t *testing.T) { }, want: gemma4RendererLarge, }, + { + name: "legacy 12b model type resolves large", + model: &Model{ + Config: testConfigWithRendererAndType(gemma4RendererLegacy, "12B"), + }, + want: gemma4RendererLarge, + }, { name: "legacy unknown defaults small", model: &Model{ diff --git a/server/renderer_resolution.go b/server/renderer_resolution.go index 870f700e4..1f348852f 100644 --- a/server/renderer_resolution.go +++ b/server/renderer_resolution.go @@ -12,10 +12,10 @@ const ( gemma4RendererSmall = "gemma4-small" gemma4RendererLarge = "gemma4-large" - // Gemma 4 small templates cover the e2b/e4b family, while 26b/31b use the - // large template. Default to the small prompt unless the model is clearly in - // the large range. - gemma4LargeMinParameterCount = 16_000_000_000 + // Gemma 4 small templates cover the e2b/e4b family, while 12b/26b/31b use + // the large template. Default to the small prompt unless the model is + // clearly in the large range. + gemma4LargeMinParameterCount = 12_000_000_000 ) func resolveRendererName(m *Model) string { @@ -67,7 +67,7 @@ func gemma4RendererFromName(name string) (string, bool) { switch { case strings.Contains(lower, "e2b"), strings.Contains(lower, "e4b"): return gemma4RendererSmall, true - case strings.Contains(lower, "26b"), strings.Contains(lower, "31b"): + case strings.Contains(lower, "12b"), strings.Contains(lower, "26b"), strings.Contains(lower, "31b"): return gemma4RendererLarge, true default: return "", false diff --git a/x/create/create.go b/x/create/create.go index d5b6ffa10..6d9f944c9 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -806,18 +806,24 @@ func (noopImportTransform) quantizationType(name string, shape []int32, quantize type tensorImportTransformFactory func(modelDir string, cfg sourceModelConfig) (tensorImportTransform, error) var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ - "Qwen3_5ForCausalLM": newQwen35ImportTransform, - "Qwen3_5ForConditionalGeneration": newQwen35ImportTransform, - "Qwen3NextForCausalLM": newQwen35ImportTransform, - "Qwen3NextForConditionalGeneration": newQwen35ImportTransform, - "Qwen3_5MoeForCausalLM": newQwen35ImportTransform, - "Qwen3_5MoeForConditionalGeneration": newQwen35ImportTransform, - "Qwen3NextMoeForCausalLM": newQwen35ImportTransform, - "Qwen3NextMoeForConditionalGeneration": newQwen35ImportTransform, - "Gemma4ForCausalLM": newGemma4ImportTransform, - "Gemma4ForConditionalGeneration": newGemma4ImportTransform, - "LagunaForCausalLM": newLagunaImportTransform, - "Gemma4AssistantForCausalLM": newGemma4ImportTransform, + "Qwen3_5ForCausalLM": newQwen35ImportTransform, + "Qwen3_5ForConditionalGeneration": newQwen35ImportTransform, + "Qwen3NextForCausalLM": newQwen35ImportTransform, + "Qwen3NextForConditionalGeneration": newQwen35ImportTransform, + "Qwen3_5MoeForCausalLM": newQwen35ImportTransform, + "Qwen3_5MoeForConditionalGeneration": newQwen35ImportTransform, + "Qwen3NextMoeForCausalLM": newQwen35ImportTransform, + "Qwen3NextMoeForConditionalGeneration": newQwen35ImportTransform, + "Gemma4ForCausalLM": newGemma4ImportTransform, + "Gemma4ForConditionalGeneration": newGemma4ImportTransform, + "Gemma4UnifiedForCausalLM": newGemma4ImportTransform, + "Gemma4UnifiedForConditionalGeneration": newGemma4ImportTransform, + "gemma4_unified": newGemma4ImportTransform, + "gemma4_unified_text": newGemma4ImportTransform, + "LagunaForCausalLM": newLagunaImportTransform, + "Gemma4AssistantForCausalLM": newGemma4ImportTransform, + "Gemma4UnifiedAssistantForCausalLM": newGemma4ImportTransform, + "gemma4_unified_assistant": newGemma4ImportTransform, } func newTensorImportTransform(modelDir string, cfg sourceModelConfig) (tensorImportTransform, error) { diff --git a/x/create/gemma4_test.go b/x/create/gemma4_test.go index 858183db0..de7b53b98 100644 --- a/x/create/gemma4_test.go +++ b/x/create/gemma4_test.go @@ -1,9 +1,52 @@ package create import ( + "os" + "path/filepath" "testing" ) +func TestGemma4UnifiedImportTransformRegistration(t *testing.T) { + tests := []struct { + name string + configJSON string + cfg sourceModelConfig + }{ + { + name: "unified conditional generation architecture", + configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"text_config":{"num_hidden_layers":48}}`, + cfg: sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}}, + }, + { + name: "unified model type fallback", + configJSON: `{"model_type":"gemma4_unified","text_config":{"num_hidden_layers":48}}`, + cfg: sourceModelConfig{ModelType: "gemma4_unified"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(tt.configJSON), 0o644); err != nil { + t.Fatal(err) + } + + transform, err := newTensorImportTransform(dir, tt.cfg) + if err != nil { + t.Fatalf("newTensorImportTransform() error = %v", err) + } + + gemmaTransform, ok := transform.(gemma4ImportTransform) + if !ok { + t.Fatalf("newTensorImportTransform() = %T, want gemma4ImportTransform", transform) + } + if gemmaTransform.numLayers != 48 { + t.Fatalf("numLayers = %d, want 48", gemmaTransform.numLayers) + } + }) + } +} + func TestGemma4QuantizationType(t *testing.T) { // 26B MoE: 30 layers, 128 experts transform26B := gemma4ImportTransform{numLayers: 30, numExperts: 128} diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index adfc7ab55..9f8029580 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -18,8 +18,13 @@ import ( func init() { base.Register("Gemma4ForCausalLM", newModel) base.Register("Gemma4ForConditionalGeneration", newModel) + base.Register("Gemma4UnifiedForCausalLM", newModel) + base.Register("Gemma4UnifiedForConditionalGeneration", newModel) + base.Register("gemma4_unified", newModel) base.RegisterDraft("Gemma4AssistantForCausalLM", newAssistantModel) + base.RegisterDraft("Gemma4UnifiedAssistantForCausalLM", newAssistantModel) base.RegisterDraft("gemma4_assistant", newAssistantModel) + base.RegisterDraft("gemma4_unified_assistant", newAssistantModel) } // Compile-time interface checks. @@ -93,6 +98,10 @@ type TextConfig struct { KVDonors map[int32]bool `json:"-"` } +type generationConfig struct { + SuppressTokens []int32 `json:"suppress_tokens"` +} + // sharedHistory carries a donor layer's K/V to donees that share // it. Exactly one of history or (k, v) is populated: history when // the donor had a cache (donees feed it to SDPA via WithKVHistory; @@ -342,7 +351,8 @@ type Model struct { tok *tokenizer.Tokenizer *TextConfig - weightPrefix string + SuppressLogitBias *mlx.Array + weightPrefix string } func parseTextConfig(configData []byte) (TextConfig, error) { @@ -467,6 +477,14 @@ func parseTextConfig(configData []byte) (TextConfig, error) { return cfg, nil } +func parseSuppressTokens(configData []byte) []int32 { + var cfg generationConfig + if err := json.Unmarshal(configData, &cfg); err != nil { + return nil + } + return cfg.SuppressTokens +} + func (m *Model) EnableCompile() bool { return true } @@ -591,8 +609,10 @@ func newModel(root *model.Root) (base.Model, error) { } tokConfig := &tokenizer.TokenizerConfig{ConfigJSON: configData} + var suppressTokens []int32 if genConfigData, err := root.Manifest.ReadConfig("generation_config.json"); err == nil { tokConfig.GenerationConfigJSON = genConfigData + suppressTokens = parseSuppressTokens(genConfigData) } if tokConfigData, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil { tokConfig.TokenizerConfigJSON = tokConfigData @@ -604,9 +624,10 @@ func newModel(root *model.Root) (base.Model, error) { } m := &Model{ - Layers: make([]*DecoderLayer, cfg.NumHiddenLayers), - TextConfig: &cfg, - tok: tok, + Layers: make([]*DecoderLayer, cfg.NumHiddenLayers), + TextConfig: &cfg, + tok: tok, + SuppressLogitBias: makeSuppressLogitBias(suppressTokens, cfg.VocabSize), } for i := range m.Layers { @@ -1076,7 +1097,40 @@ func (m *Model) Unembed(x *mlx.Array) *mlx.Array { logits = mlx.LogitSoftcap(logits, cap) } - return logits + return suppressTokenLogits(logits, m.SuppressLogitBias) +} + +func makeSuppressLogitBias(tokenIDs []int32, vocabSize int32) *mlx.Array { + if len(tokenIDs) == 0 || vocabSize <= 0 { + return nil + } + + bias := make([]float32, int(vocabSize)) + any := false + for _, tokenID := range tokenIDs { + if tokenID >= 0 && tokenID < vocabSize { + bias[tokenID] = float32(math.Inf(-1)) + any = true + } + } + if !any { + return nil + } + + return mlx.FromValues(bias, 1, 1, int(vocabSize)) +} + +func suppressTokenLogits(logits, bias *mlx.Array) *mlx.Array { + if bias == nil { + return logits + } + + dims := logits.Dims() + if len(dims) != 3 { + return logits + } + + return logits.Add(bias.AsType(logits.DType())) } func (m *Model) NumLayers() int { diff --git a/x/models/gemma4/gemma4_test.go b/x/models/gemma4/gemma4_test.go index d00864e10..432342524 100644 --- a/x/models/gemma4/gemma4_test.go +++ b/x/models/gemma4/gemma4_test.go @@ -1,11 +1,29 @@ package gemma4 import ( + "encoding/json" + "slices" "testing" "github.com/ollama/ollama/x/mlxrunner/mlx" ) +func TestParseSuppressTokens(t *testing.T) { + got := parseSuppressTokens([]byte(`{"suppress_tokens":[258883,258882]}`)) + want := []int32{258883, 258882} + if !slices.Equal(got, want) { + t.Fatalf("parseSuppressTokens() = %v, want %v", got, want) + } + + if got := parseSuppressTokens([]byte(`{"eos_token_id":[1,106,50]}`)); got != nil { + t.Fatalf("parseSuppressTokens() without suppress_tokens = %v, want nil", got) + } + + if got := parseSuppressTokens([]byte(`{`)); got != nil { + t.Fatalf("parseSuppressTokens() with invalid JSON = %v, want nil", got) + } +} + func TestParseTextConfigE2B(t *testing.T) { skipIfNoMLX(t) data := []byte(`{ @@ -300,6 +318,115 @@ func TestParseTextConfig31B(t *testing.T) { } } +func TestParseTextConfig12BUnified(t *testing.T) { + skipIfNoMLX(t) + + layerTypes := make([]string, 0, 48) + for i := range 48 { + if i%6 == 5 { + layerTypes = append(layerTypes, "full_attention") + } else { + layerTypes = append(layerTypes, "sliding_attention") + } + } + + data, err := json.Marshal(map[string]any{ + "architectures": []string{"Gemma4UnifiedForConditionalGeneration"}, + "model_type": "gemma4_unified", + "text_config": map[string]any{ + "hidden_size": 3840, + "num_hidden_layers": 48, + "intermediate_size": 15360, + "num_attention_heads": 16, + "num_key_value_heads": 8, + "num_global_key_value_heads": 1, + "head_dim": 256, + "global_head_dim": 512, + "vocab_size": 262144, + "rms_norm_eps": 1e-6, + "max_position_embeddings": 131072, + "sliding_window": 1024, + "final_logit_softcapping": 30.0, + "use_double_wide_mlp": false, + "num_kv_shared_layers": 0, + "hidden_size_per_layer_input": 0, + "vocab_size_per_layer_input": 262144, + "attention_k_eq_v": true, + "enable_moe_block": false, + "tie_word_embeddings": true, + "layer_types": layerTypes, + "rope_parameters": map[string]any{ + "full_attention": map[string]any{ + "partial_rotary_factor": 0.25, + "rope_theta": 1000000.0, + "rope_type": "proportional", + }, + "sliding_attention": map[string]any{ + "rope_theta": 10000.0, + "rope_type": "default", + }, + }, + }, + }) + if err != nil { + t.Fatalf("json.Marshal failed: %v", err) + } + + cfg, err := parseTextConfig(data) + if err != nil { + t.Fatalf("parseTextConfig failed: %v", err) + } + + if cfg.HiddenSize != 3840 { + t.Errorf("HiddenSize = %d, want 3840", cfg.HiddenSize) + } + if cfg.NumHiddenLayers != 48 { + t.Errorf("NumHiddenLayers = %d, want 48", cfg.NumHiddenLayers) + } + if cfg.IntermediateSize != 15360 { + t.Errorf("IntermediateSize = %d, want 15360", cfg.IntermediateSize) + } + if cfg.NumAttentionHeads != 16 { + t.Errorf("NumAttentionHeads = %d, want 16", cfg.NumAttentionHeads) + } + if cfg.NumKeyValueHeads != 8 { + t.Errorf("NumKeyValueHeads = %d, want 8", cfg.NumKeyValueHeads) + } + if cfg.NumGlobalKeyValueHeads != 1 { + t.Errorf("NumGlobalKeyValueHeads = %d, want 1", cfg.NumGlobalKeyValueHeads) + } + if !cfg.AttentionKEqV { + t.Error("AttentionKEqV should be true") + } + if cfg.EnableMoeBlock { + t.Error("EnableMoeBlock should be false") + } + if cfg.HiddenSizePerLayer != 0 { + t.Errorf("HiddenSizePerLayer = %d, want 0", cfg.HiddenSizePerLayer) + } + if cfg.NumKVSharedLayers != 0 { + t.Errorf("NumKVSharedLayers = %d, want 0", cfg.NumKVSharedLayers) + } + if len(cfg.KVShareMap) != 0 { + t.Errorf("KVShareMap should be empty, got %d entries", len(cfg.KVShareMap)) + } + if cfg.FullRopeDims != 512 { + t.Errorf("FullRopeDims = %d, want 512", cfg.FullRopeDims) + } + if cfg.FullRopeFreqs == nil { + t.Error("FullRopeFreqs should be precomputed for proportional RoPE") + } + if isLayerSliding(5, &cfg) { + t.Error("layer 5 should be full attention") + } + if !isLayerSliding(6, &cfg) { + t.Error("layer 6 should be sliding") + } + if isLayerSliding(47, &cfg) { + t.Error("layer 47 should be full attention") + } +} + func TestParseTextConfigE4B(t *testing.T) { skipIfNoMLX(t) data := []byte(`{