mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 05:28:00 -05:00
mlxrunner: lay out model by contract, checkpoint and construction
model is one package with three jobs: the contract between the runner and the architectures, the opened checkpoint, and building nn layers from checkpoint tensors. Its files did not say which was which. base.go carried the folded package's name over the interfaces and the registry, root.go held the safetensors header scan next to Root, and quant.go mixed the nvfp4 global-scale helpers with quant parameter resolution. base.go becomes model.go, named for what it holds. root.go keeps Root and Open; TensorQuantInfo and the header scan join quant.go, so everything the checkpoint says about quantization is read and resolved in one file. The global-scale helpers move to globalscale.go with their tests. Root.Close, a no-op with one caller, goes. No code changes otherwise.
This commit is contained in:
Jesse Gross
@@ -0,0 +1,89 @@
|
||||
package model
|
||||
|
||||
import "github.com/ollama/ollama/mlx"
|
||||
|
||||
// Import rewrites every vendor spelling to ".global_scale"; "_scale_2" is
|
||||
// ModelOpt's own name, reached when a checkpoint skips import.
|
||||
var globalScaleSuffixes = []string{".global_scale", "_scale_2"}
|
||||
|
||||
// These scale the activations, never the weight, but are freed alongside it.
|
||||
var activationScaleSuffixes = []string{".input_global_scale", ".input_scale"}
|
||||
|
||||
// ReadGlobalScale returns a weight's NVFP4 global scale in MLX's
|
||||
// representation, and the companion keys the caller should release. Candidate
|
||||
// keys are tried in order, so pass the resolved tensor key before any base.
|
||||
func ReadGlobalScale(tensors map[string]*mlx.Array, weightKeys ...string) (*mlx.Array, []string) {
|
||||
var found *mlx.Array
|
||||
var consumed []string
|
||||
for _, key := range weightKeys {
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
for _, suffix := range globalScaleSuffixes {
|
||||
scale, ok := tensors[key+suffix]
|
||||
if !ok || scale == nil {
|
||||
continue
|
||||
}
|
||||
if found == nil {
|
||||
found = scale
|
||||
}
|
||||
consumed = append(consumed, key+suffix)
|
||||
}
|
||||
for _, suffix := range activationScaleSuffixes {
|
||||
if _, ok := tensors[key+suffix]; ok {
|
||||
consumed = append(consumed, key+suffix)
|
||||
}
|
||||
}
|
||||
}
|
||||
return ToMLXGlobalScale(found), consumed
|
||||
}
|
||||
|
||||
// ToMLXGlobalScale converts a checkpoint multiplier into the representation
|
||||
// every global scale is held in once loaded. Shape is flattened too: a scalar
|
||||
// ships as either [] or [1], and stacking a mix of the two fails.
|
||||
func ToMLXGlobalScale(globalScale *mlx.Array) *mlx.Array {
|
||||
if globalScale == nil {
|
||||
return nil
|
||||
}
|
||||
flat := mlx.Reshape(globalScale.AsType(mlx.DTypeFloat32), int32(globalScale.Size()))
|
||||
return mlx.MulScalar(flat, mlx.Nvfp4MaxProduct)
|
||||
}
|
||||
|
||||
// PrepareGatherQMMGlobalScale broadcasts an already-converted global scale
|
||||
// into the one-entry-per-expert bank gather_qmm wants. Materialized dense: the
|
||||
// kernel indexes it by raw offset, and a broadcast view is one element of
|
||||
// storage.
|
||||
func PrepareGatherQMMGlobalScale(globalScale *mlx.Array, numExperts int) *mlx.Array {
|
||||
if globalScale == nil {
|
||||
return nil
|
||||
}
|
||||
return mlx.Contiguous(mlx.BroadcastTo(globalScale, int32(numExperts)), false)
|
||||
}
|
||||
|
||||
// GatherQMMIdentityScale is the scale that leaves an expert bank unscaled,
|
||||
// for rows folded into a scaled bank without a scale of their own.
|
||||
func GatherQMMIdentityScale() *mlx.Array {
|
||||
return mlx.FromValues([]float32{mlx.Nvfp4MaxProduct}, 1)
|
||||
}
|
||||
|
||||
// SameGlobalScales reports whether two prepared banks hold the same scale for
|
||||
// every expert, which is what lets two projections share one fused bank.
|
||||
func SameGlobalScales(a, b *mlx.Array) bool {
|
||||
if a == nil || b == nil {
|
||||
return a == nil && b == nil
|
||||
}
|
||||
if a == b {
|
||||
return true
|
||||
}
|
||||
if a.Size() != b.Size() {
|
||||
return false
|
||||
}
|
||||
mlx.Eval(a, b)
|
||||
aValues, bValues := a.Floats(), b.Floats()
|
||||
for i := range aValues {
|
||||
if aValues[i] != bValues[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
+160
-62
@@ -1,94 +1,192 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlx/quant"
|
||||
)
|
||||
|
||||
// Import rewrites every vendor spelling to ".global_scale"; "_scale_2" is
|
||||
// ModelOpt's own name, reached when a checkpoint skips import.
|
||||
var globalScaleSuffixes = []string{".global_scale", "_scale_2"}
|
||||
// TensorQuantInfo describes per-tensor quantization metadata.
|
||||
type TensorQuantInfo struct {
|
||||
QuantType string
|
||||
GroupSize int
|
||||
}
|
||||
|
||||
// These scale the activations, never the weight, but are freed alongside it.
|
||||
var activationScaleSuffixes = []string{".input_global_scale", ".input_scale"}
|
||||
func readBlobTensorQuantInfo(path string) (map[string]*TensorQuantInfo, string, int, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
// ReadGlobalScale returns a weight's NVFP4 global scale in MLX's
|
||||
// representation, and the companion keys the caller should release. Candidate
|
||||
// keys are tried in order, so pass the resolved tensor key before any base.
|
||||
func ReadGlobalScale(tensors map[string]*mlx.Array, weightKeys ...string) (*mlx.Array, []string) {
|
||||
var found *mlx.Array
|
||||
var consumed []string
|
||||
for _, key := range weightKeys {
|
||||
if key == "" {
|
||||
var headerSize uint64
|
||||
if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
if headerSize > 100*1024*1024 {
|
||||
return nil, "", 0, fmt.Errorf("header too large: %d", headerSize)
|
||||
}
|
||||
|
||||
data := make([]byte, headerSize)
|
||||
if _, err := io.ReadFull(f, data); err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
|
||||
var header map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &header); err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
|
||||
globalQuantType, globalGroupSize := parseGlobalQuantMetadata(header)
|
||||
globalQuantType = strings.ToUpper(globalQuantType)
|
||||
|
||||
// Parse full metadata for per-tensor quant info
|
||||
var metaMap map[string]string
|
||||
if metaRaw, ok := header["__metadata__"]; ok {
|
||||
json.Unmarshal(metaRaw, &metaMap)
|
||||
}
|
||||
|
||||
mainNames := mainTensorNames(header)
|
||||
infos := make(map[string]*TensorQuantInfo)
|
||||
for _, name := range mainNames {
|
||||
if _, ok := header[name+".scale"]; !ok {
|
||||
continue
|
||||
}
|
||||
for _, suffix := range globalScaleSuffixes {
|
||||
scale, ok := tensors[key+suffix]
|
||||
if !ok || scale == nil {
|
||||
continue
|
||||
|
||||
quantType := globalQuantType
|
||||
groupSize := globalGroupSize
|
||||
|
||||
// Check per-tensor metadata (e.g. from packed expert blobs with mixed precision)
|
||||
if metaMap != nil {
|
||||
if qt, ok := metaMap[name+".quant_type"]; ok && qt != "" {
|
||||
quantType = strings.ToUpper(qt)
|
||||
}
|
||||
if found == nil {
|
||||
found = scale
|
||||
}
|
||||
consumed = append(consumed, key+suffix)
|
||||
}
|
||||
for _, suffix := range activationScaleSuffixes {
|
||||
if _, ok := tensors[key+suffix]; ok {
|
||||
consumed = append(consumed, key+suffix)
|
||||
if gs, ok := metaMap[name+".group_size"]; ok && gs != "" {
|
||||
if v, err := strconv.Atoi(gs); err == nil {
|
||||
groupSize = v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
inferredType, inferredGroup := inferQuantTypeFromShapes(header, name, quantType)
|
||||
if quantType == "" {
|
||||
quantType = inferredType
|
||||
}
|
||||
if groupSize == 0 {
|
||||
groupSize = inferredGroup
|
||||
}
|
||||
if quantType == "" {
|
||||
continue
|
||||
}
|
||||
if groupSize == 0 {
|
||||
groupSize = defaultGroupSize(quantType)
|
||||
}
|
||||
|
||||
infos[name] = &TensorQuantInfo{QuantType: quantType, GroupSize: groupSize}
|
||||
}
|
||||
return ToMLXGlobalScale(found), consumed
|
||||
|
||||
return infos, globalQuantType, globalGroupSize, nil
|
||||
}
|
||||
|
||||
// ToMLXGlobalScale converts a checkpoint multiplier into the representation
|
||||
// every global scale is held in once loaded. Shape is flattened too: a scalar
|
||||
// ships as either [] or [1], and stacking a mix of the two fails.
|
||||
func ToMLXGlobalScale(globalScale *mlx.Array) *mlx.Array {
|
||||
if globalScale == nil {
|
||||
return nil
|
||||
func parseGlobalQuantMetadata(header map[string]json.RawMessage) (quantType string, groupSize int) {
|
||||
metaRaw, ok := header["__metadata__"]
|
||||
if !ok {
|
||||
return "", 0
|
||||
}
|
||||
flat := mlx.Reshape(globalScale.AsType(mlx.DTypeFloat32), int32(globalScale.Size()))
|
||||
return mlx.MulScalar(flat, mlx.Nvfp4MaxProduct)
|
||||
|
||||
var meta map[string]string
|
||||
if err := json.Unmarshal(metaRaw, &meta); err != nil {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
quantType = meta["quant_type"]
|
||||
if gs := meta["group_size"]; gs != "" {
|
||||
groupSize, _ = strconv.Atoi(gs)
|
||||
}
|
||||
return quantType, groupSize
|
||||
}
|
||||
|
||||
// PrepareGatherQMMGlobalScale broadcasts an already-converted global scale
|
||||
// into the one-entry-per-expert bank gather_qmm wants. Materialized dense: the
|
||||
// kernel indexes it by raw offset, and a broadcast view is one element of
|
||||
// storage.
|
||||
func PrepareGatherQMMGlobalScale(globalScale *mlx.Array, numExperts int) *mlx.Array {
|
||||
if globalScale == nil {
|
||||
return nil
|
||||
func mainTensorNames(header map[string]json.RawMessage) []string {
|
||||
names := make([]string, 0, len(header))
|
||||
for name := range header {
|
||||
if name == "__metadata__" || strings.HasSuffix(name, ".scale") || strings.HasSuffix(name, ".bias") {
|
||||
continue
|
||||
}
|
||||
names = append(names, name)
|
||||
}
|
||||
return mlx.Contiguous(mlx.BroadcastTo(globalScale, int32(numExperts)), false)
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
// GatherQMMIdentityScale is the scale that leaves an expert bank unscaled,
|
||||
// for rows folded into a scaled bank without a scale of their own.
|
||||
func GatherQMMIdentityScale() *mlx.Array {
|
||||
return mlx.FromValues([]float32{mlx.Nvfp4MaxProduct}, 1)
|
||||
}
|
||||
func inferQuantTypeFromShapes(header map[string]json.RawMessage, tensorName string, hintQuantType string) (string, int) {
|
||||
type tensorShape struct {
|
||||
Shape []int64 `json:"shape"`
|
||||
}
|
||||
|
||||
// SameGlobalScales reports whether two prepared banks hold the same scale for
|
||||
// every expert, which is what lets two projections share one fused bank.
|
||||
func SameGlobalScales(a, b *mlx.Array) bool {
|
||||
if a == nil || b == nil {
|
||||
return a == nil && b == nil
|
||||
mainRaw, ok := header[tensorName]
|
||||
if !ok {
|
||||
return "", 0
|
||||
}
|
||||
if a == b {
|
||||
return true
|
||||
scaleRaw, ok := header[tensorName+".scale"]
|
||||
if !ok {
|
||||
return "", 0
|
||||
}
|
||||
if a.Size() != b.Size() {
|
||||
return false
|
||||
|
||||
var mainInfo tensorShape
|
||||
if err := json.Unmarshal(mainRaw, &mainInfo); err != nil || len(mainInfo.Shape) == 0 {
|
||||
return "", 0
|
||||
}
|
||||
mlx.Eval(a, b)
|
||||
aValues, bValues := a.Floats(), b.Floats()
|
||||
for i := range aValues {
|
||||
if aValues[i] != bValues[i] {
|
||||
return false
|
||||
|
||||
var scaleInfo tensorShape
|
||||
if err := json.Unmarshal(scaleRaw, &scaleInfo); err != nil || len(scaleInfo.Shape) == 0 {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
weightCols := int(mainInfo.Shape[len(mainInfo.Shape)-1])
|
||||
scalesCols := int(scaleInfo.Shape[len(scaleInfo.Shape)-1])
|
||||
if weightCols <= 0 || scalesCols <= 0 {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
groupSize4 := weightCols * 8 / scalesCols
|
||||
groupSize8 := weightCols * 4 / scalesCols
|
||||
|
||||
switch {
|
||||
case groupSize4 == 32:
|
||||
return "INT4", 32
|
||||
case groupSize8 == 64:
|
||||
return "INT8", 64
|
||||
case groupSize4 == 64 && groupSize8 == 32:
|
||||
h := strings.ToUpper(hintQuantType)
|
||||
if strings.Contains(h, "8") {
|
||||
return "INT8", 32
|
||||
}
|
||||
if strings.Contains(h, "4") {
|
||||
return "INT4", 64
|
||||
}
|
||||
}
|
||||
return true
|
||||
|
||||
if isCommonGroupSize(groupSize4) && !isCommonGroupSize(groupSize8) {
|
||||
return "INT4", groupSize4
|
||||
}
|
||||
if isCommonGroupSize(groupSize8) && !isCommonGroupSize(groupSize4) {
|
||||
return "INT8", groupSize8
|
||||
}
|
||||
|
||||
return "", 0
|
||||
}
|
||||
|
||||
func defaultGroupSize(quantType string) int {
|
||||
groupSize, _, _ := QuantizationParams(quantType)
|
||||
return groupSize
|
||||
}
|
||||
|
||||
// QuantizationParams returns default groupSize, bits, and mode for a
|
||||
|
||||
@@ -1,25 +1,14 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/manifest"
|
||||
modeltypes "github.com/ollama/ollama/types/model"
|
||||
)
|
||||
|
||||
// TensorQuantInfo describes per-tensor quantization metadata.
|
||||
type TensorQuantInfo struct {
|
||||
QuantType string
|
||||
GroupSize int
|
||||
}
|
||||
|
||||
// Root wraps a model's manifest with pre-scanned quantization metadata.
|
||||
type Root struct {
|
||||
Manifest *manifest.Manifest
|
||||
@@ -106,9 +95,6 @@ func readDraftConfig(m *manifest.Manifest) *modeltypes.Draft {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close is a no-op for now (future: release resources).
|
||||
func (r *Root) Close() {}
|
||||
|
||||
// QuantType returns the quantization type detected from the first tensor blob metadata.
|
||||
func (r *Root) QuantType() string { return r.quantType }
|
||||
|
||||
@@ -135,172 +121,3 @@ func (r *Root) AllTensorQuant() map[string]*TensorQuantInfo {
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func defaultGroupSize(quantType string) int {
|
||||
groupSize, _, _ := QuantizationParams(quantType)
|
||||
return groupSize
|
||||
}
|
||||
|
||||
func readBlobTensorQuantInfo(path string) (map[string]*TensorQuantInfo, string, int, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
var headerSize uint64
|
||||
if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
if headerSize > 100*1024*1024 {
|
||||
return nil, "", 0, fmt.Errorf("header too large: %d", headerSize)
|
||||
}
|
||||
|
||||
data := make([]byte, headerSize)
|
||||
if _, err := io.ReadFull(f, data); err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
|
||||
var header map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &header); err != nil {
|
||||
return nil, "", 0, err
|
||||
}
|
||||
|
||||
globalQuantType, globalGroupSize := parseGlobalQuantMetadata(header)
|
||||
globalQuantType = strings.ToUpper(globalQuantType)
|
||||
|
||||
// Parse full metadata for per-tensor quant info
|
||||
var metaMap map[string]string
|
||||
if metaRaw, ok := header["__metadata__"]; ok {
|
||||
json.Unmarshal(metaRaw, &metaMap)
|
||||
}
|
||||
|
||||
mainNames := mainTensorNames(header)
|
||||
infos := make(map[string]*TensorQuantInfo)
|
||||
for _, name := range mainNames {
|
||||
if _, ok := header[name+".scale"]; !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
quantType := globalQuantType
|
||||
groupSize := globalGroupSize
|
||||
|
||||
// Check per-tensor metadata (e.g. from packed expert blobs with mixed precision)
|
||||
if metaMap != nil {
|
||||
if qt, ok := metaMap[name+".quant_type"]; ok && qt != "" {
|
||||
quantType = strings.ToUpper(qt)
|
||||
}
|
||||
if gs, ok := metaMap[name+".group_size"]; ok && gs != "" {
|
||||
if v, err := strconv.Atoi(gs); err == nil {
|
||||
groupSize = v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
inferredType, inferredGroup := inferQuantTypeFromShapes(header, name, quantType)
|
||||
if quantType == "" {
|
||||
quantType = inferredType
|
||||
}
|
||||
if groupSize == 0 {
|
||||
groupSize = inferredGroup
|
||||
}
|
||||
if quantType == "" {
|
||||
continue
|
||||
}
|
||||
if groupSize == 0 {
|
||||
groupSize = defaultGroupSize(quantType)
|
||||
}
|
||||
|
||||
infos[name] = &TensorQuantInfo{QuantType: quantType, GroupSize: groupSize}
|
||||
}
|
||||
|
||||
return infos, globalQuantType, globalGroupSize, nil
|
||||
}
|
||||
|
||||
func parseGlobalQuantMetadata(header map[string]json.RawMessage) (quantType string, groupSize int) {
|
||||
metaRaw, ok := header["__metadata__"]
|
||||
if !ok {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
var meta map[string]string
|
||||
if err := json.Unmarshal(metaRaw, &meta); err != nil {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
quantType = meta["quant_type"]
|
||||
if gs := meta["group_size"]; gs != "" {
|
||||
groupSize, _ = strconv.Atoi(gs)
|
||||
}
|
||||
return quantType, groupSize
|
||||
}
|
||||
|
||||
func mainTensorNames(header map[string]json.RawMessage) []string {
|
||||
names := make([]string, 0, len(header))
|
||||
for name := range header {
|
||||
if name == "__metadata__" || strings.HasSuffix(name, ".scale") || strings.HasSuffix(name, ".bias") {
|
||||
continue
|
||||
}
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
func inferQuantTypeFromShapes(header map[string]json.RawMessage, tensorName string, hintQuantType string) (string, int) {
|
||||
type tensorShape struct {
|
||||
Shape []int64 `json:"shape"`
|
||||
}
|
||||
|
||||
mainRaw, ok := header[tensorName]
|
||||
if !ok {
|
||||
return "", 0
|
||||
}
|
||||
scaleRaw, ok := header[tensorName+".scale"]
|
||||
if !ok {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
var mainInfo tensorShape
|
||||
if err := json.Unmarshal(mainRaw, &mainInfo); err != nil || len(mainInfo.Shape) == 0 {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
var scaleInfo tensorShape
|
||||
if err := json.Unmarshal(scaleRaw, &scaleInfo); err != nil || len(scaleInfo.Shape) == 0 {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
weightCols := int(mainInfo.Shape[len(mainInfo.Shape)-1])
|
||||
scalesCols := int(scaleInfo.Shape[len(scaleInfo.Shape)-1])
|
||||
if weightCols <= 0 || scalesCols <= 0 {
|
||||
return "", 0
|
||||
}
|
||||
|
||||
groupSize4 := weightCols * 8 / scalesCols
|
||||
groupSize8 := weightCols * 4 / scalesCols
|
||||
|
||||
switch {
|
||||
case groupSize4 == 32:
|
||||
return "INT4", 32
|
||||
case groupSize8 == 64:
|
||||
return "INT8", 64
|
||||
case groupSize4 == 64 && groupSize8 == 32:
|
||||
h := strings.ToUpper(hintQuantType)
|
||||
if strings.Contains(h, "8") {
|
||||
return "INT8", 32
|
||||
}
|
||||
if strings.Contains(h, "4") {
|
||||
return "INT4", 64
|
||||
}
|
||||
}
|
||||
|
||||
if isCommonGroupSize(groupSize4) && !isCommonGroupSize(groupSize8) {
|
||||
return "INT4", groupSize4
|
||||
}
|
||||
if isCommonGroupSize(groupSize8) && !isCommonGroupSize(groupSize4) {
|
||||
return "INT8", groupSize8
|
||||
}
|
||||
|
||||
return "", 0
|
||||
}
|
||||
|
||||
@@ -78,7 +78,6 @@ func (r *Runner) loadModel(modelName string) (weights []*mlx.Array, err error) {
|
||||
err = e
|
||||
return nil
|
||||
}
|
||||
defer root.Close()
|
||||
|
||||
m, e := model.New(root)
|
||||
if e != nil {
|
||||
|
||||
Reference in New Issue
Block a user