2d870385f9
Add LoRA field to Linear for transparent adapter injection via model's Forward() path. ApplyLoRA() on Qwen3/Gemma3 wraps target projections. Deterministic param ordering for adapter save/load consistency. MaskedCrossEntropyLoss for training on assistant tokens only. Co-Authored-By: Virgil <virgil@lethean.io>
78 lines
2.1 KiB
Go
78 lines
2.1 KiB
Go
//go:build darwin && arm64
|
|
|
|
// Package model provides transformer model architectures for MLX inference.
|
|
package model
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
|
|
"forge.lthn.ai/core/go-ai/mlx"
|
|
"forge.lthn.ai/core/go-ai/mlx/cache"
|
|
"forge.lthn.ai/core/go-ai/mlx/tokenizer"
|
|
)
|
|
|
|
// Model is the common interface for all transformer model architectures.
|
|
type Model interface {
|
|
// Forward runs the model forward pass on token IDs with KV caches.
|
|
Forward(tokens *mlx.Array, caches []cache.Cache) *mlx.Array
|
|
|
|
// NewCache creates per-layer KV caches for generation.
|
|
NewCache() []cache.Cache
|
|
|
|
// NumLayers returns the number of transformer layers.
|
|
NumLayers() int
|
|
|
|
// Tokenizer returns the model's tokenizer.
|
|
Tokenizer() *tokenizer.Tokenizer
|
|
|
|
// ModelType returns the architecture identifier (e.g. "gemma3", "qwen3").
|
|
ModelType() string
|
|
|
|
// ApplyLoRA wraps target projection layers with LoRA adapters for training.
|
|
// Returns the adapter which holds references to all LoRA layers.
|
|
ApplyLoRA(cfg mlx.LoRAConfig) *mlx.LoRAAdapter
|
|
}
|
|
|
|
// QuantizationConfig holds quantization parameters from config.json.
|
|
type QuantizationConfig struct {
|
|
GroupSize int `json:"group_size"`
|
|
Bits int `json:"bits"`
|
|
}
|
|
|
|
// resolveWeight looks up a weight with optional "language_model." prefix.
|
|
func resolveWeight(weights map[string]*mlx.Array, name string) *mlx.Array {
|
|
if w, ok := weights[name]; ok {
|
|
return w
|
|
}
|
|
if w, ok := weights["language_model."+name]; ok {
|
|
return w
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// LoadModel auto-detects the model architecture from config.json and loads it.
|
|
func LoadModel(modelPath string) (Model, error) {
|
|
data, err := os.ReadFile(filepath.Join(modelPath, "config.json"))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("model: load config: %w", err)
|
|
}
|
|
|
|
var probe struct {
|
|
ModelType string `json:"model_type"`
|
|
}
|
|
if err := json.Unmarshal(data, &probe); err != nil {
|
|
return nil, fmt.Errorf("model: parse model_type: %w", err)
|
|
}
|
|
|
|
switch probe.ModelType {
|
|
case "qwen3":
|
|
return LoadQwen3(modelPath)
|
|
case "gemma3", "gemma2":
|
|
return LoadGemma3(modelPath)
|
|
default:
|
|
return nil, fmt.Errorf("model: unsupported architecture %q", probe.ModelType)
|
|
}
|
|
}
|