ollama source for Momentry Core verification
This commit is contained in:
Accusys
200
x/create/laguna_test.go
Normal file
200
x/create/laguna_test.go
Normal file
@@ -0,0 +1,200 @@
|
||||
package create
|
||||
|
||||
import (
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
st "github.com/ollama/ollama/x/safetensors"
|
||||
)
|
||||
|
||||
func TestCreateSafetensorsModel_LagunaHFFP8RespectsSourceTensorPrecision(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
requested string
|
||||
wantFP8Gate string
|
||||
wantFP8Up string
|
||||
wantFP8Down string
|
||||
wantBF16QProj string
|
||||
}{
|
||||
{
|
||||
name: "default mxfp8 import keeps source bf16 tensors",
|
||||
requested: "",
|
||||
wantFP8Gate: "mxfp8",
|
||||
wantFP8Up: "mxfp8",
|
||||
wantFP8Down: "mxfp8",
|
||||
wantBF16QProj: "",
|
||||
},
|
||||
{
|
||||
name: "nvfp4 import keeps source bf16 tensors and preserves down_proj at mxfp8",
|
||||
requested: "nvfp4",
|
||||
wantFP8Gate: "nvfp4",
|
||||
wantFP8Up: "nvfp4",
|
||||
wantFP8Down: "mxfp8",
|
||||
wantBF16QProj: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
configJSON := `{
|
||||
"model_type": "laguna",
|
||||
"architectures": ["LagunaForCausalLM"],
|
||||
"quantization_config": {"quant_method": "fp8", "weight_block_size": [128, 128]}
|
||||
}`
|
||||
if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil {
|
||||
t.Fatalf("failed to write config.json: %v", err)
|
||||
}
|
||||
|
||||
createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.gate_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.gate_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.up_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.up_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("lm_head.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.gate.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
})
|
||||
|
||||
quantizeByName := make(map[string]string)
|
||||
|
||||
createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
|
||||
if _, err := io.ReadAll(r); err != nil {
|
||||
return LayerInfo{}, err
|
||||
}
|
||||
return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil
|
||||
}
|
||||
createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
|
||||
if _, err := io.ReadAll(r); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
quantizeByName[name] = quantize
|
||||
return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil
|
||||
}
|
||||
writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil }
|
||||
|
||||
if err := CreateSafetensorsModel("test-model", dir, tt.requested, createLayer, createTensorLayer, writeManifest, func(string) {}); err != nil {
|
||||
t.Fatalf("CreateSafetensorsModel failed: %v", err)
|
||||
}
|
||||
|
||||
if got := quantizeByName["model.layers.0.mlp.experts.0.gate_proj.weight"]; got != tt.wantFP8Gate {
|
||||
t.Fatalf("gate_proj quantization = %q, want %q", got, tt.wantFP8Gate)
|
||||
}
|
||||
if got := quantizeByName["model.layers.0.mlp.experts.0.up_proj.weight"]; got != tt.wantFP8Up {
|
||||
t.Fatalf("up_proj quantization = %q, want %q", got, tt.wantFP8Up)
|
||||
}
|
||||
if got := quantizeByName["model.layers.0.mlp.experts.0.down_proj.weight"]; got != tt.wantFP8Down {
|
||||
t.Fatalf("down_proj quantization = %q, want %q", got, tt.wantFP8Down)
|
||||
}
|
||||
for _, name := range []string{
|
||||
"model.layers.0.self_attn.q_proj.weight",
|
||||
"model.embed_tokens.weight",
|
||||
"lm_head.weight",
|
||||
"model.layers.0.mlp.gate.weight",
|
||||
} {
|
||||
if got := quantizeByName[name]; got != tt.wantBF16QProj {
|
||||
t.Fatalf("%s quantization = %q, want %q", name, got, tt.wantBF16QProj)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateSafetensorsModel_LagunaBF16QuantizesOnlyRoutedExperts(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
requested string
|
||||
want map[string]string
|
||||
}{
|
||||
{
|
||||
name: "int8 quantizes only routed experts",
|
||||
requested: "int8",
|
||||
want: map[string]string{
|
||||
"model.layers.0.mlp.experts.0.gate_proj.weight": "int8",
|
||||
"model.layers.0.mlp.experts.0.up_proj.weight": "int8",
|
||||
"model.layers.0.mlp.experts.0.down_proj.weight": "int8",
|
||||
"model.layers.0.mlp.shared_experts.gate_proj.weight": "",
|
||||
"model.layers.0.mlp.shared_experts.down_proj.weight": "",
|
||||
"model.layers.0.self_attn.q_proj.weight": "",
|
||||
"model.layers.0.mlp.down_proj.weight": "",
|
||||
"model.embed_tokens.weight": "",
|
||||
"lm_head.weight": "",
|
||||
"model.layers.0.mlp.gate.weight": "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "int4 keeps routed down_proj at int8 and leaves others bf16",
|
||||
requested: "int4",
|
||||
want: map[string]string{
|
||||
"model.layers.0.mlp.experts.0.gate_proj.weight": "int4",
|
||||
"model.layers.0.mlp.experts.0.up_proj.weight": "int4",
|
||||
"model.layers.0.mlp.experts.0.down_proj.weight": "int8",
|
||||
"model.layers.0.mlp.shared_experts.gate_proj.weight": "",
|
||||
"model.layers.0.mlp.shared_experts.down_proj.weight": "",
|
||||
"model.layers.0.self_attn.q_proj.weight": "",
|
||||
"model.layers.0.mlp.down_proj.weight": "",
|
||||
"model.embed_tokens.weight": "",
|
||||
"lm_head.weight": "",
|
||||
"model.layers.0.mlp.gate.weight": "",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
configJSON := `{
|
||||
"model_type": "laguna",
|
||||
"architectures": ["LagunaForCausalLM"]
|
||||
}`
|
||||
if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil {
|
||||
t.Fatalf("failed to write config.json: %v", err)
|
||||
}
|
||||
|
||||
createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.gate_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.up_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.shared_experts.gate_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.shared_experts.down_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.down_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("lm_head.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
st.NewTensorDataFromBytes("model.layers.0.mlp.gate.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)),
|
||||
})
|
||||
|
||||
quantizeByName := make(map[string]string)
|
||||
|
||||
createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) {
|
||||
if _, err := io.ReadAll(r); err != nil {
|
||||
return LayerInfo{}, err
|
||||
}
|
||||
return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil
|
||||
}
|
||||
createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) {
|
||||
if _, err := io.ReadAll(r); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
quantizeByName[name] = quantize
|
||||
return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil
|
||||
}
|
||||
writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil }
|
||||
|
||||
if err := CreateSafetensorsModel("test-model", dir, tt.requested, createLayer, createTensorLayer, writeManifest, func(string) {}); err != nil {
|
||||
t.Fatalf("CreateSafetensorsModel failed: %v", err)
|
||||
}
|
||||
|
||||
for name, want := range tt.want {
|
||||
if got := quantizeByName[name]; got != want {
|
||||
t.Fatalf("%s quantization = %q, want %q", name, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user