feat(ml): comprehensive GPU saturation audit — 58 fixes across all 10 models
Phase 1 — Fix broken models (P0): - Diffusion: wire optimizer_step to actually apply gradients (was no-op) - TLOB: connect forward pass to projection layers (was Tensor::zeros) - Mamba2: F64→F32 migration across 5 files (~30x faster on L40S tensor cores) Phase 2 — Eliminate hot-path GPU sync stalls: - Mamba2: keep dt on GPU in discretize_ssm (4 functions, no CPU round-trip) - TFT: gate attention weight logging to eval only (8 syncs/forward eliminated) - Mamba2: defer loss scalar after backward (pipeline stall removed) - Mamba2: delete dead gradient clipping (4N wasted GPU syncs removed) Phase 3 — Enable BF16 for supervised models: - Flip mixed_precision defaults to true in 4 config locations - Fix cuda_layer_norm to support BF16/F16 via F32 intermediate Phase 4 — Raise hyperopt bounds for datacenter GPUs: - 7 adapters with VRAM-aware tiers (TFT, Liquid, TGGN, KAN, xLSTM, Diffusion, TLOB) — L40S gets full hidden_dim range - Fix L40S tier boundary (was excluded at <48000, now >=40000) Phase 5 — Update memory estimates: - 10 param_count estimates updated (DQN 200K→12M, TFT 2M→50M, etc.) - Fix power-of-two rounding (was wasting up to 49% of budget) - Correct MODEL_OVERHEAD_MB in DQN/PPO/TFT adapters Phase 6 — Fix per-epoch CPU bottlenecks: - PPO: deduplicate double advantage normalization (correctness fix) - PPO: GPU tensor reward normalization + explained variance - Fuse per-parameter grad norm to single GPU sync (xLSTM, KAN, TGGN) Phase 7 — Data pipeline: - GpuBufferPool: use from_slice (eliminate staging buffer copy) Phase 8 — Correctness: - TFT: remove broken .detach() in forward_checkpointed (restore gradients) - Update stale RTX 3050 Ti doc references 33 files changed, 2451 tests pass, 0 clippy warnings. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -29,6 +29,8 @@ pub struct DiffusionTrainableAdapter {
|
||||
learning_rate: f64,
|
||||
loss_history: Vec<f64>,
|
||||
config: DiffusionConfig,
|
||||
/// Gradients from the last backward pass, consumed by optimizer_step.
|
||||
last_grads: Option<candle_core::backprop::GradStore>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for DiffusionTrainableAdapter {
|
||||
@@ -84,6 +86,7 @@ impl DiffusionTrainableAdapter {
|
||||
learning_rate: lr,
|
||||
loss_history: Vec::new(),
|
||||
config,
|
||||
last_grads: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -155,10 +158,18 @@ impl UnifiedTrainable for DiffusionTrainableAdapter {
|
||||
|
||||
let loss_val = loss.to_scalar::<f32>().unwrap_or(f32::NAN) as f64;
|
||||
self.loss_history.push(loss_val);
|
||||
|
||||
// Store gradients for optimizer_step to consume
|
||||
self.last_grads = Some(grads);
|
||||
|
||||
Ok(total_norm.sqrt())
|
||||
}
|
||||
|
||||
fn optimizer_step(&mut self) -> Result<(), MLError> {
|
||||
if let Some(grads) = self.last_grads.take() {
|
||||
self.optimizer.step(&grads)
|
||||
.map_err(|e| MLError::ModelError(e.to_string()))?;
|
||||
}
|
||||
self.step += 1;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user