entrenar/hf_pipeline/fine_tune/
memory.rs

1//! Memory estimation and requirements for fine-tuning
2//!
3//! Provides utilities for estimating memory usage during training.
4
5/// Mixed precision training options
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum MixedPrecision {
8    /// FP16 mixed precision
9    Fp16,
10    /// BF16 mixed precision (better for training)
11    Bf16,
12}
13
14/// Memory requirements breakdown
15#[derive(Debug, Clone, Copy)]
16pub struct MemoryRequirement {
17    /// Model weights
18    pub model: u64,
19    /// Optimizer states
20    pub optimizer: u64,
21    /// Gradients
22    pub gradients: u64,
23    /// Activations
24    pub activations: u64,
25}
26
27impl MemoryRequirement {
28    /// Total memory required
29    #[must_use]
30    pub fn total(&self) -> u64 {
31        self.model + self.optimizer + self.gradients + self.activations
32    }
33
34    /// Check if fits in available memory
35    #[must_use]
36    pub fn fits_in(&self, available: u64) -> bool {
37        self.total() <= available
38    }
39
40    /// Memory savings compared to full fine-tuning
41    #[must_use]
42    pub fn savings_vs_full(&self, full_params: u64) -> f32 {
43        let full_memory = full_params * 4 + full_params * 4 * 2 + full_params * 4;
44        1.0 - (self.total() as f32 / full_memory as f32)
45    }
46
47    /// Format as human-readable string
48    #[must_use]
49    pub fn format_human(&self) -> String {
50        format!(
51            "Model: {:.1}GB, Optimizer: {:.1}GB, Gradients: {:.1}GB, Activations: {:.1}GB, Total: {:.1}GB",
52            self.model as f64 / 1e9,
53            self.optimizer as f64 / 1e9,
54            self.gradients as f64 / 1e9,
55            self.activations as f64 / 1e9,
56            self.total() as f64 / 1e9
57        )
58    }
59}
entrenar/hf_pipeline/fine_tune/memory.rs

entrenar/hf_pipeline/fine_tune/
memory.rs
