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/imagegen/safetensors/loader.go
+++ b/x/imagegen/safetensors/loader.go
@@ -0,0 +1,429 @@
+package safetensors
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+)
+
+// WeightSource is an interface for loading weights.
+// Both ModelWeights (directory-based) and ManifestWeights (blob-based) implement this.
+type WeightSource interface {
+	GetTensor(name string) (*mlx.Array, error)
+	ListTensors() []string
+	HasTensor(name string) bool
+	Quantization() string // Returns "NVFP4", "INT4", "INT8", or ""
+	GroupSize() int       // Returns quantization group size, or 0 if not specified
+}
+
+// QuantizationParams returns groupSize, bits, mode for a quantization type.
+// MLX quantization modes:
+//   - "affine": scale + zero-point bias, group_size=32/64/128
+//   - "nvfp4": NVIDIA FP4 with E4M3 scales, group_size=16 (no bias)
+//   - "mxfp8": Microsoft MX FP8 with E4M3 scales, group_size=32 (no bias)
+func QuantizationParams(quantization string) (groupSize, bits int, mode string) {
+	switch strings.ToUpper(quantization) {
+	case "NVFP4":
+		// NVIDIA FP4: group_size=16, bits=4, E4M3 scales (no qbias)
+		return 16, 4, "nvfp4"
+	case "FP4", "Q4", "INT4":
+		// 4-bit quantization with affine mode (scale + qbias)
+		return 32, 4, "affine"
+	case "MXFP8":
+		// Microsoft MX FP8: group_size=32, bits=8, E4M3 scales (no qbias)
+		return 32, 8, "mxfp8"
+	case "FP8", "Q8", "INT8":
+		// 8-bit quantization with affine mode (default for quantized models)
+		return 64, 8, "affine"
+	case "":
+		return 0, 0, ""
+	default:
+		return 32, 8, "affine" // Default to affine
+	}
+}
+
+// Transformer allows structs to transform weight arrays before assignment.
+// Implement this to apply operations like transpose during loading.
+type Transformer interface {
+	Transform(field string, arr *mlx.Array) *mlx.Array
+}
+
+// LoadModule loads weights into a struct using reflection and struct tags.
+//
+// Struct tags use the format: `weight:"path[,optional]"`
+//   - path: the weight name suffix (appended to prefix)
+//   - optional: if present, missing weights don't cause errors
+//   - "-": skip this field entirely
+//   - no tag on struct pointer: recurse with current prefix
+//   - no tag on *mlx.Array: skip (computed fields don't need loading)
+//
+// For slices of struct pointers, the loader iterates with .0, .1, .2... suffixes.
+// The slice must be pre-allocated to the correct length.
+//
+// Example:
+//
+//	type Attention struct {
+//	    QProj *nn.Linear  `weight:"self_attn.q_proj"`
+//	    KProj *nn.Linear  `weight:"self_attn.k_proj"`
+//	    Cache *mlx.Array  // no tag = skipped (computed field)
+//	}
+//
+//	err := LoadModule(&attn, weights, "model.layers.0")
+func LoadModule(dst any, weights WeightSource, prefix string) error {
+	v := reflect.ValueOf(dst)
+	if v.Kind() != reflect.Ptr || v.IsNil() {
+		return fmt.Errorf("LoadModule: dst must be a non-nil pointer")
+	}
+	v = v.Elem()
+	if v.Kind() != reflect.Struct {
+		return fmt.Errorf("LoadModule: dst must be a pointer to struct, got %v", v.Kind())
+	}
+
+	var errs []string
+	loadStruct(v, weights, prefix, &errs, false)
+
+	if len(errs) > 0 {
+		return fmt.Errorf("LoadModule: missing weights:\n  %s", strings.Join(errs, "\n  "))
+	}
+	return nil
+}
+
+// loadStruct recursively loads weights into a struct value.
+func loadStruct(v reflect.Value, weights WeightSource, prefix string, errs *[]string, parentOptional bool) {
+	t := v.Type()
+
+	for i := 0; i < t.NumField(); i++ {
+		field := t.Field(i)
+		fieldVal := v.Field(i)
+
+		// Skip unexported fields
+		if !fieldVal.CanSet() {
+			continue
+		}
+
+		// Parse tag
+		tag, hasTag := field.Tag.Lookup("weight")
+		if tag == "-" {
+			continue
+		}
+
+		// Parse tag options
+		optional := parentOptional
+		weightPath := tag
+		if idx := strings.Index(tag, ","); idx != -1 {
+			weightPath = tag[:idx]
+			if strings.Contains(tag[idx+1:], "optional") {
+				optional = true
+			}
+		}
+
+		// Build full path
+		fullPath := joinPath(prefix, weightPath)
+
+		// For struct pointers without a tag, recurse with current prefix
+		if !hasTag && fieldVal.Kind() == reflect.Ptr {
+			elemType := fieldVal.Type().Elem()
+			if elemType.Kind() == reflect.Struct && elemType != reflect.TypeOf(mlx.Array{}) {
+				if fieldVal.IsNil() {
+					fieldVal.Set(reflect.New(elemType))
+				}
+				loadStruct(fieldVal.Elem(), weights, prefix, errs, optional)
+				continue
+			}
+		}
+
+		// Handle nn.LinearLayer interface fields specially
+		linearLayerType := reflect.TypeOf((*nn.LinearLayer)(nil)).Elem()
+		if field.Type == linearLayerType {
+			if !hasTag {
+				continue // no tag = skip
+			}
+			layer, err := LoadLinearLayer(weights, fullPath)
+			if err != nil {
+				if !optional {
+					*errs = append(*errs, fullPath+": "+err.Error())
+				}
+				continue
+			}
+			fieldVal.Set(reflect.ValueOf(layer))
+			continue
+		}
+
+		// Handle nn.MultiLinearLayer interface fields specially
+		multiLinearLayerType := reflect.TypeOf((*nn.MultiLinearLayer)(nil)).Elem()
+		if field.Type == multiLinearLayerType {
+			if !hasTag {
+				continue // no tag = skip
+			}
+			layer, err := LoadMultiLinearLayer(weights, fullPath)
+			if err != nil {
+				if !optional {
+					*errs = append(*errs, fullPath+": "+err.Error())
+				}
+				continue
+			}
+			fieldVal.Set(reflect.ValueOf(layer))
+			continue
+		}
+
+		// Handle by kind
+		switch fieldVal.Kind() {
+		case reflect.Ptr:
+			elemType := fieldVal.Type().Elem()
+
+			// *mlx.Array - load directly (but skip if no tag - computed fields)
+			if fieldVal.Type() == reflect.TypeOf((*mlx.Array)(nil)) {
+				if !hasTag {
+					continue // no tag on *mlx.Array = computed field, skip
+				}
+				arr, err := weights.GetTensor(fullPath)
+				if err != nil {
+					if !optional {
+						*errs = append(*errs, fullPath)
+					}
+					continue
+				}
+				// Transform before assigning if parent implements Transformer
+				if t, ok := v.Addr().Interface().(Transformer); ok {
+					arr = t.Transform(field.Name, arr)
+				}
+				fieldVal.Set(reflect.ValueOf(arr))
+				continue
+			}
+
+			// Pointer to struct - allocate and recurse
+			if elemType.Kind() == reflect.Struct {
+				if optional && !hasWeightsWithPrefix(weights, fullPath) {
+					continue
+				}
+				if fieldVal.IsNil() {
+					fieldVal.Set(reflect.New(elemType))
+				}
+				loadStruct(fieldVal.Elem(), weights, fullPath, errs, optional)
+			}
+
+		case reflect.Slice:
+			elemType := fieldVal.Type().Elem()
+			if elemType.Kind() == reflect.Ptr && elemType.Elem().Kind() == reflect.Struct {
+				loadSlice(fieldVal, weights, fullPath, errs)
+			}
+		}
+	}
+}
+
+// hasWeightsWithPrefix checks if any weights exist with the given prefix.
+func hasWeightsWithPrefix(weights WeightSource, prefix string) bool {
+	for _, name := range weights.ListTensors() {
+		if strings.HasPrefix(name, prefix+".") || name == prefix {
+			return true
+		}
+	}
+	return false
+}
+
+// loadSlice loads weights into each element of a slice of struct pointers.
+func loadSlice(v reflect.Value, weights WeightSource, prefix string, errs *[]string) {
+	elemStructType := v.Type().Elem().Elem()
+
+	for i := 0; i < v.Len(); i++ {
+		elem := v.Index(i)
+		if elem.IsNil() {
+			elem.Set(reflect.New(elemStructType))
+		}
+		loadStruct(elem.Elem(), weights, fmt.Sprintf("%s.%d", prefix, i), errs, false)
+	}
+}
+
+// joinPath joins path segments with dots, handling empty segments.
+func joinPath(prefix, suffix string) string {
+	if prefix == "" {
+		return suffix
+	}
+	if suffix == "" {
+		return prefix
+	}
+	return prefix + "." + suffix
+}
+
+// LoadMultiLinearLayer loads a per-head linear layer from weights.
+// Weight shape should be [num_heads, output_dims, input_dims].
+// If quantized, always dequantizes since batched quantized matmul isn't supported.
+func LoadMultiLinearLayer(weights WeightSource, path string) (nn.MultiLinearLayer, error) {
+	// Check if this is a quantized layer by looking for scale tensor
+	scalePath := path + ".weight_scale"
+	hasScale := weights.HasTensor(scalePath)
+
+	weight, err := weights.GetTensor(path + ".weight")
+	if err != nil {
+		return nil, fmt.Errorf("failed to load weight %s: %w", path, err)
+	}
+
+	if hasScale {
+		scales, err := weights.GetTensor(scalePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load scales %s: %w", scalePath, err)
+		}
+
+		var qbiases *mlx.Array
+		qbiasPath := path + ".weight_qbias"
+		if weights.HasTensor(qbiasPath) {
+			qbiases, _ = weights.GetTensor(qbiasPath)
+		}
+
+		// Always dequantize for MultiLinear - no batched quantized matmul support
+		// Detect bits from tensor shapes (supports mixed-precision Q4/Q8)
+		weightShape := weight.Shape()
+		scalesShape := scales.Shape()
+		weightCols := int(weightShape[len(weightShape)-1])
+		scalesCols := int(scalesShape[len(scalesShape)-1])
+
+		// Detect quantization from tensor shapes
+		// groupSize = weightCols * packFactor / scalesCols
+		// Note: groupSize4 = 2 * groupSize8 always, so ambiguous cases need metadata
+		groupSize4 := weightCols * 8 / scalesCols
+		groupSize8 := weightCols * 4 / scalesCols
+
+		var bits, groupSize int
+		// Use metadata to help disambiguate when shapes are ambiguous
+		// (e.g., Q4 with group_size=64 has same shapes as Q8 with group_size=32)
+		quantType := strings.ToUpper(weights.Quantization())
+		isQ8Type := quantType == "Q8" || quantType == "FP8" || quantType == "INT8"
+
+		if groupSize4 == 32 {
+			// Unambiguous: Q4 with group_size=32
+			bits = 4
+			groupSize = 32
+		} else if groupSize8 == 64 {
+			// Unambiguous: Q8 with group_size=64
+			bits = 8
+			groupSize = 64
+		} else if groupSize4 == 64 && groupSize8 == 32 {
+			// Ambiguous: could be Q4/gs=64 or Q8/gs=32, use metadata
+			if isQ8Type {
+				bits = 8
+				groupSize = 32
+			} else {
+				bits = 4
+				groupSize = 64
+			}
+		} else {
+			// Fallback: use global quantization params
+			_, bits, _ = QuantizationParams(weights.Quantization())
+			packFactor := 32 / bits
+			groupSize = weightCols * packFactor / scalesCols
+		}
+		weight = mlx.Dequantize(weight, scales, qbiases, groupSize, bits, "affine")
+	}
+
+	return nn.NewMultiLinear(weight), nil
+}
+
+// LoadLinearLayer loads a linear layer from weights, automatically detecting if it's quantized.
+// If {path}.weight_scale exists, creates a QuantizedLinear layer (or dequantizes if no kernel support).
+func LoadLinearLayer(weights WeightSource, path string) (nn.LinearLayer, error) {
+	// Check if this is a quantized layer by looking for scale tensor
+	scalePath := path + ".weight_scale"
+	hasScale := weights.HasTensor(scalePath)
+	if hasScale {
+		weight, err := weights.GetTensor(path + ".weight")
+		if err != nil {
+			return nil, fmt.Errorf("failed to load quantized weight %s: %w", path, err)
+		}
+
+		scales, err := weights.GetTensor(scalePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load scales %s: %w", scalePath, err)
+		}
+
+		// Bias is optional
+		var bias *mlx.Array
+		biasPath := path + ".bias"
+		if weights.HasTensor(biasPath) {
+			bias, _ = weights.GetTensor(biasPath)
+		}
+
+		var qbiases *mlx.Array
+		qbiasPath := path + ".weight_qbias"
+		if weights.HasTensor(qbiasPath) {
+			qbiases, _ = weights.GetTensor(qbiasPath)
+		}
+
+		// Detect bits from tensor shapes (supports mixed-precision Q4/Q8)
+		weightShape := weight.Shape()
+		scalesShape := scales.Shape()
+		weightCols := int(weightShape[len(weightShape)-1])
+		scalesCols := int(scalesShape[len(scalesShape)-1])
+
+		// Detect quantization from tensor shapes
+		// groupSize = weightCols * packFactor / scalesCols
+		// Note: groupSize4 = 2 * groupSize8 always, so ambiguous cases need metadata
+		groupSize4 := weightCols * 8 / scalesCols
+		groupSize8 := weightCols * 4 / scalesCols
+
+		var bits, groupSize int
+		mode := "affine"
+		// Use metadata to help disambiguate when shapes are ambiguous
+		// (e.g., Q4 with group_size=64 has same shapes as Q8 with group_size=32)
+		quantType := strings.ToUpper(weights.Quantization())
+		isQ8Type := quantType == "Q8" || quantType == "FP8" || quantType == "INT8"
+
+		if groupSize4 == 32 {
+			// Unambiguous: Q4 with group_size=32
+			bits = 4
+			groupSize = 32
+		} else if groupSize8 == 64 {
+			// Unambiguous: Q8 with group_size=64
+			bits = 8
+			groupSize = 64
+		} else if groupSize4 == 64 && groupSize8 == 32 {
+			// Ambiguous: could be Q4/gs=64 or Q8/gs=32, use metadata
+			if isQ8Type {
+				bits = 8
+				groupSize = 32
+			} else {
+				bits = 4
+				groupSize = 64
+			}
+		} else {
+			// Fallback: use global quantization params
+			_, bits, mode = QuantizationParams(weights.Quantization())
+			packFactor := 32 / bits
+			groupSize = weightCols * packFactor / scalesCols
+		}
+
+		// NVFP4 and MXFP8 don't have native quantized matmul kernels in MLX,
+		// so we always dequantize at load time. Affine modes (FP4, FP8) have kernel support.
+		if mlx.MetalIsAvailable() && mode != "nvfp4" && mode != "mxfp8" {
+			return &nn.QuantizedLinear{
+				Weight:    weight,
+				Scales:    scales,
+				QBiases:   qbiases,
+				Bias:      bias,
+				GroupSize: groupSize,
+				Bits:      bits,
+				Mode:      mode,
+			}, nil
+		}
+
+		dequantized := mlx.Dequantize(weight, scales, qbiases, groupSize, bits, mode)
+		return nn.NewLinear(dequantized, bias), nil
+	}
+
+	// Load as regular Linear
+	weight, err := weights.GetTensor(path + ".weight")
+	if err != nil {
+		return nil, fmt.Errorf("failed to load weight %s: %w", path, err)
+	}
+
+	// Bias is optional
+	var bias *mlx.Array
+	biasPath := path + ".bias"
+	if weights.HasTensor(biasPath) {
+		bias, _ = weights.GetTensor(biasPath)
+	}
+
+	return nn.NewLinear(weight, bias), nil
+}