Gradient Clipping and Scaling

GPU kernels for:

1. Global gradient norm — L2 norm across all LoRA parameter gradients
2. Gradient clipping — scale gradients if norm exceeds threshold
3. Gradient scaling — divide gradients by token count (loss normalization)

Standard Values (matches PyTorch defaults)

• max_grad_norm = 1.0
• Loss normalization: divide gradients by number of output tokens

structure Hesper.Optimizer.GradientClip.ClipBuffers : Type

Buffers needed for gradient clipping

• normSqBuf : WebGPU.Buffer

  Accumulator for sum of squared gradients [1]

• partialBuf : WebGPU.Buffer

  Temporary for per-buffer partial sums [1]

def Hesper.Optimizer.GradientClip.createClipBuffers (device : WebGPU.Device) : IO ClipBuffers

Create clip buffers

Equations

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

Sum of Squares Kernel

def Hesper.Optimizer.GradientClip.sumSquaredKernel (numElements : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Compute sum of squares of a buffer, write result to accumulator (ADD to existing value). Uses single workgroup with shared memory reduction.

Equations

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

def Hesper.Optimizer.GradientClip.executeSumSquared (device : WebGPU.Device) (gradBuf accumBuf : WebGPU.Buffer) (numElements : Nat) : IO Unit

Execute sum of squares and add to accumulator

Equations

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

Gradient Clip Kernel

def Hesper.Optimizer.GradientClip.clipKernel (numElements : Nat) (maxNorm : Float) : WGSL.Monad.ShaderM Unit

Scale gradient buffer by clip_factor = maxNorm / globalNorm (if norm > maxNorm). Reads globalNormSq[0], computes norm = sqrt(normSq), clips if needed.

Equations

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

def Hesper.Optimizer.GradientClip.executeClip (device : WebGPU.Device) (gradBuf normSqBuf : WebGPU.Buffer) (numElements : Nat) (maxNorm : Float) : IO Unit

Execute gradient clipping on a single buffer

Equations

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

In-place Gradient Scale Kernel

def Hesper.Optimizer.GradientClip.scaleKernel (numElements : Nat) (scaleFactor : Float) : WGSL.Monad.ShaderM Unit

Scale gradient in-place: grad[i] *= scaleFactor

Equations

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

def Hesper.Optimizer.GradientClip.executeScale (device : WebGPU.Device) (gradBuf : WebGPU.Buffer) (numElements : Nat) (scaleFactor : Float) : IO Unit

Equations

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

High-Level API

def Hesper.Optimizer.GradientClip.clipGradNorm (device : WebGPU.Device) (adapter : LoRA.Adapter) (grads : LoRA.AdapterGrad) (maxNorm : Float) (clipBufs : ClipBuffers) : IO Float

Clip gradients of all LoRA parameters to maxNorm (global L2 norm). Returns the gradient norm before clipping (for logging).

Equations

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

def Hesper.Optimizer.GradientClip.scaleGrads (device : WebGPU.Device) (adapter : LoRA.Adapter) (grads : LoRA.AdapterGrad) (scaleFactor : Float) : IO Unit

Scale all gradients by a factor (e.g., 1/numTokens for loss normalization)

Equations

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