mirror of
https://github.com/ollama/ollama.git
synced 2026-11-12 03:46:14 -06:00
Jesse Gross
2e036e7cdf
The MLX runner is the only Go inference runner left and is no longer experimental, so its packages leave x/. The bindings become a top-level mlx package beside the carried patches in mlx/compat, mirroring how llama/ holds the llama.cpp integration, and the runner becomes mlxrunner with the architectures nested under the package they implement. Subpackages move with their parent unless listed. x/mlxrunner/mlx mlx x/internal/mlxthread mlx/mlxthread x/internal/mlxthreadtest mlx/mlxthread/mlxthreadtest x/internal/mlxtest mlx/mlxtest x/quant mlx/quant mlx/compat/*.patch mlx/compat/mlx-c (MLX patches go in mlx/compat/mlx) x/mlxrunner mlxrunner x/models/nn mlxrunner/nn x/models/<arch> mlxrunner/model/<arch> x/mlxrunner/imports.go mlxrunner/model/architectures (new package) x/create create x/safetensors fs/safetensors x/tokenizer mlxrunner/tokenizer Every package keeps its name, so the Go changes are the import path rewrites the moves force, and the CMake, Dockerfile, CI cache keys, drift check and Darwin payload script follow the new paths. Four edits are not paths: the runner's blank architecture imports become the package mlxrunner/model/architectures, so the list to extend for a new model sits beside the architecture directories; a depguard rule keeps the two test harnesses out of non-test code, as the x/internal placement used to; the CI change filter's two entries for the long-deleted x/imagegen/mlx now name the bindings' CMake project and the carried patches, so a change to either builds the payload; and the tokenizer parity test reads its fixtures from its own testdata instead of walking out of x/. x/server and x/imagegen/manifest stay for the next two commits.
221 lines
12 KiB
Go
221 lines
12 KiB
Go
package create
|
|
|
|
import (
|
|
"encoding/json"
|
|
"testing"
|
|
)
|
|
|
|
func TestGemma4UnifiedImportTransformRegistration(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
configJSON string
|
|
cfg sourceModelConfig
|
|
wantErr bool
|
|
wantLayers int
|
|
}{
|
|
{
|
|
name: "unified conditional generation architecture",
|
|
configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"text_config":{"num_hidden_layers":48}}`,
|
|
cfg: sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}},
|
|
wantLayers: 48,
|
|
},
|
|
{
|
|
name: "unified model type fallback",
|
|
configJSON: `{"model_type":"gemma4_unified","text_config":{"num_hidden_layers":48}}`,
|
|
cfg: sourceModelConfig{ModelType: "gemma4_unified"},
|
|
wantLayers: 48,
|
|
},
|
|
{
|
|
name: "malformed config is a hard error",
|
|
configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"num_hidden_layers":"oops"}`,
|
|
cfg: sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}},
|
|
wantErr: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
inv := Inventory{Config: tt.cfg, RawConfig: json.RawMessage(tt.configJSON)}
|
|
|
|
transform, err := newTensorImportTransform(inv)
|
|
if tt.wantErr {
|
|
if err == nil {
|
|
t.Fatalf("newTensorImportTransform() error = nil, want error")
|
|
}
|
|
return
|
|
}
|
|
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 != tt.wantLayers {
|
|
t.Fatalf("numLayers = %d, want %d", gemmaTransform.numLayers, tt.wantLayers)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGemma4QuantizationType(t *testing.T) {
|
|
// 26B MoE: 30 layers, 128 experts
|
|
transform26B := gemma4ImportTransform{numLayers: 30, numExperts: 128}
|
|
// 8-expert model (hypothetical)
|
|
transform8E := gemma4ImportTransform{numLayers: 30, numExperts: 8}
|
|
|
|
aligned := []int32{2816, 2816} // divisible by 64 (int4/int8 group size) and 16 (nvfp4)
|
|
|
|
tests := []struct {
|
|
name string
|
|
transform gemma4ImportTransform
|
|
tensor string
|
|
shape []int32
|
|
quantize string
|
|
want string
|
|
}{
|
|
// === embed_tokens: quantize to 8-bit variant (serves as both embed and lm_head) ===
|
|
{"embed_tokens int4", transform26B, "model.embed_tokens.weight", aligned, "int4", "int8"},
|
|
{"embed_tokens nvfp4", transform26B, "model.embed_tokens.weight", aligned, "nvfp4", "mxfp8"},
|
|
{"embed_tokens mxfp4", transform26B, "model.embed_tokens.weight", aligned, "mxfp4", "mxfp8"},
|
|
{"embed_tokens int8", transform26B, "model.embed_tokens.weight", aligned, "int8", "int8"},
|
|
{"embed_tokens mxfp8", transform26B, "model.embed_tokens.weight", aligned, "mxfp8", "mxfp8"},
|
|
|
|
// === v_proj: layer-position heuristic for int4/nvfp4 ===
|
|
// Layer 0 is in first 1/8 (30/8=3) → promoted
|
|
{"v_proj int4 promoted layer", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "int4", "int8"},
|
|
// Layer 4 is NOT in useMoreBits → base quant
|
|
{"v_proj int4 non-promoted layer", transform26B, "model.layers.4.self_attn.v_proj.weight", aligned, "int4", "int4"},
|
|
// Layer 29 is in last 1/8 → promoted
|
|
{"v_proj int4 last layer promoted", transform26B, "model.layers.29.self_attn.v_proj.weight", aligned, "int4", "int8"},
|
|
// nvfp4: promote to mxfp8 (cross-family, validated by MLX quantized_matmul)
|
|
{"v_proj nvfp4 promoted layer", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "nvfp4", "mxfp8"},
|
|
{"v_proj nvfp4 non-promoted layer", transform26B, "model.layers.4.self_attn.v_proj.weight", aligned, "nvfp4", "nvfp4"},
|
|
// mxfp4: promoted to mxfp8 at promoted layers (same mxfp family)
|
|
{"v_proj mxfp4 promoted layer", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "mxfp4", "mxfp8"},
|
|
{"v_proj mxfp4 non-promoted layer", transform26B, "model.layers.4.self_attn.v_proj.weight", aligned, "mxfp4", "mxfp4"},
|
|
// int8/mxfp8: no promotion (already 8-bit)
|
|
{"v_proj int8 base", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "int8", "int8"},
|
|
{"v_proj mxfp8 base", transform26B, "model.layers.0.self_attn.v_proj.weight", aligned, "mxfp8", "mxfp8"},
|
|
|
|
// === down_proj (dense MLP): same heuristic as v_proj ===
|
|
{"dense down_proj int4 promoted", transform26B, "model.layers.0.mlp.down_proj.weight", aligned, "int4", "int8"},
|
|
{"dense down_proj int4 non-promoted", transform26B, "model.layers.4.mlp.down_proj.weight", aligned, "int4", "int4"},
|
|
{"dense down_proj nvfp4 promoted", transform26B, "model.layers.0.mlp.down_proj.weight", aligned, "nvfp4", "mxfp8"},
|
|
{"dense down_proj nvfp4 non-promoted", transform26B, "model.layers.4.mlp.down_proj.weight", aligned, "nvfp4", "nvfp4"},
|
|
{"dense down_proj mxfp4 promoted", transform26B, "model.layers.0.mlp.down_proj.weight", aligned, "mxfp4", "mxfp8"},
|
|
{"dense down_proj mxfp4 non-promoted", transform26B, "model.layers.4.mlp.down_proj.weight", aligned, "mxfp4", "mxfp4"},
|
|
|
|
// === Expert down_proj: int4→int8, nvfp4→nvfp8 at promoted layers ===
|
|
{"expert down_proj int4 promoted", transform26B, "model.layers.0.moe.experts.42.down_proj.weight", aligned, "int4", "int8"},
|
|
{"expert down_proj int4 non-promoted", transform26B, "model.layers.4.moe.experts.42.down_proj.weight", aligned, "int4", "int4"},
|
|
// nvfp4 experts: promote to mxfp8 (all experts at a layer get same treatment,
|
|
// so GatherQMM sees uniform quant per projection per layer)
|
|
{"expert down_proj nvfp4 promoted layer", transform26B, "model.layers.0.moe.experts.42.down_proj.weight", aligned, "nvfp4", "mxfp8"},
|
|
{"expert down_proj nvfp4 non-promoted layer", transform26B, "model.layers.4.moe.experts.42.down_proj.weight", aligned, "nvfp4", "nvfp4"},
|
|
// mxfp4 experts: promote to mxfp8 (same mxfp family, GatherQMM compatible)
|
|
{"expert down_proj mxfp4 promoted layer", transform26B, "model.layers.0.moe.experts.42.down_proj.weight", aligned, "mxfp4", "mxfp8"},
|
|
{"expert down_proj mxfp4 non-promoted layer", transform26B, "model.layers.4.moe.experts.42.down_proj.weight", aligned, "mxfp4", "mxfp4"},
|
|
|
|
// === Expert gate_up_proj: always base quant (not a sensitive tensor) ===
|
|
{"expert gate_up int4", transform26B, "model.layers.0.moe.experts.42.gate_up_proj.weight", aligned, "int4", "int4"},
|
|
{"expert gate_up nvfp4", transform26B, "model.layers.0.moe.experts.42.gate_up_proj.weight", aligned, "nvfp4", "nvfp4"},
|
|
{"expert gate_up mxfp4", transform26B, "model.layers.0.moe.experts.42.gate_up_proj.weight", aligned, "mxfp4", "mxfp4"},
|
|
|
|
// === Router projection: expert selection is sensitive; keep source precision ===
|
|
{"router proj int4", transform26B, "model.layers.0.router.proj.weight", aligned, "int4", ""},
|
|
{"router proj nvfp4", transform26B, "model.layers.0.router.proj.weight", aligned, "nvfp4", ""},
|
|
{"router proj mxfp4", transform26B, "model.layers.0.router.proj.weight", aligned, "mxfp4", ""},
|
|
|
|
// === k_proj: promoted only for 8-expert models ===
|
|
{"k_proj 128 experts int4", transform26B, "model.layers.0.self_attn.k_proj.weight", aligned, "int4", "int4"},
|
|
{"k_proj 8 experts int4", transform8E, "model.layers.0.self_attn.k_proj.weight", aligned, "int4", "int8"},
|
|
{"k_proj 8 experts nvfp4", transform8E, "model.layers.0.self_attn.k_proj.weight", aligned, "nvfp4", "mxfp8"},
|
|
{"k_proj 8 experts mxfp4", transform8E, "model.layers.0.self_attn.k_proj.weight", aligned, "mxfp4", "mxfp8"},
|
|
|
|
// === q_proj, o_proj, gate_proj, up_proj: always base quant ===
|
|
{"q_proj int4", transform26B, "model.layers.0.self_attn.q_proj.weight", aligned, "int4", "int4"},
|
|
{"o_proj int4", transform26B, "model.layers.0.self_attn.o_proj.weight", aligned, "int4", "int4"},
|
|
{"gate_proj int4", transform26B, "model.layers.0.mlp.gate_proj.weight", aligned, "int4", "int4"},
|
|
{"up_proj int4", transform26B, "model.layers.0.mlp.up_proj.weight", aligned, "int4", "int4"},
|
|
|
|
// === Non-quantizable tensors: always bf16 ===
|
|
{"embed_tokens per_layer skip", transform26B, "model.embed_tokens_per_layer.weight", aligned, "int4", ""},
|
|
{"norm", transform26B, "model.layers.0.input_layernorm.weight", []int32{2816}, "int4", ""},
|
|
{"router scale", transform26B, "model.layers.0.router.scale", []int32{2816}, "int4", ""},
|
|
|
|
// === Audio/vision tower tensors: must pass through unquantized for all quant types ===
|
|
// These contain .v_proj and down_proj but should NOT be intercepted by
|
|
// the sensitive-tensor promotion logic.
|
|
{"audio norm int4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "int4", ""},
|
|
{"audio norm nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "nvfp4", ""},
|
|
{"audio norm int8", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "int8", ""},
|
|
{"audio norm mxfp8", transform26B, "model.audio_tower.subsample_conv_projection.layer0.norm.weight", []int32{128}, "mxfp8", ""},
|
|
{"audio conv int4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.conv.weight", []int32{128, 1, 3, 3}, "int4", ""},
|
|
{"audio conv nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.layer0.conv.weight", []int32{128, 1, 3, 3}, "nvfp4", ""},
|
|
{"audio linear int4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "int4", ""},
|
|
{"audio linear nvfp4", transform26B, "model.audio_tower.subsample_conv_projection.input_proj_linear.weight", aligned, "nvfp4", ""},
|
|
// Audio tower v_proj — must NOT be promoted despite containing .v_proj
|
|
{"audio v_proj int4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", ""},
|
|
{"audio v_proj nvfp4", transform26B, "model.audio_tower.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", ""},
|
|
// Vision tower — source precision for every quant family.
|
|
{"vision v_proj int4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", ""},
|
|
{"vision v_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", ""},
|
|
{"vision q_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight", aligned, "nvfp4", ""},
|
|
{"unified vision embedder nvfp4", transform26B, "model.vision_embedder.patch_dense.weight", aligned, "nvfp4", ""},
|
|
{"vision projection nvfp4", transform26B, "model.embed_vision.embedding_projection.linear.weight", aligned, "nvfp4", ""},
|
|
// Audio tower down_proj
|
|
{"audio down_proj int4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "int4", ""},
|
|
{"audio down_proj nvfp4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "nvfp4", ""},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := tt.transform.quantizationType(tt.tensor, tt.shape, tt.quantize)
|
|
if got != tt.want {
|
|
t.Errorf("quantizationType(%q, %v, %q) = %q, want %q",
|
|
tt.tensor, tt.shape, tt.quantize, got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestUseMoreBits(t *testing.T) {
|
|
// 30 layers: first 1/8 = layers 0-2, last 1/8 = layers 27-29
|
|
// In between: every 3rd from offset (i - n/8) % 3 == 2
|
|
n := 30
|
|
promoted := map[int]bool{}
|
|
for i := range n {
|
|
if useMoreBits(i, n) {
|
|
promoted[i] = true
|
|
}
|
|
}
|
|
|
|
// First 1/8 (30/8 = 3): layers 0, 1, 2
|
|
for _, i := range []int{0, 1, 2} {
|
|
if !promoted[i] {
|
|
t.Errorf("layer %d should be promoted (first 1/8)", i)
|
|
}
|
|
}
|
|
|
|
// Last 1/8: layers 26, 27, 28, 29 (>= 7*30/8 = 26)
|
|
for _, i := range []int{26, 27, 28, 29} {
|
|
if !promoted[i] {
|
|
t.Errorf("layer %d should be promoted (last 1/8)", i)
|
|
}
|
|
}
|
|
|
|
// Some middle layers should NOT be promoted
|
|
for _, i := range []int{3, 4, 6, 7} {
|
|
if promoted[i] {
|
|
t.Errorf("layer %d should NOT be promoted", i)
|
|
}
|
|
}
|
|
|
|
// Layer 5 should be promoted: (5 - 3) % 3 == 2
|
|
if !promoted[5] {
|
|
t.Errorf("layer 5 should be promoted (periodic)")
|
|
}
|
|
}
|