lora

Documentation

Overview

Package lora implements LoRA and QLoRA fine-tuning adapters.

Stability: beta

Package lora provides Low-Rank Adaptation layers for parameter-efficient fine-tuning.

Index

• func InjectLoRA[T tensor.Numeric](m Model[T], rank int, alpha float32, targetModules []string, ...) error
• func LoadAdapter[T tensor.Numeric](path string, m Model[T]) error
• func SaveAdapter[T tensor.Numeric](path string, m Model[T]) error
• func TotalParamCount[T tensor.Numeric](m Model[T]) int
• func TrainableParamCount[T tensor.Numeric](m Model[T]) int
• type Layer
• type LinearInfo
• type LoraLinear
  • func NewLoraLinear[T tensor.Numeric](name string, base graph.Node[T], rank int, alpha float32, ...) (*LoraLinear[T], error)
  • func (l *LoraLinear[T]) Alpha() float32
  • func (l *LoraLinear[T]) Attributes() map[string]interface{}
  • func (l *LoraLinear[T]) Backward(ctx context.Context, mode types.BackwardMode, ...) ([]*tensor.TensorNumeric[T], error)
  • func (l *LoraLinear[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)
  • func (l *LoraLinear[T]) Name() string
  • func (l *LoraLinear[T]) OpType() string
  • func (l *LoraLinear[T]) OutputShape() []int
  • func (l *LoraLinear[T]) Parameters() []*graph.Parameter[T]
  • func (l *LoraLinear[T]) Rank() int
  • func (l *LoraLinear[T]) SetName(name string)
• type Model
• type Named
• type QLoRATrainer
  • func NewQLoRATrainer[T tensor.Numeric](model Model[T], rank int, alpha float32, targetModules []string, ...) (*QLoRATrainer[T], error)
  • func (q *QLoRATrainer[T]) Alpha() float32
  • func (q *QLoRATrainer[T]) Rank() int
  • func (q *QLoRATrainer[T]) Step(ctx context.Context, input *tensor.TensorNumeric[T], ...) (float32, error)
  • func (q *QLoRATrainer[T]) TrainableParams() []*graph.Parameter[T]

Functions

func InjectLoRA

func InjectLoRA[T tensor.Numeric](
	m Model[T],
	rank int,
	alpha float32,
	targetModules []string,
	engine compute.Engine[T],
) error

InjectLoRA walks the model, finds all Linear layers whose names match targetModules, and wraps them with LoraLinear. Base model parameters are frozen (only LoRA A and B matrices are trainable).

func LoadAdapter

func LoadAdapter[T tensor.Numeric](path string, m Model[T]) error

LoadAdapter reads a LoRA adapter GGUF checkpoint and restores the A and B matrices into matching LoraLinear layers in the model. Layer names are matched by stripping the "lora." prefix and ".weight_a"/".weight_b" suffix.

func SaveAdapter

func SaveAdapter[T tensor.Numeric](path string, m Model[T]) error

SaveAdapter writes all LoRA adapter matrices (A, B) from the given model to a GGUF file. Tensors are stored as float32 with naming convention lora.{layer}.weight_a and lora.{layer}.weight_b. Metadata includes the LoRA rank and alpha from the first adapter found.

func TotalParamCount

func TotalParamCount[T tensor.Numeric](m Model[T]) int

TotalParamCount returns the total number of parameters in the model.

func TrainableParamCount

func TrainableParamCount[T tensor.Numeric](m Model[T]) int

TrainableParamCount returns the count of trainable parameters (those from LoRA layers). After InjectLoRA, only LoRA A and B matrices are trainable.

Types

type Layer

type Layer[T tensor.Numeric] interface {
	graph.Node[T]
	Named
}

Layer is a named node with dimension information, implemented by Linear layers.

type LinearInfo

type LinearInfo interface {
	InputFeatures() int
	OutputFeatures() int
}

LinearInfo provides the dimensions of a Linear layer for LoRA wrapping.

type LoraLinear

type LoraLinear[T tensor.Numeric] struct {
	A *graph.Parameter[T] // r x d_in, initialized N(0,1)
	B *graph.Parameter[T] // d_out x r, initialized zero
	// contains filtered or unexported fields
}

LoraLinear wraps an existing linear layer with low-rank adaptation matrices. During forward: y = base(x) + (alpha/rank) * B @ A @ x Only A and B are trainable; the base layer weights are frozen.

func NewLoraLinear

func NewLoraLinear[T tensor.Numeric](
	name string,
	base graph.Node[T],
	rank int,
	alpha float32,
	engine compute.Engine[T],
	dIn, dOut int,
) (*LoraLinear[T], error)

NewLoraLinear creates a LoRA adapter around an existing base layer. rank is the low-rank dimension; alpha is the scaling factor. A is initialized with N(0,1) random values; B is initialized to zero.

func (*LoraLinear[T]) Alpha

func (l *LoraLinear[T]) Alpha() float32

Alpha returns the LoRA alpha scaling factor.

func (*LoraLinear[T]) Attributes

func (l *LoraLinear[T]) Attributes() map[string]interface{}

Attributes returns layer attributes.

func (*LoraLinear[T]) Backward

func (l *LoraLinear[T]) Backward(ctx context.Context, mode types.BackwardMode, outputGradient *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes gradients for A and B (not the base layer weights).

func (*LoraLinear[T]) Forward

func (l *LoraLinear[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward computes y = base(x) + (alpha/rank) * B @ A @ x.

func (*LoraLinear[T]) Name

func (l *LoraLinear[T]) Name() string

Name returns the name of the layer.

func (*LoraLinear[T]) OpType

func (l *LoraLinear[T]) OpType() string

OpType returns the operation type.

func (*LoraLinear[T]) OutputShape

func (l *LoraLinear[T]) OutputShape() []int

OutputShape returns the output shape of the base layer.

func (*LoraLinear[T]) Parameters

func (l *LoraLinear[T]) Parameters() []*graph.Parameter[T]

Parameters returns only the LoRA parameters A and B (not the base layer weights).

func (*LoraLinear[T]) Rank

func (l *LoraLinear[T]) Rank() int

Rank returns the LoRA rank.

func (*LoraLinear[T]) SetName

func (l *LoraLinear[T]) SetName(name string)

SetName sets the name of the layer.

type Model

type Model[T tensor.Numeric] interface {
	// Layers returns all named layers in the model.
	Layers() []Layer[T]

	// ReplaceLayer replaces the layer with the given name.
	// Returns an error if the name is not found.
	ReplaceLayer(name string, replacement Layer[T]) error
}

Model provides access to named layers for LoRA injection. Implementations must support replacing layers by name.

type Named

type Named interface {
	Name() string
}

Named is implemented by graph nodes that have a name.

type QLoRATrainer

type QLoRATrainer[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

QLoRATrainer wraps a model whose base Linear weights are NF4-quantized. Only LoRA A and B matrices are trained; base weights stay frozen. During forward, NF4 weights are dequantized to the working precision on the fly.

func NewQLoRATrainer

func NewQLoRATrainer[T tensor.Numeric](
	model Model[T],
	rank int,
	alpha float32,
	targetModules []string,
	engine compute.Engine[T],
	opt optimizer.Optimizer[T],
) (*QLoRATrainer[T], error)

NewQLoRATrainer creates a QLoRA trainer. It injects LoRA adapters into the model's target modules. Base weights are expected to already be NF4-quantized (via NF4Storage). Only the LoRA A and B matrices are trainable.

func (*QLoRATrainer[T]) Alpha

func (q *QLoRATrainer[T]) Alpha() float32

Alpha returns the LoRA alpha scaling factor.

func (*QLoRATrainer[T]) Rank

func (q *QLoRATrainer[T]) Rank() int

Rank returns the LoRA rank.

func (*QLoRATrainer[T]) Step

func (q *QLoRATrainer[T]) Step(
	ctx context.Context,
	input *tensor.TensorNumeric[T],
	target *tensor.TensorNumeric[T],
) (float32, error)

Step performs a single training step: forward pass through the model, compute loss (MSE against target), backward pass through LoRA layers, and optimizer step on LoRA parameters only.

Returns the scalar loss value.

func (*QLoRATrainer[T]) TrainableParams

func (q *QLoRATrainer[T]) TrainableParams() []*graph.Parameter[T]

TrainableParams returns the LoRA A and B parameters from all injected layers.