From 14d405f112bcfa0b5348a06662001c1e757e68bb Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 28 Feb 2026 20:58:04 +0100 Subject: [PATCH] refactor(ml): remove legacy base_model_memory_mb from BatchSizeConfig Single model_memory_mb field now serves as the sole source of truth for base model memory. Removes the redundant base_model_memory_mb field, legacy dual-path branches in calculate_optimal_batch_size() and max_safe_batch_size(), and the associated legacy test. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/gpu/memory_profile.rs | 1 - .../memory_optimization/auto_batch_size.rs | 91 +++---------------- crates/ml/src/trainers/tft/trainer.rs | 3 +- 3 files changed, 16 insertions(+), 79 deletions(-) diff --git a/crates/ml/src/gpu/memory_profile.rs b/crates/ml/src/gpu/memory_profile.rs index d792d1d4f..175151fa3 100644 --- a/crates/ml/src/gpu/memory_profile.rs +++ b/crates/ml/src/gpu/memory_profile.rs @@ -167,7 +167,6 @@ pub fn resolve_batch_size( let batch_config = BatchSizeConfig { model_memory_mb: model_size_mb, model_precision: ModelPrecision::FP32, - base_model_memory_mb: model_size_mb, sequence_length: estimate.default_seq_len, feature_dim: estimate.default_feature_dim, gradient_checkpointing: estimate.supports_checkpointing, diff --git a/crates/ml/src/memory_optimization/auto_batch_size.rs b/crates/ml/src/memory_optimization/auto_batch_size.rs index 58c161332..2b46c4a3f 100644 --- a/crates/ml/src/memory_optimization/auto_batch_size.rs +++ b/crates/ml/src/memory_optimization/auto_batch_size.rs @@ -100,24 +100,20 @@ impl ModelPrecision { /// Batch size configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BatchSizeConfig { - /// Model memory in MB (parameters only) - /// DEPRECATED: Use model_precision and base_model_memory_mb instead + /// Base model memory in MB (scaled by precision multiplier for actual usage). + /// For TFT-225: ~125MB (INT8) → ~500MB (FP32). pub model_memory_mb: f64, - /// Model precision (FP32 or INT8) + /// Model precision (FP32, INT8, or QAT) pub model_precision: ModelPrecision, - /// Base model memory in MB for INT8 (will be scaled by precision multiplier) - /// For TFT-225: ~125MB (INT8) -> ~500MB (FP32) - pub base_model_memory_mb: f64, - /// Sequence length (lookback window) pub sequence_length: usize, /// Feature dimension (number of input features) pub feature_dim: usize, - /// Enable gradient checkpointing (reduces activation memory by ~50%) + /// Enable gradient checkpointing (reduces activation memory by ~35%) pub gradient_checkpointing: bool, /// Optimizer type (affects memory overhead) @@ -136,9 +132,8 @@ pub struct BatchSizeConfig { impl Default for BatchSizeConfig { fn default() -> Self { Self { - model_memory_mb: 125.0, // TFT default (INT8) + model_memory_mb: 125.0, model_precision: ModelPrecision::INT8, - base_model_memory_mb: 125.0, // TFT-225 base size (INT8) sequence_length: 60, feature_dim: 225, gradient_checkpointing: false, @@ -219,28 +214,8 @@ impl AutoBatchSizer { usable_memory_mb ); - // Calculate actual model memory based on precision - // Use new precision-aware fields if available, fall back to model_memory_mb for backward compatibility - let model_mb = if config.base_model_memory_mb.abs() < 0.01 { - // Legacy path: base_model_memory_mb is zero/unset, use model_memory_mb directly - debug!( - "Model memory (legacy): {:.1}MB (using model_memory_mb field)", - config.model_memory_mb - ); - config.model_memory_mb - } else { - // New path: Scale base memory by precision multiplier - let scaled_mb = - config.base_model_memory_mb * config.model_precision.memory_multiplier(); - debug!( - "Model memory (precision-aware): {:.1}MB (base: {:.1}MB × {:.1}x for {:?})", - scaled_mb, - config.base_model_memory_mb, - config.model_precision.memory_multiplier(), - config.model_precision - ); - scaled_mb - }; + // Scale base memory by precision multiplier + let model_mb = config.model_memory_mb * config.model_precision.memory_multiplier(); // Calculate fixed memory overhead (model + optimizer states) let optimizer_mb = model_mb * config.optimizer_type.memory_multiplier(); @@ -358,11 +333,7 @@ impl AutoBatchSizer { .max(config.safety_margin); let usable_memory_mb = self.free_memory_mb * (1.0 - effective_safety_margin); - let model_mb = if config.base_model_memory_mb.abs() < 0.01 { - config.model_memory_mb - } else { - config.base_model_memory_mb * config.model_precision.memory_multiplier() - }; + let model_mb = config.model_memory_mb * config.model_precision.memory_multiplier(); let optimizer_mb = model_mb * config.optimizer_type.memory_multiplier(); let gradient_mb = model_mb; @@ -527,7 +498,7 @@ mod tests { fn test_batch_size_config_default() { let config = BatchSizeConfig::default(); assert_eq!(config.model_memory_mb, 125.0); - assert_eq!(config.base_model_memory_mb, 125.0); + assert_eq!(config.model_memory_mb, 125.0); assert_eq!(config.model_precision, ModelPrecision::INT8); assert_eq!(config.sequence_length, 60); assert_eq!(config.feature_dim, 225); @@ -544,7 +515,7 @@ mod tests { let config = BatchSizeConfig { model_memory_mb: 125.0, // TFT model (INT8) model_precision: ModelPrecision::INT8, - base_model_memory_mb: 125.0, + sequence_length: 60, feature_dim: 225, gradient_checkpointing: false, @@ -583,7 +554,7 @@ mod tests { let config = BatchSizeConfig { model_memory_mb: 125.0, model_precision: ModelPrecision::INT8, - base_model_memory_mb: 125.0, + sequence_length: 60, feature_dim: 225, gradient_checkpointing: false, @@ -692,7 +663,7 @@ mod tests { let config_fp32 = BatchSizeConfig { model_memory_mb: 80.0, // Legacy field (backward compat) model_precision: ModelPrecision::FP32, - base_model_memory_mb: 80.0, // Small base for testing (scaled to 320MB) + sequence_length: 60, feature_dim: 150, // Fewer features for testing gradient_checkpointing: true, // Enable to fit in memory @@ -706,7 +677,7 @@ mod tests { let config_int8 = BatchSizeConfig { model_memory_mb: 125.0, model_precision: ModelPrecision::INT8, - base_model_memory_mb: 125.0, + sequence_length: 60, feature_dim: 225, gradient_checkpointing: false, @@ -768,7 +739,7 @@ mod tests { let config_fp32 = BatchSizeConfig { model_memory_mb: 125.0, model_precision: ModelPrecision::FP32, - base_model_memory_mb: 125.0, + sequence_length: 60, feature_dim: 225, gradient_checkpointing: false, @@ -804,7 +775,7 @@ mod tests { let config_int8 = BatchSizeConfig { model_memory_mb: 125.0, model_precision: ModelPrecision::INT8, - base_model_memory_mb: 125.0, + sequence_length: 60, feature_dim: 225, gradient_checkpointing: false, @@ -824,38 +795,6 @@ mod tests { ); } - #[test] - fn test_legacy_model_memory_mb_still_works() { - // RTX 3050 Ti: 4GB total, ~3.2GB free (realistic) - let sizer = AutoBatchSizer::with_manual_memory(4096.0, 3200.0, "RTX 3050 Ti".to_owned()); - - // Legacy config (no precision fields, only model_memory_mb) - let config_legacy = BatchSizeConfig { - model_memory_mb: 500.0, // Manually specify FP32 size - model_precision: ModelPrecision::INT8, // Ignored when base=0 - base_model_memory_mb: 0.0, // Zero triggers legacy path - sequence_length: 60, - feature_dim: 225, - gradient_checkpointing: true, // Enable to fit in memory - optimizer_type: OptimizerType::Adam, - safety_margin: 0.20, - min_batch_size: 1, - max_batch_size: 64, - }; - - let batch_size = sizer.calculate_optimal_batch_size(&config_legacy).unwrap(); - - // Should get small batch size similar to FP32 (since model_memory_mb=500) - // With 3200MB free, 20% margin = 2560MB usable: - // 500 (model) + 1000 (optimizer) + 500 (gradients) + 250 (activations w/ checkpointing) = 2250MB fixed overhead - // → 310MB for batches → batch_size ~8-32 (with gradient checkpointing) - assert!( - batch_size >= 1 && batch_size <= 64, - "Legacy config with 500MB model should get batch_size 1-64, got {}", - batch_size - ); - } - #[test] fn test_reduce_batch_size() { // Test exponential backoff: 64 → 32 → 16 → 8 → 4 → 2 → 1 diff --git a/crates/ml/src/trainers/tft/trainer.rs b/crates/ml/src/trainers/tft/trainer.rs index dcbff72df..1cab17bfb 100644 --- a/crates/ml/src/trainers/tft/trainer.rs +++ b/crates/ml/src/trainers/tft/trainer.rs @@ -229,9 +229,8 @@ impl TFTTrainer { }; let batch_config = BatchSizeConfig { - model_memory_mb: base_model_memory_mb, // Legacy field (for backward compatibility) + model_memory_mb: base_model_memory_mb, model_precision, - base_model_memory_mb, sequence_length: config.lookback_window, feature_dim: 225, // Wave C (201) + Wave D (24) gradient_checkpointing: config.use_gradient_checkpointing,