gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(ml): remove legacy base_model_memory_mb from BatchSizeConfig

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-28 20:58:04 +01:00
parent 6804f0facf
commit 14d405f112
3 changed files with 16 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml/src/gpu/memory_profile.rs
View File

@@ -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,

91
crates/ml/src/memory_optimization/auto_batch_size.rs
View File

@@ -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

3
crates/ml/src/trainers/tft/trainer.rs
View File

@@ -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,