Mirror_Repos/ollama_ollama
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

models: add support for gemma4-12b (#16457)

Browse Source
This commit is contained in:
Patrick Devine
2026-06-03 07:44:57 -07:00
committed by GitHub
parent 4b5bdd3b25
commit 50bbda5660
6 changed files with 268 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/gemma4_test.go
+16
View File
@@ -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{
server/renderer_resolution.go
+5 -5
View File
@@ -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
x/create/create.go
+18 -12
View File
@@ -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) {
x/create/gemma4_test.go
+43
View File
@@ -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}
x/models/gemma4/gemma4.go
+59 -5
View File
@@ -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 {
x/models/gemma4/gemma4_test.go
+127
View File
@@ -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(`{