LoRA Training Loop

Teacher-forcing training loop for Alpaca-style instruction finetuning of BitNet models with LoRA adapters.

Training Algorithm

For each example:

1. Tokenize: instruction + input → prompt tokens, output → target tokens
2. For each position t in the sequence: a. Forward: run model with LoRA to get logits b. If t >= promptLen: compute cross-entropy loss on target token c. Backward: compute LoRA gradients (dA, dB) from loss
3. Adam update on all LoRA parameters

Simplification (v1)

The backward pass only computes gradients for the LoRA parameters. The gradient signal flows through the residual stream, and LoRA gradients are computed using saved activations from the forward pass. This is standard practice in LoRA finetuning.

structure Hesper.Training.TrainLoop.TrainState : Type

Training state maintained across steps

• adapter : LoRA.Adapter

  LoRA adapter weights

• grads : LoRA.AdapterGrad

  Gradient accumulators

• adamState : LoRA.AdapterAdamState

  Adam optimizer state

• savedActs : LoRA.SavedActivations

  Saved activations for backward

• dhBuf : WebGPU.Buffer

  Temporary buffers

• dInputBuf : WebGPU.Buffer
• hBuf : WebGPU.Buffer
• yBufQ : WebGPU.Buffer
• yBufV : WebGPU.Buffer
• totalLoss : Float

  Loss tracking

• numTokens : Nat

def Hesper.Training.TrainLoop.createTrainState (device : WebGPU.Device) (adapter : LoRA.Adapter) (dim kvDim : Nat) : IO TrainState

Create training state with all necessary buffers

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.TrainLoop.zeroGrads (device : WebGPU.Device) (adapter : LoRA.Adapter) (grads : LoRA.AdapterGrad) : IO Unit

Zero all gradient buffers (call before each training step)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.TrainLoop.applyLoRAForward (device : WebGPU.Device) (layerAdapter : LoRA.LayerAdapter) (scale : Float) (inputBuf qBuf vBuf : WebGPU.Buffer) (state : TrainState) (layerIdx : Nat) : IO Unit

Apply LoRA forward pass for a single attention layer. Called after BitLinear.forward has already written the base output to qBuf/vBuf. This adds the LoRA contribution: qBuf += scale * B_Q @ (A_Q @ inputBuf)

@param device GPU device @param layerAdapter LoRA weights for this layer @param scale alpha/rank scaling factor @param inputBuf Input to attention (after RMSNorm) [dim] @param qBuf Q projection output buffer [dim] (already has base output) @param vBuf V projection output buffer [kvDim] (already has base output) @param state Training state (for temp buffers and activation saving) @param layerIdx Layer index for saving activations

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.TrainLoop.applyLoRABackward (device : WebGPU.Device) (layerAdapter : LoRA.LayerAdapter) (layerGrad : LoRA.LayerAdapterGrad) (scale : Float) (dQBuf dVBuf : WebGPU.Buffer) (state : TrainState) (layerIdx : Nat) : IO Unit

Apply LoRA backward pass for a single attention layer. Computes dA, dB for Q and V projections using saved activations.

@param device GPU device @param layerAdapter LoRA weights for this layer @param layerGrad Gradient accumulators for this layer @param scale alpha/rank scaling factor @param dQBuf Gradient w.r.t. Q output [dim] @param dVBuf Gradient w.r.t. V output [kvDim] @param state Training state (temp buffers, saved activations) @param layerIdx Layer index

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.TrainLoop.optimizerStep (device : WebGPU.Device) (state : TrainState) (config : Optimizer.AdamGPU.Config) : IO TrainState

Run a single training step on one tokenized example.

This is the main entry point for training. It:

1. Runs forward pass token-by-token with LoRA
2. Computes cross-entropy loss on output tokens
3. Runs backward pass to accumulate LoRA gradients
4. Runs Adam optimizer to update LoRA weights

Note: This function is designed to be called with the model's existing forward infrastructure. The caller is responsible for orchestrating the per-token forward pass with the model and calling applyLoRAForward at each attention layer.

@param device GPU device @param state Training state @param losses Array of per-token losses (populated during forward) @param config Optimizer config @return Updated training state

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.TrainLoop.zeroBuffer (device : WebGPU.Device) (buf : WebGPU.Buffer) (numElements : Nat) : IO Unit

Zero a GPU buffer via GPU kernel (safe to use inside batch)

Equations

• Hesper.Training.TrainLoop.zeroBuffer device buf numElements = Hesper.Optimizer.GradientClip.executeScale device buf numElements 0.0

def Hesper.Training.TrainLoop.readLoss (device : WebGPU.Device) (lossBuf : WebGPU.Buffer) : IO Float

Read loss value from GPU buffer (safe, returns 0.0 on failure)

Equations

• Hesper.Training.TrainLoop.readLoss device lossBuf = Hesper.Training.SafeBuffer.safeReadF32 device lossBuf

def Hesper.Training.TrainLoop.printProgress (epoch step : Nat) (loss : Float) (numTokens : Nat) : IO Unit

Print training progress

Equations

• One or more equations did not get rendered due to their size.