Struct WgpuTrainer 

pub struct WgpuTrainer { /* private fields */ }

wgpu-accelerated training context (zero unsafe, safe Rust API)

Implementations

impl WgpuTrainer

pub fn new() -> Result<Self>

Create a new wgpu trainer. Requests a GPU device via Vulkan/Metal/DX12.

pub fn upload(&self, data: &[f32]) -> Buffer

Upload host data to GPU buffer

pub fn zeros(&self, len: usize) -> Buffer

Allocate zero-initialized GPU buffer

pub fn download(&self, buffer: &Buffer) -> Vec<f32>

Download GPU buffer to host

pub fn matmul_forward( &self, a: &Buffer, b: &Buffer, c: &Buffer, m: u32, k: u32, n: u32, )

Matrix multiply forward: C = A @ B using WGSL tiled GEMM

pub fn matmul_backward( &self, a: &Buffer, b: &Buffer, grad_c: &Buffer, grad_a: &Buffer, grad_b: &Buffer, m: u32, k: u32, n: u32, )

Matrix multiply backward: grad_a = grad_c @ B^T, grad_b = A^T @ grad_c

Uses the SAME tiled GEMM shader with transposed arguments. This is the standard GEMM backward formula:

• ∂L/∂A = ∂L/∂C @ B^T (shape: M×N @ N×K = M×K)
• ∂L/∂B = A^T @ ∂L/∂C (shape: K×M @ M×N = K×N)

pub fn adamw_step( &mut self, params: &Buffer, grads: &Buffer, m_state: &Buffer, v_state: &Buffer, lr: f32, beta1: f32, beta2: f32, eps: f32, weight_decay: f32, )

AdamW optimizer step on GPU

pub fn clip_gradients(&self, grads: &Buffer, max_norm: f32)

Gradient clipping (downloads to compute norm, clips on GPU)

pub fn step_count(&self) -> u32

Get current step count

pub fn reset_step(&mut self)

Reset step counter

pub fn queue_ref(&self) -> &Queue

Get a reference to the wgpu queue (for buffer writes)

pub fn device_ref(&self) -> &Device

Get a reference to the wgpu device (for buffer creation)

pub fn from_device(device: Device, queue: Queue) -> Result<Self>

Create from an existing device + queue (share device with WgslForwardPass). Contract: single device for all GPU operations — no cross-device buffer access.