trustformers_optim/
traits.rs

1//! Advanced optimizer trait hierarchy for TrustformeRS.
2//!
3//! This module extends the base `Optimizer` trait from `trustformers-core` with
4//! additional specialized traits for different categories of optimizers, providing
5//! better organization and extensibility.
6//!
7//! # Trait Hierarchy
8//!
9//! ```text
10//! Optimizer (from trustformers-core)
11//!     │
12//!     ├── StatefulOptimizer
13//!     │   ├── MomentumOptimizer
14//!     │   │   ├── AdaptiveMomentumOptimizer  (Adam, AdamW, etc.)
15//!     │   │   └── ClassicalMomentumOptimizer (SGD with momentum)
16//!     │   └── SecondOrderOptimizer (L-BFGS, Newton-CG, etc.)
17//!     │
18//!     ├── DistributedOptimizer
19//!     │   ├── GradientCompressionOptimizer
20//!     │   ├── FederatedOptimizer
21//!     │   └── AsyncOptimizer
22//!     │
23//!     ├── HardwareOptimizer
24//!     │   ├── SIMDOptimizer
25//!     │   ├── GPUOptimizer
26//!     │   └── EdgeOptimizer
27//!     │
28//!     └── MetaOptimizer
29//!         ├── LookaheadOptimizer
30//!         ├── ScheduledOptimizer
31//!         └── CompositeOptimizer
32//! ```
33
34use crate::common::StateMemoryStats;
35use std::collections::HashMap;
36use trustformers_core::errors::{Result, TrustformersError};
37use trustformers_core::tensor::Tensor;
38use trustformers_core::traits::Optimizer;
39
40/// Extended optimizer trait with state management capabilities.
41///
42/// This trait builds on the base `Optimizer` trait to provide standardized
43/// state management, serialization, and configuration access.
44pub trait StatefulOptimizer: Optimizer {
45    /// The configuration type for this optimizer.
46    type Config: Clone + Send + Sync;
47
48    /// The state type used by this optimizer.
49    type State: Send + Sync;
50
51    /// Gets a reference to the optimizer's configuration.
52    fn config(&self) -> &Self::Config;
53
54    /// Gets a reference to the optimizer's internal state.
55    fn state(&self) -> &Self::State;
56
57    /// Gets a mutable reference to the optimizer's internal state.
58    fn state_mut(&mut self) -> &mut Self::State;
59
60    /// Saves the optimizer state to a dictionary for checkpointing.
61    fn state_dict(&self) -> Result<HashMap<String, Tensor>>;
62
63    /// Loads optimizer state from a dictionary during checkpoint restoration.
64    fn load_state_dict(&mut self, state: HashMap<String, Tensor>) -> Result<()>;
65
66    /// Gets memory usage statistics for this optimizer.
67    fn memory_usage(&self) -> StateMemoryStats;
68
69    /// Resets the optimizer state (useful for training restarts).
70    fn reset_state(&mut self);
71
72    /// Returns the number of parameters being optimized.
73    fn num_parameters(&self) -> usize;
74
75    /// Saves the optimizer state to `path`.
76    ///
77    /// The default implementation serialises [`Self::state_dict`] with `oxicode`, so
78    /// every implementor gets checkpointing for free and all implementors share one
79    /// on-disk format. Override only to add a format of your own.
80    ///
81    /// # Errors
82    ///
83    /// Returns an error when the state cannot be produced, encoded, or written.
84    fn save_state(&self, path: &std::path::Path) -> Result<()> {
85        let state = self.state_dict()?;
86        let encoded = encode_state_dict(&state)?;
87        std::fs::write(path, encoded).map_err(|error| {
88            TrustformersError::io_error(format!(
89                "failed to write optimizer state to {}: {error}",
90                path.display()
91            ))
92        })
93    }
94
95    /// Loads the optimizer state written by [`Self::save_state`].
96    ///
97    /// # Errors
98    ///
99    /// Returns an error when the file cannot be read, is not a state dictionary this
100    /// crate wrote, or does not match this optimizer's expectations.
101    fn load_state(&mut self, path: &std::path::Path) -> Result<()> {
102        let bytes = std::fs::read(path).map_err(|error| {
103            TrustformersError::io_error(format!(
104                "failed to read optimizer state from {}: {error}",
105                path.display()
106            ))
107        })?;
108        let state = decode_state_dict(&bytes)?;
109        self.load_state_dict(state)
110    }
111}
112
113/// Wire format of a serialised optimizer state dictionary.
114///
115/// Tensors are stored as `(name, shape, f32 payload)` triples. `f32` is the only
116/// dtype optimizer state uses in this crate.
117type WireStateDict = Vec<(String, Vec<usize>, Vec<f32>)>;
118
119/// Encodes a state dictionary for [`StatefulOptimizer::save_state`].
120///
121/// # Errors
122///
123/// Returns an error when a tensor is not `f32`-readable or encoding fails.
124pub fn encode_state_dict(state: &HashMap<String, Tensor>) -> Result<Vec<u8>> {
125    let mut wire: WireStateDict = Vec::with_capacity(state.len());
126    for (name, tensor) in state {
127        wire.push((name.clone(), tensor.shape().to_vec(), tensor.data_f32()?));
128    }
129    // Deterministic order keeps checkpoints byte-reproducible.
130    wire.sort_by(|a, b| a.0.cmp(&b.0));
131
132    oxicode::serde::encode_to_vec(&wire, oxicode::config::standard()).map_err(|error| {
133        TrustformersError::invalid_state(format!("failed to encode optimizer state: {error}"))
134    })
135}
136
137/// Decodes a state dictionary written by [`encode_state_dict`].
138///
139/// # Errors
140///
141/// Returns an error when the bytes are not a state dictionary or a payload length
142/// disagrees with its shape.
143pub fn decode_state_dict(bytes: &[u8]) -> Result<HashMap<String, Tensor>> {
144    let (wire, _): (WireStateDict, usize) =
145        oxicode::serde::decode_from_slice(bytes, oxicode::config::standard()).map_err(|error| {
146            TrustformersError::invalid_state(format!("failed to decode optimizer state: {error}"))
147        })?;
148
149    let mut state = HashMap::with_capacity(wire.len());
150    for (name, shape, values) in wire {
151        let expected: usize = shape.iter().product();
152        if values.len() != expected {
153            return Err(TrustformersError::invalid_state(format!(
154                "optimizer state entry '{name}' has {} values but shape {shape:?} needs {expected}",
155                values.len()
156            )));
157        }
158        state.insert(name, Tensor::from_vec(values, &shape)?);
159    }
160    Ok(state)
161}
162
163/// Trait for optimizers that use momentum-based updates.
164///
165/// This includes both classical momentum (SGD) and adaptive momentum (Adam family).
166pub trait MomentumOptimizer: StatefulOptimizer {
167    /// Gets the momentum decay coefficient (β1 in Adam, momentum in SGD).
168    fn momentum_coeff(&self) -> f32;
169
170    /// Sets the momentum decay coefficient.
171    fn set_momentum_coeff(&mut self, coeff: f32);
172
173    /// Gets the current momentum buffers (for debugging/analysis).
174    fn momentum_buffers(&self) -> &HashMap<String, Vec<f32>>;
175
176    /// Clears all momentum buffers (useful for fine-tuning).
177    fn clear_momentum(&mut self);
178}
179
180/// Trait for adaptive momentum optimizers (Adam, AdamW, RAdam, etc.).
181///
182/// These optimizers maintain both first and second moment estimates.
183pub trait AdaptiveMomentumOptimizer: MomentumOptimizer {
184    /// Gets the second moment decay coefficient (β2 in Adam).
185    fn variance_coeff(&self) -> f32;
186
187    /// Sets the second moment decay coefficient.
188    fn set_variance_coeff(&mut self, coeff: f32);
189
190    /// Gets the epsilon value for numerical stability.
191    fn epsilon(&self) -> f32;
192
193    /// Sets the epsilon value.
194    fn set_epsilon(&mut self, eps: f32);
195
196    /// Gets the current variance buffers (for debugging/analysis).
197    fn variance_buffers(&self) -> &HashMap<String, Vec<f32>>;
198
199    /// Clears variance buffers.
200    fn clear_variance(&mut self);
201
202    /// Applies bias correction to momentum and variance estimates.
203    fn apply_bias_correction(&self, momentum: f32, variance: f32, step: usize) -> (f32, f32);
204}
205
206/// Trait for classical momentum optimizers (SGD variants).
207pub trait ClassicalMomentumOptimizer: MomentumOptimizer {
208    /// Gets the dampening factor.
209    fn dampening(&self) -> f32;
210
211    /// Sets the dampening factor.
212    fn set_dampening(&mut self, dampening: f32);
213
214    /// Whether Nesterov momentum is enabled.
215    fn nesterov(&self) -> bool;
216
217    /// Enables or disables Nesterov momentum.
218    fn set_nesterov(&mut self, nesterov: bool);
219}
220
221/// Trait for second-order optimization methods.
222///
223/// These optimizers use curvature information (Hessian approximations).
224pub trait SecondOrderOptimizer: StatefulOptimizer {
225    /// The type used to represent curvature information.
226    type CurvatureInfo;
227
228    /// Updates the curvature approximation with new gradient information.
229    fn update_curvature(&mut self, gradients: &[Tensor]) -> Result<()>;
230
231    /// Gets the current curvature approximation.
232    fn curvature_info(&self) -> &Self::CurvatureInfo;
233
234    /// Applies the inverse Hessian approximation to compute search direction.
235    fn apply_inverse_hessian(&self, gradient: &Tensor) -> Result<Tensor>;
236
237    /// Gets the maximum number of curvature pairs stored (for L-BFGS).
238    fn history_size(&self) -> usize;
239}
240
241/// Trait for distributed optimization capabilities.
242///
243/// Provides interfaces for gradient synchronization and distributed training.
244pub trait DistributedOptimizer: Optimizer {
245    /// The communicator type used for distributed operations.
246    type Communicator;
247
248    /// Performs all-reduce operation on gradients.
249    fn all_reduce_gradients(&mut self, gradients: &mut [Tensor]) -> Result<()>;
250
251    /// Broadcasts parameters from rank 0 to all other ranks.
252    fn broadcast_parameters(&mut self, parameters: &mut [Tensor]) -> Result<()>;
253
254    /// Gets the current rank in the distributed group.
255    fn rank(&self) -> usize;
256
257    /// Gets the total number of ranks in the distributed group.
258    fn world_size(&self) -> usize;
259
260    /// Synchronizes optimizer state across all ranks.
261    fn sync_state(&mut self) -> Result<()>;
262}
263
264/// Trait for optimizers with gradient compression capabilities.
265pub trait GradientCompressionOptimizer: DistributedOptimizer {
266    /// The compression method used.
267    type CompressionMethod;
268
269    /// Compresses gradients before communication.
270    fn compress_gradients(&self, gradients: &[Tensor]) -> Result<Vec<u8>>;
271
272    /// Decompresses received gradient data.
273    fn decompress_gradients(&self, data: &[u8]) -> Result<Vec<Tensor>>;
274
275    /// Gets the compression ratio achieved.
276    fn compression_ratio(&self) -> f32;
277
278    /// Sets the compression parameters.
279    fn set_compression_config(&mut self, config: Self::CompressionMethod);
280}
281
282/// Trait for federated learning optimizers.
283pub trait FederatedOptimizer: DistributedOptimizer {
284    /// Client information type.
285    type ClientInfo;
286
287    /// Aggregates model updates from multiple clients.
288    fn aggregate_updates(
289        &mut self,
290        updates: &[Tensor],
291        clients: &[Self::ClientInfo],
292    ) -> Result<Tensor>;
293
294    /// Selects clients for the next round of training.
295    fn select_clients(
296        &self,
297        available_clients: &[Self::ClientInfo],
298        num_clients: usize,
299    ) -> Vec<usize>;
300
301    /// Applies differential privacy to updates.
302    fn apply_differential_privacy(&mut self, update: &mut Tensor) -> Result<()>;
303}
304
305/// Trait for asynchronous optimization methods.
306pub trait AsyncOptimizer: DistributedOptimizer {
307    /// Applies delayed gradients with staleness compensation.
308    fn apply_delayed_gradients(&mut self, gradients: &[Tensor], staleness: usize) -> Result<()>;
309
310    /// Gets the maximum allowed staleness.
311    fn max_staleness(&self) -> usize;
312
313    /// Sets the staleness compensation method.
314    fn set_staleness_compensation(&mut self, method: StalenessCompensation);
315}
316
317/// Staleness compensation methods for asynchronous optimization.
318#[derive(Debug, Clone, Copy)]
319pub enum StalenessCompensation {
320    /// No compensation for staleness.
321    None,
322    /// Linear scaling by staleness factor.
323    Linear,
324    /// Exponential decay based on staleness.
325    Exponential,
326    /// Polynomial scaling with configurable degree.
327    Polynomial(f32),
328}
329
330/// Trait for hardware-specific optimizer optimizations.
331pub trait HardwareOptimizer: Optimizer {
332    /// The target hardware type.
333    type HardwareTarget;
334
335    /// Optimizes the optimizer for specific hardware.
336    fn optimize_for_hardware(&mut self, target: Self::HardwareTarget) -> Result<()>;
337
338    /// Gets hardware utilization statistics.
339    fn hardware_utilization(&self) -> HardwareStats;
340
341    /// Checks if the optimizer is compatible with the current hardware.
342    fn is_hardware_compatible(&self) -> bool;
343}
344
345/// Hardware utilization statistics.
346#[derive(Debug, Clone)]
347pub struct HardwareStats {
348    /// Memory bandwidth utilization (0.0 to 1.0).
349    pub memory_bandwidth_utilization: f32,
350    /// Compute utilization (0.0 to 1.0).
351    pub compute_utilization: f32,
352    /// Cache hit rate (0.0 to 1.0).
353    pub cache_hit_rate: f32,
354    /// FLOPS per second achieved.
355    pub flops_per_second: f64,
356}
357
358/// Trait for SIMD-optimized operations.
359pub trait SIMDOptimizer: HardwareOptimizer {
360    /// The SIMD instruction set being used.
361    type SIMDType;
362
363    /// Checks if SIMD operations are available.
364    fn simd_available(&self) -> bool;
365
366    /// Gets the SIMD vector width.
367    fn vector_width(&self) -> usize;
368
369    /// Applies SIMD-optimized parameter updates.
370    fn simd_update(&mut self, parameters: &mut [Tensor], gradients: &[Tensor]) -> Result<()>;
371}
372
373/// Trait for GPU-accelerated optimizers.
374pub trait GPUOptimizer: HardwareOptimizer {
375    /// The GPU compute capability.
376    type ComputeCapability;
377
378    /// Transfers optimizer state to GPU.
379    fn to_gpu(&mut self) -> Result<()>;
380
381    /// Transfers optimizer state to CPU.
382    fn to_cpu(&mut self) -> Result<()>;
383
384    /// Launches GPU kernels for parameter updates.
385    fn gpu_update(&mut self, parameters: &mut [Tensor], gradients: &[Tensor]) -> Result<()>;
386
387    /// Gets GPU memory usage.
388    fn gpu_memory_usage(&self) -> GPUMemoryStats;
389}
390
391/// GPU memory usage statistics.
392#[derive(Debug, Clone)]
393pub struct GPUMemoryStats {
394    /// Total GPU memory in bytes.
395    pub total_memory: usize,
396    /// Used GPU memory in bytes.
397    pub used_memory: usize,
398    /// Available GPU memory in bytes.
399    pub available_memory: usize,
400    /// Memory usage by optimizer state.
401    pub optimizer_memory: usize,
402}
403
404/// Trait for edge device optimized optimizers.
405pub trait EdgeOptimizer: HardwareOptimizer {
406    /// Power consumption statistics.
407    type PowerStats;
408
409    /// Optimizes for low power consumption.
410    fn optimize_for_power(&mut self) -> Result<()>;
411
412    /// Gets current power consumption statistics.
413    fn power_stats(&self) -> Self::PowerStats;
414
415    /// Reduces precision to save memory and power.
416    fn reduce_precision(&mut self, bits: u8) -> Result<()>;
417}
418
419/// Trait for meta-optimizers that wrap other optimizers.
420pub trait MetaOptimizer: Optimizer {
421    /// The base optimizer type being wrapped.
422    type BaseOptimizer: Optimizer;
423
424    /// Gets a reference to the base optimizer.
425    fn base_optimizer(&self) -> &Self::BaseOptimizer;
426
427    /// Gets a mutable reference to the base optimizer.
428    fn base_optimizer_mut(&mut self) -> &mut Self::BaseOptimizer;
429
430    /// Applies the meta-optimization strategy.
431    fn apply_meta_strategy(
432        &mut self,
433        parameters: &mut [Tensor],
434        gradients: &[Tensor],
435    ) -> Result<()>;
436}
437
438/// Trait for lookahead meta-optimizers.
439pub trait LookaheadOptimizer: MetaOptimizer {
440    /// Gets the lookahead step size (α).
441    fn lookahead_alpha(&self) -> f32;
442
443    /// Sets the lookahead step size.
444    fn set_lookahead_alpha(&mut self, alpha: f32);
445
446    /// Gets the lookahead update frequency (k).
447    fn lookahead_k(&self) -> usize;
448
449    /// Sets the lookahead update frequency.
450    fn set_lookahead_k(&mut self, k: usize);
451
452    /// Gets the slow weights (for debugging).
453    fn slow_weights(&self) -> &HashMap<String, Vec<f32>>;
454}
455
456/// Trait for scheduled optimizers with learning rate scheduling.
457pub trait ScheduledOptimizer: Optimizer {
458    /// The scheduler type.
459    type Scheduler;
460
461    /// Gets a reference to the scheduler.
462    fn scheduler(&self) -> &Self::Scheduler;
463
464    /// Gets a mutable reference to the scheduler.
465    fn scheduler_mut(&mut self) -> &mut Self::Scheduler;
466
467    /// Updates the learning rate based on the scheduler.
468    fn update_lr(&mut self) -> Result<()>;
469
470    /// Gets the current scheduled learning rate.
471    fn current_lr(&self) -> f32;
472}
473
474/// Trait for composite optimizers that combine multiple optimization strategies.
475pub trait CompositeOptimizer: Optimizer {
476    /// The component optimizer types.
477    type Components;
478
479    /// Gets references to all component optimizers.
480    fn components(&self) -> &Self::Components;
481
482    /// Gets mutable references to all component optimizers.
483    fn components_mut(&mut self) -> &mut Self::Components;
484
485    /// Applies updates from all component optimizers.
486    fn apply_composite_update(
487        &mut self,
488        parameters: &mut [Tensor],
489        gradients: &[Tensor],
490    ) -> Result<()>;
491
492    /// Gets the weight assigned to each component.
493    fn component_weights(&self) -> Vec<f32>;
494
495    /// Sets the weights for each component.
496    fn set_component_weights(&mut self, weights: Vec<f32>) -> Result<()>;
497}
498
499/// Optimizer factory trait for creating optimizers with different configurations.
500pub trait OptimizerFactory {
501    /// The optimizer type produced by this factory.
502    type Optimizer: Optimizer;
503
504    /// The configuration type for the optimizer.
505    type Config;
506
507    /// Creates a new optimizer with the given configuration.
508    fn create(&self, config: Self::Config) -> Result<Self::Optimizer>;
509
510    /// Lists all available optimizer variants.
511    fn available_variants(&self) -> Vec<&'static str>;
512
513    /// Creates an optimizer by name with default configuration.
514    fn create_by_name(&self, name: &str) -> Result<Self::Optimizer>;
515}
516
517/// Trait for optimizers that can be serialized and restored.
518pub trait SerializableOptimizer: Optimizer {
519    /// Serializes the optimizer to bytes.
520    fn serialize(&self) -> Result<Vec<u8>>;
521
522    /// Deserializes an optimizer from bytes.
523    fn deserialize(data: &[u8]) -> Result<Self>
524    where
525        Self: Sized;
526
527    /// Gets the serialization format version.
528    fn version(&self) -> u32;
529}
530
531#[cfg(test)]
532mod tests {
533    use super::*;
534
535    #[test]
536    fn test_staleness_compensation() {
537        let compensation = StalenessCompensation::Linear;
538        assert!(
539            matches!(compensation, StalenessCompensation::Linear),
540            "Expected Linear staleness compensation"
541        );
542    }
543
544    #[test]
545    fn test_hardware_stats() {
546        let stats = HardwareStats {
547            memory_bandwidth_utilization: 0.8,
548            compute_utilization: 0.9,
549            cache_hit_rate: 0.95,
550            flops_per_second: 1e12,
551        };
552
553        assert_eq!(stats.memory_bandwidth_utilization, 0.8);
554        assert_eq!(stats.compute_utilization, 0.9);
555        assert_eq!(stats.cache_hit_rate, 0.95);
556        assert_eq!(stats.flops_per_second, 1e12);
557    }
558
559    #[test]
560    fn test_gpu_memory_stats() {
561        let stats = GPUMemoryStats {
562            total_memory: 16 * 1024 * 1024 * 1024,    // 16 GB
563            used_memory: 8 * 1024 * 1024 * 1024,      // 8 GB
564            available_memory: 8 * 1024 * 1024 * 1024, // 8 GB
565            optimizer_memory: 1024 * 1024 * 1024,     // 1 GB
566        };
567
568        assert_eq!(stats.total_memory, 16 * 1024 * 1024 * 1024);
569        assert_eq!(
570            stats.used_memory + stats.available_memory,
571            stats.total_memory
572        );
573        assert!(stats.optimizer_memory <= stats.used_memory);
574    }
575
576    #[test]
577    fn test_staleness_compensation_none() {
578        let comp = StalenessCompensation::None;
579        assert!(matches!(comp, StalenessCompensation::None));
580    }
581
582    #[test]
583    fn test_staleness_compensation_exponential() {
584        let comp = StalenessCompensation::Exponential;
585        assert!(matches!(comp, StalenessCompensation::Exponential));
586    }
587
588    #[test]
589    fn test_staleness_compensation_polynomial() {
590        let comp = StalenessCompensation::Polynomial(2.0);
591        if let StalenessCompensation::Polynomial(degree) = comp {
592            assert_eq!(degree, 2.0);
593        } else {
594            panic!("Expected Polynomial variant");
595        }
596    }
597
598    #[test]
599    fn test_hardware_stats_all_zero() {
600        let stats = HardwareStats {
601            memory_bandwidth_utilization: 0.0,
602            compute_utilization: 0.0,
603            cache_hit_rate: 0.0,
604            flops_per_second: 0.0,
605        };
606        assert_eq!(stats.memory_bandwidth_utilization, 0.0);
607        assert_eq!(stats.flops_per_second, 0.0);
608    }
609
610    #[test]
611    fn test_hardware_stats_max_utilization() {
612        let stats = HardwareStats {
613            memory_bandwidth_utilization: 1.0,
614            compute_utilization: 1.0,
615            cache_hit_rate: 1.0,
616            flops_per_second: 1e15,
617        };
618        assert!(stats.memory_bandwidth_utilization <= 1.0);
619        assert!(stats.compute_utilization <= 1.0);
620        assert!(stats.cache_hit_rate <= 1.0);
621    }
622
623    #[test]
624    fn test_gpu_memory_stats_zero_usage() {
625        let stats = GPUMemoryStats {
626            total_memory: 16 * 1024 * 1024 * 1024,
627            used_memory: 0,
628            available_memory: 16 * 1024 * 1024 * 1024,
629            optimizer_memory: 0,
630        };
631        assert_eq!(stats.used_memory, 0);
632        assert_eq!(stats.total_memory, stats.available_memory);
633    }
634
635    #[test]
636    fn test_gpu_memory_stats_full_usage() {
637        let total = 8 * 1024 * 1024 * 1024_usize;
638        let stats = GPUMemoryStats {
639            total_memory: total,
640            used_memory: total,
641            available_memory: 0,
642            optimizer_memory: total / 4,
643        };
644        assert_eq!(stats.available_memory, 0);
645        assert!(stats.optimizer_memory <= stats.used_memory);
646    }
647
648    #[test]
649    fn test_gpu_memory_stats_optimizer_fraction() {
650        let total = 16 * 1024 * 1024 * 1024_usize;
651        let used = 12 * 1024 * 1024 * 1024_usize;
652        let optimizer = 3 * 1024 * 1024 * 1024_usize;
653        let stats = GPUMemoryStats {
654            total_memory: total,
655            used_memory: used,
656            available_memory: total - used,
657            optimizer_memory: optimizer,
658        };
659        assert_eq!(stats.available_memory, 4 * 1024 * 1024 * 1024);
660        assert!(stats.optimizer_memory < stats.used_memory);
661    }
662
663    #[test]
664    fn test_hardware_stats_clone() {
665        let stats = HardwareStats {
666            memory_bandwidth_utilization: 0.5,
667            compute_utilization: 0.7,
668            cache_hit_rate: 0.9,
669            flops_per_second: 5e11,
670        };
671        let cloned = stats.clone();
672        assert_eq!(cloned.memory_bandwidth_utilization, 0.5);
673        assert_eq!(cloned.compute_utilization, 0.7);
674    }
675
676    #[test]
677    fn test_gpu_memory_stats_clone() {
678        let stats = GPUMemoryStats {
679            total_memory: 1000,
680            used_memory: 500,
681            available_memory: 500,
682            optimizer_memory: 100,
683        };
684        let cloned = stats.clone();
685        assert_eq!(cloned.total_memory, 1000);
686    }
687
688    #[test]
689    fn test_staleness_compensation_copy() {
690        let comp = StalenessCompensation::Linear;
691        let copied = comp;
692        assert!(matches!(copied, StalenessCompensation::Linear));
693    }
694
695    #[test]
696    fn test_hardware_stats_realistic_gpu() {
697        let stats = HardwareStats {
698            memory_bandwidth_utilization: 0.75,
699            compute_utilization: 0.85,
700            cache_hit_rate: 0.92,
701            flops_per_second: 1.2e13,
702        };
703        assert!(
704            stats.memory_bandwidth_utilization > 0.0 && stats.memory_bandwidth_utilization <= 1.0
705        );
706        assert!(stats.compute_utilization > 0.0 && stats.compute_utilization <= 1.0);
707        assert!(stats.flops_per_second > 1e12);
708    }
709
710    #[test]
711    fn test_hardware_stats_edge_device() {
712        let stats = HardwareStats {
713            memory_bandwidth_utilization: 0.3,
714            compute_utilization: 0.4,
715            cache_hit_rate: 0.6,
716            flops_per_second: 1e9,
717        };
718        assert!(stats.flops_per_second < 1e10);
719        assert!(stats.compute_utilization < 0.5);
720    }
721
722    #[test]
723    fn test_gpu_memory_stats_consistency() {
724        let stats = GPUMemoryStats {
725            total_memory: 8 * 1024 * 1024 * 1024,
726            used_memory: 6 * 1024 * 1024 * 1024,
727            available_memory: 2 * 1024 * 1024 * 1024,
728            optimizer_memory: 2 * 1024 * 1024 * 1024,
729        };
730        assert_eq!(
731            stats.used_memory + stats.available_memory,
732            stats.total_memory
733        );
734    }
735
736    #[test]
737    fn test_staleness_polynomial_fractional() {
738        let comp = StalenessCompensation::Polynomial(0.5);
739        if let StalenessCompensation::Polynomial(degree) = comp {
740            assert!(degree > 0.0 && degree < 1.0);
741        }
742    }
743
744    #[test]
745    fn test_staleness_polynomial_high_degree() {
746        let comp = StalenessCompensation::Polynomial(10.0);
747        if let StalenessCompensation::Polynomial(degree) = comp {
748            assert!(degree > 5.0);
749        }
750    }
751
752    #[test]
753    fn test_hardware_stats_debug_format() {
754        let stats = HardwareStats {
755            memory_bandwidth_utilization: 0.5,
756            compute_utilization: 0.5,
757            cache_hit_rate: 0.5,
758            flops_per_second: 1.0,
759        };
760        let debug_str = format!("{:?}", stats);
761        assert!(debug_str.contains("HardwareStats"));
762    }
763
764    #[test]
765    fn test_gpu_memory_stats_debug_format() {
766        let stats = GPUMemoryStats {
767            total_memory: 100,
768            used_memory: 50,
769            available_memory: 50,
770            optimizer_memory: 10,
771        };
772        let debug_str = format!("{:?}", stats);
773        assert!(debug_str.contains("GPUMemoryStats"));
774    }
775}
776
777#[cfg(test)]
778mod state_persistence_tests {
779    use super::*;
780    use crate::adam::Adam;
781
782    /// Regression: `StatefulOptimizer` declared `state_dict`/`load_state_dict` as
783    /// required methods with no file-level counterpart, so every implementor had to
784    /// hand-roll checkpointing. The trait now ships a default round trip.
785    #[test]
786    fn save_state_and_load_state_round_trip() {
787        let mut optimizer = Adam::new(0.01, (0.9, 0.999), 1e-8, 0.0);
788        let mut param = Tensor::from_vec(vec![1.0_f32, 2.0], &[2]).expect("tensor");
789        let grad = Tensor::from_vec(vec![0.5_f32, -0.5], &[2]).expect("grad");
790        optimizer.update_named("w", &mut param, &grad).expect("step 1");
791        Optimizer::step(&mut optimizer);
792        optimizer.update_named("w", &mut param, &grad).expect("step 2");
793
794        let path = std::env::temp_dir().join(format!(
795            "trustformers-optim-state-{}.bin",
796            std::process::id()
797        ));
798        optimizer.save_state(&path).expect("save_state");
799
800        let mut restored = Adam::new(0.5, (0.1, 0.1), 1e-2, 0.9);
801        restored.load_state(&path).expect("load_state");
802        let _ = std::fs::remove_file(&path);
803
804        let original = optimizer.state_dict().expect("state_dict");
805        let round_trip = restored.state_dict().expect("state_dict");
806        assert_eq!(original.len(), round_trip.len(), "every entry must survive");
807        for (key, tensor) in &original {
808            let other = round_trip.get(key).unwrap_or_else(|| panic!("missing '{key}'"));
809            assert_eq!(other.shape(), tensor.shape(), "shape of '{key}'");
810            assert_eq!(
811                other.data_f32().expect("data"),
812                tensor.data_f32().expect("data"),
813                "payload of '{key}'"
814            );
815        }
816    }
817
818    /// Corrupt bytes must be reported, not silently ignored.
819    #[test]
820    fn decoding_rejects_corrupt_state() {
821        assert!(decode_state_dict(&[0xff, 0x00, 0x13, 0x37]).is_err());
822    }
823
824    /// A payload whose length disagrees with its shape must be rejected.
825    #[test]
826    fn decoding_rejects_a_shape_payload_mismatch() {
827        let wire: Vec<(String, Vec<usize>, Vec<f32>)> =
828            vec![("w".to_string(), vec![4], vec![1.0, 2.0])];
829        let bytes =
830            oxicode::serde::encode_to_vec(&wire, oxicode::config::standard()).expect("encode");
831        assert!(decode_state_dict(&bytes).is_err());
832    }
833
834    /// Loading must fail loudly when the file does not exist.
835    #[test]
836    fn load_state_reports_a_missing_file() {
837        let mut optimizer = Adam::new(0.01, (0.9, 0.999), 1e-8, 0.0);
838        let path = std::env::temp_dir().join("trustformers-optim-definitely-absent.bin");
839        let _ = std::fs::remove_file(&path);
840        assert!(optimizer.load_state(&path).is_err());
841    }
842}
trustformers_optim/traits.rs

trustformers_optim/
traits.rs
