Parameter-Efficient Fine-Tuning (PEFT) methods—LoRA, QLoRA, DoRA, and IA³—allow large language models to be adapted for specific tasks by training only a small fraction of the total parameters. Instead of fine-tuning an entire multi-billion-parameter model, PEFT injects small trainable matrices (adapters) into specific layers while keeping the base weights frozen. This reduces memory requirements by 10–100x and enables serving multiple task-specific variants from a single base model.
A Go PEFT library implements a thread-safe adapter manager that layers on top of GoMLX contexts. Adapters are organized under ctx.In("adapters"), where each adapter is a named collection of low-rank weight matrices that can be activated, deactivated, or swapped per request. The manager freezes base model weights, counts trainable parameters, and routes layer outputs through the currently active adapter—all without reloading the base model from disk.
In a production AI platform, the library enables serving multiple fine-tuned variants of embedding and classification models from a single base model instance. A search-optimized LoRA adapter can be active for document indexing requests while a classification-focused adapter handles categorization queries—both sharing the same base model weights in memory. Hot-swapping between adapters happens in microseconds because only the adapter weight pointers change, not the base model.
PEFT adapters work by decomposing weight updates into low-rank matrices. For a weight matrix W of shape (d, k), LoRA introduces two small matrices: A of shape (d, r) and B of shape (r, k), where r << min(d, k) is the rank. The effective weight becomes W + α/r · A×B, where α is a scaling factor. Only A and B are trained; W stays frozen.
The library organizes adapter weights using GoMLX’s hierarchical context system. Base model weights live at the root context, while adapter weights are scoped under ctx.In("adapters").In(adapterName). This scoping means the optimizer only sees adapter parameters during training, base weights are automatically excluded from gradient computation, and multiple adapters can coexist in the same context tree without name collisions. Activation simply redirects layer output through the scoped adapter path.
The library supports multiple PEFT methods, each with different trade-offs. LoRA (Low-Rank Adaptation) adds trainable low-rank matrices to attention layers. QLoRA combines LoRA with 4-bit quantized base weights, reducing memory by another 4x. DoRA (Weight-Decomposed Low-Rank Adaptation) separates weight magnitude and direction, applying LoRA only to the direction component for better training stability. IA³ (Infused Adapter by Inhibiting and Amplifying Inner Activations) multiplies activations by learned scaling vectors instead of adding low-rank matrices, requiring even fewer parameters than LoRA.
The adapter manager provides concurrent access to adapter operations. Multiple goroutines can read the active adapter simultaneously (for inference), while adapter swaps acquire a write lock. The manager uses sync.RWMutex to allow concurrent reads with exclusive writes, ensuring that no request sees a partially-swapped adapter state. Adapter activation is an atomic pointer swap—the old adapter remains valid until all in-flight requests complete.
Hot-swapping means changing the active adapter while the base model stays loaded in memory. Because adapters are small (typically 0.1–1% of base model parameters), loading a new adapter from disk takes milliseconds. The swap itself is an atomic pointer update: set the new adapter as active, and subsequent requests route through it. In-flight requests on the old adapter complete normally because the old adapter weights aren’t deallocated until their reference count drops to zero. This enables zero-downtime adapter updates in production.
When an adapter is attached, the manager marks all base model parameters as frozen by setting their Trainable flag to false in the GoMLX context. Only adapter parameters retain Trainable: true. The CountTrainableParams method walks the context tree, summing the element counts of all trainable tensors. For a 7B-parameter base model with a rank-16 LoRA adapter on attention layers, trainable parameters are typically around 4–8 million (0.05–0.1% of total).
In multi-tenant serving scenarios, different requests may need different adapters. The adapter manager supports per-request routing by accepting an adapter name in the inference context. Instead of globally activating one adapter, each request specifies which adapter to use. The manager resolves the adapter weights from the context tree and injects them into the forward pass. This enables a single model server to handle requests for dozens of fine-tuned variants simultaneously without any adapter being “globally active.”
Create an adapter manager that wraps a GoMLX context and manages adapter lifecycle:
package peft
import (
"fmt"
"sync"
"github.com/gomlx/gomlx/ml/context"
)
// AdapterType enumerates supported PEFT methods.
type AdapterType int
const (
AdapterLoRA AdapterType = iota
AdapterQLoRA
AdapterDoRA
AdapterIA3
)
// AdapterConfig defines the configuration for a PEFT adapter.
type AdapterConfig struct {
Name string
Type AdapterType
Rank int // Low-rank dimension (LoRA/QLoRA/DoRA)
Alpha float32 // Scaling factor: effective scale = alpha/rank
Dropout float32 // Dropout probability during training
TargetModules []string // Layer names to apply adapter to
}
// AdapterManager manages multiple PEFT adapters on a shared base model.
type AdapterManager struct {
mu sync.RWMutex
ctx *context.Context
adapters map[string]*Adapter
active string // currently active adapter name
}
// NewAdapterManager creates a manager rooted at the given GoMLX context.
func NewAdapterManager(ctx *context.Context) *AdapterManager {
return &AdapterManager{
ctx: ctx,
adapters: make(map[string]*Adapter),
}
}
// Adapter holds the state for a single PEFT adapter.
type Adapter struct {
Config AdapterConfig
Scope *context.Context // ctx.In("adapters").In(name)
TrainableCount int64
Loaded bool
}Register a new adapter, freeze base weights, and inject low-rank matrices into target layers:
// Register creates a new adapter with the given configuration.
// It freezes all base model weights and initializes adapter parameters.
func (m *AdapterManager) Register(cfg AdapterConfig) error {
m.mu.Lock()
defer m.mu.Unlock()
if _, exists := m.adapters[cfg.Name]; exists {
return fmt.Errorf("adapter %q already registered", cfg.Name)
}
// Scope adapter weights under ctx/adapters/{name}
adapterCtx := m.ctx.In("adapters").In(cfg.Name)
// Freeze all base model parameters
m.ctx.EnumerateVariables(func(v *context.Variable) {
if !v.InScope("adapters") {
v.SetTrainable(false)
}
})
// Initialize low-rank matrices for each target module
for _, module := range cfg.TargetModules {
moduleCtx := adapterCtx.In(module)
// Get the original weight shape from the base model
baseVar := m.ctx.In(module).Variable("weights")
d, k := baseVar.Shape().Dimensions[0], baseVar.Shape().Dimensions[1]
// A matrix: (d, rank) initialized with Kaiming uniform
moduleCtx.VariableWithShape("lora_a",
context.Shape{d, cfg.Rank}).
SetInitializer(context.KaimingUniformInit)
// B matrix: (rank, k) initialized with zeros
moduleCtx.VariableWithShape("lora_b",
context.Shape{cfg.Rank, k}).
SetInitializer(context.ZerosInit)
}
adapter := &Adapter{
Config: cfg,
Scope: adapterCtx,
Loaded: true,
}
adapter.TrainableCount = m.countTrainable(adapterCtx)
m.adapters[cfg.Name] = adapter
return nil
}
// Activate sets the named adapter as the active one for inference.
func (m *AdapterManager) Activate(name string) error {
m.mu.Lock()
defer m.mu.Unlock()
adapter, exists := m.adapters[name]
if !exists {
return fmt.Errorf("adapter %q not found", name)
}
if !adapter.Loaded {
return fmt.Errorf("adapter %q not loaded", name)
}
m.active = name
return nil
}Route layer outputs through the active adapter during inference:
// Forward computes the adapter-augmented output for a linear layer.
// base is the original layer output (x @ W), adapterName selects which
// adapter to apply. If adapterName is empty, uses the globally active one.
func (m *AdapterManager) Forward(base *Node, layerName string,
input *Node, adapterName string) *Node {
m.mu.RLock()
name := adapterName
if name == "" {
name = m.active
}
adapter, exists := m.adapters[name]
m.mu.RUnlock()
if !exists || !adapter.Loaded {
return base // No adapter: return base output unchanged
}
// Retrieve low-rank matrices from the adapter scope
moduleCtx := adapter.Scope.In(layerName)
loraA := moduleCtx.Variable("lora_a").Value()
loraB := moduleCtx.Variable("lora_b").Value()
// Compute adapter contribution: scale * (input @ A @ B)
scale := adapter.Config.Alpha / float32(adapter.Config.Rank)
delta := MatMul(MatMul(input, loraA), loraB)
delta = MulScalar(delta, scale)
// Apply dropout during training
if adapter.Config.Dropout > 0 && IsTraining() {
delta = Dropout(delta, adapter.Config.Dropout)
}
// Add adapter contribution to base output
return Add(base, delta)
}
// countTrainable sums trainable parameters under a context scope.
func (m *AdapterManager) countTrainable(scope *context.Context) int64 {
var count int64
scope.EnumerateVariables(func(v *context.Variable) {
if v.Trainable() {
count += v.Shape().Size()
}
})
return count
}The LoRA scaling factor α/r controls how much the adapter modifies the base model’s behavior. A higher ratio amplifies the adapter’s effect; a lower ratio keeps outputs closer to the base model. The common practice is to set α = 2r (so the effective scale is 2.0) for initial experiments, then tune based on validation loss. During inference, the scaling is applied as a simple scalar multiply on the adapter’s output—it adds zero computational overhead beyond the low-rank matrix multiplications.
A single embedding model server hosts the base model plus tenant-specific LoRA adapters. Tenant A’s adapter is fine-tuned on legal documents for improved retrieval in their domain, while Tenant B’s adapter is tuned on medical literature. Both share the same base model weights in GPU/CPU memory. The adapter manager routes each embedding request to the correct adapter based on the tenant ID in the request context, providing customized embedding quality without the cost of hosting separate model instances per tenant.
During model evaluation, multiple adapter variants can be loaded simultaneously. An A/B testing framework randomly assigns incoming requests to different adapters and logs quality metrics per adapter. Because adapter loading is cheap (milliseconds) and switching is atomic, new variants can be deployed and tested in production without downtime or traffic disruption.
In a typical search pipeline, adapters enable incremental specialization of embedding models. The base model provides general-purpose embeddings, while task-specific adapters improve quality for particular query types. A code-search adapter is trained on programming Q&A pairs, a documentation-search adapter on technical writing, and a conversation-search adapter on chat logs. The system selects the appropriate adapter based on query classification, routing code questions through the code adapter and documentation queries through the docs adapter.
The adapter manager integrates with GoMLX’s training loop. When training an adapter, the manager ensures that only adapter parameters receive gradients, the base model’s memory footprint stays constant (no gradient buffers allocated for frozen weights), and checkpoints save only adapter weights (a few MB instead of GB). This makes fine-tuning feasible on machines with limited memory that couldn’t hold full gradient state for the base model.
Save and load adapter weights independently from the base model:
// SaveAdapter serializes only the adapter weights to disk.
func (m *AdapterManager) SaveAdapter(name, path string) error {
m.mu.RLock()
adapter, exists := m.adapters[name]
m.mu.RUnlock()
if !exists {
return fmt.Errorf("adapter %q not found", name)
}
// Collect adapter variables
vars := make(map[string][]float32)
adapter.Scope.EnumerateVariables(func(v *context.Variable) {
vars[v.ScopedName()] = v.Value().Float32s()
})
// Serialize to safetensors format
return safetensors.Save(path, vars)
}
// LoadAdapter loads adapter weights from disk into the context.
func (m *AdapterManager) LoadAdapter(name, path string,
cfg AdapterConfig) error {
// Register the adapter structure first
if err := m.Register(cfg); err != nil {
return err
}
// Load weights from safetensors
vars, err := safetensors.Load(path)
if err != nil {
return fmt.Errorf("loading adapter weights: %w", err)
}
m.mu.Lock()
defer m.mu.Unlock()
adapter := m.adapters[name]
adapter.Scope.EnumerateVariables(func(v *context.Variable) {
if data, ok := vars[v.ScopedName()]; ok {
v.SetValue(data)
}
})
adapter.Loaded = true
return nil
}DoRA decomposes weight updates into magnitude and direction components for better training dynamics:
// DoRAForward applies Weight-Decomposed Low-Rank Adaptation.
// It separates the weight update into magnitude (scalar) and
// direction (unit vector) components.
func (m *AdapterManager) DoRAForward(base *Node, layerName string,
input *Node, adapterName string) *Node {
adapter := m.getAdapter(adapterName)
moduleCtx := adapter.Scope.In(layerName)
// Retrieve DoRA-specific parameters
loraA := moduleCtx.Variable("lora_a").Value()
loraB := moduleCtx.Variable("lora_b").Value()
magnitude := moduleCtx.Variable("magnitude").Value() // learned scalar
// Compute adapted weight direction
baseWeight := m.ctx.In(layerName).Variable("weights").Value()
scale := adapter.Config.Alpha / float32(adapter.Config.Rank)
adaptedWeight := Add(baseWeight, MulScalar(MatMul(loraA, loraB), scale))
// Normalize to unit direction, then scale by learned magnitude
direction := Normalize(adaptedWeight, -1) // L2 norm along last dim
weightFinal := Mul(magnitude, direction)
return MatMul(input, weightFinal)
}IA³ (Infused Adapter by Inhibiting and Amplifying Inner Activations) uses learned scaling vectors instead of low-rank matrices. Each target layer gets a single vector of the same dimension as its output, which element-wise multiplies the layer’s activations. This requires orders of magnitude fewer parameters than LoRA (one vector vs. two matrices) and adds negligible inference latency. IA³ is ideal when you need many adapters and memory is the primary constraint, trading some adaptation quality for extreme parameter efficiency.
CountTrainableParams returns a value much smaller than total parameters after adapter registration.α/r must be applied. Merging without scaling produces a model that behaves as if the adapter contribution was divided by alpha—subtle degradation that’s easy to miss in quick evaluations.