319 lines
9.6 KiB
Rust
319 lines
9.6 KiB
Rust
//! GPU-Native Adam Optimizer
|
|
//!
|
|
//! This module provides an Adam optimizer that keeps ALL state on GPU,
|
|
//! eliminating CPU-GPU synchronization during training steps.
|
|
//!
|
|
//! ## Design Goals
|
|
//!
|
|
//! 1. **Zero CPU Sync**: All optimizer state (m, v, parameters) stays on GPU
|
|
//! 2. **Fused Operations**: Minimize kernel launches per step
|
|
//! 3. **In-Place Updates**: Modify weights in-place for graph capture compatibility
|
|
//!
|
|
//! ## Adam Update Rule
|
|
//!
|
|
//! For each parameter θ with gradient g:
|
|
//!
|
|
//! ```text
|
|
//! m = β₁ * m + (1 - β₁) * g // Update first moment
|
|
//! v = β₂ * v + (1 - β₂) * g² // Update second moment
|
|
//! m̂ = m / (1 - β₁^t) // Bias-corrected first moment
|
|
//! v̂ = v / (1 - β₂^t) // Bias-corrected second moment
|
|
//! θ = θ - lr * m̂ / (√v̂ + ε) // Update parameters
|
|
//! ```
|
|
|
|
use rtx_tensor::{Tensor, Device};
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use crate::PinnStreamContext;
|
|
|
|
/// Result type for optimizer operations
|
|
pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
|
|
|
|
/// GPU-native Adam optimizer
|
|
///
|
|
/// All optimizer state lives on GPU. No CPU-GPU synchronization during training.
|
|
#[derive(Debug)]
|
|
pub struct GpuAdam {
|
|
/// First moment estimates (one per parameter tensor)
|
|
pub m: Vec<Tensor>,
|
|
|
|
/// Second moment estimates (one per parameter tensor)
|
|
pub v: Vec<Tensor>,
|
|
|
|
/// Timestep counter
|
|
pub t: u64,
|
|
|
|
/// Learning rate
|
|
pub lr: f32,
|
|
|
|
/// Decay rate for first moment estimate
|
|
pub beta1: f32,
|
|
|
|
/// Decay rate for second moment estimate
|
|
pub beta2: f32,
|
|
|
|
/// Small constant for numerical stability
|
|
pub eps: f32,
|
|
|
|
/// Device for allocation
|
|
device: Device,
|
|
}
|
|
|
|
impl GpuAdam {
|
|
/// Create a new GPU Adam optimizer
|
|
///
|
|
/// # Arguments
|
|
/// * `param_shapes` - Shapes of parameter tensors
|
|
/// * `device` - GPU device for state allocation
|
|
/// * `lr` - Learning rate (default: 1e-3)
|
|
/// * `beta1` - First moment decay (default: 0.9)
|
|
/// * `beta2` - Second moment decay (default: 0.999)
|
|
/// * `eps` - Numerical stability constant (default: 1e-8)
|
|
pub fn new(
|
|
param_shapes: &[Vec<usize>],
|
|
device: &Device,
|
|
lr: f32,
|
|
beta1: f32,
|
|
beta2: f32,
|
|
eps: f32,
|
|
) -> Result<Self> {
|
|
let mut m = Vec::with_capacity(param_shapes.len());
|
|
let mut v = Vec::with_capacity(param_shapes.len());
|
|
|
|
for shape in param_shapes {
|
|
// Initialize first and second moments to zero on GPU
|
|
let shape_i64: Vec<usize> = shape.iter().map(|&x| x).collect();
|
|
m.push(Tensor::zeros(&shape_i64, device)?);
|
|
v.push(Tensor::zeros(&shape_i64, device)?);
|
|
}
|
|
|
|
Ok(Self {
|
|
m,
|
|
v,
|
|
t: 0,
|
|
lr,
|
|
beta1,
|
|
beta2,
|
|
eps,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
/// Create with default hyperparameters
|
|
///
|
|
/// lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8
|
|
pub fn with_defaults(param_shapes: &[Vec<usize>], device: &Device) -> Result<Self> {
|
|
Self::new(param_shapes, device, 1e-3, 0.9, 0.999, 1e-8)
|
|
}
|
|
|
|
/// Perform one optimization step
|
|
///
|
|
/// Updates parameters in-place using the provided gradients.
|
|
/// All operations stay on GPU - no CPU-GPU synchronization.
|
|
///
|
|
/// # Arguments
|
|
/// * `params` - Parameter tensors to update (modified in-place)
|
|
/// * `grads` - Gradient tensors (same shapes as params)
|
|
pub fn step(&mut self, params: &mut [Tensor], grads: &[Tensor]) -> Result<()> {
|
|
if params.len() != grads.len() || params.len() != self.m.len() {
|
|
return Err(format!(
|
|
"Mismatch: {} params, {} grads, {} optimizer states",
|
|
params.len(),
|
|
grads.len(),
|
|
self.m.len()
|
|
).into());
|
|
}
|
|
|
|
self.t += 1;
|
|
|
|
// Compute bias corrections
|
|
let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
|
|
let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);
|
|
|
|
// Adjusted learning rate with bias correction
|
|
let lr_adj = self.lr * (bias_correction2.sqrt()) / bias_correction1;
|
|
|
|
for i in 0..params.len() {
|
|
let g = &grads[i];
|
|
|
|
// m = β₁ * m + (1 - β₁) * g
|
|
let m_scaled = self.m[i].mul_scalar(self.beta1)?;
|
|
let g_scaled = g.mul_scalar(1.0 - self.beta1)?;
|
|
self.m[i] = m_scaled.add(&g_scaled)?;
|
|
|
|
// v = β₂ * v + (1 - β₂) * g²
|
|
let g_sq = g.mul(g)?;
|
|
let v_scaled = self.v[i].mul_scalar(self.beta2)?;
|
|
let g_sq_scaled = g_sq.mul_scalar(1.0 - self.beta2)?;
|
|
self.v[i] = v_scaled.add(&g_sq_scaled)?;
|
|
|
|
// θ = θ - lr_adj * m / (√v + ε)
|
|
// Using bias-corrected formulation with lr_adj
|
|
let v_sqrt = self.v[i].sqrt()?;
|
|
let v_sqrt_eps = v_sqrt.add_scalar(self.eps)?;
|
|
let update = self.m[i].div(&v_sqrt_eps)?;
|
|
let update_scaled = update.mul_scalar(lr_adj)?;
|
|
params[i] = params[i].sub(&update_scaled)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get current timestep
|
|
pub fn timestep(&self) -> u64 {
|
|
self.t
|
|
}
|
|
|
|
/// Reset optimizer state (first and second moments)
|
|
pub fn reset(&mut self) -> Result<()> {
|
|
for i in 0..self.m.len() {
|
|
self.m[i] = Tensor::zeros(self.m[i].shape(), &self.device)?;
|
|
self.v[i] = Tensor::zeros(self.v[i].shape(), &self.device)?;
|
|
}
|
|
self.t = 0;
|
|
Ok(())
|
|
}
|
|
|
|
/// Set learning rate
|
|
pub fn set_lr(&mut self, lr: f32) {
|
|
self.lr = lr;
|
|
}
|
|
|
|
/// Get learning rate
|
|
pub fn lr(&self) -> f32 {
|
|
self.lr
|
|
}
|
|
|
|
/// Fused optimizer step using a single CUDA kernel per parameter.
|
|
///
|
|
/// This is the zero-allocation version that uses the PinnStreamContext
|
|
/// for launching the fused Adam kernel. Each parameter is updated in-place
|
|
/// with a single kernel launch that computes:
|
|
/// - m update
|
|
/// - v update
|
|
/// - parameter update
|
|
///
|
|
/// # Arguments
|
|
/// * `params` - Mutable references to parameter tensors (updated in-place)
|
|
/// * `grads` - Gradient tensors (read-only)
|
|
/// * `ctx` - CUDA stream context for kernel launch
|
|
#[cfg(feature = "cuda")]
|
|
pub fn step_fused(
|
|
&mut self,
|
|
params: &mut [&mut Tensor],
|
|
grads: &[&Tensor],
|
|
ctx: &PinnStreamContext,
|
|
) -> Result<()> {
|
|
if params.len() != grads.len() || params.len() != self.m.len() {
|
|
return Err(format!(
|
|
"Mismatch: {} params, {} grads, {} optimizer states",
|
|
params.len(),
|
|
grads.len(),
|
|
self.m.len()
|
|
).into());
|
|
}
|
|
|
|
self.t += 1;
|
|
|
|
// Compute bias corrections
|
|
let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
|
|
let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);
|
|
|
|
// Adjusted learning rate with bias correction
|
|
let lr_adj = self.lr * (bias_correction2.sqrt()) / bias_correction1;
|
|
|
|
for i in 0..params.len() {
|
|
// One kernel launch per parameter: updates param, m, v in-place
|
|
ctx.adam_step(
|
|
params[i],
|
|
&mut self.m[i],
|
|
&mut self.v[i],
|
|
grads[i],
|
|
lr_adj,
|
|
self.beta1,
|
|
self.beta2,
|
|
self.eps,
|
|
).map_err(|e| format!("adam_step failed for param {}: {}", i, e))?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_gpu_adam_creation() {
|
|
let device = Device::Cpu;
|
|
let shapes = vec![
|
|
vec![64, 32], // W0
|
|
vec![64], // b0
|
|
vec![2, 64], // Wout
|
|
vec![2], // bout
|
|
];
|
|
|
|
let opt = GpuAdam::with_defaults(&shapes, &device).unwrap();
|
|
assert_eq!(opt.m.len(), 4);
|
|
assert_eq!(opt.v.len(), 4);
|
|
assert_eq!(opt.timestep(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_adam_step() {
|
|
let device = Device::Cpu;
|
|
let shapes = vec![vec![2, 2]];
|
|
|
|
let mut opt = GpuAdam::with_defaults(&shapes, &device).unwrap();
|
|
|
|
// Create parameter and gradient tensors
|
|
let mut params = vec![
|
|
Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap()
|
|
];
|
|
let grads = vec![
|
|
Tensor::from_vec(vec![0.1f32, 0.2, 0.3, 0.4], &[2, 2], &device).unwrap()
|
|
];
|
|
|
|
// Perform optimization step
|
|
opt.step(&mut params, &grads).unwrap();
|
|
|
|
assert_eq!(opt.timestep(), 1);
|
|
|
|
// Parameters should have changed
|
|
let params_cpu = params[0].to_cpu().unwrap();
|
|
// After one Adam step with lr=1e-3, params should decrease slightly
|
|
assert!(params_cpu[0] < 1.0);
|
|
assert!(params_cpu[1] < 2.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_adam_reset() {
|
|
let device = Device::Cpu;
|
|
let shapes = vec![vec![2, 2]];
|
|
|
|
let mut opt = GpuAdam::with_defaults(&shapes, &device).unwrap();
|
|
|
|
// Do a step
|
|
let mut params = vec![
|
|
Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap()
|
|
];
|
|
let grads = vec![
|
|
Tensor::from_vec(vec![0.1f32, 0.2, 0.3, 0.4], &[2, 2], &device).unwrap()
|
|
];
|
|
opt.step(&mut params, &grads).unwrap();
|
|
|
|
assert_eq!(opt.timestep(), 1);
|
|
|
|
// Reset
|
|
opt.reset().unwrap();
|
|
assert_eq!(opt.timestep(), 0);
|
|
|
|
// Check moments are zeroed
|
|
let m0_cpu = opt.m[0].to_cpu().unwrap();
|
|
for val in m0_cpu.iter() {
|
|
assert!(*val == 0.0);
|
|
}
|
|
}
|
|
}
|