Implementation:Ollama Ollama MLXRunner Model Base

Knowledge Sources	Ollama
Domains	MLX Runtime, Model Architecture
Last Updated	2025-02-15 00:00 GMT

Overview

Defines the model registry and Model interface for the MLX runner, enabling architecture-specific model implementations to register and be instantiated by name.

Description

The Model interface requires Forward, Unembed, NumLayers, Tokenizer, and LoadWeights methods. Register stores model constructor functions by architecture name (called from init() in model packages). New reads config.json from the manifest, extracts the architecture string, and dispatches to the registered constructor. Weights returns a model's LoadWeights method for the runner to call after loading tensors.

Usage

Central registry that decouples model loading from model architecture. New architectures register themselves via init().

Code Reference

Source Location

• Repository: Ollama
• File: x/mlxrunner/model/base/base.go
• Lines: 1-85

Signature

type Model interface {
    Forward(inputs *mlx.Array, cache []cache.Cache) *mlx.Array
    Unembed(x *mlx.Array) *mlx.Array
    NumLayers() int
    Tokenizer() *tokenizer.Tokenizer
    LoadWeights(tensors map[string]*mlx.Array) error
}

func Register(arch string, fn func(root *model.Root) (Model, error))
func New(root *model.Root) (Model, error)
func Weights(m Model) func(map[string]*mlx.Array) error

Import

import "github.com/ollama/ollama/x/mlxrunner/model/base"

I/O Contract

Inputs

Name	Type	Required	Description
arch	string	Yes	Architecture name (e.g. "Glm4MoeLiteForCausalLM")
fn	func(*model.Root) (Model, error)	Yes	Constructor function for the model

Outputs

Name	Type	Description
model	Model	Instantiated model with config parsed
error	error	Non-nil if architecture is unsupported or config is invalid

Usage Examples

// In model package init()
func init() {
    base.Register("Glm4MoeLiteForCausalLM", func(root *model.Root) (base.Model, error) {
        // Parse config, create model struct
        return &Model{...}, nil
    })
}

// In runner
model, err := base.New(root)
loadWeights := base.Weights(model)
err = loadWeights(tensors)

Related Pages

• Principle:Ollama_Ollama_MLXRunner_Architecture