mirror of
https://github.com/ollama/ollama.git
synced 2026-07-07 00:08:13 +00:00
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.
83 lines
3 KiB
Go
83 lines
3 KiB
Go
package create
|
|
|
|
import "fmt"
|
|
|
|
// CreateDraftLayers imports a draft (speculative-decoding / MTP assistant)
|
|
// safetensors model into prefixed tensor and config blobs and returns the
|
|
// layers WITHOUT writing a manifest — the caller folds them into the target
|
|
// model's manifest. A draft never stands alone; it always accompanies a target
|
|
// model named on the Modelfile's FROM line.
|
|
//
|
|
// It runs the same read → classify → plan → write pipeline as Create. Output
|
|
// tensor names keep their source form, namespaced by tensorPrefix (e.g.
|
|
// "draft.") so they cannot collide with the target's tensors; config blobs are
|
|
// named under configPrefix (e.g. "draft/").
|
|
func CreateDraftLayers(modelDir, tensorPrefix, configPrefix, quantize string, store BlobStore, fn func(status string)) ([]LayerInfo, error) {
|
|
if tensorPrefix == "" {
|
|
return nil, fmt.Errorf("draft tensor prefix must not be empty")
|
|
}
|
|
if configPrefix == "" {
|
|
return nil, fmt.Errorf("draft config prefix must not be empty")
|
|
}
|
|
defer sweepMLX()
|
|
|
|
inv, err := ReadInventory(modelDir)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read draft model: %w", err)
|
|
}
|
|
class, err := Classify(inv, quantize)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
policy, err := newTensorImportTransform(inv)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("build draft quantization policy for %q: %w", inv.Config.Architecture(), err)
|
|
}
|
|
specs, err := Plan(inv, class, draftPolicy{policy})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("plan draft model: %w", err)
|
|
}
|
|
specs = prefixSpecs(specs, tensorPrefix)
|
|
|
|
fn(fmt.Sprintf("importing draft (%d tensors%s)", len(inv.Tensors), quantizeStatus(class)))
|
|
layers, err := WriteBlobs(specs, modelDir, store)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
configLayers, _, err := importConfigBlobs(modelDir, configPrefix, store, fn)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return append(layers, configLayers...), nil
|
|
}
|
|
|
|
// prefixSpecs returns specs with prefix prepended to every output blob name and
|
|
// output tensor name, leaving the source references (which point at the source
|
|
// files) untouched. Scale/bias keys derive from the tensor name, so they inherit
|
|
// the prefix automatically.
|
|
func prefixSpecs(specs []BlobSpec, prefix string) []BlobSpec {
|
|
out := make([]BlobSpec, len(specs))
|
|
for i, spec := range specs {
|
|
tensors := make([]TensorSpec, len(spec.Tensors))
|
|
for j, ts := range spec.Tensors {
|
|
ts.Name = prefix + ts.Name
|
|
tensors[j] = ts
|
|
}
|
|
out[i] = BlobSpec{Name: prefix + spec.Name, Tensors: tensors, Metadata: spec.Metadata}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// draftPolicy wraps an architecture policy to keep a draft model's token
|
|
// embedding at source precision (drafts start with unquantized embeddings; this
|
|
// may change later). Every other tensor follows the wrapped policy. It is given
|
|
// unprefixed source names, since planning runs before prefixSpecs.
|
|
type draftPolicy struct{ inner quantizePolicy }
|
|
|
|
func (p draftPolicy) quantizationType(name string, shape []int32, requested string) string {
|
|
if isEmbedTokensWeight(name) {
|
|
return ""
|
|
}
|
|
return p.inner.quantizationType(name, shape, requested)
|
|
}
|