Struct KernelFusedAdam 

pub struct KernelFusedAdam { /* private fields */ }

Kernel fusion optimized Adam optimizer.

Implementations

impl KernelFusedAdam

pub fn new(lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32) -> Self

Creates a new kernel fused Adam optimizer.

pub fn with_config( lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32, config: KernelFusionConfig, ) -> Self

Creates optimizer with specific GPU configuration.

pub fn for_a100(lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32) -> Self

Creates A100-optimized variant.

pub fn for_h100(lr: f32, betas: (f32, f32), eps: f32, weight_decay: f32) -> Self

Creates H100-optimized variant.

pub fn update_fused( &mut self, params: Vec<(&str, &mut [f32], &[f32])>, ) -> Result<()>

Updates multiple parameters using fused kernels.

pub fn gpu_stats(&self) -> GPUMemoryStats

Gets GPU performance statistics.

Trait Implementations

impl Debug for KernelFusedAdam

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Optimizer for KernelFusedAdam

fn update(&mut self, parameter: &mut Tensor, grad: &Tensor) -> Result<()>

Updates a parameter based on its gradient.

fn zero_grad(&mut self)

Clears all accumulated gradients.

fn step(&mut self)

Performs a single optimization step.

fn get_lr(&self) -> f32

Gets the current learning rate.

fn set_lr(&mut self, lr: f32)

Sets a new learning rate.

fn accumulate_grad( &mut self, parameter: &mut Tensor, grad: &Tensor, ) -> Result<(), TrustformersError>

Accumulates gradients for gradient accumulation.

fn apply_accumulated_grads( &mut self, accumulation_steps: usize, ) -> Result<(), TrustformersError>

Applies accumulated gradients after gradient accumulation.