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.
72 lines
1.6 KiB
Go
72 lines
1.6 KiB
Go
package create
|
|
|
|
import (
|
|
"fmt"
|
|
"runtime"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
|
)
|
|
|
|
var (
|
|
mlxThreadOnce sync.Once
|
|
mlxThreadStarted atomic.Bool
|
|
mlxWork chan func()
|
|
mlxInitErr error
|
|
)
|
|
|
|
// runOnMLXThread runs f on the MLX thread and returns its error. The thread is
|
|
// started (and MLX initialized) on first use. A panic in f is recovered and
|
|
// returned as an error so a kernel failure cannot kill the pinned thread.
|
|
//
|
|
// TODO(pdevine): This method should be revisited when the `ollama create` is
|
|
// instead run on the ollama server process instead of the client.
|
|
func runOnMLXThread(f func() error) error {
|
|
mlxThreadOnce.Do(func() {
|
|
mlxWork = make(chan func())
|
|
ready := make(chan error)
|
|
go func() {
|
|
runtime.LockOSThread() // pinned for the process lifetime; never unlocked
|
|
err := mlx.CheckInit()
|
|
if err == nil && mlx.GPUIsAvailable() {
|
|
mlx.SetDefaultDeviceGPU()
|
|
}
|
|
ready <- err
|
|
if err != nil {
|
|
return
|
|
}
|
|
for work := range mlxWork {
|
|
work()
|
|
}
|
|
}()
|
|
mlxInitErr = <-ready
|
|
mlxThreadStarted.Store(mlxInitErr == nil)
|
|
})
|
|
if mlxInitErr != nil {
|
|
return fmt.Errorf("MLX init failed: %w", mlxInitErr)
|
|
}
|
|
|
|
done := make(chan error, 1)
|
|
mlxWork <- func() {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
done <- fmt.Errorf("mlx: %v", r)
|
|
}
|
|
}()
|
|
done <- f()
|
|
}
|
|
return <-done
|
|
}
|
|
|
|
// sweepMLX releases the MLX buffer cache. It is a no-op if no MLX work has run.
|
|
func sweepMLX() {
|
|
if !mlxThreadStarted.Load() {
|
|
return
|
|
}
|
|
_ = runOnMLXThread(func() error {
|
|
mlx.ClearCache()
|
|
mlx.Sweep()
|
|
return nil
|
|
})
|
|
}
|