mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 13:38:14 -05: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.
449 lines
13 KiB
Go
449 lines
13 KiB
Go
package create
|
|
|
|
import (
|
|
"encoding/json"
|
|
"strconv"
|
|
"testing"
|
|
)
|
|
|
|
func TestLagunaImportTransformRegistration(t *testing.T) {
|
|
inv := Inventory{
|
|
Config: sourceModelConfig{Architectures: []string{"LagunaForCausalLM"}},
|
|
RawConfig: json.RawMessage(`{"num_hidden_layers":40,"mlp_layer_types":["dense","sparse","sparse"]}`),
|
|
}
|
|
|
|
policy, err := newTensorImportTransform(inv)
|
|
if err != nil {
|
|
t.Fatalf("newTensorImportTransform() error = %v", err)
|
|
}
|
|
|
|
transform, ok := policy.(lagunaImportTransform)
|
|
if !ok {
|
|
t.Fatalf("newTensorImportTransform() = %T, want lagunaImportTransform", policy)
|
|
}
|
|
if !transform.denseMLPLayers[0] || transform.denseMLPLayers[1] {
|
|
t.Fatalf("denseMLPLayers = %v, want only layer 0", transform.denseMLPLayers)
|
|
}
|
|
if transform.numLayers != 40 {
|
|
t.Fatalf("numLayers = %d, want 40", transform.numLayers)
|
|
}
|
|
}
|
|
|
|
func TestLagunaImportTransformSameRecipeAcrossConfigs(t *testing.T) {
|
|
configs := map[string]json.RawMessage{
|
|
"laguna xs.2": json.RawMessage(`{
|
|
"num_hidden_layers": 40,
|
|
"mlp_layer_types": ["dense", "sparse", "sparse"]
|
|
}`),
|
|
"laguna xs 2.1": json.RawMessage(`{
|
|
"num_hidden_layers": 40,
|
|
"mlp_only_layers": [0],
|
|
"gating_types": ["per_head"]
|
|
}`),
|
|
}
|
|
|
|
for name, rawConfig := range configs {
|
|
t.Run(name, func(t *testing.T) {
|
|
policy, err := newLagunaImportTransform(rawConfig)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
testLagunaQuantizationRecipe(t, policy)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestLagunaPlanQuantizesEmbeddingAndHeadAtMXFP8(t *testing.T) {
|
|
policy, err := newLagunaImportTransform(json.RawMessage(`{
|
|
"num_hidden_layers": 40,
|
|
"mlp_only_layers": [0]
|
|
}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
const (
|
|
embedding = "model.embed_tokens.weight"
|
|
head = "lm_head.weight"
|
|
)
|
|
inv := Inventory{Tensors: map[string]SourceTensor{
|
|
embedding: {
|
|
Name: embedding,
|
|
Dtype: "BF16",
|
|
Shape: []int32{100352, 2048},
|
|
},
|
|
head: {
|
|
Name: head,
|
|
Dtype: "BF16",
|
|
Shape: []int32{100352, 2048},
|
|
},
|
|
}}
|
|
|
|
specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "nvfp4"}, policy)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
for _, name := range []string{embedding, head} {
|
|
if got := quantizeForPlannedTensor(specs, name); got != "mxfp8" {
|
|
t.Errorf("planned quantization for %s = %q, want mxfp8", name, got)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestLagunaPlanExpertGroupUsesStackedDownProjectionPolicy(t *testing.T) {
|
|
policy, err := newLagunaImportTransform(json.RawMessage(`{
|
|
"num_hidden_layers": 40,
|
|
"mlp_layer_types": ["dense", "sparse"]
|
|
}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
inv := Inventory{Tensors: map[string]SourceTensor{}}
|
|
for _, layer := range []int{1, 5} {
|
|
for expert := range 2 {
|
|
for _, projection := range []string{"gate_proj", "down_proj"} {
|
|
name := "model.layers." + strconv.Itoa(layer) + ".mlp.experts." + strconv.Itoa(expert) + "." + projection + ".weight"
|
|
inv.Tensors[name] = SourceTensor{
|
|
Name: name,
|
|
Dtype: "BF16",
|
|
Shape: []int32{2048, 512},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "mxfp8"}, policy)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
tests := map[string]string{
|
|
"model.layers.1.mlp.experts.down_proj.weight": "",
|
|
"model.layers.1.mlp.experts.gate_proj.weight": "mxfp8",
|
|
"model.layers.5.mlp.experts.down_proj.weight": "mxfp8",
|
|
}
|
|
for tensor, want := range tests {
|
|
if got := quantizeForPlannedTensor(specs, tensor); got != want {
|
|
t.Fatalf("planned quantization for %s = %q, want %q", tensor, got, want)
|
|
}
|
|
}
|
|
|
|
if _, ok := specByName(specs, "model.layers.1.mlp.experts"); ok {
|
|
t.Fatal("mixed layer 1 expert projections should use separate blobs")
|
|
}
|
|
for _, name := range []string{
|
|
"model.layers.1.mlp.experts.down_proj.weight",
|
|
"model.layers.1.mlp.experts.gate_proj.weight",
|
|
} {
|
|
spec, ok := specByName(specs, name)
|
|
if !ok || len(spec.Tensors) != 1 {
|
|
t.Fatalf("missing homogeneous blob %s; got %v", name, specNames(specs))
|
|
}
|
|
}
|
|
|
|
spec, ok := specByName(specs, "model.layers.5.mlp.experts")
|
|
if !ok || len(spec.Tensors) != 2 {
|
|
t.Fatalf("uniform layer 5 projections should share one blob; got %v", specNames(specs))
|
|
}
|
|
}
|
|
|
|
func quantizeForPlannedTensor(specs []BlobSpec, name string) string {
|
|
for _, spec := range specs {
|
|
for _, ts := range spec.Tensors {
|
|
if ts.Name == name {
|
|
return ts.Quantize
|
|
}
|
|
}
|
|
}
|
|
return "<missing>"
|
|
}
|
|
|
|
func testLagunaQuantizationRecipe(t *testing.T, policy quantizePolicy) {
|
|
t.Helper()
|
|
|
|
tests := []struct {
|
|
name string
|
|
tensor string
|
|
shape []int32
|
|
quantize string
|
|
want string
|
|
}{
|
|
{
|
|
name: "attention q projection uses requested fp4",
|
|
tensor: "model.layers.1.self_attn.q_proj.weight",
|
|
shape: []int32{8192, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "attention v projection uses requested fp4 before promotion layer",
|
|
tensor: "model.layers.0.self_attn.v_proj.weight",
|
|
shape: []int32{1024, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "attention v projection uses requested fp4 on layer 4",
|
|
tensor: "model.layers.4.self_attn.v_proj.weight",
|
|
shape: []int32{1024, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "attention v projection uses requested fp4 after promotion layer",
|
|
tensor: "model.layers.5.self_attn.v_proj.weight",
|
|
shape: []int32{1024, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "attention k projection uses requested fp4 on input layer",
|
|
tensor: "model.layers.0.self_attn.k_proj.weight",
|
|
shape: []int32{1024, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "attention k projection uses requested fp4 past layer 0",
|
|
tensor: "model.layers.4.self_attn.k_proj.weight",
|
|
shape: []int32{1024, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "attention k projection stays source precision for mxfp8",
|
|
tensor: "model.layers.4.self_attn.k_proj.weight",
|
|
shape: []int32{1024, 2048},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "attention v projection stays source precision for mxfp8",
|
|
tensor: "model.layers.4.self_attn.v_proj.weight",
|
|
shape: []int32{1024, 2048},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "attention q projection stays source precision for mxfp8",
|
|
tensor: "model.layers.4.self_attn.q_proj.weight",
|
|
shape: []int32{8192, 2048},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "attention o projection stays source precision for mxfp8",
|
|
tensor: "model.layers.4.self_attn.o_proj.weight",
|
|
shape: []int32{2048, 8192},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "attention gate projection stays source precision for mxfp8",
|
|
tensor: "model.layers.4.self_attn.g_proj.weight",
|
|
shape: []int32{64, 2048},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "attention gate projection uses requested fp4",
|
|
tensor: "model.layers.1.self_attn.g_proj.weight",
|
|
shape: []int32{64, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "dense gate projection uses requested fp4",
|
|
tensor: "model.layers.0.mlp.gate_proj.weight",
|
|
shape: []int32{8192, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "dense down projection uses requested fp4",
|
|
tensor: "model.layers.0.mlp.down_proj.weight",
|
|
shape: []int32{2048, 8192},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "unsupported dense projection in sparse layer stays source precision",
|
|
tensor: "model.layers.1.mlp.down_proj.weight",
|
|
shape: []int32{2048, 8192},
|
|
quantize: "nvfp4",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "routed expert gate uses requested fp4",
|
|
tensor: "model.layers.1.mlp.experts.gate_proj.weight",
|
|
shape: []int32{256, 512, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "routed expert down uses requested fp4 on cadence layer",
|
|
tensor: "model.layers.1.mlp.experts.down_proj.weight",
|
|
shape: []int32{256, 2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "routed expert down uses requested fp4 on later layer",
|
|
tensor: "model.layers.5.mlp.experts.down_proj.weight",
|
|
shape: []int32{256, 2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "shared expert gate uses requested fp4",
|
|
tensor: "model.layers.1.mlp.shared_expert.gate_proj.weight",
|
|
shape: []int32{512, 2048},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "shared expert down promotes to mxfp8 on input-side layer",
|
|
tensor: "model.layers.1.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "shared expert down uses requested fp4 before middle cadence",
|
|
tensor: "model.layers.5.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "shared expert down promotes to mxfp8 on first selected middle layer",
|
|
tensor: "model.layers.7.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "shared expert down promotes to mxfp8 on early middle layer",
|
|
tensor: "model.layers.10.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "shared expert down promotes to mxfp8 on last selected middle layer",
|
|
tensor: "model.layers.13.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "shared expert down uses requested fp4 after selected middle layers",
|
|
tensor: "model.layers.16.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "shared expert down uses requested fp4 on late middle layer",
|
|
tensor: "model.layers.19.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "nvfp4",
|
|
},
|
|
{
|
|
name: "shared expert down promotes to mxfp8 on final layers",
|
|
tensor: "model.layers.39.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "nvfp4",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "shared expert down stays source precision for mxfp8 on selected layer",
|
|
tensor: "model.layers.1.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "shared expert down stays source precision for mxfp8 off selected layers",
|
|
tensor: "model.layers.5.mlp.shared_expert.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "routed expert down stays source precision for mxfp8 on selected layer",
|
|
tensor: "model.layers.1.mlp.experts.down_proj.weight",
|
|
shape: []int32{256, 2048, 512},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "per-expert routed down stays source precision for mxfp8 on selected layer",
|
|
tensor: "model.layers.1.mlp.experts.0.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "mxfp8",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "routed expert down uses mxfp8 off selected layers",
|
|
tensor: "model.layers.5.mlp.experts.down_proj.weight",
|
|
shape: []int32{256, 2048, 512},
|
|
quantize: "mxfp8",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "per-expert routed down uses mxfp8 off selected layers",
|
|
tensor: "model.layers.5.mlp.experts.0.down_proj.weight",
|
|
shape: []int32{2048, 512},
|
|
quantize: "mxfp8",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "router gate stays source precision",
|
|
tensor: "model.layers.1.mlp.gate.weight",
|
|
shape: []int32{256, 2048},
|
|
quantize: "nvfp4",
|
|
want: "",
|
|
},
|
|
{
|
|
name: "embedding promotes to mxfp8 for nvfp4",
|
|
tensor: "model.embed_tokens.weight",
|
|
shape: []int32{100352, 2048},
|
|
quantize: "nvfp4",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "embedding uses requested mxfp8",
|
|
tensor: "model.embed_tokens.weight",
|
|
shape: []int32{100352, 2048},
|
|
quantize: "mxfp8",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "lm head promotes to mxfp8 for nvfp4",
|
|
tensor: "lm_head.weight",
|
|
shape: []int32{100352, 2048},
|
|
quantize: "nvfp4",
|
|
want: "mxfp8",
|
|
},
|
|
{
|
|
name: "lm head uses requested mxfp8",
|
|
tensor: "lm_head.weight",
|
|
shape: []int32{100352, 2048},
|
|
quantize: "mxfp8",
|
|
want: "mxfp8",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
if got := policy.quantizationType(tt.tensor, tt.shape, tt.quantize); got != tt.want {
|
|
t.Fatalf("quantizationType(%q, %v, %q) = %q, want %q", tt.tensor, tt.shape, tt.quantize, got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|