osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
341 lines
12 KiB
Rust
341 lines
12 KiB
Rust
//! Optimizers for transformer training
|
|
//!
|
|
//! This module provides high-performance optimizers including Adam and `AdamW`
|
|
//! with support for distributed training and revolutionary enhancements.
|
|
|
|
use crate::Result;
|
|
use rtx_tensor::Tensor;
|
|
use std::collections::HashMap;
|
|
use tracing::debug;
|
|
|
|
// ============================================================================
|
|
// PHASE 1: CORE OPTIMIZERS (Essential functionality)
|
|
// ============================================================================
|
|
|
|
pub mod adam;
|
|
pub mod adamw;
|
|
|
|
#[cfg(all(test, feature = "disabled_tests"))]
|
|
pub mod adam_test;
|
|
|
|
// Re-export main optimizer types
|
|
pub use adam::AdamState;
|
|
|
|
// ============================================================================
|
|
// PHASE 2: ADVANCED OPTIMIZERS (Temporarily disabled)
|
|
// ============================================================================
|
|
// These will be re-enabled once core optimizers are stable
|
|
|
|
// pub mod ademamix; // TODO: Re-enable in Phase 2
|
|
// pub mod adabound; // TODO: Re-enable in Phase 2
|
|
// pub mod lion; // TODO: Re-enable in Phase 2
|
|
// pub mod sophia; // TODO: Re-enable in Phase 2
|
|
// pub mod shampoo; // TODO: Re-enable in Phase 2
|
|
// pub mod kfac; // TODO: Re-enable in Phase 2
|
|
// pub mod kfac_advanced; // TODO: Re-enable in Phase 2
|
|
// pub mod kfac_example; // TODO: Re-enable in Phase 2
|
|
// pub mod ranger; // TODO: Re-enable in Phase 2
|
|
// pub mod novograd; // TODO: Re-enable in Phase 2
|
|
// pub mod natural_gradient; // TODO: Re-enable in Phase 2
|
|
// pub mod trust_region; // TODO: Re-enable in Phase 2
|
|
// pub mod lbfgs; // TODO: Re-enable in Phase 2
|
|
|
|
pub mod matrix_utils;
|
|
|
|
// Re-export tensor bridge traits for all optimizer modules
|
|
pub use crate::tensor_bridge::{TensorBridge, TensorBridgeStatic, TensorCompat};
|
|
|
|
// Test modules temporarily disabled due to missing dependencies
|
|
// #[cfg(test)]
|
|
// pub mod parameter_update_tests;
|
|
|
|
// All test modules temporarily disabled
|
|
// Will be re-enabled after implementing missing optimizer types
|
|
|
|
// ============================================================================
|
|
// PHASE 1: CORE OPTIMIZER EXPORTS (Only enabled optimizers)
|
|
// ============================================================================
|
|
|
|
pub use adam::AdamOptimizer;
|
|
pub use adamw::AdamWOptimizer;
|
|
|
|
// ============================================================================
|
|
// PHASE 2: ADVANCED OPTIMIZER EXPORTS (Temporarily disabled)
|
|
// ============================================================================
|
|
// These will be re-enabled when the modules are uncommented
|
|
|
|
// pub use ademamix::{AdEMAMixOptimizer, AdEMAMixConfig};
|
|
// pub use adabound::{AdaBoundOptimizer, AdaBoundConfig};
|
|
// pub use lion::{LionOptimizer, LionConfig};
|
|
// pub use sophia::{SophiaOptimizer, SophiaConfig};
|
|
// pub use shampoo::{ShampooOptimizer, ShampooConfig};
|
|
// pub use kfac::{KFacOptimizer, KFacConfig, LayerType};
|
|
// pub use kfac_advanced::{AdvancedKFacOptimizer, AdvancedKFacConfig, ConvergenceMetrics};
|
|
// pub use ranger::{RangerOptimizer, RangerConfig};
|
|
// pub use novograd::{NovoGradOptimizer, NovoGradConfig};
|
|
// pub use natural_gradient::{NaturalGradientOptimizer, NaturalGradientConfig, FisherApproximation};
|
|
// pub use trust_region::{TrustRegionOptimizer, TrustRegionConfig, SubproblemSolver};
|
|
// pub use lbfgs::{LBFGSOptimizer, LBFGSConfig};
|
|
|
|
/// Base optimizer implementation with gradient storage
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct BaseOptimizer {
|
|
/// Stored gradients for batch processing
|
|
#[serde(skip)]
|
|
pub stored_gradients: HashMap<String, Tensor>,
|
|
/// Current learning rate
|
|
pub learning_rate: f64,
|
|
}
|
|
|
|
impl BaseOptimizer {
|
|
/// Create a new base optimizer
|
|
#[must_use]
|
|
pub fn new(learning_rate: f64) -> Self {
|
|
Self {
|
|
stored_gradients: HashMap::new(),
|
|
learning_rate,
|
|
}
|
|
}
|
|
|
|
/// Get stored gradients
|
|
#[must_use]
|
|
pub fn stored_gradients(&self) -> &HashMap<String, Tensor> {
|
|
&self.stored_gradients
|
|
}
|
|
|
|
/// Clear stored gradients
|
|
pub fn clear_gradients(&mut self) {
|
|
self.stored_gradients.clear();
|
|
}
|
|
|
|
/// Set learning rate
|
|
pub fn set_learning_rate(&mut self, lr: f64) {
|
|
self.learning_rate = lr;
|
|
}
|
|
|
|
/// Store gradients
|
|
pub fn store_gradients(&mut self, gradients: HashMap<String, Tensor>) {
|
|
self.stored_gradients.clear();
|
|
self.stored_gradients.extend(gradients);
|
|
}
|
|
}
|
|
|
|
/// Configuration for different optimizer types (Phase 1: Core optimizers only)
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub enum OptimizerConfig {
|
|
/// Adam optimizer configuration
|
|
Adam(AdamConfig),
|
|
/// `AdamW` optimizer configuration
|
|
AdamW(AdamWConfig),
|
|
// ========================================================================
|
|
// PHASE 2: ADVANCED OPTIMIZERS (Temporarily disabled)
|
|
// ========================================================================
|
|
// These will be re-enabled once the corresponding modules are uncommented
|
|
|
|
// /// AdEMAMix optimizer configuration
|
|
// AdEMAMix(AdEMAMixConfig),
|
|
// /// AdaBound optimizer configuration
|
|
// AdaBound(AdaBoundConfig),
|
|
// /// Lion optimizer configuration
|
|
// Lion(LionConfig),
|
|
// /// Sophia optimizer configuration
|
|
// Sophia(SophiaConfig),
|
|
// /// Shampoo optimizer configuration
|
|
// Shampoo(ShampooConfig),
|
|
// /// K-FAC optimizer configuration
|
|
// KFac(KFacConfig),
|
|
// /// Advanced K-FAC optimizer configuration
|
|
// AdvancedKFac(AdvancedKFacConfig),
|
|
// /// Ranger optimizer configuration (RAdam + Lookahead)
|
|
// Ranger(RangerConfig),
|
|
// /// NovoGrad optimizer configuration (layer-wise gradient normalization)
|
|
// NovoGrad(NovoGradConfig),
|
|
// /// Natural Gradient optimizer configuration (Fisher Information Matrix preconditioning)
|
|
// NaturalGradient(NaturalGradientConfig),
|
|
// /// Trust Region optimizer configuration (second-order with adaptive radius)
|
|
// TrustRegion(TrustRegionConfig),
|
|
// /// L-BFGS optimizer configuration (quasi-Newton with limited memory)
|
|
// LBFGS(LBFGSConfig),
|
|
}
|
|
|
|
/// Adam optimizer configuration
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct AdamConfig {
|
|
/// Learning rate
|
|
pub learning_rate: f64,
|
|
/// Beta1 parameter (momentum)
|
|
pub beta1: f64,
|
|
/// Beta2 parameter (`RMSprop`)
|
|
pub beta2: f64,
|
|
/// Epsilon for numerical stability
|
|
pub epsilon: f64,
|
|
/// Weight decay coefficient
|
|
pub weight_decay: f64,
|
|
/// Whether to use `AMSGrad` variant
|
|
pub amsgrad: bool,
|
|
}
|
|
|
|
impl Default for AdamConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
learning_rate: 1e-3,
|
|
beta1: 0.9,
|
|
beta2: 0.999,
|
|
epsilon: 1e-8,
|
|
weight_decay: 0.0,
|
|
amsgrad: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// `AdamW` optimizer configuration
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct AdamWConfig {
|
|
/// Learning rate
|
|
pub learning_rate: f64,
|
|
/// Beta1 parameter (momentum)
|
|
pub beta1: f64,
|
|
/// Beta2 parameter (`RMSprop`)
|
|
pub beta2: f64,
|
|
/// Epsilon for numerical stability
|
|
pub epsilon: f64,
|
|
/// Weight decay coefficient
|
|
pub weight_decay: f64,
|
|
/// Whether to use bias correction
|
|
pub amsgrad: bool,
|
|
}
|
|
|
|
impl Default for AdamWConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
learning_rate: 1e-3,
|
|
beta1: 0.9,
|
|
beta2: 0.999,
|
|
epsilon: 1e-8,
|
|
weight_decay: 1e-2,
|
|
amsgrad: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Create an optimizer from configuration (Phase 1: Core optimizers only)
|
|
pub fn create_optimizer(
|
|
config: OptimizerConfig,
|
|
_parameters: HashMap<String, Tensor>,
|
|
) -> Result<Box<dyn Optimizer>> {
|
|
match config {
|
|
OptimizerConfig::Adam(adam_config) => {
|
|
debug!("Creating Adam optimizer with config: {:?}", adam_config);
|
|
Ok(Box::new(AdamOptimizer::new(adam_config)?))
|
|
}
|
|
OptimizerConfig::AdamW(adamw_config) => {
|
|
debug!("Creating AdamW optimizer with config: {:?}", adamw_config);
|
|
Ok(Box::new(AdamWOptimizer::new(
|
|
adamw_config.learning_rate,
|
|
adamw_config.beta1,
|
|
adamw_config.beta2,
|
|
adamw_config.epsilon,
|
|
adamw_config.weight_decay,
|
|
adamw_config.amsgrad, // Using amsgrad field as decoupled flag for now
|
|
)?))
|
|
} // ====================================================================
|
|
// PHASE 2: ADVANCED OPTIMIZERS (Temporarily disabled)
|
|
// ====================================================================
|
|
// These cases will be re-enabled when the optimizer modules are uncommented
|
|
|
|
// OptimizerConfig::AdEMAMix(ademamix_config) => {
|
|
// debug!("Creating AdEMAMix optimizer with config: {:?}", ademamix_config);
|
|
// Ok(Box::new(AdEMAMixOptimizer::new(ademamix_config)?))
|
|
// }
|
|
|
|
// All other optimizer cases are temporarily disabled since the modules are commented out
|
|
}
|
|
}
|
|
|
|
/// Trait for all optimizers in the transformer training infrastructure
|
|
pub trait Optimizer: Send + Sync {
|
|
/// Perform a single optimization step for a parameter
|
|
fn step_param(&mut self, param_name: &str, param: &Tensor, grad: &Tensor) -> Result<Tensor>;
|
|
|
|
/// Get the current learning rate
|
|
fn learning_rate(&self) -> f64;
|
|
|
|
/// Set the learning rate
|
|
fn set_learning_rate(&mut self, lr: f64) -> Result<()>;
|
|
|
|
/// Check if optimizer has state for a parameter
|
|
fn has_state(&self, param_name: &str) -> bool;
|
|
|
|
/// Reset state for a specific parameter
|
|
fn reset_state(&mut self, param_name: &str) -> Result<()>;
|
|
|
|
/// Reset all optimizer state
|
|
fn reset_all_state(&mut self);
|
|
|
|
/// Get step count for a parameter
|
|
fn get_step_count(&self, param_name: &str) -> Result<i64>;
|
|
|
|
/// Get optimizer type name
|
|
fn optimizer_type(&self) -> &'static str;
|
|
|
|
/// Get mutable reference as Any for downcasting
|
|
fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
|
|
|
|
/// Zero all gradients (clear stored gradients)
|
|
fn zero_grad(&mut self) {
|
|
// Default implementation clears stored gradients
|
|
if let Ok(()) = self.store_gradients_internal(HashMap::new()) {
|
|
// Gradients cleared successfully
|
|
}
|
|
}
|
|
|
|
/// Store gradients for batch processing (used by trainer)
|
|
fn set_gradients(&mut self, gradients: HashMap<String, Tensor>) -> Result<()> {
|
|
// Default implementation stores gradients in internal map
|
|
self.store_gradients_internal(gradients)
|
|
}
|
|
|
|
/// Perform optimization step on all stored gradients with given learning rate
|
|
fn step(&mut self, learning_rate: f64) -> Result<HashMap<String, Tensor>> {
|
|
// Default implementation processes stored gradients
|
|
self.process_stored_gradients(learning_rate)
|
|
}
|
|
|
|
/// Internal method to store gradients (default implementation)
|
|
fn store_gradients_internal(&mut self, gradients: HashMap<String, Tensor>) -> Result<()>;
|
|
|
|
/// Internal method to process stored gradients (default implementation)
|
|
fn process_stored_gradients(&mut self, learning_rate: f64) -> Result<HashMap<String, Tensor>>;
|
|
}
|
|
|
|
/// Optimizer type enumeration
|
|
#[derive(Debug, Clone)]
|
|
pub enum OptimizerType {
|
|
/// Adam optimizer
|
|
Adam,
|
|
/// `AdamW` optimizer with decoupled weight decay
|
|
AdamW,
|
|
/// `AdEMAMix` optimizer with adaptive EMA mixture
|
|
AdEMAMix,
|
|
/// `AdaBound` optimizer with dynamic learning rate bounds
|
|
AdaBound,
|
|
/// Lion optimizer with evolved sign momentum
|
|
Lion,
|
|
/// Sophia optimizer with second-order clipped stochastic optimization
|
|
Sophia,
|
|
/// Shampoo optimizer with full matrix preconditioning
|
|
Shampoo,
|
|
/// K-FAC optimizer with Kronecker-factored approximation
|
|
KFac,
|
|
/// Ranger optimizer combining `RAdam` with Lookahead
|
|
Ranger,
|
|
/// `NovoGrad` optimizer with layer-wise gradient normalization
|
|
NovoGrad,
|
|
/// Natural Gradient optimizer with Fisher Information Matrix preconditioning
|
|
NaturalGradient,
|
|
/// Trust Region optimizer with adaptive radius and second-order approximation
|
|
TrustRegion,
|
|
/// L-BFGS optimizer with limited-memory quasi-Newton approximation
|
|
LBFGS,
|
|
}
|