diff --git a/ml/src/trainers/tft.rs b/ml/src/trainers/tft.rs index e60d29db5..2b55e9ed3 100644 --- a/ml/src/trainers/tft.rs +++ b/ml/src/trainers/tft.rs @@ -945,7 +945,7 @@ impl TFTTrainer { // Apply QAT-specific learning rate schedule (if enabled) if self.use_qat { - self.apply_qat_lr_schedule(epoch); + self.apply_qat_lr_schedule(epoch)?; } // Training phase with OOM retry logic (note: data loader recreation not yet supported) @@ -1110,8 +1110,21 @@ impl TFTTrainer { } } + // Drop optimizer to free 1100MB AdamW state (momentum + velocity buffers) + // This prevents OOM during validation by freeing GPU memory + let optimizer_backup = self.optimizer.take(); + info!( + "[MEMORY] Dropped optimizer before validation to free ~1100MB AdamW state" + ); + let result = self.validate_epoch(&mut val_loader, epoch).await?; + // Restore optimizer for next training epoch + self.optimizer = optimizer_backup; + info!( + "[MEMORY] Restored optimizer after validation" + ); + // Memory profiling: Log GPU memory after validation #[cfg(feature = "cuda")] if self.device.is_cuda() { @@ -2286,7 +2299,7 @@ impl TFTTrainer { /// Epoch 90: 0.1 * base_lr (cooldown start) /// Epoch 99: 0.1 * base_lr (cooldown end) /// ``` - fn apply_qat_lr_schedule(&mut self, epoch: usize) { + fn apply_qat_lr_schedule(&mut self, epoch: usize) -> MLResult<()> { let total_epochs = self.training_config.epochs; let base_lr = self.training_config.learning_rate; @@ -2310,17 +2323,34 @@ impl TFTTrainer { self.state.learning_rate = new_lr; // Apply to optimizer (if initialized) - if let Some(ref mut _opt) = self.optimizer { - // Update optimizer learning rate - // Note: candle_nn::AdamW doesn't have a direct set_learning_rate method - // In practice, we recreate the optimizer with new LR or use parameter groups - debug!( - "QAT LR Schedule - Epoch {}: {:.2e} (warmup: {}, cooldown: {})", - epoch, - new_lr, - epoch < self.qat_warmup_epochs, - epoch >= cooldown_start_epoch - ); + // Check if LR actually changed before recreating optimizer + if let Some(ref opt) = self.optimizer { + let current_lr = opt.learning_rate(); + + // Only recreate optimizer if LR changed by more than epsilon (1e-10) + if (current_lr - new_lr).abs() > 1e-10 { + info!( + "🔄 QAT LR Schedule - Recreating optimizer: {:.2e} → {:.2e} (epoch {})", + current_lr, new_lr, epoch + ); + + // Drop old optimizer to free memory (~1100MB) + drop(self.optimizer.take()); + + // Update config with new LR + self.training_config.learning_rate = new_lr; + + // Recreate optimizer with new LR (allocates ~1100MB) + self.initialize_optimizer()?; + } else { + debug!( + "QAT LR Schedule - Epoch {}: {:.2e} (unchanged, warmup: {}, cooldown: {})", + epoch, + new_lr, + epoch < self.qat_warmup_epochs, + epoch >= cooldown_start_epoch + ); + } } // Log major phase transitions @@ -2331,6 +2361,8 @@ impl TFTTrainer { } else if epoch == cooldown_start_epoch { info!("🔽 QAT Cooldown Phase: Reducing LR to {:.2e} ({:.1}x reduction) at epoch {}", new_lr, self.qat_cooldown_factor, epoch); } + + Ok(()) } } @@ -2538,14 +2570,14 @@ mod tests { let mut trainer = TFTTrainer::new(config, storage).expect("Failed to create trainer"); // Test warmup phase - trainer.apply_qat_lr_schedule(0); + trainer.apply_qat_lr_schedule(0).expect("Failed to apply LR schedule"); assert!( (trainer.state.learning_rate - 1e-4).abs() < 1e-9, "Epoch 0: Expected 1e-4 (10% of 1e-3), got {}", trainer.state.learning_rate ); - trainer.apply_qat_lr_schedule(5); + trainer.apply_qat_lr_schedule(5).expect("Failed to apply LR schedule"); let expected_mid_warmup = 1e-3 * 0.55; // 55% progress assert!( (trainer.state.learning_rate - expected_mid_warmup).abs() < 1e-9, @@ -2554,7 +2586,7 @@ mod tests { trainer.state.learning_rate ); - trainer.apply_qat_lr_schedule(10); + trainer.apply_qat_lr_schedule(10).expect("Failed to apply LR schedule"); assert!( (trainer.state.learning_rate - 1e-3).abs() < 1e-9, "Epoch 10: Expected 1e-3 (full LR), got {}", @@ -2562,7 +2594,7 @@ mod tests { ); // Test normal training phase - trainer.apply_qat_lr_schedule(50); + trainer.apply_qat_lr_schedule(50).expect("Failed to apply LR schedule"); assert!( (trainer.state.learning_rate - 1e-3).abs() < 1e-9, "Epoch 50: Expected 1e-3 (full LR), got {}", @@ -2570,7 +2602,7 @@ mod tests { ); // Test cooldown phase (starts at epoch 90 for 100 total epochs) - trainer.apply_qat_lr_schedule(90); + trainer.apply_qat_lr_schedule(90).expect("Failed to apply LR schedule"); assert!( (trainer.state.learning_rate - 1e-4).abs() < 1e-9, "Epoch 90: Expected 1e-4 (10% of 1e-3), got {}",