ollama source for Momentry Core verification
This commit is contained in:
Accusys
450
convert/convert_laguna_test.go
Normal file
450
convert/convert_laguna_test.go
Normal file
@@ -0,0 +1,450 @@
|
||||
package convert
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
|
||||
"github.com/ollama/ollama/fs/ggml"
|
||||
)
|
||||
|
||||
type lagunaTestTensor struct {
|
||||
tensorBase
|
||||
}
|
||||
|
||||
func newLagunaTestTensor(name string, shape ...uint64) Tensor {
|
||||
return &lagunaTestTensor{tensorBase: tensorBase{name: name, shape: shape}}
|
||||
}
|
||||
|
||||
func (t *lagunaTestTensor) WriteTo(io.Writer) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (t *lagunaTestTensor) Clone() Tensor {
|
||||
return &lagunaTestTensor{tensorBase: tensorBase{
|
||||
name: t.name,
|
||||
shape: append([]uint64(nil), t.shape...),
|
||||
}}
|
||||
}
|
||||
|
||||
func TestLagunaReplacements(t *testing.T) {
|
||||
p := lagunaModel{}
|
||||
r := strings.NewReplacer(p.Replacements()...)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
in string
|
||||
want string
|
||||
}{
|
||||
{"embed", "model.embed_tokens.weight", "token_embd.weight"},
|
||||
{"final_norm", "model.norm.weight", "output_norm.weight"},
|
||||
{"lm_head", "lm_head.weight", "output.weight"},
|
||||
{"block prefix", "model.layers.7.input_layernorm.weight", "blk.7.attn_norm.weight"},
|
||||
{"q", "model.layers.3.self_attn.q_proj.weight", "blk.3.attn_q.weight"},
|
||||
{"k", "model.layers.3.self_attn.k_proj.weight", "blk.3.attn_k.weight"},
|
||||
{"v", "model.layers.3.self_attn.v_proj.weight", "blk.3.attn_v.weight"},
|
||||
{"o", "model.layers.3.self_attn.o_proj.weight", "blk.3.attn_output.weight"},
|
||||
{"g", "model.layers.3.self_attn.g_proj.weight", "blk.3.attn_g.weight"},
|
||||
{"q_norm", "model.layers.3.self_attn.q_norm.weight", "blk.3.attn_q_norm.weight"},
|
||||
{"k_norm", "model.layers.3.self_attn.k_norm.weight", "blk.3.attn_k_norm.weight"},
|
||||
{"post_attn_norm", "model.layers.3.post_attention_layernorm.weight", "blk.3.ffn_norm.weight"},
|
||||
{"dense gate", "model.layers.0.mlp.gate_proj.weight", "blk.0.ffn_gate.weight"},
|
||||
{"dense up", "model.layers.0.mlp.up_proj.weight", "blk.0.ffn_up.weight"},
|
||||
{"dense down", "model.layers.0.mlp.down_proj.weight", "blk.0.ffn_down.weight"},
|
||||
{"shexp gate", "model.layers.5.mlp.shared_expert.gate_proj.weight", "blk.5.ffn_gate_shexp.weight"},
|
||||
{"shexp up", "model.layers.5.mlp.shared_expert.up_proj.weight", "blk.5.ffn_up_shexp.weight"},
|
||||
{"shexp down", "model.layers.5.mlp.shared_expert.down_proj.weight", "blk.5.ffn_down_shexp.weight"},
|
||||
{"router", "model.layers.5.mlp.gate.weight", "blk.5.ffn_gate_inp.weight"},
|
||||
{"score bias", "model.layers.5.mlp.experts.e_score_correction_bias", "blk.5.exp_probs_b.bias"},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := r.Replace(tc.in); got != tc.want {
|
||||
t.Errorf("Replace(%q) = %q, want %q", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLagunaValidateRejectsUnsupportedVariants(t *testing.T) {
|
||||
base := validLagunaTestModel()
|
||||
tests := []struct {
|
||||
name string
|
||||
edit func(*lagunaModel)
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "per-element gating",
|
||||
edit: func(m *lagunaModel) {
|
||||
m.Gating = "per-element"
|
||||
},
|
||||
want: "unsupported attention gating",
|
||||
},
|
||||
{
|
||||
name: "attention sinks",
|
||||
edit: func(m *lagunaModel) {
|
||||
m.SwaAttentionSinkEnabled = true
|
||||
},
|
||||
want: "swa_attention_sink_enabled=true",
|
||||
},
|
||||
{
|
||||
name: "qk norm disabled",
|
||||
edit: func(m *lagunaModel) {
|
||||
m.QKNormType = "none"
|
||||
},
|
||||
want: "unsupported qk_norm_type",
|
||||
},
|
||||
{
|
||||
name: "softmax moe",
|
||||
edit: func(m *lagunaModel) {
|
||||
m.MoERouterUseSigmoid = false
|
||||
},
|
||||
want: "moe_router_use_sigmoid=false",
|
||||
},
|
||||
{
|
||||
name: "router weight on input",
|
||||
edit: func(m *lagunaModel) {
|
||||
m.MoEApplyRouterWeightOnInput = true
|
||||
},
|
||||
want: "moe_apply_router_weight_on_input=true",
|
||||
},
|
||||
{
|
||||
name: "unknown layer type",
|
||||
edit: func(m *lagunaModel) {
|
||||
m.LayerTypes[1] = "local_attention"
|
||||
},
|
||||
want: "unsupported layer_types[1]",
|
||||
},
|
||||
{
|
||||
name: "nonstandard dense layout",
|
||||
edit: func(m *lagunaModel) {
|
||||
m.MLPOnlyLayers = []uint32{0, 3}
|
||||
},
|
||||
want: "unsupported mlp_only_layers",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
m := base
|
||||
m.LayerTypes = append([]string(nil), base.LayerTypes...)
|
||||
m.NumAttentionHeadsPerLayer = append([]uint32(nil), base.NumAttentionHeadsPerLayer...)
|
||||
m.MLPOnlyLayers = append([]uint32(nil), base.MLPOnlyLayers...)
|
||||
tc.edit(&m)
|
||||
err := m.validate()
|
||||
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
||||
t.Fatalf("validate() error = %v, want substring %q", err, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLagunaGAConfigNormalizesBoolGatingAndNestedRope(t *testing.T) {
|
||||
var m lagunaModel
|
||||
if err := json.Unmarshal([]byte(`{
|
||||
"architectures": ["LagunaForCausalLM"],
|
||||
"num_hidden_layers": 1,
|
||||
"hidden_size": 8,
|
||||
"num_attention_heads": 2,
|
||||
"num_key_value_heads": 1,
|
||||
"head_dim": 4,
|
||||
"gating": true,
|
||||
"num_experts": 2,
|
||||
"num_experts_per_tok": 1,
|
||||
"moe_intermediate_size": 4,
|
||||
"shared_expert_intermediate_size": 4,
|
||||
"decoder_sparse_step": 1,
|
||||
"mlp_layer_types": ["dense"],
|
||||
"rope_parameters": {
|
||||
"full_attention": {
|
||||
"rope_theta": 500000,
|
||||
"rope_type": "yarn",
|
||||
"factor": 32,
|
||||
"original_max_position_embeddings": 4096,
|
||||
"beta_fast": 64,
|
||||
"beta_slow": 1,
|
||||
"attention_factor": 1,
|
||||
"partial_rotary_factor": 0.5
|
||||
},
|
||||
"sliding_attention": {
|
||||
"rope_theta": 10000,
|
||||
"rope_type": "default",
|
||||
"partial_rotary_factor": 1
|
||||
}
|
||||
}
|
||||
}`), &m); err != nil {
|
||||
t.Fatalf("json.Unmarshal() error = %v", err)
|
||||
}
|
||||
|
||||
if err := m.validate(); err != nil {
|
||||
t.Fatalf("validate() error = %v", err)
|
||||
}
|
||||
if m.Gating != "true" {
|
||||
t.Fatalf("Gating = %q, want raw true marker", m.Gating)
|
||||
}
|
||||
if !m.Gating.perHead() {
|
||||
t.Fatal("expected bool gating to normalize as per-head support")
|
||||
}
|
||||
if m.QKNormType != "rmsnorm" {
|
||||
t.Fatalf("QKNormType = %q, want rmsnorm default", m.QKNormType)
|
||||
}
|
||||
if !m.MoERouterUseSigmoid {
|
||||
t.Fatal("MoERouterUseSigmoid should default true")
|
||||
}
|
||||
if !m.NormTopKProb {
|
||||
t.Fatal("NormTopKProb should default true")
|
||||
}
|
||||
if diff := cmp.Diff(m.MLPOnlyLayers, []uint32{0}); diff != "" {
|
||||
t.Fatalf("MLPOnlyLayers mismatch (-got +want):\n%s", diff)
|
||||
}
|
||||
if m.RopeParameters.RopeTheta != 500000 || m.RopeParameters.PartialRotaryFactor != 0.5 {
|
||||
t.Fatalf("full rope = %#v, want theta=500000 partial=0.5", m.RopeParameters)
|
||||
}
|
||||
if m.SwaRopeParameters.RopeTheta != 10000 || m.SwaRopeParameters.PartialRotaryFactor != 1 {
|
||||
t.Fatalf("swa rope = %#v, want theta=10000 partial=1", m.SwaRopeParameters)
|
||||
}
|
||||
}
|
||||
|
||||
func validLagunaTestModel() lagunaModel {
|
||||
return lagunaModel{
|
||||
ModelParameters: ModelParameters{
|
||||
VocabSize: 32,
|
||||
},
|
||||
NumHiddenLayers: 2,
|
||||
HiddenSize: 8,
|
||||
IntermediateSize: 16,
|
||||
NumAttentionHeads: 2,
|
||||
NumKeyValueHeads: 1,
|
||||
HeadDim: 4,
|
||||
RMSNormEPS: 1e-6,
|
||||
MaxPositionEmbeddings: 4096,
|
||||
SlidingWindow: 512,
|
||||
Gating: "per-head",
|
||||
QKNormType: "rmsnorm",
|
||||
LayerTypes: []string{"global_attention", "sliding_attention"},
|
||||
NumAttentionHeadsPerLayer: []uint32{2, 2},
|
||||
NumExperts: 2,
|
||||
NumExpertsPerTok: 1,
|
||||
MoEIntermediateSize: 4,
|
||||
SharedExpertIntermediateSize: 4,
|
||||
NormTopKProb: true,
|
||||
MoeRoutedScalingFactor: 2.5,
|
||||
MoERouterUseSigmoid: true,
|
||||
DecoderSparseStep: 1,
|
||||
MLPOnlyLayers: []uint32{0},
|
||||
}
|
||||
}
|
||||
|
||||
func validLagunaTestTensors(m lagunaModel) []Tensor {
|
||||
ts := []Tensor{
|
||||
newLagunaTestTensor("token_embd.weight", uint64(m.VocabSize), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor("output_norm.weight", uint64(m.HiddenSize)),
|
||||
}
|
||||
|
||||
for layer := range m.NumHiddenLayers {
|
||||
prefix := fmt.Sprintf("blk.%d", layer)
|
||||
heads := uint64(m.numHeadsForLayer(layer))
|
||||
attnWidth := heads * uint64(m.HeadDim)
|
||||
kvWidth := uint64(m.NumKeyValueHeads * m.HeadDim)
|
||||
ts = append(ts,
|
||||
newLagunaTestTensor(prefix+".attn_norm.weight", uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".ffn_norm.weight", uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".attn_q.weight", attnWidth, uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".attn_k.weight", kvWidth, uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".attn_v.weight", kvWidth, uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".attn_output.weight", uint64(m.HiddenSize), attnWidth),
|
||||
newLagunaTestTensor(prefix+".attn_g.weight", heads, uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".attn_q_norm.weight", uint64(m.HeadDim)),
|
||||
newLagunaTestTensor(prefix+".attn_k_norm.weight", uint64(m.HeadDim)),
|
||||
)
|
||||
|
||||
if m.layerUsesMoE(layer) {
|
||||
ts = append(ts,
|
||||
newLagunaTestTensor(prefix+".ffn_gate_inp.weight", uint64(m.NumExperts), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".exp_probs_b.bias", uint64(m.NumExperts)),
|
||||
newLagunaTestTensor(prefix+".ffn_gate_shexp.weight", uint64(m.SharedExpertIntermediateSize), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".ffn_up_shexp.weight", uint64(m.SharedExpertIntermediateSize), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".ffn_down_shexp.weight", uint64(m.HiddenSize), uint64(m.SharedExpertIntermediateSize)),
|
||||
)
|
||||
for expert := range m.NumExperts {
|
||||
ts = append(ts,
|
||||
newLagunaTestTensor(fmt.Sprintf("%s.mlp.experts.%d.gate_proj.weight", prefix, expert), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(fmt.Sprintf("%s.mlp.experts.%d.up_proj.weight", prefix, expert), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(fmt.Sprintf("%s.mlp.experts.%d.down_proj.weight", prefix, expert), uint64(m.HiddenSize), uint64(m.MoEIntermediateSize)),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
ts = append(ts,
|
||||
newLagunaTestTensor(prefix+".ffn_gate.weight", uint64(m.IntermediateSize), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".ffn_up.weight", uint64(m.IntermediateSize), uint64(m.HiddenSize)),
|
||||
newLagunaTestTensor(prefix+".ffn_down.weight", uint64(m.HiddenSize), uint64(m.IntermediateSize)),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
return ts
|
||||
}
|
||||
|
||||
func TestLagunaTensorsMergeRoutedExperts(t *testing.T) {
|
||||
m := validLagunaTestModel()
|
||||
out := m.Tensors(validLagunaTestTensors(m))
|
||||
|
||||
tensors := make(map[string]*ggml.Tensor, len(out))
|
||||
for _, t := range out {
|
||||
tensors[t.Name] = t
|
||||
}
|
||||
|
||||
tests := map[string][]uint64{
|
||||
"blk.1.ffn_gate_exps.weight": {uint64(m.NumExperts), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)},
|
||||
"blk.1.ffn_up_exps.weight": {uint64(m.NumExperts), uint64(m.MoEIntermediateSize), uint64(m.HiddenSize)},
|
||||
"blk.1.ffn_down_exps.weight": {uint64(m.NumExperts), uint64(m.HiddenSize), uint64(m.MoEIntermediateSize)},
|
||||
}
|
||||
for name, wantShape := range tests {
|
||||
tensor, ok := tensors[name]
|
||||
if !ok {
|
||||
t.Fatalf("missing merged tensor %q", name)
|
||||
}
|
||||
if diff := cmp.Diff(wantShape, tensor.Shape); diff != "" {
|
||||
t.Fatalf("%s shape mismatch (-want +got):\n%s", name, diff)
|
||||
}
|
||||
}
|
||||
|
||||
for expert := range m.NumExperts {
|
||||
name := fmt.Sprintf("blk.1.mlp.experts.%d.gate_proj.weight", expert)
|
||||
if _, ok := tensors[name]; ok {
|
||||
t.Fatalf("unexpected unmerged expert tensor %q", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLagunaKVShape(t *testing.T) {
|
||||
m := lagunaModel{
|
||||
NumHiddenLayers: 4,
|
||||
HiddenSize: 128,
|
||||
IntermediateSize: 256,
|
||||
NumAttentionHeads: 8,
|
||||
NumKeyValueHeads: 4,
|
||||
HeadDim: 16,
|
||||
RMSNormEPS: 1e-6,
|
||||
MaxPositionEmbeddings: 4096,
|
||||
SlidingWindow: 512,
|
||||
PartialRotaryFactor: 0.5,
|
||||
Gating: "per-head",
|
||||
QKNormType: "rmsnorm",
|
||||
LayerTypes: []string{"full_attention", "sliding_attention", "sliding_attention", "sliding_attention"},
|
||||
NumAttentionHeadsPerLayer: []uint32{8, 16, 16, 16},
|
||||
NumExperts: 32,
|
||||
NumExpertsPerTok: 4,
|
||||
MoEIntermediateSize: 64,
|
||||
SharedExpertIntermediateSize: 64,
|
||||
NormTopKProb: true,
|
||||
MoeRoutedScalingFactor: 2.5,
|
||||
MoERouterUseSigmoid: true,
|
||||
DecoderSparseStep: 1,
|
||||
MLPOnlyLayers: []uint32{0},
|
||||
}
|
||||
m.RopeParameters.RopeTheta = 500000
|
||||
m.RopeParameters.RopeType = "yarn"
|
||||
m.RopeParameters.Factor = 32
|
||||
m.RopeParameters.OriginalMaxPositionEmbeddings = 4096
|
||||
m.RopeParameters.BetaFast = 64
|
||||
m.RopeParameters.BetaSlow = 1
|
||||
m.SwaRopeParameters.RopeTheta = 10000
|
||||
m.SwaRopeParameters.RopeType = "linear"
|
||||
m.SwaRopeParameters.Factor = 1
|
||||
m.SwaRopeParameters.PartialRotaryFactor = 1
|
||||
|
||||
kv := m.KV(&Tokenizer{Vocabulary: &Vocabulary{}, Template: "{% include 'chat_template.jinja' %}"})
|
||||
required := []string{
|
||||
"general.architecture",
|
||||
"tokenizer.ggml.pre",
|
||||
"laguna.block_count",
|
||||
"laguna.context_length",
|
||||
"laguna.embedding_length",
|
||||
"laguna.feed_forward_length",
|
||||
"laguna.attention.head_count",
|
||||
"laguna.attention.head_count_kv",
|
||||
"laguna.attention.key_length",
|
||||
"laguna.attention.value_length",
|
||||
"laguna.attention.layer_norm_rms_epsilon",
|
||||
"laguna.attention.sliding_window",
|
||||
"laguna.attention.layer_types",
|
||||
"laguna.attention.sliding_window_pattern",
|
||||
"laguna.attention.gating_type",
|
||||
"laguna.attention.qk_norm",
|
||||
"laguna.expert_count",
|
||||
"laguna.expert_used_count",
|
||||
"laguna.expert_feed_forward_length",
|
||||
"laguna.expert_shared_feed_forward_length",
|
||||
"laguna.expert_shared_count",
|
||||
"laguna.expert_weights_norm",
|
||||
"laguna.expert_weights_scale",
|
||||
"laguna.expert_gating_func",
|
||||
"laguna.leading_dense_block_count",
|
||||
"laguna.dense_layers",
|
||||
"laguna.rope.freq_base",
|
||||
"laguna.rope.scaling.type",
|
||||
"laguna.rope.scaling.factor",
|
||||
"laguna.rope.partial_rotary_factor",
|
||||
"laguna.rope.swa.freq_base",
|
||||
"laguna.rope.swa.scaling.type",
|
||||
"laguna.rope.dimension_count",
|
||||
"laguna.rope.swa.dimension_count",
|
||||
}
|
||||
for _, k := range required {
|
||||
if _, ok := kv[k]; !ok {
|
||||
t.Errorf("missing required KV: %s", k)
|
||||
}
|
||||
}
|
||||
|
||||
if got := kv["general.architecture"]; got != "laguna" {
|
||||
t.Errorf("architecture = %v, want laguna", got)
|
||||
}
|
||||
if got := kv["tokenizer.ggml.add_bos_token"]; got != false {
|
||||
t.Errorf("tokenizer.ggml.add_bos_token = %v, want false", got)
|
||||
}
|
||||
if _, ok := kv["tokenizer.chat_template"]; ok {
|
||||
t.Fatal("tokenizer.chat_template should be omitted for Laguna")
|
||||
}
|
||||
if got := kv["laguna.expert_gating_func"]; got != lagunaGatingFuncSigmoid {
|
||||
t.Errorf("expert_gating_func = %v, want sigmoid(%d)", got, lagunaGatingFuncSigmoid)
|
||||
}
|
||||
if got := kv["laguna.leading_dense_block_count"]; got != uint32(1) {
|
||||
t.Errorf("leading_dense_block_count = %v, want 1", got)
|
||||
}
|
||||
if got := kv["laguna.rope.dimension_count"]; got != uint32(8) {
|
||||
t.Errorf("rope.dimension_count = %v, want 8", got)
|
||||
}
|
||||
if got := kv["laguna.rope.swa.dimension_count"]; got != uint32(16) {
|
||||
t.Errorf("rope.swa.dimension_count = %v, want 16", got)
|
||||
}
|
||||
if got, ok := kv["laguna.attention.layer_types"].([]uint32); !ok || len(got) != 4 || got[0] != 0 || got[1] != 1 || got[2] != 1 || got[3] != 1 {
|
||||
t.Fatalf("layer_types = %#v, want [0 1 1 1]", kv["laguna.attention.layer_types"])
|
||||
}
|
||||
if got, ok := kv["laguna.attention.sliding_window_pattern"].([]bool); !ok || len(got) != 4 || got[0] || !got[1] || !got[2] || !got[3] {
|
||||
t.Fatalf("sliding_window_pattern = %#v, want [false true true true]", kv["laguna.attention.sliding_window_pattern"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestLagunaKVYarnAttentionFactorFallback(t *testing.T) {
|
||||
m := validLagunaTestModel()
|
||||
m.RopeParameters.RopeType = "yarn"
|
||||
m.RopeParameters.Factor = 32
|
||||
|
||||
kv := m.KV(&Tokenizer{Vocabulary: &Vocabulary{}})
|
||||
got, ok := kv["laguna.rope.scaling.attn_factor"].(float32)
|
||||
if !ok {
|
||||
t.Fatalf("attn_factor type = %T, want float32", kv["laguna.rope.scaling.attn_factor"])
|
||||
}
|
||||
|
||||
want := float32(0.1*math.Log(32) + 1)
|
||||
if diff := math.Abs(float64(got - want)); diff > 1e-6 {
|
||||
t.Fatalf("attn_factor = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user