From 864808e0568c4264c44c2529d64fbf73b911e6bc Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 3 Mar 2026 17:42:41 +0100 Subject: [PATCH] fix(ml): add ensure_training_dtype boundary casts to all model forward methods After converting all VarBuilder sites from F32 to training_dtype (BF16 on Ampere+ GPUs), input tensors from callers remain F32, causing dtype mismatches in matmul/add/mul operations. This commit adds systematic boundary casts across all 10 model architectures: - Input boundary: ensure_training_dtype() at each model forward() entry - Output boundary: to_dtype(F32) at each model forward() exit - Internal intermediates: hidden state init, gradient extraction, positional encodings, causal masks, SSM state matrices, B-spline basis values all cast to match computation dtype - Ensemble adapters: ensure_training_dtype after Tensor::from_vec 36 files, +293/-72 lines. 2640 tests pass, 0 failures, 0 clippy warnings. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/diffusion/denoiser.rs | 12 ++++++-- crates/ml/src/diffusion/trainable.rs | 2 ++ crates/ml/src/dqn/distributional_dueling.rs | 11 ++++++-- crates/ml/src/dqn/dqn.rs | 22 ++++++++++++--- crates/ml/src/dqn/dueling.rs | 11 ++++++-- crates/ml/src/dqn/mixed_precision.rs | 17 +++++++++++- crates/ml/src/dqn/network.rs | 5 ++++ crates/ml/src/dqn/quantile_regression.rs | 24 +++++++++++----- crates/ml/src/ensemble/adapters/diffusion.rs | 4 +++ crates/ml/src/ensemble/adapters/dqn.rs | 5 +++- crates/ml/src/ensemble/adapters/kan.rs | 4 +++ crates/ml/src/ensemble/adapters/mamba2.rs | 10 +++---- crates/ml/src/ensemble/adapters/tft.rs | 8 +++++- crates/ml/src/ensemble/adapters/tggn.rs | 10 +++++-- crates/ml/src/ensemble/adapters/tlob.rs | 10 +++++-- crates/ml/src/ensemble/adapters/xlstm.rs | 4 +++ crates/ml/src/gradient_utils.rs | 2 +- crates/ml/src/kan/layer.rs | 14 +++++++++- crates/ml/src/kan/network.rs | 6 +++- crates/ml/src/kan/spline.rs | 8 ++++++ crates/ml/src/kan/trainable.rs | 4 ++- crates/ml/src/liquid/candle_cfc.rs | 13 +++++++-- crates/ml/src/liquid/training.rs | 3 +- crates/ml/src/mamba/mod.rs | 29 ++++++++++++++++---- crates/ml/src/mamba/scan_algorithms.rs | 8 ++++-- crates/ml/src/ppo/lstm_networks.rs | 24 ++++++++++++++-- crates/ml/src/ppo/ppo.rs | 18 ++++++++---- crates/ml/src/tft/mod.rs | 18 +++++++++--- crates/ml/src/tft/temporal_attention.rs | 5 ++++ crates/ml/src/tft/variable_selection.rs | 2 +- crates/ml/src/tgnn/trainable_adapter.rs | 9 +++++- crates/ml/src/tlob/trainable_adapter.rs | 7 +++++ crates/ml/src/xlstm/mlstm.rs | 7 +++-- crates/ml/src/xlstm/network.rs | 17 +++++++++--- crates/ml/src/xlstm/slstm.rs | 8 ++++-- crates/ml/src/xlstm/trainable.rs | 4 ++- 36 files changed, 293 insertions(+), 72 deletions(-) diff --git a/crates/ml/src/diffusion/denoiser.rs b/crates/ml/src/diffusion/denoiser.rs index d59dd20c8..db87e5012 100644 --- a/crates/ml/src/diffusion/denoiser.rs +++ b/crates/ml/src/diffusion/denoiser.rs @@ -63,6 +63,10 @@ impl TimeEmbedding { let emb = Tensor::cat(&[&sin_emb, &cos_emb], 1) .map_err(|e| MLError::ModelError(e.to_string()))?; + // Cast to training dtype before projection through BF16 weights + let emb = crate::dqn::mixed_precision::ensure_training_dtype(&emb) + .map_err(|e| MLError::ModelError(e.to_string()))?; + // Project to hidden_dim self.proj.forward(&emb) .map_err(|e| MLError::ModelError(e.to_string())) @@ -201,13 +205,15 @@ impl Denoiser { /// /// Input x: (batch, data_dim), t: (batch,) → Output: (batch, data_dim) pub fn forward(&self, x: &Tensor, t: &Tensor) -> Result { + let x = crate::dqn::mixed_precision::ensure_training_dtype(x) + .map_err(|e| MLError::ModelError(e.to_string()))?; let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string()); // Time embedding let t_emb = self.time_embed.forward(t, &self.device)?; // Input projection → SiLU - let mut h = self.input_proj.forward(x).map_err(map_err)?; + let mut h = self.input_proj.forward(&x).map_err(map_err)?; let h_sig = candle_nn::ops::sigmoid(&h).map_err(map_err)?; h = h.mul(&h_sig).map_err(map_err)?; @@ -217,7 +223,9 @@ impl Denoiser { } // Output projection → predicted noise - self.output_proj.forward(&h).map_err(map_err) + let output = self.output_proj.forward(&h).map_err(map_err)?; + // Cast output back to F32 for API compatibility + output.to_dtype(candle_core::DType::F32).map_err(map_err) } } diff --git a/crates/ml/src/diffusion/trainable.rs b/crates/ml/src/diffusion/trainable.rs index 90f753e85..b5c61e5ea 100644 --- a/crates/ml/src/diffusion/trainable.rs +++ b/crates/ml/src/diffusion/trainable.rs @@ -113,6 +113,8 @@ impl UnifiedTrainable for DiffusionTrainableAdapter { /// Input can be (batch, seq_len * feature_dim) or (batch, seq_len, feature_dim). /// Returns predicted noise with same shape as flattened input. fn forward(&mut self, input: &Tensor) -> Result { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string()); // Flatten to (batch, data_dim) diff --git a/crates/ml/src/dqn/distributional_dueling.rs b/crates/ml/src/dqn/distributional_dueling.rs index 95f88eab6..0b1f7bbc4 100644 --- a/crates/ml/src/dqn/distributional_dueling.rs +++ b/crates/ml/src/dqn/distributional_dueling.rs @@ -39,7 +39,7 @@ //! - Dueling Network Architectures (Wang et al., 2016) //! - A Distributional Perspective on Reinforcement Learning (Bellemare et al., 2017) -use candle_core::{DType, Device, Tensor}; +use candle_core::{Device, Tensor}; use candle_nn::{Linear, Module, VarBuilder, VarMap}; use serde::{Deserialize, Serialize}; @@ -228,6 +228,8 @@ impl DistributionalDuelingQNetwork { /// - A(s,a,z_i): Advantage distribution per action (probability of atom z_i) /// - mean(A(s,·,z_i)): Mean advantage across actions (ensures identifiability) pub fn forward(&self, state: &Tensor) -> Result { + let state = crate::dqn::mixed_precision::ensure_training_dtype(state) + .map_err(|e| MLError::ModelError(e.to_string()))?; let batch_size = state .dim(0) .map_err(|e| MLError::ModelError(format!("Failed to get batch size: {}", e)))?; @@ -325,6 +327,11 @@ impl DistributionalDuelingQNetwork { MLError::ModelError(format!("Softmax over atoms failed: {}", e)) })?; + // Cast output back to F32 for API compatibility + let z_probs = z_probs.to_dtype(candle_core::DType::F32).map_err(|e| { + MLError::ModelError(format!("Output dtype cast failed: {}", e)) + })?; + Ok(z_probs) } @@ -375,7 +382,7 @@ impl DistributionalDuelingQNetwork { #[cfg(test)] mod tests { use super::*; - use candle_core::Device; + use candle_core::{DType, Device}; #[test] fn test_distributional_dueling_creation() -> anyhow::Result<()> { diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index 346c84d13..4841ae4ea 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -723,6 +723,8 @@ impl Sequential { /// Forward pass through network pub fn forward(&self, input: &Tensor) -> Result { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; let mut x = input.clone(); if self.use_noisy_nets { @@ -755,6 +757,10 @@ impl Sequential { } } + // Cast output back to F32 for API compatibility + let x = x.to_dtype(DType::F32) + .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?; + Ok(x) } @@ -1103,7 +1109,9 @@ impl DQN { /// 2. Dueling only /// 3. Standard Q-network (fallback) pub fn forward(&self, state: &Tensor) -> Result { - // Auto-convert input to correct device (Candle optimizes if already on correct device) + // Auto-convert input to correct device and dtype + let state = crate::dqn::mixed_precision::ensure_training_dtype(state) + .map_err(|e| MLError::ModelError(e.to_string()))?; let state = state .to_device(&self.device) .map_err(|e| MLError::ModelError(format!("Failed to move tensor to device: {}", e)))?; @@ -1129,10 +1137,11 @@ impl DQN { let num_actions = self.config.num_actions; // Broadcast atoms to [batch, num_actions, num_atoms] + // Note: atoms stay F32 to match z_probs output (already cast to F32 by sub-networks) let atoms_tensor = Tensor::from_vec(atoms, num_atoms, &self.device)? - .unsqueeze(0)? - .unsqueeze(0)? - .broadcast_as((batch_size, num_actions, num_atoms))?; + .unsqueeze(0)? + .unsqueeze(0)? + .broadcast_as((batch_size, num_actions, num_atoms))?; // Compute expected Q-values: Q(s,a) = Σ(z_i * p_i) let q_values_dist = (z_probs * atoms_tensor)?.sum(2)?; // Sum over atoms @@ -1215,6 +1224,9 @@ impl DQN { q_values }; + // Cast output back to F32 for API compatibility (callers expect f32 tensors) + let q_values = q_values.to_dtype(candle_core::DType::F32)?; + Ok(q_values) } @@ -1424,6 +1436,8 @@ impl DQN { /// the intermediate representation needed by IQN. fn get_state_embedding(&self, states: &Tensor) -> Result { let states = states.to_device(&self.device)?; + let states = crate::dqn::mixed_precision::ensure_training_dtype(&states) + .map_err(|e| MLError::ModelError(e.to_string()))?; if self.q_network.use_noisy_nets { let mut x = states; diff --git a/crates/ml/src/dqn/dueling.rs b/crates/ml/src/dqn/dueling.rs index fe08c228e..0e629960c 100644 --- a/crates/ml/src/dqn/dueling.rs +++ b/crates/ml/src/dqn/dueling.rs @@ -30,7 +30,7 @@ //! - A(s,a): Advantage function (per-action) //! - mean(A(s,·)): Average advantage across all actions (ensures zero mean) -use candle_core::{DType, Device, Tensor}; +use candle_core::{Device, Tensor}; use candle_nn::{Linear, Module, VarBuilder, VarMap}; use serde::{Deserialize, Serialize}; @@ -206,6 +206,8 @@ impl DuelingQNetwork { /// - A(s,a): Advantage per action /// - mean(A(s,·)): Mean advantage (ensures identifiability) pub fn forward(&self, state: &Tensor) -> Result { + let state = crate::dqn::mixed_precision::ensure_training_dtype(state) + .map_err(|e| MLError::ModelError(e.to_string()))?; // Shared feature extraction let mut h = state.clone(); for (i, layer) in self.shared_layers.iter().enumerate() { @@ -276,6 +278,11 @@ impl DuelingQNetwork { MLError::ModelError(format!("Q-value combination failed: {}", e)) })?; + // Cast output back to F32 for API compatibility + let q_values = q_values.to_dtype(candle_core::DType::F32).map_err(|e| { + MLError::ModelError(format!("Output dtype cast failed: {}", e)) + })?; + Ok(q_values) } @@ -323,7 +330,7 @@ impl DuelingQNetwork { #[cfg(test)] mod tests { use super::*; - use candle_core::Device; + use candle_core::{DType, Device}; #[test] fn test_dueling_network_creation() -> anyhow::Result<()> { diff --git a/crates/ml/src/dqn/mixed_precision.rs b/crates/ml/src/dqn/mixed_precision.rs index 774ce020c..1a8ee31f5 100644 --- a/crates/ml/src/dqn/mixed_precision.rs +++ b/crates/ml/src/dqn/mixed_precision.rs @@ -215,7 +215,22 @@ pub fn training_dtype(device: &candle_core::Device) -> candle_core::DType { candle_core::DType::F32 } } - _ => candle_core::DType::F32, + candle_core::Device::Cpu | candle_core::Device::Metal(_) => candle_core::DType::F32, + } +} + +/// Cast a tensor to the training dtype for its device if needed. +/// +/// This is the canonical "boundary cast" — call it at any model entry +/// point where the input tensor may arrive as F32 from external code +/// (tests, ensemble adapters, raw inference). When dtypes already +/// match the call is essentially free (Candle returns a clone). +pub fn ensure_training_dtype(tensor: &candle_core::Tensor) -> Result { + let target = training_dtype(tensor.device()); + if tensor.dtype() == target { + Ok(tensor.clone()) + } else { + tensor.to_dtype(target) } } diff --git a/crates/ml/src/dqn/network.rs b/crates/ml/src/dqn/network.rs index 5082d62f9..00eebd5de 100644 --- a/crates/ml/src/dqn/network.rs +++ b/crates/ml/src/dqn/network.rs @@ -190,6 +190,7 @@ impl NetworkLayers { xs: &Tensor, mixed_precision: &Option, ) -> CandleResult { + let xs = crate::dqn::mixed_precision::ensure_training_dtype(xs)?; let (x_input, use_amp) = match mixed_precision { Some(mp) if mp.enabled => { let target_dtype = mp.dtype.to_dtype(); @@ -222,6 +223,7 @@ impl NetworkLayers { impl Module for NetworkLayers { fn forward(&self, xs: &Tensor) -> CandleResult { + let xs = crate::dqn::mixed_precision::ensure_training_dtype(xs)?; let mut x = xs.clone(); // Forward through hidden layers with LeakyReLU activation and dropout @@ -236,6 +238,9 @@ impl Module for NetworkLayers { } } + // Cast output back to F32 for API compatibility + x = x.to_dtype(DType::F32)?; + Ok(x) } } diff --git a/crates/ml/src/dqn/quantile_regression.rs b/crates/ml/src/dqn/quantile_regression.rs index 1d1f86da4..fb2772434 100644 --- a/crates/ml/src/dqn/quantile_regression.rs +++ b/crates/ml/src/dqn/quantile_regression.rs @@ -144,12 +144,16 @@ impl QuantileNetwork { /// # Returns /// Quantile values [batch, num_actions, num_quantiles] pub fn forward(&self, state_embed: &Tensor, taus: &Tensor) -> Result { + let state_embed = crate::dqn::mixed_precision::ensure_training_dtype(state_embed) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let taus = crate::dqn::mixed_precision::ensure_training_dtype(taus) + .map_err(|e| MLError::ModelError(e.to_string()))?; let batch_size = state_embed.dim(0)?; let num_quantiles = taus.dim(1)?; let num_actions = self.config.num_actions; // 1. Compute cosine embedding: ψ(τ) = [cos(πi·τ) for i in 1..embedding_dim] - let cos_embed = self.cosine_embedding(taus)?; // [batch, num_quantiles, embedding_dim] + let cos_embed = self.cosine_embedding(&taus)?; // [batch, num_quantiles, embedding_dim] // 2. State embedding broadcast to match quantile dimension // [batch, state_dim] → [batch, 1, state_dim] → [batch, num_quantiles, state_dim] @@ -174,10 +178,14 @@ impl QuantileNetwork { .map_err(|e| MLError::ModelError(format!("Output layer forward failed: {}", e)))?; // 7. Transpose to [batch, num_actions, num_quantiles] - quantile_values + let output = quantile_values .reshape((batch_size, num_quantiles, num_actions))? .transpose(1, 2) - .map_err(|e| MLError::ModelError(format!("Transpose failed: {}", e))) + .map_err(|e| MLError::ModelError(format!("Transpose failed: {}", e)))?; + + // Cast output back to F32 for API compatibility + output.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e))) } /// Cosine embedding for quantile fractions τ @@ -194,17 +202,18 @@ impl QuantileNetwork { /// Cosine embeddings [batch, num_quantiles, embedding_dim] fn cosine_embedding(&self, taus: &Tensor) -> CandleResult { let device = taus.device(); + let target_dtype = taus.dtype(); let batch_size = taus.dim(0)?; let num_quantiles = taus.dim(1)?; let embed_dim = self.config.quantile_embedding_dim; - // Create indices [1, 2, 3, ..., embedding_dim] + // Create indices [1, 2, 3, ..., embedding_dim] in same dtype as taus let indices: Vec = (1..=embed_dim).map(|i| i as f32).collect(); let indices_tensor = Tensor::from_vec( indices, (1, 1, embed_dim), device, - )?; + )?.to_dtype(target_dtype)?; // Broadcast taus to [batch, num_quantiles, 1] let taus_broadcast = taus.unsqueeze(2)?; @@ -213,8 +222,9 @@ impl QuantileNetwork { let indices_broadcast = indices_tensor.broadcast_as((batch_size, num_quantiles, embed_dim))?; let taus_full = taus_broadcast.broadcast_as((batch_size, num_quantiles, embed_dim))?; - // Compute π·i·τ - let pi_tensor = Tensor::full(PI, (batch_size, num_quantiles, embed_dim), device)?; + // Compute π·i·τ (pi_tensor must match taus dtype) + let pi_tensor = Tensor::full(PI, (batch_size, num_quantiles, embed_dim), device)? + .to_dtype(target_dtype)?; let angles = (pi_tensor * indices_broadcast)? * taus_full; // cos(π·i·τ) diff --git a/crates/ml/src/ensemble/adapters/diffusion.rs b/crates/ml/src/ensemble/adapters/diffusion.rs index 91d2ec9cf..0f168bf5f 100644 --- a/crates/ml/src/ensemble/adapters/diffusion.rs +++ b/crates/ml/src/ensemble/adapters/diffusion.rs @@ -118,6 +118,8 @@ impl ModelInferenceAdapter for DiffusionInferenceAdapter { let padded = self.pad_features(&features.values); let input = Tensor::from_vec(padded, (1, self.data_dim), &self.device) .map_err(|e| MLError::ModelError(format!("Diffusion input tensor: {e}")))?; + let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("Diffusion input dtype cast: {e}")))?; // Timestep t=1 (minimal noise level for feature processing) let t = Tensor::from_vec(vec![1_u32], (1,), &self.device) @@ -160,6 +162,8 @@ impl ModelInferenceAdapter for DiffusionInferenceAdapter { let padded = self.pad_features(&features.values); let input = Tensor::from_vec(padded, (1, self.data_dim), &self.device) .map_err(|e| MLError::ModelError(format!("Diffusion input tensor: {e}")))?; + let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("Diffusion input dtype cast: {e}")))?; // Timestep t=1 (minimal noise level for feature processing) let t = Tensor::from_vec(vec![1_u32], (1,), &self.device) diff --git a/crates/ml/src/ensemble/adapters/dqn.rs b/crates/ml/src/ensemble/adapters/dqn.rs index 3a0413b54..e60804543 100644 --- a/crates/ml/src/ensemble/adapters/dqn.rs +++ b/crates/ml/src/ensemble/adapters/dqn.rs @@ -8,6 +8,7 @@ use std::sync::Mutex; use candle_core::{Device, Tensor}; use crate::dqn::dqn::{DQNConfig, DQN}; +use crate::dqn::mixed_precision::ensure_training_dtype; use crate::ensemble::inference_adapter::{ EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, }; @@ -70,9 +71,11 @@ impl ModelInferenceAdapter for DqnInferenceAdapter { let f32_values: Vec = features.values.iter().map(|&v| v as f32).collect(); let len = f32_values.len(); - // Create input tensor [1, feature_dim] + // Create input tensor [1, feature_dim] and cast to training dtype let input = Tensor::from_vec(f32_values, (1, len), &self.device) .map_err(|e| MLError::ModelError(format!("Failed to create input tensor: {e}")))?; + let input = ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("Input dtype cast failed: {e}")))?; // Run forward pass through the Q-network let model = self diff --git a/crates/ml/src/ensemble/adapters/kan.rs b/crates/ml/src/ensemble/adapters/kan.rs index 7868fa79e..23778b9f1 100644 --- a/crates/ml/src/ensemble/adapters/kan.rs +++ b/crates/ml/src/ensemble/adapters/kan.rs @@ -99,6 +99,8 @@ impl ModelInferenceAdapter for KanInferenceAdapter { let padded = self.pad_features(&features.values); let input = Tensor::from_vec(padded, (1, self.input_dim), &self.device) .map_err(|e| MLError::ModelError(format!("Failed to create KAN input tensor: {e}")))?; + let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("KAN input dtype cast: {e}")))?; let model = self .model @@ -139,6 +141,8 @@ impl ModelInferenceAdapter for KanInferenceAdapter { let padded = self.pad_features(&features.values); let input = Tensor::from_vec(padded, (1, self.input_dim), &self.device) .map_err(|e| MLError::ModelError(format!("KAN input tensor: {e}")))?; + let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("KAN input dtype cast: {e}")))?; let model = self .model diff --git a/crates/ml/src/ensemble/adapters/mamba2.rs b/crates/ml/src/ensemble/adapters/mamba2.rs index 087bfd7b0..c72d8f137 100644 --- a/crates/ml/src/ensemble/adapters/mamba2.rs +++ b/crates/ml/src/ensemble/adapters/mamba2.rs @@ -9,8 +9,9 @@ use std::collections::VecDeque; use std::sync::Mutex; -use candle_core::{DType, Device, Tensor}; +use candle_core::{Device, Tensor}; +use crate::dqn::mixed_precision::ensure_training_dtype; use crate::ensemble::inference_adapter::{ EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, }; @@ -148,10 +149,9 @@ impl ModelInferenceAdapter for Mamba2InferenceAdapter { ) .map_err(|e| MLError::ModelError(format!("Failed to create Mamba2 input tensor: {e}")))?; - // Cast to F32 (Mamba2 uses DType::F32 for GPU throughput) - let input = input - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("Failed to cast input to F32: {e}")))?; + // Cast to training dtype (BF16 on Ampere+, F32 on CPU/older GPUs) + let input = ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("Failed to cast input to training dtype: {e}")))?; // Run forward pass (needs &mut self) let mut model = self diff --git a/crates/ml/src/ensemble/adapters/tft.rs b/crates/ml/src/ensemble/adapters/tft.rs index c462f928e..28a38fb5a 100644 --- a/crates/ml/src/ensemble/adapters/tft.rs +++ b/crates/ml/src/ensemble/adapters/tft.rs @@ -247,11 +247,13 @@ impl ModelInferenceAdapter for TftInferenceAdapter { // Future (known) features: calendar-derived let future_f32 = self.build_future_features(last_ts); - // Build tensors + // Build tensors and cast to training dtype let static_tensor = Tensor::from_vec(static_f32, (1, self.num_static), &self.device).map_err(|e| { MLError::ModelError(format!("Failed to create static tensor: {e}")) })?; + let static_tensor = crate::dqn::mixed_precision::ensure_training_dtype(&static_tensor) + .map_err(|e| MLError::ModelError(format!("Static dtype cast: {e}")))?; let hist_tensor = Tensor::from_vec( hist_f32, @@ -259,6 +261,8 @@ impl ModelInferenceAdapter for TftInferenceAdapter { &self.device, ) .map_err(|e| MLError::ModelError(format!("Failed to create historical tensor: {e}")))?; + let hist_tensor = crate::dqn::mixed_precision::ensure_training_dtype(&hist_tensor) + .map_err(|e| MLError::ModelError(format!("Historical dtype cast: {e}")))?; let future_tensor = Tensor::from_vec( future_f32, @@ -266,6 +270,8 @@ impl ModelInferenceAdapter for TftInferenceAdapter { &self.device, ) .map_err(|e| MLError::ModelError(format!("Failed to create future tensor: {e}")))?; + let future_tensor = crate::dqn::mixed_precision::ensure_training_dtype(&future_tensor) + .map_err(|e| MLError::ModelError(format!("Future dtype cast: {e}")))?; // --- forward pass (model lock) --- let quantile_preds = { diff --git a/crates/ml/src/ensemble/adapters/tggn.rs b/crates/ml/src/ensemble/adapters/tggn.rs index 415bab74b..fdfda0410 100644 --- a/crates/ml/src/ensemble/adapters/tggn.rs +++ b/crates/ml/src/ensemble/adapters/tggn.rs @@ -31,16 +31,20 @@ impl TggnProjection { } fn forward(&self, input: &Tensor) -> MLResult { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(format!("TGGN input dtype cast: {e}")))?; let h = self .linear1 - .forward(input) + .forward(&input) .map_err(|e| MLError::ModelError(format!("TGGN forward linear1: {e}")))?; let h = h .relu() .map_err(|e| MLError::ModelError(format!("TGGN forward relu: {e}")))?; - self.linear2 + let out = self.linear2 .forward(&h) - .map_err(|e| MLError::ModelError(format!("TGGN forward linear2: {e}"))) + .map_err(|e| MLError::ModelError(format!("TGGN forward linear2: {e}")))?; + out.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("TGGN output dtype cast: {e}"))) } } diff --git a/crates/ml/src/ensemble/adapters/tlob.rs b/crates/ml/src/ensemble/adapters/tlob.rs index b5e1d5254..074994c59 100644 --- a/crates/ml/src/ensemble/adapters/tlob.rs +++ b/crates/ml/src/ensemble/adapters/tlob.rs @@ -43,9 +43,11 @@ impl TlobProjection { } fn forward(&self, input: &Tensor) -> MLResult { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(format!("TLOB input dtype cast: {e}")))?; let h = self .linear1 - .forward(input) + .forward(&input) .map_err(|e| MLError::ModelError(format!("TLOB forward linear1: {e}")))?; let h = h .relu() @@ -57,9 +59,11 @@ impl TlobProjection { let h = h .relu() .map_err(|e| MLError::ModelError(format!("TLOB forward relu2: {e}")))?; - self.linear3 + let out = self.linear3 .forward(&h) - .map_err(|e| MLError::ModelError(format!("TLOB forward linear3: {e}"))) + .map_err(|e| MLError::ModelError(format!("TLOB forward linear3: {e}")))?; + out.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("TLOB output dtype cast: {e}"))) } } diff --git a/crates/ml/src/ensemble/adapters/xlstm.rs b/crates/ml/src/ensemble/adapters/xlstm.rs index 1539eb726..3f9f08313 100644 --- a/crates/ml/src/ensemble/adapters/xlstm.rs +++ b/crates/ml/src/ensemble/adapters/xlstm.rs @@ -166,6 +166,8 @@ impl ModelInferenceAdapter for XlstmInferenceAdapter { &self.device, ) .map_err(|e| MLError::ModelError(format!("xLSTM input tensor: {e}")))?; + let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("xLSTM input dtype cast: {e}")))?; // XLSTMNetwork::forward takes &self (not &mut self) let model = self @@ -247,6 +249,8 @@ impl ModelInferenceAdapter for XlstmInferenceAdapter { &self.device, ) .map_err(|e| MLError::ModelError(format!("xLSTM input tensor: {e}")))?; + let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) + .map_err(|e| MLError::ModelError(format!("xLSTM input dtype cast: {e}")))?; let model = self .model diff --git a/crates/ml/src/gradient_utils.rs b/crates/ml/src/gradient_utils.rs index 6387e7769..3f8d9b159 100644 --- a/crates/ml/src/gradient_utils.rs +++ b/crates/ml/src/gradient_utils.rs @@ -44,7 +44,7 @@ pub fn clip_grad_norm( // First pass: Calculate the total L2 norm of all gradients for var in vars.iter() { if let Some(grad) = grads.get(var) { - let norm_sq = grad.sqr()?.sum_all()?.to_scalar::()? as f64; + let norm_sq = grad.sqr()?.sum_all()?.to_dtype(candle_core::DType::F32)?.to_scalar::()? as f64; total_norm_sq += norm_sq; } } diff --git a/crates/ml/src/kan/layer.rs b/crates/ml/src/kan/layer.rs index e85a3d899..0b3fa74b6 100644 --- a/crates/ml/src/kan/layer.rs +++ b/crates/ml/src/kan/layer.rs @@ -109,13 +109,25 @@ impl KANLayer { MLError::ModelError(format!("KAN basis reshape failed: {}", e)) })?; + // Cast basis_flat to match coefficients dtype (BSpline evaluates in F32 + // for precision, but coefficients are in training dtype from VarBuilder) + let coeff_dtype = self.coefficients.dtype(); + let basis_flat = basis_flat.to_dtype(coeff_dtype).map_err(|e| { + MLError::ModelError(format!("KAN basis dtype cast failed: {}", e)) + })?; + // Spline contribution: (batch, in_dim * num_bases) @ (in_dim * num_bases, out_dim) let spline_out = basis_flat.matmul(&self.coefficients).map_err(|e| { MLError::ModelError(format!("KAN spline matmul failed: {}", e)) })?; + // Cast input to match residual_weight dtype for matmul compatibility + let input_cast = input.to_dtype(coeff_dtype).map_err(|e| { + MLError::ModelError(format!("KAN input dtype cast failed: {}", e)) + })?; + // Residual contribution: (batch, in_dim) @ (in_dim, out_dim) - let residual_out = input.matmul(&self.residual_weight).map_err(|e| { + let residual_out = input_cast.matmul(&self.residual_weight).map_err(|e| { MLError::ModelError(format!("KAN residual matmul failed: {}", e)) })?; diff --git a/crates/ml/src/kan/network.rs b/crates/ml/src/kan/network.rs index aa5fb8509..cbebdfffb 100644 --- a/crates/ml/src/kan/network.rs +++ b/crates/ml/src/kan/network.rs @@ -62,11 +62,15 @@ impl KANNetwork { /// Input shape: `(batch, layer_widths[0])` /// Output shape: `(batch, layer_widths[last])` pub fn forward(&self, input: &Tensor) -> Result { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; let mut x = input.clone(); for layer in &self.layers { x = layer.forward(&x)?; } - Ok(x) + // Cast output back to F32 for API compatibility + x.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e))) } } diff --git a/crates/ml/src/kan/spline.rs b/crates/ml/src/kan/spline.rs index 177171ff7..da032b6b2 100644 --- a/crates/ml/src/kan/spline.rs +++ b/crates/ml/src/kan/spline.rs @@ -164,6 +164,10 @@ impl BSplineBasis { let x_flat = x.flatten_all().map_err(|e| { MLError::ModelError(format!("BSpline flatten failed: {e}")) })?; + // BSpline evaluation requires F32 precision for index computation + let x_flat = x_flat.to_dtype(DType::F32).map_err(|e| { + MLError::ModelError(format!("BSpline dtype cast: {e}")) + })?; let n = x_flat.dims1().map_err(|e| { MLError::ModelError(format!("BSpline dims failed: {e}")) })?; @@ -200,6 +204,10 @@ impl BSplineBasis { let x_flat = x.flatten_all().map_err(|e| { MLError::ModelError(format!("BSpline flatten failed: {e}")) })?; + // BSpline evaluation uses floor/ceil/frac which require F32 precision + let x_flat = x_flat.to_dtype(DType::F32).map_err(|e| { + MLError::ModelError(format!("BSpline dtype cast: {e}")) + })?; let n = x_flat.dims1().map_err(|e| { MLError::ModelError(format!("BSpline dims failed: {e}")) })?; diff --git a/crates/ml/src/kan/trainable.rs b/crates/ml/src/kan/trainable.rs index 419c8bf57..c1630e915 100644 --- a/crates/ml/src/kan/trainable.rs +++ b/crates/ml/src/kan/trainable.rs @@ -95,7 +95,9 @@ impl UnifiedTrainable for KANTrainableAdapter { } fn forward(&mut self, input: &Tensor) -> Result { - self.network.forward(input) + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; + self.network.forward(&input) } fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result { diff --git a/crates/ml/src/liquid/candle_cfc.rs b/crates/ml/src/liquid/candle_cfc.rs index 3bec46c43..9c8a6754d 100644 --- a/crates/ml/src/liquid/candle_cfc.rs +++ b/crates/ml/src/liquid/candle_cfc.rs @@ -4,7 +4,7 @@ //! This is the training path; the existing FixedPoint implementation in cells.rs/network.rs //! remains the production inference path. -use candle_core::{DType, Tensor}; +use candle_core::Tensor; use candle_nn::{Linear, Module, VarBuilder}; use serde::{Deserialize, Serialize}; @@ -121,6 +121,8 @@ impl BackboneMLP { /// - `f_out` is tanh-activated (bounded in [-1, 1]) /// - `tau_out` is raw (sigmoid scaling applied in the CfC cell) pub fn forward(&self, input: &Tensor) -> Result<(Tensor, Tensor), MLError> { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::InferenceError(e.to_string()))?; let mut x = input.clone(); for layer in &self.layers { x = layer @@ -285,7 +287,7 @@ impl CandleCfCNetwork { let mut h = Tensor::zeros( (batch_size, self.config.hidden_size), - DType::F32, + crate::dqn::mixed_precision::training_dtype(device), device, ) .map_err(|e| MLError::InferenceError(format!("CfC init hidden: {}", e)))?; @@ -308,7 +310,12 @@ impl CandleCfCNetwork { /// Forward compatible with UnifiedTrainable (3D input, default dt=0.01) pub fn forward(&self, input: &Tensor) -> Result { - self.forward_sequence(input, 0.01) + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::InferenceError(e.to_string()))?; + let output = self.forward_sequence(&input, 0.01)?; + // Cast output back to F32 for API compatibility + output.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::InferenceError(format!("Output dtype cast failed: {}", e))) } /// Approximate parameter count for reporting purposes diff --git a/crates/ml/src/liquid/training.rs b/crates/ml/src/liquid/training.rs index 5edea8706..e9eb0eb84 100644 --- a/crates/ml/src/liquid/training.rs +++ b/crates/ml/src/liquid/training.rs @@ -462,7 +462,7 @@ impl LiquidTrainer { // === Candle-based CfC v2 Trainer (2026 modernization) === -use candle_core::{DType, Device, Tensor}; +use candle_core::{Device, Tensor}; use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap}; use super::candle_cfc::{CandleCfCNetwork, CfCTrainConfig}; @@ -663,6 +663,7 @@ impl TrainingUtils { #[cfg(test)] mod tests { use super::*; + use candle_core::DType; use crate::liquid::candle_cfc::DeviceConfig; #[test] diff --git a/crates/ml/src/mamba/mod.rs b/crates/ml/src/mamba/mod.rs index 77deaa05e..8d2949ba2 100644 --- a/crates/ml/src/mamba/mod.rs +++ b/crates/ml/src/mamba/mod.rs @@ -779,6 +779,8 @@ impl Mamba2SSM { /// - Tensor operations fail #[instrument(skip(self, input))] pub fn forward(&mut self, input: &Tensor) -> Result { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; let start = Instant::now(); // OPTIMIZATION: Device affinity check (catch cross-device transfers early) @@ -792,7 +794,7 @@ impl Mamba2SSM { } // Input projection - let mut hidden = self.input_projection.forward(input)?; + let mut hidden = self.input_projection.forward(&input)?; // Process through each layer - collect indices first to avoid borrow conflicts let num_layers = self.ssd_layers.len(); @@ -818,6 +820,10 @@ impl Mamba2SSM { let output_raw = self.output_projection.forward(&hidden)?; let output = crate::cuda_compat::manual_sigmoid(&output_raw)?; + // Cast output back to F32 for API compatibility + let output = output.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?; + // OPTIMIZATION: Update performance metrics with VecDeque (O(1) instead of O(n)) let inference_time = start.elapsed(); self.total_inferences.fetch_add(1, Ordering::Relaxed); @@ -885,7 +891,9 @@ impl Mamba2SSM { C.dims() ); let batch_size = scanned_states.dim(0)?; - let C_t = C.t()?.contiguous()?; + // Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16) + let C_cast = C.to_dtype(scanned_states.dtype())?; + let C_t = C_cast.t()?.contiguous()?; let C_broadcasted = C_t.unsqueeze(0)? .broadcast_as((batch_size, C_t.dim(0)?, C_t.dim(1)?))?; @@ -963,7 +971,9 @@ impl Mamba2SSM { self.config.d_state ); - let B_t = B.t()?.contiguous()?; + // Cast B to match input dtype (SSM state is F32 but input may be BF16) + let B_cast = B.to_dtype(input.dtype())?; + let B_t = B_cast.t()?.contiguous()?; let d_inner = B_t.dim(0)?; let d_state = B_t.dim(1)?; let B_broadcasted = B_t @@ -1684,7 +1694,9 @@ impl Mamba2SSM { // scanned_states: [32, 60, 16] // C stored as: [d_inner, d_state] = [512, 16] // Need: [batch, d_state, d_inner] = [32, 16, 512] - let C_t = C.t()?.contiguous()?; // [512, 16] → [16, 512] + // Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16) + let C_cast = C.to_dtype(scanned_states.dtype())?; + let C_t = C_cast.t()?.contiguous()?; // [512, 16] → [16, 512] trace!("C transposed (d_state, d_inner): {:?}", C_t.dims()); // Now broadcast [16, 512] to [32, 16, 512] @@ -1744,6 +1756,9 @@ impl Mamba2SSM { let mut result = Tensor::zeros((batch_size, seq_len, d_state), input.dtype(), device)?; let mut current_state = Tensor::zeros((batch_size, d_state), input.dtype(), device)?; + // Cast A to match input dtype (SSM state matrices are F32 but computation may be BF16) + let A_cast = A.to_dtype(input.dtype())?; + // Sequential scan with state transitions (maintaining gradients) for t in 0..seq_len { let x_t = input.narrow(1, t, 1)?.squeeze(1)?; @@ -1752,7 +1767,7 @@ impl Mamba2SSM { // State transition: h_t = h_{t-1} @ A^T + x_t // current_state [batch, d_state] × A.t() [d_state, d_state] = [batch, d_state] // This is the correct way to do batch SSM state transitions - current_state = (current_state.matmul(&A.t()?)? + &x_t)?; + current_state = (current_state.matmul(&A_cast.t()?)? + &x_t)?; // Write directly to result tensor (no Vec accumulation, no Tensor::cat doubling) let current_unsqueezed = current_state.unsqueeze(1)?; @@ -1838,7 +1853,9 @@ impl Mamba2SSM { // input: [batch, seq, d_inner], B: [d_state, d_inner] // B.t(): [d_inner, d_state] → explicit repeat to [batch, d_inner, d_state] let batch_size = input.dim(0)?; - let B_t = B.t()?.contiguous()?; // [d_state, d_inner] → [d_inner, d_state] + // Cast B to match input dtype (SSM state is F32 but input may be BF16) + let B_cast = B.to_dtype(input.dtype())?; + let B_t = B_cast.t()?.contiguous()?; // [d_state, d_inner] → [d_inner, d_state] // CRITICAL FIX: Use repeat/expand instead of broadcast_as for CUDA compatibility // Create [batch, d_inner, d_state] by repeating the [d_inner, d_state] tensor diff --git a/crates/ml/src/mamba/scan_algorithms.rs b/crates/ml/src/mamba/scan_algorithms.rs index 83f387458..e13752340 100644 --- a/crates/ml/src/mamba/scan_algorithms.rs +++ b/crates/ml/src/mamba/scan_algorithms.rs @@ -327,9 +327,11 @@ impl ParallelScanEngine { let alpha_fp = FixedPoint::from_f64(0.9); // Decay factor let beta_fp = FixedPoint::from_f64(0.1); // Input weight - // Use F32 to match state/input tensor dtype (model-wide F32 for GPU throughput) - let alpha = Tensor::full(alpha_fp.to_f64() as f32, state.shape(), state.device())?; - let beta = Tensor::full(beta_fp.to_f64() as f32, input.shape(), input.device())?; + // Match state/input tensor dtype (may be BF16 during mixed-precision training) + let alpha = Tensor::full(alpha_fp.to_f64() as f32, state.shape(), state.device())? + .to_dtype(state.dtype())?; + let beta = Tensor::full(beta_fp.to_f64() as f32, input.shape(), input.device())? + .to_dtype(input.dtype())?; let decayed_state = (state * alpha)?; let input_contribution = (input * beta)?; diff --git a/crates/ml/src/ppo/lstm_networks.rs b/crates/ml/src/ppo/lstm_networks.rs index 7b69fbb73..2d0ff1cfa 100644 --- a/crates/ml/src/ppo/lstm_networks.rs +++ b/crates/ml/src/ppo/lstm_networks.rs @@ -106,10 +106,12 @@ impl LSTMPolicyNetwork { h_t: &Tensor, c_t: &Tensor, ) -> Result<(Tensor, Tensor, Tensor), MLError> { + let state = crate::dqn::mixed_precision::ensure_training_dtype(state) + .map_err(|e| MLError::ModelError(e.to_string()))?; // Project input: [batch, input_dim] → [batch, hidden_dim] let x = self .input_layer - .forward(state) + .forward(&state) .map_err(|e| MLError::ModelError(format!("Input layer forward failed: {}", e)))?; // Apply ReLU activation @@ -160,6 +162,14 @@ impl LSTMPolicyNetwork { .forward(&x) .map_err(|e| MLError::ModelError(format!("Output layer forward failed: {}", e)))?; + // Cast outputs back to F32 for API compatibility + let logits = logits.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Logits dtype cast failed: {}", e)))?; + let new_h = new_h.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Hidden state dtype cast failed: {}", e)))?; + let new_c = new_c.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Cell state dtype cast failed: {}", e)))?; + Ok((logits, new_h, new_c)) } @@ -320,10 +330,12 @@ impl LSTMValueNetwork { h_t: &Tensor, c_t: &Tensor, ) -> Result<(Tensor, Tensor, Tensor), MLError> { + let state = crate::dqn::mixed_precision::ensure_training_dtype(state) + .map_err(|e| MLError::ModelError(e.to_string()))?; // Project input: [batch, input_dim] → [batch, hidden_dim] let x = self .input_layer - .forward(state) + .forward(&state) .map_err(|e| MLError::ModelError(format!("Input layer forward failed: {}", e)))?; // Apply ReLU activation @@ -379,6 +391,14 @@ impl LSTMValueNetwork { .squeeze(1) .map_err(|e| MLError::ModelError(format!("Failed to squeeze value output: {}", e)))?; + // Cast outputs back to F32 for API compatibility + let value = value.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Value dtype cast failed: {}", e)))?; + let new_h = new_h.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Hidden state dtype cast failed: {}", e)))?; + let new_c = new_c.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Cell state dtype cast failed: {}", e)))?; + Ok((value, new_h, new_c)) } diff --git a/crates/ml/src/ppo/ppo.rs b/crates/ml/src/ppo/ppo.rs index 7c34d9ae2..7edb65027 100644 --- a/crates/ml/src/ppo/ppo.rs +++ b/crates/ml/src/ppo/ppo.rs @@ -427,7 +427,9 @@ impl PolicyNetwork { input: &Tensor, mixed_precision: &Option, ) -> Result { - let (mut x, use_amp) = match mixed_precision { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let (mut x, _use_amp) = match mixed_precision { Some(mp) if mp.enabled => { let target_dtype = mp.dtype.to_dtype(); match input.to_dtype(target_dtype) { @@ -450,9 +452,10 @@ impl PolicyNetwork { } } - if use_amp { + // Always cast output back to F32 for API compatibility + if x.dtype() != DType::F32 { x = x.to_dtype(DType::F32).map_err(|e| { - MLError::ModelError(format!("AMP cast back to F32 failed: {}", e)) + MLError::ModelError(format!("Output dtype cast to F32 failed: {}", e)) })?; } @@ -671,7 +674,9 @@ impl ValueNetwork { input: &Tensor, mixed_precision: &Option, ) -> Result { - let (mut x, use_amp) = match mixed_precision { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let (mut x, _use_amp) = match mixed_precision { Some(mp) if mp.enabled => { let target_dtype = mp.dtype.to_dtype(); match input.to_dtype(target_dtype) { @@ -697,9 +702,10 @@ impl ValueNetwork { // Squeeze the last dimension (from [batch, 1] to [batch]) x = x.squeeze(1)?; - if use_amp { + // Always cast output back to F32 for API compatibility + if x.dtype() != DType::F32 { x = x.to_dtype(DType::F32).map_err(|e| { - MLError::ModelError(format!("AMP cast back to F32 failed: {}", e)) + MLError::ModelError(format!("Output dtype cast to F32 failed: {}", e)) })?; } diff --git a/crates/ml/src/tft/mod.rs b/crates/ml/src/tft/mod.rs index b44b51402..6ef663f4e 100644 --- a/crates/ml/src/tft/mod.rs +++ b/crates/ml/src/tft/mod.rs @@ -566,10 +566,16 @@ impl TemporalFusionTransformer { future_features: &Tensor, use_checkpointing: bool, ) -> Result { + let static_features = crate::dqn::mixed_precision::ensure_training_dtype(static_features) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let historical_features = crate::dqn::mixed_precision::ensure_training_dtype(historical_features) + .map_err(|e| MLError::ModelError(e.to_string()))?; + let future_features = crate::dqn::mixed_precision::ensure_training_dtype(future_features) + .map_err(|e| MLError::ModelError(e.to_string()))?; let start_time = Instant::now(); // Validate input dimensions - self.validate_input_dimensions(static_features, historical_features, future_features)?; + self.validate_input_dimensions(&static_features, &historical_features, &future_features)?; // Log device placement for debugging debug!("Forward pass device check:"); @@ -580,7 +586,7 @@ impl TemporalFusionTransformer { // 1. Variable Selection Networks (skip absent feature paths) let static_encoded = if let Some(ref mut static_vsn) = self.static_variable_selection { - let static_selected = static_vsn.forward(static_features, None)?; + let static_selected = static_vsn.forward(&static_features, None)?; let encoder = self.static_encoder.as_mut().ok_or_else(|| { MLError::ModelError( "static_encoder must exist when static_variable_selection exists".to_owned(), @@ -597,7 +603,7 @@ impl TemporalFusionTransformer { let historical_selected = self .historical_variable_selection - .forward(historical_features, None)?; + .forward(&historical_features, None)?; let historical_encoded = if use_checkpointing { self.historical_encoder @@ -608,7 +614,7 @@ impl TemporalFusionTransformer { }; let future_encoded = if let Some(ref mut future_vsn) = self.future_variable_selection { - let future_selected = future_vsn.forward(future_features, None)?; + let future_selected = future_vsn.forward(&future_features, None)?; let encoder = self.future_encoder.as_mut().ok_or_else(|| { MLError::ModelError( "future_encoder must exist when future_variable_selection exists".to_owned(), @@ -661,6 +667,10 @@ impl TemporalFusionTransformer { debug!(" quantile_preds: {:?}", quantile_preds.device()); + // Cast output back to F32 for API compatibility + let quantile_preds = quantile_preds.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?; + // Update performance metrics let latency = start_time.elapsed().as_micros() as u64; self.update_performance_metrics(latency); diff --git a/crates/ml/src/tft/temporal_attention.rs b/crates/ml/src/tft/temporal_attention.rs index dd7227ea3..f1878defd 100644 --- a/crates/ml/src/tft/temporal_attention.rs +++ b/crates/ml/src/tft/temporal_attention.rs @@ -350,14 +350,18 @@ impl TemporalSelfAttention { // Add positional encoding (lightweight, no checkpointing needed) let pos_encoding = self.positional_encoding.forward(seq_len)?; + // Cast positional encoding to match input dtype (encoding is F32, input may be BF16) + let pos_encoding = pos_encoding.to_dtype(x.dtype())?; let pos_encoding_batch = pos_encoding .unsqueeze(0)? .broadcast_as((batch_size, seq_len, hidden_dim))?; let x_with_pos = (x + &pos_encoding_batch)?; // Create causal mask if needed (now [1, seq_len, seq_len]) + // Cast mask to match input dtype (mask is F32 but attention may be BF16) let mask = if causal_mask { let base_mask = self.create_causal_mask(seq_len)?; + let base_mask = base_mask.to_dtype(x.dtype())?; // Broadcast to [batch_size, seq_len, seq_len] Some(base_mask.broadcast_as((batch_size, seq_len, seq_len))?) } else { @@ -393,6 +397,7 @@ impl TemporalSelfAttention { for (i, weights) in attention_weights.iter().enumerate() { let mean_weight = weights .mean_all() + .and_then(|t| t.to_dtype(candle_core::DType::F32)) .and_then(|t| t.to_vec0::()) .unwrap_or(0.0) as f64; weight_stats.insert(format!("head_{}_mean", i), mean_weight); diff --git a/crates/ml/src/tft/variable_selection.rs b/crates/ml/src/tft/variable_selection.rs index 9c70ccff0..671e511e5 100644 --- a/crates/ml/src/tft/variable_selection.rs +++ b/crates/ml/src/tft/variable_selection.rs @@ -130,7 +130,7 @@ impl VariableSelectionNetwork { fn update_importance_scores(&mut self, attention_weights: &Tensor) -> Result<(), MLError> { // Compute mean attention weights across batch and time let mean_weights = attention_weights.mean_keepdim(0)?.mean_keepdim(1)?; // [1, 1, input_size] - let weights_vec = mean_weights.flatten_all()?.to_vec1::()?; + let weights_vec = mean_weights.flatten_all()?.to_dtype(candle_core::DType::F32)?.to_vec1::()?; // Clear previous scores to prevent memory growth (HashMap maintains capacity but releases entries) self.importance_scores.clear(); diff --git a/crates/ml/src/tgnn/trainable_adapter.rs b/crates/ml/src/tgnn/trainable_adapter.rs index ecd39422c..4f92f2b30 100644 --- a/crates/ml/src/tgnn/trainable_adapter.rs +++ b/crates/ml/src/tgnn/trainable_adapter.rs @@ -139,9 +139,11 @@ impl UnifiedTrainable for TGGNTrainableAdapter { } fn forward(&mut self, input: &Tensor) -> Result { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; // input: [batch, node_dim] // input_linear: node_dim -> hidden_dim - let hidden = self.input_linear.forward(input).map_err(|e| { + let hidden = self.input_linear.forward(&input).map_err(|e| { MLError::ModelError(format!("Input linear forward failed: {}", e)) })?; @@ -155,6 +157,11 @@ impl UnifiedTrainable for TGGNTrainableAdapter { MLError::ModelError(format!("Output linear forward failed: {}", e)) })?; + // Cast output back to F32 for API compatibility + let output = output.to_dtype(candle_core::DType::F32).map_err(|e| { + MLError::ModelError(format!("Output dtype cast failed: {}", e)) + })?; + Ok(output) } diff --git a/crates/ml/src/tlob/trainable_adapter.rs b/crates/ml/src/tlob/trainable_adapter.rs index ee468fd77..e801b6666 100644 --- a/crates/ml/src/tlob/trainable_adapter.rs +++ b/crates/ml/src/tlob/trainable_adapter.rs @@ -170,6 +170,8 @@ impl UnifiedTrainable for TLOBTrainableAdapter { } fn forward(&mut self, input: &Tensor) -> Result { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; // input shape: [batch, seq_len, feature_dim] or [batch, seq_len*feature_dim] // Flatten to [batch, seq_len*feature_dim] if 3D let dims = input.shape().dims(); @@ -200,6 +202,11 @@ impl UnifiedTrainable for TLOBTrainableAdapter { MLError::ModelError(format!("Output linear forward failed: {}", e)) })?; + // Cast output back to F32 for API compatibility + let output = output.to_dtype(candle_core::DType::F32).map_err(|e| { + MLError::ModelError(format!("Output dtype cast failed: {}", e)) + })?; + Ok(output) } diff --git a/crates/ml/src/xlstm/mlstm.rs b/crates/ml/src/xlstm/mlstm.rs index 560e4a7b1..11c989244 100644 --- a/crates/ml/src/xlstm/mlstm.rs +++ b/crates/ml/src/xlstm/mlstm.rs @@ -12,7 +12,7 @@ //! batch item. With default hidden=64, heads=4: head_dim=16, so C is 4×16×16=1024 //! floats per sample -- very manageable. Be cautious with hidden>128 or heads<2. -use candle_core::{DType, Tensor}; +use candle_core::Tensor; use candle_nn::{linear, Linear, Module, VarBuilder}; use crate::MLError; @@ -93,9 +93,10 @@ impl MLSTMCell { let (h_prev, c_flat_prev) = match state { Some((h, c)) => (h.clone(), c.clone()), None => { - let h_zeros = Tensor::zeros(&[batch, self.hidden_dim], DType::F32, dev) + let dtype = crate::dqn::mixed_precision::training_dtype(dev); + let h_zeros = Tensor::zeros(&[batch, self.hidden_dim], dtype, dev) .map_err(|e| MLError::ModelError(format!("mLSTM h_zeros: {e}")))?; - let c_zeros = Tensor::zeros(&[batch, c_flat_size], DType::F32, dev) + let c_zeros = Tensor::zeros(&[batch, c_flat_size], dtype, dev) .map_err(|e| MLError::ModelError(format!("mLSTM c_zeros: {e}")))?; (h_zeros, c_zeros) } diff --git a/crates/ml/src/xlstm/network.rs b/crates/ml/src/xlstm/network.rs index 0996c2076..96027d117 100644 --- a/crates/ml/src/xlstm/network.rs +++ b/crates/ml/src/xlstm/network.rs @@ -6,7 +6,7 @@ //! Block allocation: blocks 0..N*slstm_ratio are sLSTM, rest are mLSTM. //! The first block has input_dim → hidden_dim, subsequent blocks are hidden_dim → hidden_dim. -use candle_core::{DType, Tensor}; +use candle_core::Tensor; use candle_nn::{linear, Linear, Module, VarBuilder}; use crate::MLError; @@ -85,6 +85,8 @@ impl XLSTMNetwork { /// /// Returns `(batch, output_dim)`. pub fn forward(&self, input: &Tensor) -> Result { + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; let dims = input.dims(); let (batch, seq_len, _features) = match dims.len() { 2 => (dims[0], 1_usize, dims[1]), @@ -101,7 +103,11 @@ impl XLSTMNetwork { // State per block: (h, c_flat) let mut states: Vec> = vec![None; self.blocks.len()]; - let mut final_h = Tensor::zeros(&[batch, self.hidden_dim], DType::F32, input.device()) + let mut final_h = Tensor::zeros( + &[batch, self.hidden_dim], + crate::dqn::mixed_precision::training_dtype(input.device()), + input.device(), + ) .map_err(|e| MLError::ModelError(format!("xLSTM zeros: {e}")))?; for t in 0..seq_len { @@ -139,8 +145,11 @@ impl XLSTMNetwork { } // Output projection - self.output_head.forward(&final_h) - .map_err(|e| MLError::ModelError(format!("xLSTM output: {e}"))) + let output = self.output_head.forward(&final_h) + .map_err(|e| MLError::ModelError(format!("xLSTM output: {e}")))?; + // Cast output back to F32 for API compatibility + output.to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("xLSTM output dtype cast: {e}"))) } } diff --git a/crates/ml/src/xlstm/slstm.rs b/crates/ml/src/xlstm/slstm.rs index 20d15daf5..903810720 100644 --- a/crates/ml/src/xlstm/slstm.rs +++ b/crates/ml/src/xlstm/slstm.rs @@ -5,7 +5,7 @@ //! This prevents gradient vanishing over long sequences. A normalizer state `n_t` //! stabilizes the exponential values. -use candle_core::{DType, Tensor}; +use candle_core::Tensor; use candle_nn::{linear, Linear, Module, VarBuilder}; use crate::MLError; @@ -57,7 +57,11 @@ impl SLSTMCell { let (h_prev, c_prev) = match state { Some((h, c)) => (h.clone(), c.clone()), None => { - let zeros = Tensor::zeros(&[batch, self.hidden_dim], DType::F32, dev) + let zeros = Tensor::zeros( + &[batch, self.hidden_dim], + crate::dqn::mixed_precision::training_dtype(dev), + dev, + ) .map_err(|e| MLError::ModelError(format!("sLSTM init zeros: {e}")))?; (zeros.clone(), zeros) } diff --git a/crates/ml/src/xlstm/trainable.rs b/crates/ml/src/xlstm/trainable.rs index bea8a6452..23597e73a 100644 --- a/crates/ml/src/xlstm/trainable.rs +++ b/crates/ml/src/xlstm/trainable.rs @@ -88,7 +88,9 @@ impl UnifiedTrainable for XLSTMTrainableAdapter { } fn forward(&mut self, input: &Tensor) -> Result { - self.network.forward(input) + let input = crate::dqn::mixed_precision::ensure_training_dtype(input) + .map_err(|e| MLError::ModelError(e.to_string()))?; + self.network.forward(&input) } fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result {