mirror of
https://github.com/ollama/ollama.git
synced 2026-07-06 07:41:05 +00:00
03aee88186
* mlx: Support NVIDIA TensorRT Model Optimizer import * x/create: support FP8 safetensors import Decode HF F8_E4M3 safetensors with block scale companions into MLX-importable tensor blobs, including compressed-tensors weight_scale metadata, packed NVFP4 layouts, and mixed-precision tensor headers. Use that source-precision metadata during create quantization: default FP8-sourced imports to mxfp8, allow source FP8 to target MLX low-bit formats, preserve source-quantized NVFP4 layouts, selectively keep or promote tensors based on their source precision, and detect quantized dtype from mixed-precision safetensors manifests. * review comments
99 lines
2.4 KiB
Go
99 lines
2.4 KiB
Go
package model
|
|
|
|
import (
|
|
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
|
"github.com/ollama/ollama/x/models/nn"
|
|
)
|
|
|
|
// LinearFactory builds linear layers using shared tensor maps and quant defaults.
|
|
type LinearFactory struct {
|
|
tensors map[string]*mlx.Array
|
|
defaultGroupSize int
|
|
defaultBits int
|
|
defaultMode string
|
|
tensorQuant map[string]*TensorQuantInfo
|
|
}
|
|
|
|
// NewLinearFactory creates a reusable constructor for model linear layers.
|
|
func NewLinearFactory(
|
|
tensors map[string]*mlx.Array,
|
|
defaultGroupSize, defaultBits int,
|
|
defaultMode string,
|
|
tensorQuant map[string]*TensorQuantInfo,
|
|
) LinearFactory {
|
|
return LinearFactory{
|
|
tensors: tensors,
|
|
defaultGroupSize: defaultGroupSize,
|
|
defaultBits: defaultBits,
|
|
defaultMode: defaultMode,
|
|
tensorQuant: tensorQuant,
|
|
}
|
|
}
|
|
|
|
// Make constructs a linear layer at path.
|
|
func (f LinearFactory) Make(path string) nn.LinearLayer {
|
|
return MakeLinearLayer(
|
|
f.tensors,
|
|
path,
|
|
f.defaultGroupSize,
|
|
f.defaultBits,
|
|
f.defaultMode,
|
|
f.tensorQuant,
|
|
)
|
|
}
|
|
|
|
// MakeLinearLayer constructs a linear layer from a tensor map.
|
|
//
|
|
// For quantized tensors (path.weight + path.weight_scale), it resolves per-tensor
|
|
// quant params via TensorQuant metadata (with shape-based affine fallback).
|
|
// For non-quantized tensors, it returns a standard nn.Linear.
|
|
func MakeLinearLayer(
|
|
tensors map[string]*mlx.Array,
|
|
path string,
|
|
defaultGroupSize, defaultBits int,
|
|
defaultMode string,
|
|
tensorQuant map[string]*TensorQuantInfo,
|
|
) nn.LinearLayer {
|
|
w := tensors[path+".weight"]
|
|
if w == nil {
|
|
return nil
|
|
}
|
|
|
|
scales := tensors[path+".weight_scale"]
|
|
if scales != nil {
|
|
qbiases := tensors[path+".weight_qbias"]
|
|
bias := tensors[path+".bias"]
|
|
|
|
groupSize, bits, mode := ResolveLinearQuantParams(
|
|
defaultGroupSize,
|
|
defaultBits,
|
|
defaultMode,
|
|
tensorQuant,
|
|
path+".weight",
|
|
w,
|
|
scales,
|
|
)
|
|
|
|
// Check for per-tensor global scale (NVIDIA double-scale nvfp4).
|
|
// NVIDIA ModelOpt stores this as "weight_scale_2"; our import
|
|
// pipeline maps it to "weight.global_scale".
|
|
globalScale := tensors[path+".weight.global_scale"]
|
|
if globalScale == nil {
|
|
globalScale = tensors[path+".weight_scale_2"]
|
|
}
|
|
|
|
return &nn.QuantizedLinear{
|
|
Weight: w,
|
|
Scales: scales,
|
|
QBiases: qbiases,
|
|
Bias: bias,
|
|
GlobalScale: globalScale,
|
|
GroupSize: groupSize,
|
|
Bits: bits,
|
|
Mode: mode,
|
|
}
|
|
}
|
|
|
|
bias := tensors[path+".bias"]
|
|
return nn.NewLinear(w, bias)
|
|
}
|