ollama source for Momentry Core verification

2026-05-22 17:19:10 +08:00
commit 0b31ff9135
2020 changed files with 1413145 additions and 0 deletions
--- a/x/create/qwen35.go
+++ b/x/create/qwen35.go
@@ -0,0 +1,351 @@
+package create
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/ollama/ollama/x/safetensors"
+)
+
+type qwen35ImportTransform struct {
+	shouldShiftNormWeights bool
+	rewriteLanguageModel   bool
+}
+
+type qwen35SourceInfo struct {
+	hasPrequantizedWeights bool
+	shouldShiftNormWeights bool
+}
+
+func newQwen35ImportTransform(modelDir string, cfg sourceModelConfig) (tensorImportTransform, error) {
+	sourceInfo, err := qwen35InspectSource(modelDir)
+	if err != nil {
+		return qwen35ImportTransform{}, err
+	}
+	if sourceInfo.hasPrequantizedWeights {
+		return noopImportTransform{}, nil
+	}
+
+	return qwen35ImportTransform{
+		shouldShiftNormWeights: sourceInfo.shouldShiftNormWeights,
+		rewriteLanguageModel:   strings.Contains(cfg.Architecture(), "ConditionalGeneration"),
+	}, nil
+}
+
+func qwen35InspectSource(modelDir string) (qwen35SourceInfo, error) {
+	entries, err := os.ReadDir(modelDir)
+	if err != nil {
+		return qwen35SourceInfo{}, err
+	}
+
+	var info qwen35SourceInfo
+	for _, entry := range entries {
+		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".safetensors") {
+			continue
+		}
+
+		extractor, err := safetensors.OpenForExtraction(filepath.Join(modelDir, entry.Name()))
+		if err != nil {
+			return qwen35SourceInfo{}, err
+		}
+
+		for _, name := range extractor.ListTensors() {
+			if strings.HasSuffix(name, ".scales") {
+				extractor.Close()
+				info.hasPrequantizedWeights = true
+				return info, nil
+			}
+			// This should change when MTP is supported
+			if strings.Contains(name, "mtp.") {
+				info.shouldShiftNormWeights = true
+				continue
+			}
+			if info.shouldShiftNormWeights || !strings.Contains(name, "conv1d.weight") {
+				continue
+			}
+
+			td, err := extractor.GetTensor(name)
+			if err != nil {
+				extractor.Close()
+				return qwen35SourceInfo{}, err
+			}
+			if len(td.Shape) == 3 && td.Shape[2] != 1 {
+				info.shouldShiftNormWeights = true
+			}
+		}
+
+		extractor.Close()
+	}
+
+	return info, nil
+}
+
+func (t qwen35ImportTransform) skipTensor(name string) bool {
+	return strings.Contains(name, "mtp.")
+}
+
+func qwen35ShouldKeepBF16ForDirectNonAffine(name string) bool {
+	switch {
+	case strings.HasSuffix(name, "embed_tokens.weight"):
+		return true
+	case strings.HasSuffix(name, "lm_head.weight"):
+		return true
+	case strings.HasSuffix(name, ".linear_attn.in_proj_a.weight"):
+		return true
+	case strings.HasSuffix(name, ".linear_attn.in_proj_b.weight"):
+		return true
+	case strings.HasSuffix(name, ".linear_attn.in_proj_ba.weight"):
+		return true
+	case strings.HasSuffix(name, ".mlp.gate.weight") && !strings.Contains(name, "_proj"):
+		return true
+	case strings.HasSuffix(name, ".mlp.shared_expert_gate.weight"):
+		return true
+	default:
+		return false
+	}
+}
+
+func (t qwen35ImportTransform) quantizationType(name string, shape []int32, quantize string) string {
+	if strings.HasPrefix(name, "vision_tower.") {
+		return ""
+	}
+
+	stackedExpert := isStackedExpertWeight(name)
+	if strings.HasSuffix(name, ".bias") || strings.HasSuffix(name, ".scale") || strings.HasSuffix(name, ".qbias") ||
+		strings.HasSuffix(name, ".biases") || strings.HasSuffix(name, ".scales") {
+		return ""
+	}
+	if !stackedExpert && !strings.HasSuffix(name, ".weight") {
+		return ""
+	}
+	if strings.Contains(name, "norm") || strings.Contains(name, "ln_") || strings.Contains(name, "layernorm") {
+		return ""
+	}
+	if len(shape) != 2 && !(len(shape) == 3 && stackedExpert) {
+		return ""
+	}
+
+	var elems int64 = 1
+	for _, d := range shape {
+		elems *= int64(d)
+	}
+	if elems < 1024 {
+		return ""
+	}
+
+	quantNorm := normalizeQuantType(quantize)
+	groupSize := int32(32)
+	switch quantNorm {
+	case "nvfp4":
+		groupSize = 16
+	case "int4", "int8":
+		groupSize = 64
+	}
+	if shape[len(shape)-1]%groupSize != 0 {
+		return ""
+	}
+
+	// Match the working HF-FP8 import policy for direct NVFP4/MXFP4/MXFP8 imports:
+	// keep embeddings, LM head, low-rank linear_attn projections, and routing
+	// gates in BF16 rather than forcing them into a non-affine quantized format.
+	if (quantNorm == "nvfp4" || quantNorm == "mxfp4" || quantNorm == "mxfp8") && qwen35ShouldKeepBF16ForDirectNonAffine(name) {
+		return ""
+	}
+
+	return quantNorm
+}
+
+func (t qwen35ImportTransform) rewriteTensorData(td *safetensors.TensorData) (*safetensors.TensorData, error) {
+	if td == nil {
+		return td, nil
+	}
+
+	shiftNorm := t.shouldShiftNormWeights && qwen35ShouldShiftNormKey(td.Name)
+	transposeConv := strings.Contains(td.Name, "conv1d.weight") && len(td.Shape) == 3 && td.Shape[2] != 1
+	castToBF16 := qwen35NeedsCastToBF16(td.Name, td.Dtype)
+	if !shiftNorm && !transposeConv && !castToBF16 {
+		return td, nil
+	}
+
+	raw, err := io.ReadAll(td.Reader())
+	if err != nil {
+		return nil, fmt.Errorf("failed to read tensor %s: %w", td.Name, err)
+	}
+
+	values, err := DecodeFloatTensor(td.Dtype, raw)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decode tensor %s: %w", td.Name, err)
+	}
+
+	shape := append([]int32(nil), td.Shape...)
+	if transposeConv {
+		values, shape = qwen35TransposeConv1D(values, shape)
+	}
+	if shiftNorm {
+		for i := range values {
+			values[i] += 1.0
+		}
+	}
+
+	targetDtype := td.Dtype
+	if castToBF16 {
+		targetDtype = "BF16"
+	}
+
+	out, err := EncodeFloatTensor(targetDtype, values)
+	if err != nil {
+		return nil, fmt.Errorf("failed to encode tensor %s: %w", td.Name, err)
+	}
+
+	return safetensors.NewTensorDataFromBytes(td.Name, targetDtype, shape, out), nil
+}
+
+func (t qwen35ImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) {
+	if td == nil {
+		return nil, nil
+	}
+
+	name := t.canonicalTensorName(td.Name)
+
+	// Phase 1: rename/split into intermediate tensors
+	var intermediates []*safetensors.TensorData
+	stripped := strings.TrimSuffix(name, ".weight")
+	switch {
+	case strings.HasSuffix(stripped, ".mlp.experts.gate_up_proj"):
+		prefix := strings.TrimSuffix(stripped, ".mlp.experts.gate_up_proj")
+		raw, err := io.ReadAll(td.Reader())
+		if err != nil {
+			return nil, fmt.Errorf("failed to read tensor %s: %w", td.Name, err)
+		}
+		gateRaw, upRaw, splitShape, err := qwen35SplitAxis1Raw(raw, td.Dtype, td.Shape)
+		if err != nil {
+			return nil, fmt.Errorf("failed to split tensor %s: %w", td.Name, err)
+		}
+		intermediates = []*safetensors.TensorData{
+			safetensors.NewTensorDataFromBytes(prefix+".mlp.switch_mlp.gate_proj.weight", td.Dtype, splitShape, gateRaw),
+			safetensors.NewTensorDataFromBytes(prefix+".mlp.switch_mlp.up_proj.weight", td.Dtype, splitShape, upRaw),
+		}
+	case strings.HasSuffix(stripped, ".mlp.experts.down_proj"):
+		newName := strings.TrimSuffix(stripped, ".mlp.experts.down_proj") + ".mlp.switch_mlp.down_proj.weight"
+		intermediates = []*safetensors.TensorData{td.WithName(newName)}
+	default:
+		intermediates = []*safetensors.TensorData{td.WithName(name)}
+	}
+
+	// Phase 2: rewrite all intermediates
+	results := make([]*safetensors.TensorData, 0, len(intermediates))
+	for _, inter := range intermediates {
+		rewritten, err := t.rewriteTensorData(inter)
+		if err != nil {
+			return nil, err
+		}
+		results = append(results, rewritten)
+	}
+	return results, nil
+}
+
+func (t qwen35ImportTransform) canonicalTensorName(name string) string {
+	// Vision tensors: normalize to vision_tower.* prefix
+	switch {
+	case strings.HasPrefix(name, "model.visual."):
+		return "vision_tower." + strings.TrimPrefix(name, "model.visual.")
+	case strings.HasPrefix(name, "vision_tower."):
+		return name
+	}
+
+	// Language model tensors: normalize to language_model.model.* prefix
+	if !t.rewriteLanguageModel {
+		return name
+	}
+	switch {
+	case strings.HasPrefix(name, "model.language_model"):
+		return "language_model.model" + strings.TrimPrefix(name, "model.language_model")
+	case strings.HasPrefix(name, "language_model."):
+		return name
+	default:
+		return "language_model." + name
+	}
+}
+
+func qwen35ShouldShiftNormKey(key string) bool {
+	for _, suffix := range []string{
+		".input_layernorm.weight",
+		".post_attention_layernorm.weight",
+		"model.norm.weight",
+		".q_norm.weight",
+		".k_norm.weight",
+	} {
+		if strings.HasSuffix(key, suffix) {
+			return true
+		}
+	}
+	return false
+}
+
+func qwen35NeedsCastToBF16(name, dtype string) bool {
+	if strings.HasSuffix(name, "A_log") {
+		return false
+	}
+	switch strings.ToUpper(dtype) {
+	case "F16", "F32", "F64":
+		return true
+	default:
+		return false
+	}
+}
+
+func qwen35TransposeConv1D(values []float32, shape []int32) ([]float32, []int32) {
+	if len(shape) != 3 {
+		return values, shape
+	}
+
+	d0, d1, d2 := int(shape[0]), int(shape[1]), int(shape[2])
+	out := make([]float32, len(values))
+	for i := range d0 {
+		for j := range d1 {
+			for k := range d2 {
+				inIdx := (i*d1+j)*d2 + k
+				outIdx := (i*d2+k)*d1 + j
+				out[outIdx] = values[inIdx]
+			}
+		}
+	}
+
+	return out, []int32{shape[0], shape[2], shape[1]}
+}
+
+func qwen35SplitAxis1Raw(raw []byte, dtype string, shape []int32) ([]byte, []byte, []int32, error) {
+	if len(shape) != 3 {
+		return nil, nil, nil, fmt.Errorf("expected 3D tensor, got shape %v", shape)
+	}
+	if shape[1]%2 != 0 {
+		return nil, nil, nil, fmt.Errorf("axis 1 dim %d is not even", shape[1])
+	}
+
+	elemSize, err := DTypeSize(dtype)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	d0, d1, d2 := int(shape[0]), int(shape[1]), int(shape[2])
+	perExpertBytes := d1 * d2 * elemSize
+	if len(raw) != d0*perExpertBytes {
+		return nil, nil, nil, fmt.Errorf("raw byte length %d does not match shape %v and dtype %s", len(raw), shape, dtype)
+	}
+
+	halfD1 := d1 / 2
+	halfExpertBytes := halfD1 * d2 * elemSize
+	gateRaw := make([]byte, d0*halfExpertBytes)
+	upRaw := make([]byte, d0*halfExpertBytes)
+	for e := range d0 {
+		src := e * perExpertBytes
+		dst := e * halfExpertBytes
+		copy(gateRaw[dst:dst+halfExpertBytes], raw[src:src+halfExpertBytes])
+		copy(upRaw[dst:dst+halfExpertBytes], raw[src+halfExpertBytes:src+perExpertBytes])
+	}
+
+	return gateRaw, upRaw, []int32{shape[0], int32(halfD1), shape[2]}, nil
+}