Patrick Devine
964ea42c09
This is a rewrite of the create functionality for the MLX engine. The core idea behind the create functionality is to break the import/convert into a pipeline of distinct phases: * Read (scan the safetensors directory for the various bits of metadata) * Classify (determine what the import type) * Plan (determine any transforms that need to be done) * Write (transform any data as necessary and write out the blobs) * Create the manifest Each architecture has a "policy" which determines how to convert the model correctly. A number of different formats for safetensors are supported including: * nvfp4 (two formats: model optimized, torch) * fp8 datatypes (convert to mxfp8) * standard bf16 based weights A number of cleanups/simplifications have been done including: * using the baked in names for the tensors instead of munging them into something else * unified 3d expert tensors (instead of separate per expert tensors) * fewer unnecessary transforms to the various tensors in a model (keep a model as close to the source as possible) * unified capability checking * draft model handling (for MTP) is done on the same path Image generation has been intentionally removed.
96 lines
3.0 KiB
Go
96 lines
3.0 KiB
Go
package create
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
)
|
|
|
|
type gemma4ImportTransform struct {
|
|
numLayers int
|
|
numExperts int
|
|
}
|
|
|
|
// gemma4Config is a minimal subset of the Gemma 4 config.json used for quant decisions.
|
|
type gemma4Config struct {
|
|
NumHiddenLayers int `json:"num_hidden_layers"`
|
|
NumExperts int `json:"num_experts"`
|
|
TextConfig struct {
|
|
NumHiddenLayers int `json:"num_hidden_layers"`
|
|
NumExperts int `json:"num_experts"`
|
|
} `json:"text_config"`
|
|
}
|
|
|
|
func newGemma4ImportTransform(rawConfig json.RawMessage) (quantizePolicy, error) {
|
|
var cfg gemma4Config
|
|
if err := json.Unmarshal(rawConfig, &cfg); err != nil {
|
|
return nil, fmt.Errorf("gemma4: parse config.json: %w", err)
|
|
}
|
|
|
|
numLayers := cfg.NumHiddenLayers
|
|
if numLayers == 0 {
|
|
numLayers = cfg.TextConfig.NumHiddenLayers
|
|
}
|
|
numExperts := cfg.NumExperts
|
|
if numExperts == 0 {
|
|
numExperts = cfg.TextConfig.NumExperts
|
|
}
|
|
|
|
return gemma4ImportTransform{numLayers: numLayers, numExperts: numExperts}, nil
|
|
}
|
|
|
|
func (t gemma4ImportTransform) quantizationType(name string, shape []int32, quantize string) string {
|
|
base := normalizeQuantType(quantize)
|
|
switch {
|
|
case isEmbedTokensWeight(name):
|
|
// The embedding doubles as the lm_head projection; an 8-bit type keeps
|
|
// quality close to bf16 (matching GGUF Q6_K) while saving bandwidth.
|
|
// Fall back to the base type when 8-bit does not fit the vocab shape.
|
|
if e := promoteEmbedding(shape, base); e != "" {
|
|
return e
|
|
}
|
|
if isAligned(shape, base) {
|
|
return base
|
|
}
|
|
return ""
|
|
case t.isSensitiveProjection(name) && eightBit(base) != base:
|
|
return sensitiveType(t.promoteSensitive(name), shape, base)
|
|
default:
|
|
// Routing gates, norms, embeddings, etc. are handled by the generic
|
|
// policy; everything else quantizes at the requested type.
|
|
return GetTensorQuantization(name, shape, quantize)
|
|
}
|
|
}
|
|
|
|
// isSensitiveProjection reports the value/key/down projections whose precision
|
|
// most affects quality — attention output (v/k) and the residual stream
|
|
// (down). Audio and vision tower tensors are excluded and follow the generic
|
|
// policy.
|
|
func (t gemma4ImportTransform) isSensitiveProjection(name string) bool {
|
|
if isVisionTower(name) || isAudioTower(name) {
|
|
return false
|
|
}
|
|
return strings.Contains(name, ".v_proj") ||
|
|
strings.Contains(name, ".k_proj") ||
|
|
strings.Contains(name, "down_proj")
|
|
}
|
|
|
|
// promoteSensitive decides whether a sensitive projection uses 8-bit precision.
|
|
// 8-expert models share very few KV heads, so their k/v projections are always
|
|
// promoted; otherwise v/down projections are promoted at the input and output
|
|
// layers and periodically between (useMoreBits), where residual-stream error
|
|
// accumulates most.
|
|
func (t gemma4ImportTransform) promoteSensitive(name string) bool {
|
|
if t.numLayers == 0 {
|
|
return false
|
|
}
|
|
if t.numExperts == 8 && (strings.Contains(name, ".v_proj") || strings.Contains(name, ".k_proj")) {
|
|
return true
|
|
}
|
|
if strings.Contains(name, ".k_proj") {
|
|
return false // k_proj is promoted only via the 8-expert path
|
|
}
|
|
layer := layerIndex(name)
|
|
return layer >= 0 && useMoreBits(layer, t.numLayers)
|
|
}
|