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:
jgrusewski
2026-03-02 12:38:02 +01:00
parent b4dc9766f9
commit 2aedc2ae1a
33 changed files with 605 additions and 421 deletions

View File

@@ -23,6 +23,15 @@ pub struct StartTrainingRequest {
::prost::alloc::string::String,
::prost::alloc::string::String,
>,
/// FULL (default) or FINE_TUNE
#[prost(enumeration = "TrainingMode", tag = "7")]
pub mode: i32,
/// Path to checkpoint for fine-tune
#[prost(string, tag = "8")]
pub resume_checkpoint_path: ::prost::alloc::string::String,
/// Override epoch count (0 = use default)
#[prost(uint32, tag = "9")]
pub max_epochs: u32,
}
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
pub struct StartTrainingResponse {
@@ -880,6 +889,35 @@ pub struct RejectPromotionResponse {
#[prost(string, tag = "2")]
pub message: ::prost::alloc::string::String,
}
/// Training mode selection
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
#[repr(i32)]
pub enum TrainingMode {
/// Full training from scratch (default, backward-compatible)
Full = 0,
/// Fine-tune from existing checkpoint
FineTune = 1,
}
impl TrainingMode {
/// String value of the enum field names used in the ProtoBuf definition.
///
/// The values are not transformed in any way and thus are considered stable
/// (if the ProtoBuf definition does not change) and safe for programmatic use.
pub fn as_str_name(&self) -> &'static str {
match self {
Self::Full => "TRAINING_MODE_FULL",
Self::FineTune => "TRAINING_MODE_FINE_TUNE",
}
}
/// Creates an enum from field names used in the ProtoBuf definition.
pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
match value {
"TRAINING_MODE_FULL" => Some(Self::Full),
"TRAINING_MODE_FINE_TUNE" => Some(Self::FineTune),
_ => None,
}
}
}
/// Type of progress update
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
#[repr(i32)]