GPU-Accelerated Adam Optimizer

Implements the Adam optimizer (Kingma & Ba, 2014) as a GPU compute kernel for efficient parameter updates on LoRA weights.

m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
v_t = β₂ * v_{t-1} + (1 - β₂) * g_t²
m̂_t = m_t / (1 - β₁^t)
v̂_t = v_t / (1 - β₂^t)
θ_t = θ_{t-1} - lr * m̂_t / (√v̂_t + ε)

All updates happen in-place on GPU buffers (param, m, v, grad).

Reference

CPU implementation: Hesper/Optimizer/Adam.lean

structure Hesper.Optimizer.AdamGPU.Config : Type

AdamW hyperparameters (matches PyTorch defaults)

• lr : Float
• beta1 : Float
• beta2 : Float
• eps : Float
• weightDecay : Float

def Hesper.Optimizer.AdamGPU.instReprConfig.repr : Config → Nat → Std.Format

Equations

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

instance Hesper.Optimizer.AdamGPU.instReprConfig : Repr Config

Equations

• Hesper.Optimizer.AdamGPU.instReprConfig = { reprPrec := Hesper.Optimizer.AdamGPU.instReprConfig.repr }

def Hesper.Optimizer.AdamGPU.adamUpdateKernel (numElements : Nat) (lr beta1 beta2 eps weightDecay biasCorrection1 biasCorrection2 : Float) : WGSL.Monad.ShaderM Unit

GPU kernel: Adam parameter update.

For each element i: m[i] = beta1 * m[i] + (1 - beta1) * grad[i] v[i] = beta2 * v[i] + (1 - beta2) * grad[i]^2 m_hat = m[i] / (1 - beta1^step) v_hat = v[i] / (1 - beta2^step) param[i] -= lr * m_hat / (sqrt(v_hat) + eps) grad[i] = 0 (zero gradient for next step)

Buffers: param, grad, m, v (all read-write, [numElements] FP32)

Equations

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

def Hesper.Optimizer.AdamGPU.executeAdamUpdate (device : WebGPU.Device) (paramBuf gradBuf mBuf vBuf : WebGPU.Buffer) (numElements : Nat) (config : Config) (step : Nat) : IO Unit

Execute Adam update on a single parameter buffer

Equations

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

def Hesper.Optimizer.AdamGPU.updateLoRAAdapter (device : WebGPU.Device) (adapter : LoRA.Adapter) (grads : LoRA.AdapterGrad) (adamState : LoRA.AdapterAdamState) (config : Config) : IO LoRA.AdapterAdamState

Execute Adam update on all LoRA parameters in the adapter

Equations

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