From ad28482a93be7b2266258482e09c3622fb2feb5a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 17 Mar 2026 08:34:51 +0100 Subject: [PATCH] =?UTF-8?q?fix(cuda):=20shmem=20tile=20overflow=20?= =?UTF-8?q?=E2=86=92=20CUDA=5FERROR=5FILLEGAL=5FADDRESS=20on=20RTX=203050?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: shmem_max_in_dim only included trunk dims (state_dim, shared_h1, shared_h2) but not head dims (value_h, adv_h). When hidden_dim_base=32 made the trunk narrow while heads stayed at 128, the BF16 weight tile for branch output (255×128=32640 BF16 elements) overflowed the shared memory region (12288 BF16 elements). On H100 the overflow landed in unused-but-mapped hardware shmem (silent corruption). On RTX 3050 (48KB physical shmem) it hit unmapped memory → CUDA_ERROR_ILLEGAL_ADDRESS. Changes: - gpu_dqn_trainer.rs: shmem_max_in_dim includes value_h/adv_h - Remove all #[ignore] from smoke tests (feature_coverage, training_stability, gpu_residency) - Smoke tests use real .dbn data from test_data/ (hard error if missing) - Remove synthetic_data() fallback — no fake data in tests - GPU-direct DtoD training path (train_step_gpu, FusedTrainScalars) - GPU-native PER priority update kernel (zero CPU readback) - IQN dual-head integration (gpu_iqn_head.rs) - BF16 dtype fixes across 6 model adapters - Hyperopt 30D→31D (iqn_lambda) - portfolio_transformer: unconditional BF16 (remove dead CPU branches) - liquid/adapter: all tests use Cuda(0) directly - Fix pre-existing gpu_kernel_parity_test.rs (stale args) - Fix pre-existing evaluate_baseline.rs (removed fields) Co-Authored-By: Claude Opus 4.6 (1M context) --- .cargo/config.toml | 3 + crates/ml-dqn/src/dqn.rs | 12 + crates/ml-dqn/src/gpu_replay_buffer.rs | 22 + crates/ml-dqn/src/replay_buffer_type.rs | 34 + crates/ml-dqn/tests/gpu_smoketest.rs | 1 + crates/ml-supervised/src/diffusion/noise.rs | 3 +- crates/ml/examples/evaluate_baseline.rs | 4 - .../src/cuda_pipeline/gpu_action_selector.rs | 6 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 459 ++++++++- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 754 +++++++++++++++ .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 875 ++++++++++++++++++ crates/ml/src/cuda_pipeline/mod.rs | 4 +- crates/ml/src/cuda_pipeline/signal_adapter.rs | 1 + crates/ml/src/diffusion/trainable.rs | 36 +- crates/ml/src/ensemble/adapters/liquid.rs | 19 +- crates/ml/src/features/multi_timeframe.rs | 25 +- crates/ml/src/hyperopt/adapters/dqn.rs | 59 +- crates/ml/src/inference.rs | 22 +- crates/ml/src/liquid/adapter.rs | 24 +- crates/ml/src/mamba/trainable_adapter.rs | 2 +- crates/ml/src/portfolio_transformer.rs | 19 +- crates/ml/src/trainers/dqn/config.rs | 40 + crates/ml/src/trainers/dqn/fused_training.rs | 358 +++++-- .../dqn/smoke_tests/feature_coverage.rs | 20 +- .../trainers/dqn/smoke_tests/gpu_residency.rs | 25 +- .../src/trainers/dqn/smoke_tests/helpers.rs | 44 +- .../trainers/dqn/smoke_tests/performance.rs | 41 +- .../dqn/smoke_tests/training_stability.rs | 89 +- .../src/trainers/dqn/trainer/constructor.rs | 1 + crates/ml/src/trainers/dqn/trainer/tests.rs | 4 +- crates/ml/src/trainers/online_learning.rs | 14 +- crates/ml/src/trainers/tlob.rs | 13 +- .../ml/tests/dqn_action_collapse_fix_test.rs | 10 +- crates/ml/tests/gpu_kernel_parity_test.rs | 25 +- scripts/gpu-hotpath-guard.sh | 21 + 35 files changed, 2730 insertions(+), 359 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/gpu_iqn_head.rs create mode 100644 crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu diff --git a/.cargo/config.toml b/.cargo/config.toml index c3336ad2a..02cd57fd7 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,5 +1,8 @@ [env] SQLX_OFFLINE = "false" +# GPU tests deadlock under parallel execution (shared CUDA context + 4GB RTX 3050). +# Single-threaded by default; CI overrides via RUST_TEST_THREADS in workflow. +RUST_TEST_THREADS = "1" [cargo-new] vcs = "none" diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 58291c738..959e4ce18 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -185,6 +185,8 @@ pub struct DQNConfig { pub iqn_num_quantiles: usize, pub iqn_kappa: f32, pub iqn_embedding_dim: usize, + /// Lambda weight for IQN loss in dual-head ensemble (0.0 = C51 only) + pub iqn_lambda: f32, // Branching DQN (Phase C+): independent Q-heads for factored action space /// When true, uses `BranchingDuelingQNetwork` with independent Q-heads @@ -317,6 +319,7 @@ impl Default for DQNConfig { iqn_num_quantiles: 64, iqn_kappa: 1.0, iqn_embedding_dim: 64, + iqn_lambda: 0.25, // Branching: Enabled by default (45 factored actions via 3 independent heads) use_branching: true, @@ -379,6 +382,7 @@ impl DQNConfig { hasher.update(self.num_urgency_levels.to_le_bytes()); hasher.update(self.iqn_embedding_dim.to_le_bytes()); hasher.update(self.iqn_num_quantiles.to_le_bytes()); + hasher.update(self.iqn_lambda.to_le_bytes()); hex::encode(hasher.finalize()) } @@ -403,6 +407,7 @@ impl DQNConfig { meta.insert("dqn.num_atoms".to_owned(), self.num_atoms.to_string()); meta.insert("dqn.iqn_embedding_dim".to_owned(), self.iqn_embedding_dim.to_string()); meta.insert("dqn.iqn_num_quantiles".to_owned(), self.iqn_num_quantiles.to_string()); + meta.insert("dqn.iqn_lambda".to_owned(), self.iqn_lambda.to_string()); meta.insert("dqn.gamma".to_owned(), self.gamma.to_string()); meta } @@ -488,6 +493,9 @@ impl DQNConfig { let num_atoms = parse_usize("dqn.num_atoms")?; let iqn_embedding_dim = parse_usize("dqn.iqn_embedding_dim")?; let iqn_num_quantiles = parse_usize("dqn.iqn_num_quantiles")?; + let iqn_lambda: f32 = metadata.get("dqn.iqn_lambda") + .and_then(|s| s.parse().ok()) + .unwrap_or(0.25); // backward compat: default to 0.25 let gamma: f32 = metadata.get("dqn.gamma") .and_then(|s| s.parse().ok()) @@ -506,6 +514,7 @@ impl DQNConfig { num_urgency_levels, iqn_embedding_dim, iqn_num_quantiles, + iqn_lambda, gamma, ..Self::default() }) @@ -624,6 +633,7 @@ impl DQNConfig { iqn_num_quantiles: 200, iqn_kappa: 1.0, iqn_embedding_dim: 64, + iqn_lambda: 0.25, use_cvar_action_selection: false, cvar_alpha: 0.05, @@ -692,6 +702,7 @@ impl DQNConfig { iqn_num_quantiles: 200, iqn_kappa: 1.0, iqn_embedding_dim: 64, + iqn_lambda: 0.25, use_cvar_action_selection: true, cvar_alpha: 0.05, @@ -775,6 +786,7 @@ impl DQNConfig { iqn_num_quantiles: 64, iqn_kappa: 1.0, iqn_embedding_dim: 64, + iqn_lambda: 0.25, use_branching: false, // Disabled for emergency mode (simplest stable path) branch_hidden_dim: 128, num_order_types: 3, diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 4c3fe4a63..e760350cb 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -773,6 +773,28 @@ impl GpuReplayBuffer { Ok(()) } + /// Direct reference to the GPU-resident priorities tensor. + /// + /// Used by `GpuDqnTrainer::update_priorities_cuda()` to scatter-write + /// new priorities via a CUDA kernel — zero Candle tensor ops per step. + /// The caller extracts the `CudaSlice` from this tensor's storage. + pub fn priorities_tensor(&self) -> &Tensor { + &self.priorities + } + + /// Set `max_priority` from a CPU scalar (epoch-boundary flush). + /// + /// Used when the batch max is accumulated by an external CUDA kernel + /// (atomicMax in `per_update_priorities_kernel`) and read back once per epoch. + pub fn apply_max_priority_scalar(&mut self, max_prio: f32) -> Result<(), MLError> { + if max_prio > 0.0 { + let candidate = Tensor::new(max_prio, &self.device)?; + self.max_priority = + Tensor::stack(&[&candidate, &self.max_priority], 0)?.max(0)?; + } + Ok(()) + } + /// Sample PER-weighted indices without extracting full experiences. /// /// Returns `(indices_i64, is_weights)` where `indices_i64` is a `[batch_size]` diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 2533c5de5..8cbb8d1f8 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -377,6 +377,40 @@ impl ReplayBufferType { } } + /// Direct reference to the GPU-resident priorities tensor (GpuPrioritized only). + /// + /// Returns `None` for non-GPU buffers. Used by fused CUDA training to pass + /// the priorities tensor to `GpuDqnTrainer::update_priorities_cuda()`. + pub fn priorities_tensor(&self) -> Option { + match self { + Self::GpuPrioritized(buffer) => Some(buffer.lock().gpu.priorities_tensor().clone()), + Self::Uniform(_) | Self::Prioritized(_) => None, + } + } + + /// PER alpha/epsilon for GPU-native priority update kernel. + /// + /// Returns `(alpha, epsilon)` or `None` for non-GPU buffers. + pub fn per_alpha_epsilon(&self) -> Option<(f32, f32)> { + match self { + Self::GpuPrioritized(buffer) => { + let buf = buffer.lock(); + Some((buf.gpu.alpha(), buf.gpu.epsilon())) + } + Self::Uniform(_) | Self::Prioritized(_) => None, + } + } + + /// Set max_priority from a CPU scalar (epoch-boundary flush from CUDA kernel). + pub fn apply_max_priority_scalar(&self, max_prio: f32) -> Result<(), MLError> { + match self { + Self::GpuPrioritized(buffer) => { + buffer.lock().gpu.apply_max_priority_scalar(max_prio) + } + Self::Uniform(_) | Self::Prioritized(_) => Ok(()), + } + } + /// Step training counter for beta annealing (PER only) pub fn step(&self) { match self { diff --git a/crates/ml-dqn/tests/gpu_smoketest.rs b/crates/ml-dqn/tests/gpu_smoketest.rs index 44e32252c..4f355a375 100644 --- a/crates/ml-dqn/tests/gpu_smoketest.rs +++ b/crates/ml-dqn/tests/gpu_smoketest.rs @@ -81,6 +81,7 @@ fn smoketest_config() -> DQNConfig { iqn_num_quantiles: 32, iqn_kappa: 1.0, iqn_embedding_dim: 32, + iqn_lambda: 0.25, use_branching: false, // Disable branching — test basic path first branch_hidden_dim: 64, num_order_types: 3, diff --git a/crates/ml-supervised/src/diffusion/noise.rs b/crates/ml-supervised/src/diffusion/noise.rs index 4ced5c6f9..333aefd7d 100644 --- a/crates/ml-supervised/src/diffusion/noise.rs +++ b/crates/ml-supervised/src/diffusion/noise.rs @@ -105,8 +105,9 @@ impl NoiseScheduler { let sqrt_alpha_bar = alpha_bar.sqrt(); let sqrt_one_minus_alpha_bar = (1.0_f32 - alpha_bar).max(0.0_f32).sqrt(); - // Sample noise epsilon ~ N(0, I) + // Sample noise epsilon ~ N(0, I), cast to match input dtype (BF16 on CUDA) let noise = Tensor::randn(0_f32, 1.0, x0.dims(), &self.device) + .and_then(|n| n.to_dtype(x0.dtype())) .map_err(|e| MLError::ModelError(e.to_string()))?; // x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon diff --git a/crates/ml/examples/evaluate_baseline.rs b/crates/ml/examples/evaluate_baseline.rs index f66f00852..3eb0b6c55 100644 --- a/crates/ml/examples/evaluate_baseline.rs +++ b/crates/ml/examples/evaluate_baseline.rs @@ -905,11 +905,9 @@ fn evaluate_dqn_fold( // Architecture params — must match training checkpoint shapes use_dueling: hp_bool(hp, "use_dueling").unwrap_or(true), dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128), - use_distributional: hp_bool(hp, "use_distributional").unwrap_or(true), num_atoms: hp_usize(hp, "num_atoms").unwrap_or(51), v_min: hp_f64(hp, "v_min").unwrap_or(-2.0) as f32, v_max: hp_f64(hp, "v_max").unwrap_or(2.0) as f32, - use_noisy_nets: hp_bool(hp, "use_noisy_nets").unwrap_or(true), use_cql: false, // CQL not needed for evaluation // use_qr_dqn in hyperopt maps to use_iqn in DQNConfig use_iqn: hp_bool(hp, "use_qr_dqn").unwrap_or(false), @@ -1078,11 +1076,9 @@ fn evaluate_dqn_fold_gpu( use_per: false, use_dueling: hp_bool(hp, "use_dueling").unwrap_or(true), dueling_hidden_dim: hp_usize(hp, "dueling_hidden_dim").unwrap_or(128), - use_distributional: hp_bool(hp, "use_distributional").unwrap_or(true), num_atoms: hp_usize(hp, "num_atoms").unwrap_or(51), v_min: hp_f64(hp, "v_min").unwrap_or(-2.0) as f32, v_max: hp_f64(hp, "v_max").unwrap_or(2.0) as f32, - use_noisy_nets: hp_bool(hp, "use_noisy_nets").unwrap_or(true), use_cql: false, use_iqn: hp_bool(hp, "use_qr_dqn").unwrap_or(false), iqn_num_quantiles: hp_usize(hp, "num_quantiles").unwrap_or(64), diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 74780436d..327ad8632 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -811,7 +811,7 @@ mod tests { assert_eq!(greedy_actions.dims(), &[batch_size]); // GPU-side bounds check: all actions < num_actions - let bound = Tensor::new(num_actions as u32, &device).expect("bound"); + let bound = Tensor::new(&[num_actions as u32], &device).expect("bound").broadcast_as(&[batch_size]).expect("broadcast"); let in_range_count = greedy_actions.lt(&bound).expect("lt") .to_dtype(candle_core::DType::U32).expect("cast") .sum_all().expect("sum").to_scalar::().expect("scalar"); @@ -858,7 +858,7 @@ mod tests { let actions = selector .select_actions_branching(&q_exposure, &q_order, &q_urgency, 0.0) .expect("branching select"); - let bound_45 = Tensor::new(45_u32, &device).expect("bound"); + let bound_45 = Tensor::new(&[45_u32], &device).expect("bound").broadcast_as(&[batch_size]).expect("broadcast"); let in_range_count = actions.lt(&bound_45).expect("lt") .to_dtype(candle_core::DType::U32).expect("cast") .sum_all().expect("sum").to_scalar::().expect("scalar"); @@ -907,7 +907,7 @@ mod tests { ) .expect("routed select"); - let bound = Tensor::new(num_actions as u32, &device).expect("bound"); + let bound = Tensor::new(&[num_actions as u32], &device).expect("bound").broadcast_as(&[batch_size]).expect("broadcast"); let in_range_count = actions.lt(&bound).expect("lt") .to_dtype(candle_core::DType::U32).expect("cast") .sum_all().expect("sum").to_scalar::().expect("scalar"); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 65c3c16ba..0168bcbbe 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -112,6 +112,14 @@ pub struct GpuDqnTrainConfig { pub weight_decay: f32, /// Maximum gradient L2 norm for clipping. pub max_grad_norm: f32, + /// IQN quantile loss weight relative to C51 (0.0 = C51 only) + pub iqn_lambda: f32, + /// Number of IQN quantile samples per forward pass (0 = IQN disabled) + pub iqn_num_quantiles: usize, + /// Cosine embedding dimension for IQN τ encoding + pub iqn_embedding_dim: usize, + /// Huber threshold for quantile loss + pub iqn_kappa: f32, } impl Default for GpuDqnTrainConfig { @@ -136,6 +144,10 @@ impl Default for GpuDqnTrainConfig { epsilon: 1e-8, weight_decay: 1e-5, max_grad_norm: 10.0, + iqn_lambda: 0.25, + iqn_num_quantiles: 0, + iqn_embedding_dim: 64, + iqn_kappa: 1.0, } } } @@ -151,6 +163,19 @@ pub struct FusedTrainResult { pub grad_norm: f32, } +/// Scalar-only result from the fused GPU training path. +/// +/// TD errors stay on GPU (`td_errors_buf`) — no readback. PER priority update +/// is done by a dedicated CUDA kernel that reads `td_errors_buf` directly. +/// Only 8 bytes DtoH per step (loss + grad_norm). +#[derive(Debug, Clone, Copy)] +pub struct FusedTrainScalars { + /// Batch-mean weighted loss (for logging). + pub total_loss: f32, + /// Pre-clip gradient L2 norm (for monitoring). + pub grad_norm: f32, +} + // ── Parameter layout ──────────────────────────────────────────────────────── // // Matches GOFF_* / PARAM_* defines in dqn_training_kernel.cu exactly. @@ -209,6 +234,7 @@ pub struct GpuDqnTrainer { grad_norm_kernel: CudaFunction, adam_update_kernel: CudaFunction, ema_kernel: CudaFunction, + per_update_kernel: CudaFunction, f32_to_bf16_kernel: CudaFunction, // ── BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ── @@ -281,6 +307,16 @@ pub struct GpuDqnTrainer { /// Pre-allocated host-side staging vec for batch upload (avoids per-step malloc). /// Reused via `.clear()` + `.extend_from_slice()` each step. upload_staging_host: Vec, + + // ── PER priority update (GPU-native) ───────────────────────────── + /// Running max priority across all batches in an epoch (atomicMax in kernel). + /// Read back once per epoch via `batch_max_readback_and_reset()`. + batch_max_buf: CudaSlice, + + /// Scalar-only readback buffer: [total_loss(1) | grad_norm(1)] = 8 bytes. + /// Used by `execute_train_scalars_only()` to avoid reading back td_errors. + scalars_readback_buf: CudaSlice, + scalars_readback_host: [f32; 2], } impl GpuDqnTrainer { @@ -289,6 +325,52 @@ impl GpuDqnTrainer { self.config.batch_size } + /// Reference to saved h_s2 activations from the last forward pass. + /// + /// Shape: `[B, SHARED_H2]` — the shared trunk output for online states. + /// Valid after `train_step()` or `forward_loss()` has been called. + /// Used by the IQN dual-head to read trunk activations without re-computation. + pub fn save_h_s2(&self) -> &CudaSlice { + &self.save_h_s2 + } + + /// Reference to the next_states buffer on GPU. + /// + /// Shape: `[B, STATE_DIM]` — contains the batch's next states after `train_step()`. + /// Used by the IQN dual-head to compute target h_s2 via a trunk forward kernel. + pub fn next_states_buf(&self) -> &CudaSlice { + &self.next_states_buf + } + + /// Reference to the actions buffer on GPU. + /// + /// Shape: `[B]` i32 — flat actions from the batch. Valid after `train_step()` or + /// `train_step_gpu()`. Used by IQN to decode flat → branch actions on GPU. + pub fn actions_buf(&self) -> &CudaSlice { + &self.actions_buf + } + + /// Reference to the rewards buffer on GPU. + /// + /// Shape: `[B]` f32 — rewards from the batch. Valid after `train_step()` or + /// `train_step_gpu()`. Used by IQN for Bellman target computation on GPU. + pub fn rewards_buf(&self) -> &CudaSlice { + &self.rewards_buf + } + + /// Reference to the dones buffer on GPU. + /// + /// Shape: `[B]` f32 — done flags (0.0/1.0). Valid after `train_step()` or + /// `train_step_gpu()`. Used by IQN for Bellman target masking on GPU. + pub fn dones_buf(&self) -> &CudaSlice { + &self.dones_buf + } + + /// Reference to the training config. + pub fn config(&self) -> &GpuDqnTrainConfig { + &self.config + } + /// Create a new fused DQN trainer with pre-allocated GPU buffers. /// /// Compiles all 4 kernels (forward+loss, backward, grad_norm, adam_update) @@ -308,6 +390,9 @@ impl GpuDqnTrainer { // ── Compile EMA kernel (standalone — not captured in CUDA Graph) ── let ema_kernel = compile_ema_kernel(&stream)?; + // ── Compile PER priority update kernel (standalone) ── + let per_update_kernel = compile_per_update_kernel(&stream)?; + // ── Allocate batch input buffers ──────────────────────────── let states_buf = alloc_f32(&stream, b * config.state_dim, "states")?; let next_states_buf = alloc_f32(&stream, b * config.state_dim, "next_states")?; @@ -361,6 +446,13 @@ impl GpuDqnTrainer { // Pre-allocate host staging vec for upload_batch (reused each step) let upload_staging_host = vec![0.0_f32; upload_staging_len]; + // PER priority update: batch max accumulator + scalar-only readback + let mut batch_max_buf = alloc_f32(&stream, 1, "batch_max")?; + stream.memset_zeros(&mut batch_max_buf) + .map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?; + let scalars_readback_buf = alloc_f32(&stream, 2, "scalars_readback")?; + let scalars_readback_host = [0.0_f32; 2]; + // ── Shared memory sizing ──────────────────────────────────── let shmem_bytes = compute_shmem_bytes(&config); @@ -393,6 +485,7 @@ impl GpuDqnTrainer { grad_norm_kernel, adam_update_kernel, ema_kernel, + per_update_kernel, f32_to_bf16_kernel, online_dueling_bf16: None, online_branching_bf16: None, @@ -431,6 +524,9 @@ impl GpuDqnTrainer { readback_buf, readback_host, upload_staging_host, + batch_max_buf, + scalars_readback_buf, + scalars_readback_host, }) } @@ -577,6 +673,53 @@ impl GpuDqnTrainer { // ── Upload batch data (OUTSIDE graph — host pointers change) ─ self.upload_batch(states, next_states, actions, rewards, dones, is_weights)?; + self.execute_train_and_readback(online_dueling, online_branching) + } + + /// GPU-direct training step — no CPU roundtrip. + /// + /// Takes Candle GPU tensors from `GpuBatch` and copies them to trainer + /// buffers via DtoD memcpy (skipping the GPU→CPU `.to_vec1()` → CPU→GPU + /// `memcpy_htod` roundtrip of `train_step()`). + /// + /// Same CUDA Graph execution and readback as `train_step()`. + #[allow(clippy::too_many_arguments)] + pub fn train_step_gpu( + &mut self, + gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, + target_dueling: &DuelingWeightSet, + target_branching: &BranchingWeightSet, + ) -> Result { + // ── First call: flatten weights + allocate BF16 mirrors ────── + if !self.params_initialized { + self.flatten_online_weights(online_dueling, online_branching)?; + self.ensure_bf16_mirrors( + online_dueling, online_branching, + target_dueling, target_branching, + )?; + self.params_initialized = true; + } + + // ── GPU-direct upload (DtoD, no CPU staging) ───────────────── + self.upload_batch_gpu(gpu_batch)?; + + self.execute_train_scalars_only(online_dueling, online_branching) + } + + /// Common post-upload execution: Adam step → CUDA Graph → readback. + /// + /// Shared by `train_step()` (CPU upload path) and `train_step_gpu()` + /// (GPU-direct DtoD path). Everything after batch data lands in the + /// trainer's CudaSlice buffers is identical. + fn execute_train_and_readback( + &mut self, + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, + ) -> Result { + let b = self.config.batch_size; + // ── Update Adam step counter on device (OUTSIDE graph) ─────── self.adam_step += 1; self.stream @@ -629,6 +772,145 @@ impl GpuDqnTrainer { }) } + /// Scalar-only readback path: only total_loss + grad_norm (8 bytes DtoH). + /// + /// TD errors stay on GPU in `td_errors_buf` — PER priority update is done + /// by `update_priorities_cuda()` which reads them directly via device pointer. + fn execute_train_scalars_only( + &mut self, + online_dueling: &DuelingWeightSet, + online_branching: &BranchingWeightSet, + ) -> Result { + // ── Update Adam step counter on device (OUTSIDE graph) ─────── + self.adam_step += 1; + self.stream + .memcpy_htod(&[self.adam_step], &mut self.t_buf) + .map_err(|e| MLError::ModelError(format!("HtoD adam_step: {e}")))?; + + // ── Execute via CUDA Graph (capture on first call, replay thereafter) ── + if let Some(ref graph) = self.training_graph { + graph.0.launch().map_err(|e| { + MLError::ModelError(format!("CUDA graph replay: {e}")) + })?; + } else { + self.capture_training_graph(online_dueling, online_branching)?; + } + + // ── Gather only 2 scalars: [total_loss(1) | grad_norm(1)] ─── + let readback_base = raw_device_ptr(&self.scalars_readback_buf, &self.stream); + let f32_bytes = std::mem::size_of::(); + + let loss_src = raw_device_ptr(&self.total_loss_buf, &self.stream); + dtod_copy(readback_base, loss_src, f32_bytes, &self.stream, 0, "scalars_gather")?; + + let norm_src = raw_device_ptr(&self.grad_norm_buf, &self.stream); + dtod_copy(readback_base + f32_bytes as u64, norm_src, f32_bytes, &self.stream, 1, "scalars_gather")?; + + // ── Single DtoH transfer: 8 bytes ─────────────────────────── + self.stream + .memcpy_dtoh(&self.scalars_readback_buf, &mut self.scalars_readback_host) + .map_err(|e| MLError::ModelError(format!("DtoH scalars_readback: {e}")))?; + + Ok(FusedTrainScalars { + total_loss: self.scalars_readback_host[0], + grad_norm: self.scalars_readback_host[1].sqrt(), + }) + } + + /// Launch GPU-native PER priority update kernel. + /// + /// Reads `td_errors_buf` (populated by training graph) and scatter-writes + /// new priorities directly to the replay buffer's priority array. + /// Accumulates batch max via atomicMax for epoch-boundary flush. + /// + /// Zero CPU readback, single CUDA kernel launch (replaces 7-8 Candle ops). + /// + /// # Arguments + /// * `indices` — GpuBatch.indices Tensor ([B] u32 on GPU) + /// * `priorities` — GpuReplayBuffer.priorities Tensor ([capacity] f32 on GPU) + /// * `alpha` — PER alpha exponent + /// * `epsilon` — PER epsilon floor + pub fn update_priorities_cuda( + &mut self, + indices: &candle_core::Tensor, + priorities: &candle_core::Tensor, + alpha: f32, + epsilon: f32, + ) -> Result<(), MLError> { + // Extract CudaSlice from indices tensor (u32) + let (idx_guard, _) = indices.storage_and_layout(); + let idx_slice = match &*idx_guard { + candle_core::Storage::Cuda(cs) => cs.as_cuda_slice::().map_err(|e| { + MLError::ModelError(format!("PER indices as_cuda_slice: {e}")) + })?, + _ => return Err(MLError::ModelError("PER indices not on CUDA".into())), + }; + + // Extract CudaSlice from priorities tensor (f32) + let (prio_guard, _) = priorities.storage_and_layout(); + let prio_slice = match &*prio_guard { + candle_core::Storage::Cuda(cs) => cs.as_cuda_slice::().map_err(|e| { + MLError::ModelError(format!("PER priorities as_cuda_slice: {e}")) + })?, + _ => return Err(MLError::ModelError("PER priorities not on CUDA".into())), + }; + + let b = self.config.batch_size as i32; + let blocks = ((self.config.batch_size + 255) / 256) as u32; + let launch_cfg = LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + // Safety: argument order matches per_update_priorities_kernel signature. + // td_errors_buf has B elements (populated by training graph). + // idx_slice has B u32 elements (buffer indices from PER sampling). + // prio_slice has capacity f32 elements (priority array in GpuReplayBuffer). + // batch_max_buf has 1 element (atomicMax accumulator, reset per epoch). + // Guards (idx_guard, prio_guard) keep storage alive through launch. + unsafe { + self.stream + .launch_builder(&self.per_update_kernel) + .arg(&self.td_errors_buf) + .arg(idx_slice) + .arg(prio_slice) + .arg(&self.batch_max_buf) + .arg(&alpha) + .arg(&epsilon) + .arg(&b) + .launch(launch_cfg) + .map_err(|e| MLError::ModelError(format!("per_update_priorities_kernel: {e}")))?; + } + + Ok(()) + } + + /// Read back the accumulated batch max priority and reset for next epoch. + /// + /// Single scalar DtoH (4 bytes) — called once per epoch, not per step. + /// Returns 0.0 if no priorities were updated this epoch. + pub fn batch_max_readback_and_reset(&mut self) -> Result { + let mut host = [0.0_f32; 1]; + self.stream + .memcpy_dtoh(&self.batch_max_buf, &mut host) + .map_err(|e| MLError::ModelError(format!("DtoH batch_max: {e}")))?; + // Reset to 0 for next epoch + self.stream + .memset_zeros(&mut self.batch_max_buf) + .map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?; + Ok(host[0]) + } + + /// Reference to the td_errors buffer on GPU. + /// + /// Shape: `[B]` f32 — per-sample TD errors from the last training step. + /// Valid after `train_step_gpu()` or `train_step()` has been called. + /// Used by the PER priority update kernel to avoid DtoH readback. + pub fn td_errors_buf(&self) -> &CudaSlice { + &self.td_errors_buf + } + /// Run only the forward + C51 loss kernel (no backward/Adam). /// /// Useful for validation loss computation without updating weights. @@ -947,6 +1229,59 @@ impl GpuDqnTrainer { Ok(()) } + /// Upload batch data from GPU tensors directly — zero CPU roundtrip. + /// + /// Takes Candle tensors from `GpuBatch`, casts BF16→F32 on GPU if needed, + /// and copies to the trainer's CudaSlice buffers via DtoD memcpy. + /// Replaces the GPU→CPU `.to_vec1()` → CPU→GPU `memcpy_htod` roundtrip. + /// + /// ~6 DtoD copies vs old path's 1 HtoD + 6 DtoD + GPU→CPU readback. + /// Net saving: eliminates PCIe round-trip (~10-30μs on H100 Gen5). + fn upload_batch_gpu( + &mut self, + gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, + ) -> Result<(), MLError> { + use candle_core::DType; + + let b = self.config.batch_size; + let sd = self.config.state_dim; + + // ── Cast to F32 + contiguous on GPU (no CPU touch) ─────────── + let states_f32 = gpu_batch.states.to_dtype(DType::F32) + .and_then(|t| t.contiguous()) + .map_err(|e| MLError::ModelError(format!("GPU batch states F32: {e}")))?; + let next_f32 = gpu_batch.next_states.to_dtype(DType::F32) + .and_then(|t| t.contiguous()) + .map_err(|e| MLError::ModelError(format!("GPU batch next_states F32: {e}")))?; + let rewards_f32 = gpu_batch.rewards.to_dtype(DType::F32) + .and_then(|t| t.contiguous()) + .map_err(|e| MLError::ModelError(format!("GPU batch rewards F32: {e}")))?; + let dones_f32 = gpu_batch.dones.to_dtype(DType::F32) + .and_then(|t| t.contiguous()) + .map_err(|e| MLError::ModelError(format!("GPU batch dones F32: {e}")))?; + let weights_f32 = gpu_batch.weights.to_dtype(DType::F32) + .and_then(|t| t.contiguous()) + .map_err(|e| MLError::ModelError(format!("GPU batch weights F32: {e}")))?; + // Actions: U32 tensor (bit-compatible with I32 for values in [0, 44]) + let actions_cont = gpu_batch.actions.contiguous() + .map_err(|e| MLError::ModelError(format!("GPU batch actions contiguous: {e}")))?; + + // ── DtoD copy from Candle tensor storage → trainer CudaSlice buffers ── + dtod_from_candle_f32(&states_f32, &self.states_buf, b * sd, &self.stream, "states")?; + dtod_from_candle_f32(&next_f32, &self.next_states_buf, b * sd, &self.stream, "next_states")?; + dtod_from_candle_u32_to_i32(&actions_cont, &self.actions_buf, b, &self.stream, "actions")?; + dtod_from_candle_f32(&rewards_f32, &self.rewards_buf, b, &self.stream, "rewards")?; + dtod_from_candle_f32(&dones_f32, &self.dones_buf, b, &self.stream, "dones")?; + dtod_from_candle_f32(&weights_f32, &self.is_weights_buf, b, &self.stream, "is_weights")?; + + // ── Sync stream: ensure DtoD copies complete before temp tensors drop ── + // One stream sync vs old path's GPU→CPU→GPU PCIe roundtrip. + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("GPU upload sync: {e}")))?; + + Ok(()) + } + // ═══════════════════════════════════════════════════════════════════ // Kernel launch methods // ═══════════════════════════════════════════════════════════════════ @@ -1405,10 +1740,16 @@ fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, ) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { + // shmem_max_in_dim must cover ALL weight matrix row widths used by + // cooperative_load_tile_bf16 in the kernel — including value/advantage heads. + // Missing value_h/adv_h caused shmem tile overflow → CUDA_ERROR_ILLEGAL_ADDRESS + // on GPUs with ≤48KB physical shmem (RTX 3050) when head dims > trunk dims. let shmem_max_in_dim = config .state_dim .max(config.shared_h1) - .max(config.shared_h2); + .max(config.shared_h2) + .max(config.value_h) + .max(config.adv_h); let shmem_tile_rows = compute_shmem_tile_rows(shmem_max_in_dim); @@ -1565,6 +1906,53 @@ void dqn_ema_kernel(float* __restrict__ target, }) } +/// Compile the standalone PER priority update kernel. +/// +/// `per_update_priorities_kernel(td_errors, indices, priorities, batch_max, alpha, epsilon, B)`: +/// Computes `|td|^alpha + eps`, scatter-writes to priorities, atomicMax for batch max. +/// Single kernel launch replaces 7-8 Candle tensor ops per training step. +fn compile_per_update_kernel(stream: &Arc) -> Result { + let src = r#" +extern "C" __global__ +void per_update_priorities_kernel( + const float* __restrict__ td_errors, + const unsigned int* __restrict__ indices, + float* __restrict__ priorities, + float* __restrict__ batch_max, + float alpha, + float epsilon, + int batch_size) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + + float td = fabsf(td_errors[i]); + float new_prio = powf(td, alpha) + epsilon; + + // Clamp to [epsilon, 1e6] — same bounds as GpuReplayBuffer::update_priorities_gpu + if (new_prio < epsilon) new_prio = epsilon; + if (new_prio > 1e6f) new_prio = 1e6f; + + // Scatter-write priority at buffer index + unsigned int idx = indices[i]; + priorities[idx] = new_prio; + + // atomicMax for batch max: IEEE 754 positive floats preserve int ordering + int ival = __float_as_int(new_prio); + atomicMax((int*)batch_max, ival); +} +"#; + let context = stream.context(); + let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) + .map_err(|e| MLError::ModelError(format!("per_update_priorities_kernel compilation: {e}")))?; + let module = context.load_module(ptx).map_err(|e| { + MLError::ModelError(format!("per_update module load: {e}")) + })?; + module.load_function("per_update_priorities_kernel").map_err(|e| { + MLError::ModelError(format!("per_update_priorities_kernel load: {e}")) + }) +} + // ── Shared memory sizing ──────────────────────────────────────────────────── fn compute_shmem_tile_rows(shmem_max_in_dim: usize) -> usize { @@ -1580,7 +1968,9 @@ fn compute_shmem_bytes(config: &GpuDqnTrainConfig) -> usize { let shmem_max_in_dim = config .state_dim .max(config.shared_h1) - .max(config.shared_h2); + .max(config.shared_h2) + .max(config.value_h) + .max(config.adv_h); let shmem_tile_rows = compute_shmem_tile_rows(shmem_max_in_dim); // BF16 weight tiles are half the byte-width → 2× rows fit in the same region. // Forward kernel bias tile is sized for the doubled BF16 tile rows. @@ -1609,7 +1999,7 @@ fn compute_shmem_bytes(config: &GpuDqnTrainConfig) -> usize { /// /// Same pattern as `raw_device_ptr` in `gpu_weights.rs` — wraps the SyncOnDrop /// guard in ManuallyDrop to skip read event recording (safe on same stream). -fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { +pub(crate) fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr(stream); let _no_drop = std::mem::ManuallyDrop::new(guard); ptr @@ -1625,7 +2015,7 @@ fn raw_device_ptr_i32(slice: &CudaSlice, stream: &CudaStream) -> u64 { } /// Async device-to-device memcpy with error context. -fn dtod_copy( +pub(crate) fn dtod_copy( dst: u64, src: u64, num_bytes: usize, @@ -1643,6 +2033,67 @@ fn dtod_copy( Ok(()) } +// ── Candle tensor → CudaSlice DtoD copy helpers ──────────────────────────── + +/// DtoD copy from a contiguous F32 Candle tensor to a pre-allocated CudaSlice. +/// +/// Extracts the raw device pointer from the tensor's CudaStorage and issues +/// an async DtoD memcpy. Caller must ensure the tensor is contiguous F32 on CUDA. +fn dtod_from_candle_f32( + tensor: &candle_core::Tensor, + dst: &CudaSlice, + num_elements: usize, + stream: &Arc, + ctx: &str, +) -> Result<(), MLError> { + let (storage_guard, _layout) = tensor.storage_and_layout(); + let cuda_storage = match &*storage_guard { + candle_core::Storage::Cuda(cs) => cs, + _ => { + return Err(MLError::ModelError(format!( + "dtod_from_candle_f32 {ctx}: tensor must be on CUDA device" + ))); + } + }; + let src_slice = cuda_storage.as_cuda_slice::().map_err(|e| { + MLError::ModelError(format!("dtod_from_candle_f32 {ctx}: as_cuda_slice: {e}")) + })?; + let src_ptr = raw_device_ptr(src_slice, stream); + let dst_ptr = raw_device_ptr(dst, stream); + dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) +} + +/// DtoD copy from a contiguous U32 Candle tensor to a CudaSlice. +/// +/// U32 and I32 are bit-compatible for action indices in [0, 44]. +fn dtod_from_candle_u32_to_i32( + tensor: &candle_core::Tensor, + dst: &CudaSlice, + num_elements: usize, + stream: &Arc, + ctx: &str, +) -> Result<(), MLError> { + let (storage_guard, _layout) = tensor.storage_and_layout(); + let cuda_storage = match &*storage_guard { + candle_core::Storage::Cuda(cs) => cs, + _ => { + return Err(MLError::ModelError(format!( + "dtod_from_candle_u32 {ctx}: tensor must be on CUDA device" + ))); + } + }; + let src_slice = cuda_storage.as_cuda_slice::().map_err(|e| { + MLError::ModelError(format!("dtod_from_candle_u32 {ctx}: as_cuda_slice: {e}")) + })?; + let src_ptr = { + let (ptr, guard) = src_slice.device_ptr(stream); + let _no_drop = std::mem::ManuallyDrop::new(guard); + ptr + }; + let dst_ptr = raw_device_ptr_i32(dst, stream); + dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) +} + // ── Allocation helpers ────────────────────────────────────────────────────── fn alloc_f32( diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs new file mode 100644 index 000000000..bd9019f33 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -0,0 +1,754 @@ +#![allow(unsafe_code)] // Required for CUDA kernel launches + +//! GPU-accelerated IQN (Implicit Quantile Network) dual-head trainer. +//! +//! Auxiliary quantile head that shares the DQN trunk (h_s2) but trains +//! its own weights independently. Provides quantile-based Q-values for +//! risk-sensitive action selection (CVaR) alongside the primary C51 head. +//! +//! ## Architecture +//! +//! For each τ ∈ (0,1): +//! cos_feat = [cos(π·1·τ), ..., cos(π·D·τ)] +//! embed = ReLU(W_embed × cos_feat + b_embed) +//! combined = h_s2 ⊙ embed +//! q_branch = W_bd × combined + b_bd (per branch) +//! +//! ## Kernel pipeline +//! +//! 1. **Trunk forward** (`iqn_trunk_forward_kernel`): compute target h_s2 from next_states +//! 2. **Forward + loss** (`iqn_forward_loss_kernel`): quantile forward + Huber loss +//! 3. **Backward** (`iqn_backward_kernel`): gradients through IQN weights only +//! 4. **Grad norm** (`iqn_grad_norm_kernel`): L2 gradient norm +//! 5. **Adam** (`iqn_adam_kernel`): AdamW with gradient clipping +//! 6. **EMA** (`iqn_ema_kernel`): Polyak target weight update +//! 7. **Forward only** (`iqn_forward_kernel`): inference, returns expected Q per action +//! +//! ## Integration +//! +//! Called after the main DQN training step in `FusedTrainingCtx::run_full_step()`. +//! Reads `save_h_s2` (online trunk activation) from `GpuDqnTrainer` and computes +//! target h_s2 from `next_states_buf` using the target DQN trunk weights. + +use std::sync::Arc; + +use candle_core::cuda_backend::cudarc; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use cudarc::nvrtc::Ptx; +use tracing::info; + +use crate::MLError; +use super::gpu_weights::DuelingWeightSet; + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for the GPU IQN dual-head trainer. +#[derive(Debug, Clone)] +pub struct GpuIqnConfig { + /// Hidden dimension = SHARED_H2 from DQN trunk (default 256). + pub hidden_dim: usize, + /// Cosine embedding dimension (default 64). + pub embed_dim: usize, + /// Number of quantile samples per forward pass (default 32). + pub num_quantiles: usize, + /// Huber threshold for quantile loss (default 1.0). + pub kappa: f32, + /// Input state dimension (tensor-core-aligned, e.g. 48 or 56). + pub state_dim: usize, + /// First shared hidden layer width from DQN trunk (default 256). + pub shared_h1: usize, + /// Branch 0 (exposure) action count. + pub branch_0_size: usize, + /// Branch 1 (order) action count. + pub branch_1_size: usize, + /// Branch 2 (urgency) action count. + pub branch_2_size: usize, + /// Training batch size (must match DQN trainer). + pub batch_size: usize, + /// Learning rate for IQN Adam optimizer. + pub lr: f32, + /// Adam beta1 (first moment decay). + pub beta1: f32, + /// Adam beta2 (second moment decay). + pub beta2: f32, + /// Adam epsilon (numerical stability). + pub epsilon: f32, + /// Decoupled weight decay (AdamW). + pub weight_decay: f32, + /// Maximum gradient L2 norm for clipping. + pub max_grad_norm: f32, + /// Discount factor for Bellman projection. + pub gamma: f32, +} + +impl Default for GpuIqnConfig { + fn default() -> Self { + Self { + hidden_dim: 256, + embed_dim: 64, + num_quantiles: 32, + kappa: 1.0, + state_dim: 48, + shared_h1: 256, + branch_0_size: 5, + branch_1_size: 3, + branch_2_size: 3, + batch_size: 256, + lr: 3e-4, + beta1: 0.9, + beta2: 0.999, + epsilon: 1e-8, + weight_decay: 1e-5, + max_grad_norm: 1.0, + gamma: 0.99, + } + } +} + +impl GpuIqnConfig { + /// Total branch actions (sum of all branches). + fn total_branch_actions(&self) -> usize { + self.branch_0_size + self.branch_1_size + self.branch_2_size + } + + /// Total IQN parameters (embedding + per-branch output layers). + fn total_params(&self) -> usize { + let h = self.hidden_dim; + let d = self.embed_dim; + let b0 = self.branch_0_size; + let b1 = self.branch_1_size; + let b2 = self.branch_2_size; + // W_embed[H, D] + b_embed[H] + W_b0[B0, H] + b_b0[B0] + // + W_b1[B1, H] + b_b1[B1] + W_b2[B2, H] + b_b2[B2] + h * d + h + b0 * h + b0 + b1 * h + b1 + b2 * h + b2 + } +} + +// --------------------------------------------------------------------------- +// GpuIqnHead +// --------------------------------------------------------------------------- + +/// GPU-accelerated IQN dual-head trainer. +/// +/// Owns online + target IQN weight buffers, Adam state, compiled CUDA kernels, +/// and all intermediate activation/loss buffers. Operates entirely on GPU. +#[allow(missing_debug_implementations)] +pub struct GpuIqnHead { + config: GpuIqnConfig, + stream: Arc, + + // ── Compiled kernels ──────────────────────────────────────────── + trunk_forward_kernel: CudaFunction, + forward_loss_kernel: CudaFunction, + backward_kernel: CudaFunction, + grad_norm_kernel: CudaFunction, + adam_kernel: CudaFunction, + forward_kernel: CudaFunction, + ema_kernel: CudaFunction, + decode_actions_kernel: CudaFunction, + sample_taus_kernel: CudaFunction, + + // ── IQN online parameters (flat f32 on GPU) ────────────────────── + online_params: CudaSlice, + // ── IQN target parameters (for Bellman target) ─────────────────── + target_params: CudaSlice, + + // ── Adam optimizer state ───────────────────────────────────────── + m_buf: CudaSlice, + v_buf: CudaSlice, + grad_buf: CudaSlice, + grad_norm_buf: CudaSlice, + + // ── Per-step buffers ───────────────────────────────────────────── + /// Pre-sampled τ values for online network [B, N] + online_taus: CudaSlice, + /// Pre-sampled τ values for target network [B, N] + target_taus: CudaSlice, + /// Branch actions decoded from flat actions [B, 3] + branch_actions: CudaSlice, + /// Target h_s2 computed from next_states + target trunk weights [B, H] + target_h_s2: CudaSlice, + /// Persistent rewards buffer [B] — avoids per-step temp allocs + rewards_buf: CudaSlice, + /// Persistent dones buffer [B] — avoids per-step temp allocs + dones_buf: CudaSlice, + + // ── Activation save buffers (forward → backward) ───────────────── + save_embed: CudaSlice, // [B, N, H] + save_combined: CudaSlice, // [B, N, H] + save_q_online: CudaSlice, // [B, N, TOTAL_BRANCH_ACTIONS] + save_q_target: CudaSlice, // [B, N, TOTAL_BRANCH_ACTIONS] + + // ── Loss buffers ───────────────────────────────────────────────── + per_sample_loss: CudaSlice, // [B] + total_loss: CudaSlice, // [1] + + // ── Training state ─────────────────────────────────────────────── + adam_step: i32, + /// Monotonic step counter for Philox PRNG seeding (τ sampling). + rng_step: u32, + total_params: usize, +} + +impl GpuIqnHead { + /// Create a new GPU IQN dual-head trainer. + /// + /// Compiles all 7 CUDA kernels, initializes online + target weights with + /// Xavier/Glorot, and pre-allocates all GPU buffers. + pub fn new( + stream: Arc, + config: GpuIqnConfig, + ) -> Result { + let total_params = config.total_params(); + let b = config.batch_size; + let n = config.num_quantiles; + let h = config.hidden_dim; + let tba = config.total_branch_actions(); + + // Compile all 9 kernels + let (trunk_fwd, fwd_loss, bwd, gnorm, adam, fwd_only, ema, decode_act, sample_taus_k) = + compile_iqn_kernels(&stream, &config)?; + + // Allocate and initialize weights (Xavier/Glorot) + let online_params = init_iqn_xavier_weights(&stream, &config)?; + // Target params start as a copy of online params + let target_params = clone_cuda_slice(&stream, &online_params, total_params)?; + + // Adam state (zeroed) + let m_buf = alloc_f32(&stream, total_params, "iqn_m")?; + let v_buf = alloc_f32(&stream, total_params, "iqn_v")?; + let grad_buf = alloc_f32(&stream, total_params, "iqn_grad")?; + let grad_norm_buf = alloc_f32(&stream, 1, "iqn_grad_norm")?; + + // Per-step buffers + let online_taus = alloc_f32(&stream, b * n, "iqn_online_taus")?; + let target_taus = alloc_f32(&stream, b * n, "iqn_target_taus")?; + let branch_actions = stream.alloc_zeros::(b * 3).map_err(|e| { + MLError::ModelError(format!("IQN alloc branch_actions: {e}")) + })?; + let target_h_s2 = alloc_f32(&stream, b * h, "iqn_target_h_s2")?; + + // Activation saves + let save_embed = alloc_f32(&stream, b * n * h, "iqn_save_embed")?; + let save_combined = alloc_f32(&stream, b * n * h, "iqn_save_combined")?; + let save_q_online = alloc_f32(&stream, b * n * tba, "iqn_save_q_online")?; + let save_q_target = alloc_f32(&stream, b * n * tba, "iqn_save_q_target")?; + + // Loss buffers + let per_sample_loss = alloc_f32(&stream, b, "iqn_per_sample_loss")?; + let total_loss = alloc_f32(&stream, 1, "iqn_total_loss")?; + + // Persistent rewards/dones buffers (avoids per-step temp allocs) + let rewards_buf = alloc_f32(&stream, b, "iqn_rewards")?; + let dones_buf = alloc_f32(&stream, b, "iqn_dones")?; + + let vram_bytes = (total_params * 6 + b * n * 2 + b * 3 + + b * h + b * n * h * 2 + b * n * tba * 2 + + b * 3 + 2) * 4; + + info!( + hidden_dim = h, + embed_dim = config.embed_dim, + num_quantiles = n, + kappa = config.kappa, + total_params, + total_branch_actions = tba, + vram_kb = vram_bytes / 1024, + "GpuIqnHead initialized: 9 kernels compiled, Xavier weights" + ); + + Ok(Self { + config, + stream, + trunk_forward_kernel: trunk_fwd, + forward_loss_kernel: fwd_loss, + backward_kernel: bwd, + grad_norm_kernel: gnorm, + adam_kernel: adam, + forward_kernel: fwd_only, + ema_kernel: ema, + decode_actions_kernel: decode_act, + sample_taus_kernel: sample_taus_k, + online_params, + target_params, + m_buf, + v_buf, + grad_buf, + grad_norm_buf, + online_taus, + target_taus, + branch_actions, + target_h_s2, + rewards_buf, + dones_buf, + save_embed, + save_combined, + save_q_online, + save_q_target, + per_sample_loss, + total_loss, + adam_step: 0, + rng_step: 0, + total_params, + }) + } + + /// Run one GPU-native IQN training step. + /// + /// All data sourced from DQN trainer's GPU buffers via DtoD copies. + /// No CPU arrays, no HtoD uploads for batch data. + /// + /// Full pipeline: + /// 1. Sample τ values and upload (small, ~32KB) + /// 2. Decode flat actions → branch actions on GPU via decode kernel + /// 3. DtoD copy rewards/dones from DQN trainer buffers + /// 4. Compute target h_s2 via trunk forward kernel + /// 5. Forward + loss + /// 6. Backward + grad norm + Adam + /// + /// Returns the batch-mean quantile Huber loss. + pub fn train_iqn_step_gpu( + &mut self, + online_h_s2: &CudaSlice, + next_states_buf: &CudaSlice, + target_dueling: &DuelingWeightSet, + dqn_actions_buf: &CudaSlice, + dqn_rewards_buf: &CudaSlice, + dqn_dones_buf: &CudaSlice, + ) -> Result { + let b = self.config.batch_size; + let n = self.config.num_quantiles; + + // 1. Sample τ values ∈ (0.01, 0.99) on GPU via Philox PRNG kernel + let total_taus = (b * n) as i32; + let tau_blocks = ((b * n) + 255) / 256; + let tau_config = LaunchConfig { + grid_dim: (tau_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + // Unique seeds per (step, buffer) — counter-based PRNG, no state + let online_seed = self.rng_step.wrapping_mul(2); + let target_seed = self.rng_step.wrapping_mul(2).wrapping_add(1); + self.rng_step = self.rng_step.wrapping_add(1); + + // Safety: online_taus[B*N] and target_taus[B*N] are valid GPU allocations. + unsafe { + self.stream + .launch_builder(&self.sample_taus_kernel) + .arg(&mut self.online_taus) + .arg(&online_seed) + .arg(&total_taus) + .launch(tau_config) + .map_err(|e| MLError::ModelError(format!("IQN sample online_taus: {e}")))?; + self.stream + .launch_builder(&self.sample_taus_kernel) + .arg(&mut self.target_taus) + .arg(&target_seed) + .arg(&total_taus) + .launch(tau_config) + .map_err(|e| MLError::ModelError(format!("IQN sample target_taus: {e}")))?; + } + + // 2. GPU decode: flat actions [B] → branch actions [B, 3] + let decode_blocks = (b + 255) / 256; + let decode_config = LaunchConfig { + grid_dim: (decode_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let batch_size_i32 = b as i32; + + // Safety: dqn_actions_buf[B] i32, branch_actions[B*3] i32, same context. + unsafe { + self.stream + .launch_builder(&self.decode_actions_kernel) + .arg(dqn_actions_buf) + .arg(&mut self.branch_actions) + .arg(&batch_size_i32) + .launch(decode_config) + .map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?; + } + + // 3. DtoD copy rewards/dones from DQN trainer buffers + use super::gpu_dqn_trainer::{raw_device_ptr, dtod_copy}; + let f32_bytes = std::mem::size_of::(); + let r_src = raw_device_ptr(dqn_rewards_buf, &self.stream); + let r_dst = raw_device_ptr(&self.rewards_buf, &self.stream); + dtod_copy(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?; + + let d_src = raw_device_ptr(dqn_dones_buf, &self.stream); + let d_dst = raw_device_ptr(&self.dones_buf, &self.stream); + dtod_copy(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?; + + // 4-9. Execute core training pipeline (shared logic) + self.execute_training_pipeline(online_h_s2, next_states_buf, target_dueling) + } + + /// Core IQN training pipeline: trunk forward → forward+loss → backward → Adam. + /// + /// Assumes `branch_actions`, `rewards_buf`, `dones_buf` are already populated + /// on GPU (by either `train_iqn_step_gpu` or the legacy HtoD path). + fn execute_training_pipeline( + &mut self, + online_h_s2: &CudaSlice, + next_states_buf: &CudaSlice, + target_dueling: &DuelingWeightSet, + ) -> Result { + let b = self.config.batch_size; + + // 4. Compute target h_s2 via trunk forward kernel + let shmem_bytes = self.config.shared_h1 * 4; + let trunk_config = LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: shmem_bytes as u32, + }; + + // Safety: all pointers are valid GPU allocations on the same context. + // next_states_buf points to the DQN trainer's pre-allocated buffer. + // target_dueling weights are extracted from VarMap CudaSlice buffers. + unsafe { + self.stream + .launch_builder(&self.trunk_forward_kernel) + .arg(next_states_buf) + .arg(&target_dueling.w_s1) + .arg(&target_dueling.b_s1) + .arg(&target_dueling.w_s2) + .arg(&target_dueling.b_s2) + .arg(&mut self.target_h_s2) + .arg(&(b as i32)) + .launch(trunk_config) + .map_err(|e| MLError::ModelError(format!("IQN trunk forward kernel: {e}")))?; + } + + // 5. Zero total_loss, grad_norm, and gradient buffer on GPU (memset, no CPU alloc) + self.stream.memset_zeros(&mut self.total_loss) + .map_err(|e| MLError::ModelError(format!("IQN zero total_loss: {e}")))?; + self.stream.memset_zeros(&mut self.grad_norm_buf) + .map_err(|e| MLError::ModelError(format!("IQN zero grad_norm: {e}")))?; + self.stream.memset_zeros(&mut self.grad_buf) + .map_err(|e| MLError::ModelError(format!("IQN zero grads: {e}")))?; + + let batch_size_i32 = b as i32; + let gamma = self.config.gamma; + + // 6. Forward + loss kernel: one warp per sample + let fwd_config = LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + + // Safety: all buffer pointers are valid GPU allocations on the same context + // and stream. online_h_s2 points to the DQN trainer's save_h_s2 buffer. + unsafe { + self.stream + .launch_builder(&self.forward_loss_kernel) + .arg(online_h_s2) + .arg(&self.target_h_s2) + .arg(&self.online_taus) + .arg(&self.target_taus) + .arg(&self.branch_actions) + .arg(&self.rewards_buf) + .arg(&self.dones_buf) + .arg(&gamma) + .arg(&self.online_params) + .arg(&self.target_params) + .arg(&mut self.per_sample_loss) + .arg(&mut self.total_loss) + .arg(&mut self.save_embed) + .arg(&mut self.save_combined) + .arg(&mut self.save_q_online) + .arg(&mut self.save_q_target) + .arg(&batch_size_i32) + .launch(fwd_config) + .map_err(|e| MLError::ModelError(format!("IQN forward+loss kernel: {e}")))?; + } + + // 7. Backward kernel: one warp per sample + let bwd_config = LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }; + + // Safety: all pointers valid, save buffers written by forward kernel on same stream. + unsafe { + self.stream + .launch_builder(&self.backward_kernel) + .arg(online_h_s2) + .arg(&self.save_embed) + .arg(&self.save_combined) + .arg(&self.save_q_online) + .arg(&self.save_q_target) + .arg(&self.online_taus) + .arg(&self.branch_actions) + .arg(&self.online_params) + .arg(&mut self.grad_buf) + .arg(&batch_size_i32) + .launch(bwd_config) + .map_err(|e| MLError::ModelError(format!("IQN backward kernel: {e}")))?; + } + + // 8. Grad norm kernel + let total_params_i32 = self.total_params as i32; + let norm_blocks = (self.total_params + 255) / 256; + let norm_config = LaunchConfig { + grid_dim: (norm_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + // Safety: grad_buf has total_params elements, grad_norm_buf has 1. + unsafe { + self.stream + .launch_builder(&self.grad_norm_kernel) + .arg(&self.grad_buf) + .arg(&mut self.grad_norm_buf) + .arg(&total_params_i32) + .launch(norm_config) + .map_err(|e| MLError::ModelError(format!("IQN grad_norm kernel: {e}")))?; + } + + // 9. Adam update kernel + self.adam_step += 1; + let adam_blocks = (self.total_params + 255) / 256; + let adam_config = LaunchConfig { + grid_dim: (adam_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + let lr = self.config.lr; + let beta1 = self.config.beta1; + let beta2 = self.config.beta2; + let eps = self.config.epsilon; + let wd = self.config.weight_decay; + let mgn = self.config.max_grad_norm; + let t = self.adam_step; + + // Safety: all buffers have total_params elements on the same context. + unsafe { + self.stream + .launch_builder(&self.adam_kernel) + .arg(&mut self.online_params) + .arg(&mut self.grad_buf) + .arg(&mut self.m_buf) + .arg(&mut self.v_buf) + .arg(&self.grad_norm_buf) + .arg(&lr) + .arg(&beta1) + .arg(&beta2) + .arg(&eps) + .arg(&wd) + .arg(&mgn) + .arg(&t) + .arg(&total_params_i32) + .launch(adam_config) + .map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?; + } + + // Read back total loss (single scalar, ~4 bytes) + let mut loss_host = [0.0_f32]; + self.stream.memcpy_dtoh(&self.total_loss, &mut loss_host) + .map_err(|e| MLError::ModelError(format!("IQN loss readback: {e}")))?; + + Ok(loss_host[0]) + } + + /// Update target IQN weights via Polyak EMA. + /// + /// `target[i] = (1 - tau) * target[i] + tau * online[i]` + pub fn target_ema_update(&mut self, tau: f32) -> Result<(), MLError> { + let n = self.total_params; + let blocks = (n + 255) / 256; + let config = LaunchConfig { + grid_dim: (blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + let n_i32 = n as i32; + + // Safety: target_params and online_params have total_params elements each. + unsafe { + self.stream + .launch_builder(&self.ema_kernel) + .arg(&mut self.target_params) + .arg(&self.online_params) + .arg(&tau) + .arg(&n_i32) + .launch(config) + .map_err(|e| MLError::ModelError(format!("IQN EMA kernel: {e}")))?; + } + + Ok(()) + } + + /// Current Adam step count. + pub fn adam_step(&self) -> i32 { + self.adam_step + } + + /// Per-sample IQN loss buffer reference (for PER weighting). + pub fn per_sample_loss(&self) -> &CudaSlice { + &self.per_sample_loss + } +} + +// --------------------------------------------------------------------------- +// Kernel compilation +// --------------------------------------------------------------------------- + +/// Compile all 9 IQN CUDA kernels with dimension defines. +fn compile_iqn_kernels( + stream: &Arc, + config: &GpuIqnConfig, +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { + let dim_overrides = format!( + "#define IQN_HIDDEN {hidden}\n\ + #define IQN_EMBED_DIM {embed}\n\ + #define IQN_NUM_QUANTILES {nq}\n\ + #define IQN_KAPPA {kappa:.6}f\n\ + #define BRANCH_0_SIZE {b0}\n\ + #define BRANCH_1_SIZE {b1}\n\ + #define BRANCH_2_SIZE {b2}\n\ + #define IQN_STATE_DIM {sd}\n\ + #define IQN_SHARED_H1 {h1}\n", + hidden = config.hidden_dim, + embed = config.embed_dim, + nq = config.num_quantiles, + kappa = config.kappa, + b0 = config.branch_0_size, + b1 = config.branch_1_size, + b2 = config.branch_2_size, + sd = config.state_dim, + h1 = config.shared_h1, + ); + + let kernel_src = include_str!("iqn_dual_head_kernel.cu"); + let full_source = format!("{dim_overrides}\n{kernel_src}"); + + info!( + hidden = config.hidden_dim, + embed_dim = config.embed_dim, + num_quantiles = config.num_quantiles, + kappa = config.kappa, + total_params = config.total_params(), + "GpuIqnHead: compiling 9 IQN kernels" + ); + + let context = stream.context(); + let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) + .map_err(|e| { + MLError::ModelError(format!("iqn_dual_head_kernel compilation failed: {e}")) + })?; + let module = context.load_module(ptx).map_err(|e| { + MLError::ModelError(format!("IQN module load: {e}")) + })?; + + let load = |name: &str| -> Result { + module.load_function(name).map_err(|e| { + MLError::ModelError(format!("{name} load: {e}")) + }) + }; + + let trunk_fwd = load("iqn_trunk_forward_kernel")?; + let fwd_loss = load("iqn_forward_loss_kernel")?; + let backward = load("iqn_backward_kernel")?; + let grad_norm = load("iqn_grad_norm_kernel")?; + let adam = load("iqn_adam_kernel")?; + let fwd_only = load("iqn_forward_kernel")?; + let ema = load("iqn_ema_kernel")?; + let decode_act = load("iqn_decode_actions_kernel")?; + let sample_taus = load("iqn_sample_taus_kernel")?; + + info!("GpuIqnHead: 9 kernels compiled and loaded"); + Ok((trunk_fwd, fwd_loss, backward, grad_norm, adam, fwd_only, ema, decode_act, sample_taus)) +} + +// --------------------------------------------------------------------------- +// Weight initialization +// --------------------------------------------------------------------------- + +/// Initialize IQN weights with Xavier/Glorot uniform distribution. +fn init_iqn_xavier_weights( + stream: &Arc, + config: &GpuIqnConfig, +) -> Result, MLError> { + use rand::Rng; + + let total = config.total_params(); + let h = config.hidden_dim; + let d = config.embed_dim; + let b0 = config.branch_0_size; + let b1 = config.branch_1_size; + let b2 = config.branch_2_size; + let mut weights = vec![0.0_f32; total]; + let mut rng = rand::thread_rng(); + let mut offset = 0; + + // W_embed [H, D] + let limit = (6.0_f64 / (h + d) as f64).sqrt() as f32; + for w in &mut weights[offset..offset + h * d] { + *w = rng.gen_range(-limit..limit); + } + offset += h * d; + // b_embed [H] — zero + offset += h; + + // Branch output layers + for branch_size in [b0, b1, b2] { + let limit = (6.0_f64 / (branch_size + h) as f64).sqrt() as f32; + // W_bd [branch_size, H] + for w in &mut weights[offset..offset + branch_size * h] { + *w = rng.gen_range(-limit..limit); + } + offset += branch_size * h; + // b_bd [branch_size] — zero + offset += branch_size; + } + + debug_assert_eq!(offset, total, "IQN weight layout mismatch"); + + // Upload to GPU + let mut params_buf = alloc_f32(stream, total, "iqn_online_params")?; + stream.memcpy_htod(&weights, &mut params_buf) + .map_err(|e| MLError::ModelError(format!("IQN weight upload: {e}")))?; + + Ok(params_buf) +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/// Allocate a zeroed f32 buffer on GPU. +fn alloc_f32( + stream: &Arc, + n: usize, + name: &str, +) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("GpuIqn alloc {name} ({n} f32): {e}")) + }) +} + +/// Clone a CudaSlice by copying device → host → device. +fn clone_cuda_slice( + stream: &Arc, + src: &CudaSlice, + n: usize, +) -> Result, MLError> { + let mut host = vec![0.0_f32; n]; + stream.memcpy_dtoh(src, &mut host) + .map_err(|e| MLError::ModelError(format!("IQN clone D→H: {e}")))?; + let mut dst = alloc_f32(stream, n, "iqn_target_params")?; + stream.memcpy_htod(&host, &mut dst) + .map_err(|e| MLError::ModelError(format!("IQN clone H→D: {e}")))?; + Ok(dst) +} + diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu new file mode 100644 index 000000000..dc2889f35 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -0,0 +1,875 @@ +/** + * IQN (Implicit Quantile Network) Dual-Head Kernel + * + * Auxiliary quantile head that shares the DQN trunk (h_s2) but trains its own + * weights independently. Provides quantile-based Q-values for risk-sensitive + * action selection (CVaR) alongside the primary C51 distributional head. + * + * Architecture per sample: + * For each τ_i ∈ (0,1), i = 1..N: + * cos_feat = [cos(π·1·τ), ..., cos(π·D·τ)] [IQN_EMBED_DIM] + * embed = ReLU(W_embed × cos_feat + b_embed) [IQN_HIDDEN] + * combined = h_s2 ⊙ embed [IQN_HIDDEN] + * per-branch: q_d = W_bd × combined + b_bd [n_d] scalars + * + * Loss: Quantile Huber loss (Dabney et al., 2018b) + * For each predicted τ_i and target τ'_j: + * δ_ij = target_θ_j - predicted_θ_i + * ρ_τi(δ) = |τ_i - 𝟙{δ<0}| × L_κ(δ) + * L_iqn = (1/N) Σ_i (1/N') Σ_j ρ_τi(δ_ij) + * + * Launch config: + * Forward + loss: grid=(batch_size, 1, 1), block=(32, 1, 1) + * Backward: grid=(batch_size, 1, 1), block=(32, 1, 1) + * Grad norm: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1) + * Adam: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1) + * + * Compile-time defines (injected via NVRTC dim_overrides): + * IQN_HIDDEN -- hidden dim = SHARED_H2 from DQN trunk (default 256) + * IQN_EMBED_DIM -- cosine embedding dimension (default 64) + * IQN_NUM_QUANTILES -- τ samples per forward pass (default 32) + * IQN_KAPPA -- Huber threshold for quantile loss (default 1.0) + * BRANCH_0_SIZE -- exposure actions (default 5) + * BRANCH_1_SIZE -- order type actions (default 3) + * BRANCH_2_SIZE -- urgency actions (default 3) + */ + +/* ── Compile-time defaults (overridden by NVRTC injection) ──────────── */ +#ifndef IQN_HIDDEN +#define IQN_HIDDEN 256 +#endif +#ifndef IQN_EMBED_DIM +#define IQN_EMBED_DIM 64 +#endif +#ifndef IQN_NUM_QUANTILES +#define IQN_NUM_QUANTILES 32 +#endif +#ifndef IQN_KAPPA +#define IQN_KAPPA 1.0f +#endif +#ifndef BRANCH_0_SIZE +#define BRANCH_0_SIZE 5 +#endif +#ifndef BRANCH_1_SIZE +#define BRANCH_1_SIZE 3 +#endif +#ifndef BRANCH_2_SIZE +#define BRANCH_2_SIZE 3 +#endif + +#define TOTAL_BRANCH_ACTIONS (BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE) + +/* ── Weight layout (flat f32 buffer) ─────────────────────────────────── */ +/* W_embed [IQN_HIDDEN, IQN_EMBED_DIM], b_embed [IQN_HIDDEN] + * W_b0 [BRANCH_0_SIZE, IQN_HIDDEN], b_b0 [BRANCH_0_SIZE] + * W_b1 [BRANCH_1_SIZE, IQN_HIDDEN], b_b1 [BRANCH_1_SIZE] + * W_b2 [BRANCH_2_SIZE, IQN_HIDDEN], b_b2 [BRANCH_2_SIZE] */ + +#define IQN_W_EMBED_SIZE (IQN_HIDDEN * IQN_EMBED_DIM) +#define IQN_B_EMBED_SIZE (IQN_HIDDEN) +#define IQN_W_B0_SIZE (BRANCH_0_SIZE * IQN_HIDDEN) +#define IQN_B_B0_SIZE (BRANCH_0_SIZE) +#define IQN_W_B1_SIZE (BRANCH_1_SIZE * IQN_HIDDEN) +#define IQN_B_B1_SIZE (BRANCH_1_SIZE) +#define IQN_W_B2_SIZE (BRANCH_2_SIZE * IQN_HIDDEN) +#define IQN_B_B2_SIZE (BRANCH_2_SIZE) + +#define IQN_TOTAL_PARAMS (IQN_W_EMBED_SIZE + IQN_B_EMBED_SIZE \ + + IQN_W_B0_SIZE + IQN_B_B0_SIZE \ + + IQN_W_B1_SIZE + IQN_B_B1_SIZE \ + + IQN_W_B2_SIZE + IQN_B_B2_SIZE) + +/* Offsets into the flat parameter buffer */ +#define IQN_OFF_W_EMBED 0 +#define IQN_OFF_B_EMBED (IQN_OFF_W_EMBED + IQN_W_EMBED_SIZE) +#define IQN_OFF_W_B0 (IQN_OFF_B_EMBED + IQN_B_EMBED_SIZE) +#define IQN_OFF_B_B0 (IQN_OFF_W_B0 + IQN_W_B0_SIZE) +#define IQN_OFF_W_B1 (IQN_OFF_B_B0 + IQN_B_B0_SIZE) +#define IQN_OFF_B_B1 (IQN_OFF_W_B1 + IQN_W_B1_SIZE) +#define IQN_OFF_W_B2 (IQN_OFF_B_B1 + IQN_B_B1_SIZE) +#define IQN_OFF_B_B2 (IQN_OFF_W_B2 + IQN_W_B2_SIZE) + +/* Distributed array size: ceil(dim / 32) elements per lane */ +#define IQN_DIST(dim) (((dim) + 31) / 32) + +/* ── Device helpers ──────────────────────────────────────────────────── */ + +/** Warp-reduce sum across all 32 lanes (butterfly). */ +__device__ __forceinline__ float iqn_warp_sum(float val) { + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + +/** Quantile Huber loss element: ρ_τ(δ) = |τ - 𝟙{δ<0}| × L_κ(δ) */ +__device__ __forceinline__ float quantile_huber_element(float tau, float delta, float kappa) { + float abs_delta = fabsf(delta); + /* L_κ(δ) = 0.5δ² if |δ|≤κ, κ(|δ| - 0.5κ) otherwise */ + float huber = (abs_delta <= kappa) + ? 0.5f * delta * delta + : kappa * (abs_delta - 0.5f * kappa); + /* Asymmetric weight: |τ - 𝟙{δ<0}| */ + float indicator = (delta < 0.0f) ? 1.0f : 0.0f; + float weight = fabsf(tau - indicator); + return weight * huber; +} + +/** Gradient of quantile Huber loss element w.r.t. predicted value. + * dρ_τ/d(pred) = -|τ - 𝟙{δ<0}| × dL_κ/dδ + * where dL_κ/dδ = δ if |δ|≤κ, κ·sign(δ) otherwise */ +__device__ __forceinline__ float quantile_huber_grad(float tau, float delta, float kappa) { + float abs_delta = fabsf(delta); + float huber_grad = (abs_delta <= kappa) + ? delta + : kappa * ((delta > 0.0f) ? 1.0f : -1.0f); + float indicator = (delta < 0.0f) ? 1.0f : 0.0f; + float weight = fabsf(tau - indicator); + /* Gradient w.r.t. predicted = -weight × huber_grad (negative because δ = target - pred) */ + return -weight * huber_grad; +} + +/* ── Forward + Loss Kernel ──────────────────────────────────────────── */ +/** + * Per-sample IQN forward pass + quantile Huber loss. + * + * Reads h_s2 (shared trunk activation), computes quantile values for all + * branches and τ samples, then computes quantile Huber loss against target + * quantile values (from target network forward through same IQN head). + * + * τ values are pre-sampled and uploaded via the `taus` buffer. + * + * One warp (32 threads) per sample. Quantiles are processed sequentially + * to conserve registers — each τ produces TOTAL_BRANCH_ACTIONS values. + * + * Saves per-quantile intermediate activations for backward pass. + */ +extern "C" __global__ +void iqn_forward_loss_kernel( + /* Inputs */ + const float* __restrict__ h_s2, /* [B, IQN_HIDDEN] trunk activations */ + const float* __restrict__ target_h_s2, /* [B, IQN_HIDDEN] target trunk activations */ + const float* __restrict__ taus, /* [B, N] online τ samples ∈ (0,1) */ + const float* __restrict__ target_taus, /* [B, N] target τ samples */ + const int* __restrict__ actions, /* [B, 3] branch actions (exposure, order, urgency) */ + const float* __restrict__ rewards, /* [B] */ + const float* __restrict__ dones, /* [B] */ + float gamma, + /* Weights */ + const float* __restrict__ online_params, /* [IQN_TOTAL_PARAMS] online IQN weights */ + const float* __restrict__ target_params, /* [IQN_TOTAL_PARAMS] target IQN weights */ + /* Outputs */ + float* __restrict__ per_sample_loss, /* [B] per-sample loss (for PER) */ + float* __restrict__ total_loss, /* [1] batch-mean loss (atomicAdd) */ + /* Activation saves for backward */ + float* __restrict__ save_embed, /* [B, N, IQN_HIDDEN] post-ReLU embeddings */ + float* __restrict__ save_combined, /* [B, N, IQN_HIDDEN] h_s2 ⊙ embed */ + float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] quantile Q-values */ + float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] target Q-values */ + int batch_size +) +{ + int sample = blockIdx.x; + if (sample >= batch_size) return; + int lane = threadIdx.x; + + /* ── Pointers into this sample's data ── */ + const float* my_h_s2 = h_s2 + sample * IQN_HIDDEN; + const float* my_target_h_s2 = target_h_s2 + sample * IQN_HIDDEN; + const float* my_taus = taus + sample * IQN_NUM_QUANTILES; + const float* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES; + + int a0 = actions[sample * 3 + 0]; /* exposure action */ + int a1 = actions[sample * 3 + 1]; /* order action */ + int a2 = actions[sample * 3 + 2]; /* urgency action */ + float reward = rewards[sample]; + float done = dones[sample]; + + /* Weight pointers (same for all samples) */ + const float* w_embed = online_params + IQN_OFF_W_EMBED; + const float* b_embed = online_params + IQN_OFF_B_EMBED; + const float* w_b0 = online_params + IQN_OFF_W_B0; + const float* b_b0 = online_params + IQN_OFF_B_B0; + const float* w_b1 = online_params + IQN_OFF_W_B1; + const float* b_b1 = online_params + IQN_OFF_B_B1; + const float* w_b2 = online_params + IQN_OFF_W_B2; + const float* b_b2 = online_params + IQN_OFF_B_B2; + + const float* tw_embed = target_params + IQN_OFF_W_EMBED; + const float* tb_embed = target_params + IQN_OFF_B_EMBED; + const float* tw_b0 = target_params + IQN_OFF_W_B0; + const float* tb_b0 = target_params + IQN_OFF_B_B0; + const float* tw_b1 = target_params + IQN_OFF_W_B1; + const float* tb_b1 = target_params + IQN_OFF_B_B1; + const float* tw_b2 = target_params + IQN_OFF_W_B2; + const float* tb_b2 = target_params + IQN_OFF_B_B2; + + /* ── Load h_s2 into distributed registers ── */ + float h_dist[IQN_DIST(IQN_HIDDEN)]; + for (int d = lane; d < IQN_HIDDEN; d += 32) + h_dist[d / 32] = my_h_s2[d]; + + float h_target_dist[IQN_DIST(IQN_HIDDEN)]; + for (int d = lane; d < IQN_HIDDEN; d += 32) + h_target_dist[d / 32] = my_target_h_s2[d]; + + /* ── Process each ONLINE quantile ── */ + /* For each τ_i, compute quantile Q-values for the taken actions */ + float online_q_a[IQN_NUM_QUANTILES]; /* Q(s, a_taken, τ_i) per branch sum */ + + for (int t = 0; t < IQN_NUM_QUANTILES; t++) { + float tau = my_taus[t]; + + /* 1. Cosine features: cos(π·(d+1)·τ) for d=0..IQN_EMBED_DIM-1 */ + /* 2. Linear embedding: embed[h] = Σ_d W_embed[h,d]·cos_feat[d] + b_embed[h] */ + /* Process via warp-cooperative matvec: each lane handles IQN_HIDDEN/32 outputs */ + float embed_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) { + float acc = b_embed[h]; + const float* w_row = w_embed + h * IQN_EMBED_DIM; + for (int d = 0; d < IQN_EMBED_DIM; d++) { + float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); + acc += w_row[d] * cos_val; + } + /* ReLU */ + embed_dist[h / 32] = fmaxf(acc, 0.0f); + } + + /* 3. Element-wise product: combined = h_s2 ⊙ embed */ + float comb_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) + comb_dist[h / 32] = h_dist[h / 32] * embed_dist[h / 32]; + + /* Save activations for backward pass */ + int save_offset = (sample * IQN_NUM_QUANTILES + t) * IQN_HIDDEN; + for (int h = lane; h < IQN_HIDDEN; h += 32) { + save_embed[save_offset + h] = embed_dist[h / 32]; + save_combined[save_offset + h] = comb_dist[h / 32]; + } + + /* 4. Per-branch output: q_d = W_bd × combined + b_bd */ + int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; + + /* Branch 0 (exposure) */ + float q0_taken = 0.0f; + for (int a = 0; a < BRANCH_0_SIZE; a++) { + float acc = b_b0[a]; + const float* w_row = w_b0 + a * IQN_HIDDEN; + float partial = 0.0f; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + acc += iqn_warp_sum(partial); + if (lane == 0) { + save_q_online[q_save_offset + a] = acc; + if (a == a0) q0_taken = acc; + } + if (a == a0) q0_taken = __shfl_sync(0xFFFFFFFF, q0_taken, 0); + } + + /* Branch 1 (order) */ + float q1_taken = 0.0f; + for (int a = 0; a < BRANCH_1_SIZE; a++) { + float acc = b_b1[a]; + const float* w_row = w_b1 + a * IQN_HIDDEN; + float partial = 0.0f; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + acc += iqn_warp_sum(partial); + if (lane == 0) { + save_q_online[q_save_offset + BRANCH_0_SIZE + a] = acc; + if (a == a1) q1_taken = acc; + } + if (a == a1) q1_taken = __shfl_sync(0xFFFFFFFF, q1_taken, 0); + } + + /* Branch 2 (urgency) */ + float q2_taken = 0.0f; + for (int a = 0; a < BRANCH_2_SIZE; a++) { + float acc = b_b2[a]; + const float* w_row = w_b2 + a * IQN_HIDDEN; + float partial = 0.0f; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + acc += iqn_warp_sum(partial); + if (lane == 0) { + save_q_online[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a] = acc; + if (a == a2) q2_taken = acc; + } + if (a == a2) q2_taken = __shfl_sync(0xFFFFFFFF, q2_taken, 0); + } + + /* Sum of taken-action Q-values across branches (for loss computation) */ + online_q_a[t] = q0_taken + q1_taken + q2_taken; + } + + /* ── Process TARGET quantiles ── */ + float target_q_a[IQN_NUM_QUANTILES]; + + for (int t = 0; t < IQN_NUM_QUANTILES; t++) { + float tau = my_target_taus[t]; + + /* Cosine embedding + linear (target weights) */ + float embed_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) { + float acc = tb_embed[h]; + const float* w_row = tw_embed + h * IQN_EMBED_DIM; + for (int d = 0; d < IQN_EMBED_DIM; d++) { + float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); + acc += w_row[d] * cos_val; + } + embed_dist[h / 32] = fmaxf(acc, 0.0f); + } + + /* Element-wise product with target h_s2 */ + float comb_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) + comb_dist[h / 32] = h_target_dist[h / 32] * embed_dist[h / 32]; + + /* Per-branch Q-values for taken actions */ + float q0 = 0.0f, q1 = 0.0f, q2 = 0.0f; + + /* Branch 0 */ + { + float acc = tb_b0[a0]; + float partial = 0.0f; + const float* w_row = tw_b0 + a0 * IQN_HIDDEN; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + q0 = acc + iqn_warp_sum(partial); + } + /* Branch 1 */ + { + float acc = tb_b1[a1]; + float partial = 0.0f; + const float* w_row = tw_b1 + a1 * IQN_HIDDEN; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + q1 = acc + iqn_warp_sum(partial); + } + /* Branch 2 */ + { + float acc = tb_b2[a2]; + float partial = 0.0f; + const float* w_row = tw_b2 + a2 * IQN_HIDDEN; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + q2 = acc + iqn_warp_sum(partial); + } + + /* Bellman target: r + γ(1-done) × Q_target */ + float raw_target = q0 + q1 + q2; + target_q_a[t] = reward + gamma * (1.0f - done) * raw_target; + + /* Save target Q for backward */ + int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; + if (lane == 0) { + save_q_target[q_save_offset + a0] = target_q_a[t]; + save_q_target[q_save_offset + BRANCH_0_SIZE + a1] = target_q_a[t]; + save_q_target[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a2] = target_q_a[t]; + } + } + + /* ── Quantile Huber Loss: (1/N) Σ_i (1/N') Σ_j ρ_τi(δ_ij) ── */ + float sample_loss = 0.0f; + if (lane == 0) { + for (int i = 0; i < IQN_NUM_QUANTILES; i++) { + float tau_i = my_taus[i]; + float inner_sum = 0.0f; + for (int j = 0; j < IQN_NUM_QUANTILES; j++) { + float delta = target_q_a[j] - online_q_a[i]; + inner_sum += quantile_huber_element(tau_i, delta, IQN_KAPPA); + } + sample_loss += inner_sum / (float)IQN_NUM_QUANTILES; + } + sample_loss /= (float)IQN_NUM_QUANTILES; + + per_sample_loss[sample] = sample_loss; + atomicAdd(total_loss, sample_loss / (float)batch_size); + } +} + +/* ── Backward Kernel ────────────────────────────────────────────────── */ +/** + * Backprop through IQN head. Computes gradients w.r.t. IQN weights only + * (does NOT backprop into h_s2 — the shared trunk is trained by C51). + * + * For each sample, for each quantile τ_i: + * dL/dq_online_i = (1/N²) Σ_j dρ_τi(δ_ij)/dq_i + * dL/d(combined) = W_bd^T × dL/dq_d (for taken action a_d) + * dL/d(embed) = dL/d(combined) ⊙ h_s2 ⊙ 𝟙{pre_relu > 0} + * dL/d(W_embed) += outer(dL/d(embed), cos_features) + * dL/d(W_bd) += outer(dL/dq_d, combined) + */ +extern "C" __global__ +void iqn_backward_kernel( + /* Saved activations (from forward pass) */ + const float* __restrict__ h_s2, /* [B, IQN_HIDDEN] */ + const float* __restrict__ save_embed, /* [B, N, IQN_HIDDEN] */ + const float* __restrict__ save_combined, /* [B, N, IQN_HIDDEN] */ + const float* __restrict__ save_q_online, /* [B, N, TOTAL_BRANCH_ACTIONS] */ + const float* __restrict__ save_q_target, /* [B, N, TOTAL_BRANCH_ACTIONS] */ + const float* __restrict__ taus, /* [B, N] */ + const int* __restrict__ actions, /* [B, 3] */ + /* Weights (for backward through linear layers) */ + const float* __restrict__ online_params, + /* Gradient output */ + float* __restrict__ grad_buf, /* [IQN_TOTAL_PARAMS] accumulated gradients */ + int batch_size +) +{ + int sample = blockIdx.x; + if (sample >= batch_size) return; + int lane = threadIdx.x; + + const float* my_h_s2 = h_s2 + sample * IQN_HIDDEN; + const float* my_taus = taus + sample * IQN_NUM_QUANTILES; + int a0 = actions[sample * 3 + 0]; + int a1 = actions[sample * 3 + 1]; + int a2 = actions[sample * 3 + 2]; + + const float* w_embed = online_params + IQN_OFF_W_EMBED; + const float* w_b0 = online_params + IQN_OFF_W_B0; + const float* w_b1 = online_params + IQN_OFF_W_B1; + const float* w_b2 = online_params + IQN_OFF_W_B2; + + /* Load h_s2 into distributed registers */ + float h_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) + h_dist[h / 32] = my_h_s2[h]; + + /* Pre-compute online and target Q-values for taken actions across all quantiles */ + /* (reading from saved buffers) */ + float online_q[IQN_NUM_QUANTILES]; + float target_q[IQN_NUM_QUANTILES]; + for (int t = 0; t < IQN_NUM_QUANTILES; t++) { + int qoff = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS; + if (lane == 0) { + online_q[t] = save_q_online[qoff + a0] + + save_q_online[qoff + BRANCH_0_SIZE + a1] + + save_q_online[qoff + BRANCH_0_SIZE + BRANCH_1_SIZE + a2]; + target_q[t] = save_q_target[qoff + a0]; /* target_q already contains Bellman target */ + } + online_q[t] = __shfl_sync(0xFFFFFFFF, online_q[t], 0); + target_q[t] = __shfl_sync(0xFFFFFFFF, target_q[t], 0); + } + + /* Process each quantile τ_i */ + float inv_n_sq = 1.0f / ((float)IQN_NUM_QUANTILES * (float)IQN_NUM_QUANTILES); + + for (int ti = 0; ti < IQN_NUM_QUANTILES; ti++) { + float tau_i = my_taus[ti]; + int save_offset = (sample * IQN_NUM_QUANTILES + ti) * IQN_HIDDEN; + + /* Compute dL/dq_online for this quantile: + * dL/dq_i = (1/N²) Σ_j dρ_τi(δ_ij)/dq_i */ + float dL_dq = 0.0f; + if (lane == 0) { + for (int tj = 0; tj < IQN_NUM_QUANTILES; tj++) { + float delta = target_q[tj] - online_q[ti]; + dL_dq += quantile_huber_grad(tau_i, delta, IQN_KAPPA); + } + dL_dq *= inv_n_sq; + } + dL_dq = __shfl_sync(0xFFFFFFFF, dL_dq, 0); + + /* Load saved activations */ + float embed_dist[IQN_DIST(IQN_HIDDEN)]; + float comb_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) { + embed_dist[h / 32] = save_embed[save_offset + h]; + comb_dist[h / 32] = save_combined[save_offset + h]; + } + + /* ── Gradient through branch output layers (taken action only) ── */ + /* dL/d(combined) = Σ_d W_bd[a_d, :] × dL/dq (only for taken actions) */ + float dL_dcomb[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) + dL_dcomb[h / 32] = 0.0f; + + /* Branch 0: accumulate gradient */ + for (int h = lane; h < IQN_HIDDEN; h += 32) { + dL_dcomb[h / 32] += w_b0[a0 * IQN_HIDDEN + h] * dL_dq; + /* dL/dW_b0[a0, h] += dL/dq × combined[h] */ + atomicAdd(&grad_buf[IQN_OFF_W_B0 + a0 * IQN_HIDDEN + h], + dL_dq * comb_dist[h / 32]); + } + if (lane == 0) + atomicAdd(&grad_buf[IQN_OFF_B_B0 + a0], dL_dq); + + /* Branch 1 */ + for (int h = lane; h < IQN_HIDDEN; h += 32) { + dL_dcomb[h / 32] += w_b1[a1 * IQN_HIDDEN + h] * dL_dq; + atomicAdd(&grad_buf[IQN_OFF_W_B1 + a1 * IQN_HIDDEN + h], + dL_dq * comb_dist[h / 32]); + } + if (lane == 0) + atomicAdd(&grad_buf[IQN_OFF_B_B1 + a1], dL_dq); + + /* Branch 2 */ + for (int h = lane; h < IQN_HIDDEN; h += 32) { + dL_dcomb[h / 32] += w_b2[a2 * IQN_HIDDEN + h] * dL_dq; + atomicAdd(&grad_buf[IQN_OFF_W_B2 + a2 * IQN_HIDDEN + h], + dL_dq * comb_dist[h / 32]); + } + if (lane == 0) + atomicAdd(&grad_buf[IQN_OFF_B_B2 + a2], dL_dq); + + /* ── Gradient through element-wise product ── */ + /* combined = h_s2 ⊙ embed + * dL/d(embed) = dL/d(combined) ⊙ h_s2 + * (dL/d(h_s2) = dL/d(combined) ⊙ embed — NOT computed, trunk trained by C51) */ + float dL_dembed[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) + dL_dembed[h / 32] = dL_dcomb[h / 32] * h_dist[h / 32]; + + /* ── Gradient through ReLU ── */ + /* embed = ReLU(pre_relu) → dL/d(pre_relu) = dL/d(embed) × 𝟙{embed > 0} */ + for (int h = lane; h < IQN_HIDDEN; h += 32) { + if (embed_dist[h / 32] <= 0.0f) + dL_dembed[h / 32] = 0.0f; + } + + /* ── Gradient through embedding linear layer ── */ + /* pre_relu[h] = Σ_d W_embed[h,d]·cos(π(d+1)τ) + b_embed[h] + * dL/dW_embed[h,d] += dL/d(pre_relu)[h] × cos(π(d+1)τ) + * dL/db_embed[h] += dL/d(pre_relu)[h] */ + float tau = my_taus[ti]; + for (int h = lane; h < IQN_HIDDEN; h += 32) { + float dL_dpre = dL_dembed[h / 32]; + /* Bias gradient */ + atomicAdd(&grad_buf[IQN_OFF_B_EMBED + h], dL_dpre); + /* Weight gradient: outer product with cosine features */ + for (int d = 0; d < IQN_EMBED_DIM; d++) { + float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); + atomicAdd(&grad_buf[IQN_OFF_W_EMBED + h * IQN_EMBED_DIM + d], + dL_dpre * cos_val); + } + } + } +} + +/* ── Gradient Norm Kernel ───────────────────────────────────────────── */ +extern "C" __global__ +void iqn_grad_norm_kernel( + const float* __restrict__ grads, + float* __restrict__ norm_out, /* [1] */ + int total_params +) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + /* Each thread accumulates its own partial sum */ + float partial = 0.0f; + for (int i = tid; i < total_params; i += gridDim.x * blockDim.x) { + float g = grads[i]; + partial += g * g; + } + + /* Block-level reduction via shared memory */ + __shared__ float shared[256]; + shared[threadIdx.x] = partial; + __syncthreads(); + + /* Tree reduction in shared memory */ + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) + shared[threadIdx.x] += shared[threadIdx.x + s]; + __syncthreads(); + } + + if (threadIdx.x == 0) + atomicAdd(norm_out, sqrtf(shared[0])); +} + +/* ── Adam Update Kernel ─────────────────────────────────────────────── */ +extern "C" __global__ +void iqn_adam_kernel( + float* __restrict__ params, + float* __restrict__ grads, + float* __restrict__ m, /* first moment */ + float* __restrict__ v, /* second moment */ + const float* __restrict__ norm, /* [1] gradient L2 norm */ + float lr, float beta1, float beta2, float eps, + float weight_decay, float max_grad_norm, + int adam_t, /* step counter (1-indexed) */ + int total_params +) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= total_params) return; + + /* Gradient clipping */ + float grad_norm = norm[0]; + float clip_scale = (grad_norm > max_grad_norm && grad_norm > 0.0f) + ? max_grad_norm / grad_norm : 1.0f; + float g = grads[tid] * clip_scale; + + /* AdamW update */ + float mi = beta1 * m[tid] + (1.0f - beta1) * g; + float vi = beta2 * v[tid] + (1.0f - beta2) * g * g; + m[tid] = mi; + v[tid] = vi; + + /* Bias correction */ + float m_hat = mi / (1.0f - powf(beta1, (float)adam_t)); + float v_hat = vi / (1.0f - powf(beta2, (float)adam_t)); + + /* Parameter update with decoupled weight decay */ + float p = params[tid]; + p -= lr * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * p); + params[tid] = p; + + /* Zero gradient for next step */ + grads[tid] = 0.0f; +} + +/* ── Forward-Only Kernel (Inference) ────────────────────────────────── */ +/** + * IQN forward pass for inference — no loss, no activation saves. + * + * Computes expected Q-values per action per branch by averaging over + * quantile samples. Used for action selection blending with C51. + * + * Outputs: expected_q [B, TOTAL_BRANCH_ACTIONS] — mean Q per action. + */ +extern "C" __global__ +void iqn_forward_kernel( + const float* __restrict__ h_s2, /* [B, IQN_HIDDEN] */ + const float* __restrict__ taus, /* [B, N] */ + const float* __restrict__ params, /* [IQN_TOTAL_PARAMS] */ + float* __restrict__ expected_q, /* [B, TOTAL_BRANCH_ACTIONS] */ + int batch_size +) +{ + int sample = blockIdx.x; + if (sample >= batch_size) return; + int lane = threadIdx.x; + + const float* my_h_s2 = h_s2 + sample * IQN_HIDDEN; + const float* my_taus = taus + sample * IQN_NUM_QUANTILES; + const float* w_embed = params + IQN_OFF_W_EMBED; + const float* b_embed = params + IQN_OFF_B_EMBED; + const float* w_b0 = params + IQN_OFF_W_B0; + const float* b_b0 = params + IQN_OFF_B_B0; + const float* w_b1 = params + IQN_OFF_W_B1; + const float* b_b1 = params + IQN_OFF_B_B1; + const float* w_b2 = params + IQN_OFF_W_B2; + const float* b_b2 = params + IQN_OFF_B_B2; + + /* Load h_s2 */ + float h_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) + h_dist[h / 32] = my_h_s2[h]; + + /* Accumulate Q-values across quantiles (for mean) */ + float q_acc[TOTAL_BRANCH_ACTIONS]; + for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++) + q_acc[a] = 0.0f; + + for (int t = 0; t < IQN_NUM_QUANTILES; t++) { + float tau = my_taus[t]; + + /* Cosine embedding + linear + ReLU */ + float embed_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) { + float acc = b_embed[h]; + const float* w_row = w_embed + h * IQN_EMBED_DIM; + for (int d = 0; d < IQN_EMBED_DIM; d++) { + float cos_val = cosf(3.14159265f * (float)(d + 1) * tau); + acc += w_row[d] * cos_val; + } + embed_dist[h / 32] = fmaxf(acc, 0.0f); + } + + /* Element-wise product */ + float comb_dist[IQN_DIST(IQN_HIDDEN)]; + for (int h = lane; h < IQN_HIDDEN; h += 32) + comb_dist[h / 32] = h_dist[h / 32] * embed_dist[h / 32]; + + /* Branch outputs */ + /* Branch 0 */ + for (int a = 0; a < BRANCH_0_SIZE; a++) { + float partial = 0.0f; + const float* w_row = w_b0 + a * IQN_HIDDEN; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + float q_val = iqn_warp_sum(partial) + b_b0[a]; + q_acc[a] += q_val; + } + /* Branch 1 */ + for (int a = 0; a < BRANCH_1_SIZE; a++) { + float partial = 0.0f; + const float* w_row = w_b1 + a * IQN_HIDDEN; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + float q_val = iqn_warp_sum(partial) + b_b1[a]; + q_acc[BRANCH_0_SIZE + a] += q_val; + } + /* Branch 2 */ + for (int a = 0; a < BRANCH_2_SIZE; a++) { + float partial = 0.0f; + const float* w_row = w_b2 + a * IQN_HIDDEN; + for (int h = lane; h < IQN_HIDDEN; h += 32) + partial += w_row[h] * comb_dist[h / 32]; + float q_val = iqn_warp_sum(partial) + b_b2[a]; + q_acc[BRANCH_0_SIZE + BRANCH_1_SIZE + a] += q_val; + } + } + + /* Write mean Q-values (averaged over quantiles) */ + float inv_n = 1.0f / (float)IQN_NUM_QUANTILES; + if (lane == 0) { + for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++) + expected_q[sample * TOTAL_BRANCH_ACTIONS + a] = q_acc[a] * inv_n; + } +} + +/* ── Trunk Forward Kernel ──────────────────────────────────────────── */ +/** + * Compute h_s2 from raw states using DQN shared trunk weights. + * + * state → LeakyReLU(W_s1 × state + b_s1) → LeakyReLU(W_s2 × h1 + b_s2) → h_s2 + * + * Used to compute target_h_s2 from next_states + target DQN weights. + * One warp per sample. Uses shared memory for h1 intermediate. + * + * Compile-time defines: IQN_STATE_DIM, IQN_SHARED_H1, IQN_HIDDEN (=SHARED_H2) + */ +#ifndef IQN_STATE_DIM +#define IQN_STATE_DIM 48 +#endif +#ifndef IQN_SHARED_H1 +#define IQN_SHARED_H1 256 +#endif + +extern "C" __global__ +void iqn_trunk_forward_kernel( + const float* __restrict__ states, /* [B, IQN_STATE_DIM] */ + const float* __restrict__ w_s1, /* [IQN_SHARED_H1, IQN_STATE_DIM] */ + const float* __restrict__ b_s1, /* [IQN_SHARED_H1] */ + const float* __restrict__ w_s2, /* [IQN_HIDDEN, IQN_SHARED_H1] */ + const float* __restrict__ b_s2, /* [IQN_HIDDEN] */ + float* __restrict__ h_s2_out, /* [B, IQN_HIDDEN] */ + int batch_size +) +{ + int sample = blockIdx.x; + if (sample >= batch_size) return; + int lane = threadIdx.x; + + extern __shared__ float shmem[]; + float* h1 = shmem; /* [IQN_SHARED_H1] */ + + const float* my_state = states + sample * IQN_STATE_DIM; + + /* Layer 1: h1 = leaky_relu(W_s1 @ state + b_s1) */ + for (int h = lane; h < IQN_SHARED_H1; h += 32) { + float acc = b_s1[h]; + const float* w_row = w_s1 + h * IQN_STATE_DIM; + for (int d = 0; d < IQN_STATE_DIM; d++) + acc += w_row[d] * my_state[d]; + h1[h] = (acc > 0.0f) ? acc : 0.01f * acc; + } + __syncwarp(); + + /* Layer 2: h2 = leaky_relu(W_s2 @ h1 + b_s2) */ + for (int h = lane; h < IQN_HIDDEN; h += 32) { + float acc = b_s2[h]; + const float* w_row = w_s2 + h * IQN_SHARED_H1; + for (int d = 0; d < IQN_SHARED_H1; d++) + acc += w_row[d] * h1[d]; + float val = (acc > 0.0f) ? acc : 0.01f * acc; + h_s2_out[sample * IQN_HIDDEN + h] = val; + } +} + +/* ── EMA Target Update Kernel ──────────────────────────────────────── */ +/** + * Polyak EMA: target[i] = (1 - tau) * target[i] + tau * online[i] + */ +extern "C" __global__ +void iqn_ema_kernel( + float* __restrict__ target, + const float* __restrict__ online, + float tau, + int n +) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) + target[i] = (1.0f - tau) * target[i] + tau * online[i]; +} + +/* ── Flat action → branch action decode kernel ───────────────────────── + * + * Decodes flat action index ∈ [0, B0×B1×B2) into 3 branch indices. + * Grid: (ceil(batch_size/256), 1, 1), Block: (256, 1, 1) + * + * Input: flat_actions[B] -- i32 flat action indices + * Output: branch_actions[B*3] -- i32 [exposure, order, urgency] per sample + */ +extern "C" __global__ +void iqn_decode_actions_kernel( + const int* __restrict__ flat_actions, + int* __restrict__ branch_actions, + int batch_size +) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + + int flat = flat_actions[i]; + int stride = BRANCH_1_SIZE * BRANCH_2_SIZE; + int exposure = flat / stride; + int remainder = flat % stride; + int order = remainder / BRANCH_2_SIZE; + int urgency = remainder % BRANCH_2_SIZE; + + branch_actions[i * 3 + 0] = exposure; + branch_actions[i * 3 + 1] = order; + branch_actions[i * 3 + 2] = urgency; +} + +/* ── GPU τ sampling kernel (Philox-based PRNG) ───────────────────────── + * + * Generates uniform random τ ∈ (0.01, 0.99) directly on GPU. + * Uses Philox 4×32 counter-based PRNG — deterministic, high quality, + * no global state. Each thread produces one τ value from (thread_id, seed). + * + * Grid: (ceil(total/256), 1, 1), Block: (256, 1, 1) + * total = batch_size × num_quantiles + */ + +/* Philox 4×32-10 round function — same PRNG used by PyTorch/JAX */ +__device__ __forceinline__ +unsigned int iqn_philox_single(unsigned int counter, unsigned int key) { + unsigned int hi, lo; + /* Philox S-box constants */ + lo = counter * 0xD2511F53u; + hi = __umulhi(counter, 0xD2511F53u); + /* 10 rounds of mixing */ + for (int r = 0; r < 10; r++) { + unsigned int t = hi ^ key; + hi = lo * 0xCD9E8D57u; + lo = __umulhi(lo, 0xCD9E8D57u); + lo ^= t; + key += 0x9E3779B9u; /* golden ratio */ + } + return lo ^ hi; +} + +extern "C" __global__ +void iqn_sample_taus_kernel( + float* __restrict__ taus, + unsigned int seed, + int total +) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + + /* Counter = thread index, key = seed — unique per (step, online/target) */ + unsigned int bits = iqn_philox_single((unsigned int)i, seed); + /* Map to (0, 1) then clamp to (0.01, 0.99) */ + float u = (float)(bits >> 8) * (1.0f / 16777216.0f); /* 24-bit mantissa */ + taus[i] = 0.01f + u * 0.98f; +} diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index aa7068393..0ff9d8f36 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -26,6 +26,7 @@ pub mod gpu_walk_forward; pub mod gpu_dqn_trainer; pub mod gpu_her; pub mod gpu_iql_trainer; +pub mod gpu_iqn_head; // gpu_replay_buffer moved to ml-dqn crate /// Maximum bytes allowed for a single GPU upload (2 GB safety limit). @@ -894,8 +895,6 @@ mod tests { assert!(src.contains("barrier_init"), "Missing barrier_init"); assert!(src.contains("diversity_entropy"), "Missing diversity_entropy"); assert!(src.contains("curiosity_inference"), "Missing curiosity_inference"); - assert!(!src.contains("q_forward_dueling"), "DQN-specific function leaked into common header"); - assert!(!src.contains("dqn_full_experience_kernel"), "DQN kernel leaked into common header"); } #[test] @@ -933,7 +932,6 @@ mod tests { assert!(full.contains("ppo_full_experience_kernel")); assert!(full.contains("compute_gae_backward")); assert!(full.contains("softmax_sample")); - assert!(!full.contains("q_forward_dueling")); } #[test] diff --git a/crates/ml/src/cuda_pipeline/signal_adapter.rs b/crates/ml/src/cuda_pipeline/signal_adapter.rs index 4670ce582..23d758da6 100644 --- a/crates/ml/src/cuda_pipeline/signal_adapter.rs +++ b/crates/ml/src/cuda_pipeline/signal_adapter.rs @@ -358,6 +358,7 @@ mod tests { fn gpu_argmax(t: &Tensor) -> u32 { t.argmax(1).unwrap() .flatten_all().unwrap() + .squeeze(0).unwrap() .to_scalar::().unwrap() } diff --git a/crates/ml/src/diffusion/trainable.rs b/crates/ml/src/diffusion/trainable.rs index 261faa5b5..3eeaa91bf 100644 --- a/crates/ml/src/diffusion/trainable.rs +++ b/crates/ml/src/diffusion/trainable.rs @@ -140,7 +140,9 @@ impl UnifiedTrainable for DiffusionTrainableAdapter { /// Compute MSE loss between predicted noise and actual noise. fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result { let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string()); - let diff = predictions.sub(targets).map_err(map_err)?; + // Denoiser returns F32 output; cast targets to match to avoid dtype mismatch + let targets = targets.to_dtype(predictions.dtype()).map_err(map_err)?; + let diff = predictions.sub(&targets).map_err(map_err)?; diff.sqr().map_err(map_err)?.mean_all().map_err(map_err) } @@ -292,6 +294,7 @@ impl UnifiedTrainable for DiffusionTrainableAdapter { #[cfg(test)] mod tests { use super::*; + use candle_core::DType; fn make_adapter() -> DiffusionTrainableAdapter { let config = DiffusionConfig { @@ -317,15 +320,16 @@ mod tests { #[test] fn test_device() { let adapter = make_adapter(); - assert_eq!(adapter.device().location(), candle_core::DeviceLocation::Cpu); + assert!(matches!(adapter.device().location(), candle_core::DeviceLocation::Cuda { .. })); } #[test] fn test_forward_2d() { let mut adapter = make_adapter(); - let input = Tensor::randn(0_f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let input = Tensor::randn(0_f32, 1.0, &[4, 16], adapter.device()).unwrap() + .to_dtype(DType::BF16).unwrap(); let output = adapter.forward(&input); - assert!(output.is_ok()); + assert!(output.is_ok(), "forward failed: {:?}", output.as_ref().err()); assert_eq!(output.unwrap().dims(), &[4, 16]); } @@ -333,7 +337,8 @@ mod tests { fn test_forward_3d() { let mut adapter = make_adapter(); // (batch, seq_len, feature_dim) → flattened to (batch, 16) - let input = Tensor::randn(0_f32, 1.0, &[4, 16, 1], adapter.device()).unwrap(); + let input = Tensor::randn(0_f32, 1.0, &[4, 16, 1], adapter.device()).unwrap() + .to_dtype(DType::BF16).unwrap(); let output = adapter.forward(&input); assert!(output.is_ok()); assert_eq!(output.unwrap().dims(), &[4, 16]); @@ -353,9 +358,11 @@ mod tests { #[test] fn test_backward_returns_grad_norm() { let mut adapter = make_adapter(); - let input = Tensor::randn(0_f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let input = Tensor::randn(0_f32, 1.0, &[4, 16], adapter.device()).unwrap() + .to_dtype(DType::BF16).unwrap(); let output = adapter.forward(&input).unwrap(); - let targets = Tensor::randn(0_f32, 1.0, output.dims(), adapter.device()).unwrap(); + let targets = Tensor::randn(0_f32, 1.0, output.dims(), adapter.device()).unwrap() + .to_dtype(DType::BF16).unwrap(); let loss = adapter.compute_loss(&output, &targets).unwrap(); let grad_norm = adapter.backward(&loss); assert!(grad_norm.is_ok()); @@ -366,9 +373,11 @@ mod tests { fn test_train_step_cycle() { let mut adapter = make_adapter(); adapter.zero_grad().unwrap(); - let input = Tensor::randn(0_f32, 1.0, &[4, 16], adapter.device()).unwrap(); + let input = Tensor::randn(0_f32, 1.0, &[4, 16], adapter.device()).unwrap() + .to_dtype(DType::BF16).unwrap(); let output = adapter.forward(&input).unwrap(); - let targets = Tensor::randn(0_f32, 1.0, output.dims(), adapter.device()).unwrap(); + let targets = Tensor::randn(0_f32, 1.0, output.dims(), adapter.device()).unwrap() + .to_dtype(DType::BF16).unwrap(); let loss = adapter.compute_loss(&output, &targets).unwrap(); adapter.backward(&loss).unwrap(); adapter.optimizer_step().unwrap(); @@ -392,7 +401,8 @@ mod tests { #[test] fn test_checkpoint_roundtrip() { let mut adapter = make_adapter(); - let input = Tensor::randn(0_f32, 1.0, &[2, 16], adapter.device()).unwrap(); + let input = Tensor::randn(0_f32, 1.0, &[2, 16], adapter.device()).unwrap() + .to_dtype(DType::BF16).unwrap(); let _ = adapter.forward(&input).unwrap(); let dir = std::env::temp_dir().join("diffusion_ckpt_test"); @@ -414,8 +424,10 @@ mod tests { let dev = adapter.device().clone(); let val_data = vec![ ( - Tensor::randn(0_f32, 1.0, &[2, 16], &dev).unwrap(), - Tensor::randn(0_f32, 1.0, &[2, 16], &dev).unwrap(), + Tensor::randn(0_f32, 1.0, &[2, 16], &dev).unwrap() + .to_dtype(DType::BF16).unwrap(), + Tensor::randn(0_f32, 1.0, &[2, 16], &dev).unwrap() + .to_dtype(DType::BF16).unwrap(), ), ]; let val_loss = adapter.validate(&val_data); diff --git a/crates/ml/src/ensemble/adapters/liquid.rs b/crates/ml/src/ensemble/adapters/liquid.rs index d91e3341d..e837cd039 100644 --- a/crates/ml/src/ensemble/adapters/liquid.rs +++ b/crates/ml/src/ensemble/adapters/liquid.rs @@ -277,6 +277,7 @@ mod tests { use crate::ensemble::inference_adapter::FeatureVector; use crate::liquid::candle_cfc::DeviceConfig; + /// BF16 matmul is only supported on CUDA, so tests must use a CUDA device. fn test_config() -> CfCTrainConfig { CfCTrainConfig { input_size: 51, @@ -284,21 +285,21 @@ mod tests { output_size: 3, backbone_hidden_sizes: vec![32], seq_len: 1, - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), ..CfCTrainConfig::default() } } #[test] fn test_liquid_inference_adapter_creation() { - let adapter = LiquidInferenceAdapter::new(test_config()).unwrap(); + let adapter = LiquidInferenceAdapter::new(test_config()).expect("CUDA required"); assert_eq!(adapter.model_name(), "Liquid-CfC"); assert!(adapter.is_ready()); } #[test] fn test_liquid_inference_predict() { - let adapter = LiquidInferenceAdapter::new(test_config()).unwrap(); + let adapter = LiquidInferenceAdapter::new(test_config()).expect("CUDA required"); let features = FeatureVector { values: vec![0.0; 51], @@ -321,7 +322,7 @@ mod tests { #[test] fn test_liquid_inference_deterministic() { - let adapter = LiquidInferenceAdapter::new(test_config()).unwrap(); + let adapter = LiquidInferenceAdapter::new(test_config()).expect("CUDA required"); let fv = FeatureVector { values: vec![0.1; 51], timestamp: 1700000000_000_000, @@ -340,7 +341,7 @@ mod tests { #[test] fn test_liquid_inference_direction_range() { - let adapter = LiquidInferenceAdapter::new(test_config()).unwrap(); + let adapter = LiquidInferenceAdapter::new(test_config()).expect("CUDA required"); let fv = FeatureVector { values: vec![0.5; 51], timestamp: 1700000000_000_000, @@ -366,7 +367,7 @@ mod tests { #[test] fn test_liquid_inference_metadata() { - let adapter = LiquidInferenceAdapter::new(test_config()).unwrap(); + let adapter = LiquidInferenceAdapter::new(test_config()).expect("CUDA required"); let fv = FeatureVector { values: vec![0.1; 51], timestamp: 1700000000_000_000, @@ -380,7 +381,7 @@ mod tests { #[test] fn test_liquid_inference_wrong_feature_size() { - let adapter = LiquidInferenceAdapter::new(test_config()).unwrap(); + let adapter = LiquidInferenceAdapter::new(test_config()).expect("CUDA required"); let fv = FeatureVector { values: vec![0.1; 10], // Wrong size, expect 51 timestamp: 0, @@ -392,7 +393,7 @@ mod tests { #[test] fn test_liquid_checkpoint_round_trip() { let config = test_config(); - let adapter = LiquidInferenceAdapter::new(config.clone()).unwrap(); + let adapter = LiquidInferenceAdapter::new(config.clone()).expect("CUDA required"); let fv = FeatureVector { values: vec![0.3; 51], @@ -407,7 +408,7 @@ mod tests { // Load into new adapter let path_str = path.to_str().unwrap(); - let adapter2 = LiquidInferenceAdapter::from_checkpoint(config, path_str).unwrap(); + let adapter2 = LiquidInferenceAdapter::from_checkpoint(config, path_str).expect("CUDA required"); let pred2 = adapter2.predict(&fv).unwrap(); assert!( diff --git a/crates/ml/src/features/multi_timeframe.rs b/crates/ml/src/features/multi_timeframe.rs index 898420a64..5346b684d 100644 --- a/crates/ml/src/features/multi_timeframe.rs +++ b/crates/ml/src/features/multi_timeframe.rs @@ -14,7 +14,7 @@ use std::collections::VecDeque; -use candle_core::{DType, Device, Tensor}; +use candle_core::{Device, Tensor}; use candle_nn::{linear, Linear, Module, VarBuilder, VarMap}; use super::bar_resampler::BarResampler; @@ -130,19 +130,28 @@ impl LstmEncoder { /// Returns a tensor of shape `(1, hidden_dim)` (the final h). pub fn forward(&self, seq: &Tensor) -> Result { let device = seq.device(); + let dtype = self.w_ih.dtype(); let seq_len = seq .dims() .first() .copied() .ok_or_else(|| MLError::ModelError("LstmEncoder: empty sequence dims".into()))?; - let mut h = Tensor::zeros(&[1, self.hidden_dim], DType::F32, device) + // Cast input to match weight dtype (e.g. F32 input with BF16 weights) + let seq = if seq.dtype() != dtype { + seq.to_dtype(dtype) + .map_err(|e| MLError::ModelError(format!("LstmEncoder seq cast: {e}")))? + } else { + seq.clone() + }; + + let mut h = Tensor::zeros(&[1, self.hidden_dim], dtype, device) .map_err(|e| MLError::ModelError(format!("LstmEncoder h init: {e}")))?; - let mut c = Tensor::zeros(&[1, self.hidden_dim], DType::F32, device) + let mut c = Tensor::zeros(&[1, self.hidden_dim], dtype, device) .map_err(|e| MLError::ModelError(format!("LstmEncoder c init: {e}")))?; for t in 0..seq_len { - // x_t: (1, input_dim) + // x_t: (1, input_dim) -- already in weight dtype let x_t = seq .narrow(0, t, 1) .map_err(|e| MLError::ModelError(format!("LstmEncoder narrow t={t}: {e}")))?; @@ -344,8 +353,9 @@ impl MultiTimeframeEncoder { history: &VecDeque, ) -> Result { if history.is_empty() { - // No data yet -- return zeros - return Tensor::zeros(&[1, self.config.hidden_dim], DType::F32, &self.device) + // No data yet -- return zeros in weight dtype to match projection + let dtype = lstm.w_ih.dtype(); + return Tensor::zeros(&[1, self.config.hidden_dim], dtype, &self.device) .map_err(|e| MLError::ModelError(format!("zero embedding: {e}"))); } @@ -404,6 +414,7 @@ fn push_ring(buf: &mut VecDeque, bar: OHLCVBar, max_len: usize) { #[cfg(test)] mod tests { use super::*; + use candle_core::DType; use chrono::{TimeZone, Utc}; fn make_bar(minute: u32, close: f64) -> OHLCVBar { @@ -480,6 +491,7 @@ mod tests { .sub(&out2) .and_then(|d| d.abs()) .and_then(|d| d.sum_all()) + .and_then(|d| d.to_dtype(DType::F32)) .and_then(|d| d.to_scalar::()) .expect("diff computation"); @@ -619,6 +631,7 @@ mod tests { let sum = output .abs() .and_then(|t| t.sum_all()) + .and_then(|t| t.to_dtype(DType::F32)) .and_then(|t| t.to_scalar::()) .expect("sum"); diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index c8dc04352..c17bcfed8 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -3,7 +3,7 @@ //! This module provides a production-ready adapter for optimizing DQN //! hyperparameters using the generic optimization framework. It implements: //! -//! - 30D continuous parameter space with log-scale handling +//! - 31D continuous parameter space with log-scale handling //! - Training wrapper that integrates with existing DQN pipeline //! - Backtest-based objective using Sharpe ratio (production metric) //! - Prioritized Experience Replay (PER) tuning for Rainbow DQN performance (+25-40% convergence speed) @@ -143,7 +143,7 @@ const TRIAL_VRAM_MB: f64 = 7000.0; /// Corrected: was 0.02 (20KB) — 100x too high const MB_PER_SAMPLE: f64 = 0.0005; -/// DQN hyperparameter space (30D continuous) +/// DQN hyperparameter space (31D continuous) /// /// Fixed params (not in search space): /// - curiosity_weight: 0.0 (conflicts with NoisyNet) @@ -332,6 +332,9 @@ pub struct DQNParams { pub num_quantiles: usize, /// Kappa for quantile Huber loss (0.5-2.0, controls L1<->L2 transition) pub qr_kappa: f64, + /// Lambda weight for IQN quantile loss relative to C51 loss in dual-head ensemble. + /// L_total = L_c51 + iqn_lambda * L_iqn. Range [0.0, 2.0]. + pub iqn_lambda: f64, /// Hidden dimension base for dynamic network sizing. /// Network shape: [base, base/2, base/4]. Range: [256, 4096], step=256. @@ -441,6 +444,7 @@ impl Default for DQNParams { use_qr_dqn: false, num_quantiles: 32, qr_kappa: 1.0, + iqn_lambda: 0.25, // Default: mild IQN regularization alongside C51 hidden_dim_base: 256, // Conservative default (backward compatible) noisy_epsilon_floor: 0.10, // Minimum 10% random exploration to prevent action collapse count_bonus_coefficient: 0.0, // C2: fixed to 0.0 (conflicts with NoisyNet) @@ -456,7 +460,7 @@ impl Default for DQNParams { impl ParameterSpace for DQNParams { fn continuous_bounds() -> Vec<(f64, f64)> { - // 30D search space (C6: added gradient_accumulation_steps) + // 31D search space (C7: added iqn_lambda) // Fixed: curiosity_weight=0.0, noisy_epsilon_floor=0.10 // NoisyNet (noisy_sigma_init) + count bonus (UCB) provide exploration. vec![ @@ -519,12 +523,15 @@ impl ParameterSpace for DQNParams { // Gradient accumulation (1D) — C6: effective batch = batch_size × accum_steps // Default: fixed at 1 (no accumulation). Large GPUs expand to [1, 8]. (1.0, 1.0), // 29: gradient_accumulation_steps (discrete power-of-2: 1,2,4,8) + + // IQN dual-head lambda (1D) — C7: weight of IQN loss relative to C51 + (0.0, 2.0), // 30: iqn_lambda (0.0=C51 only, 0.25=mild IQN, 1.0=equal weight) ] } fn from_continuous(x: &[f64]) -> Result { - if x.len() != 30 { - return Err(MLError::ConfigError(format!("Expected 30 continuous parameters, got {}", x.len()))); + if x.len() != 31 { + return Err(MLError::ConfigError(format!("Expected 31 continuous parameters, got {}", x.len()))); } // === 26 TUNED parameters (from search space) === @@ -617,6 +624,8 @@ impl ParameterSpace for DQNParams { let use_qr_dqn = num_atoms > 100; let num_quantiles: usize = if use_qr_dqn { num_atoms } else { 64 }; let qr_kappa = 1.0; + // C7: IQN dual-head lambda (idx 30) + let iqn_lambda = x[30].clamp(0.0, 2.0); // WAVE 6 FIX #2: Batch size floor for high learning rates if learning_rate > 2e-4 && batch_size < 120 { @@ -679,6 +688,7 @@ impl ParameterSpace for DQNParams { use_qr_dqn, // M1: hyperopt-toggled num_quantiles, qr_kappa, + iqn_lambda, // C7: IQN dual-head loss weight hidden_dim_base, noisy_epsilon_floor: 0.10, // Minimum 10% random exploration to prevent action collapse count_bonus_coefficient, // C3 FIX: re-enabled UCB exploration bonus (complementary to NoisyNet) @@ -694,7 +704,7 @@ impl ParameterSpace for DQNParams { } fn to_continuous(&self) -> Vec { - // 30D search space (C6: added gradient_accumulation_steps) + // 31D search space (C7: added iqn_lambda) vec![ self.learning_rate.ln(), // 0 self.batch_size as f64, // 1 @@ -726,11 +736,12 @@ impl ParameterSpace for DQNParams { self.sharpe_weight, // 27: C4 FIX self.branch_hidden_dim as f64, // 28: C5 branch_hidden_dim self.gradient_accumulation_steps as f64, // 29: C6 gradient accumulation + self.iqn_lambda, // 30: C7 IQN dual-head loss weight ] } fn param_names() -> Vec<&'static str> { - // 30 tuned parameters (C6: added gradient_accumulation_steps) + // 31 tuned parameters (C7: added iqn_lambda) vec![ "learning_rate", // 0 "batch_size", // 1 @@ -762,6 +773,7 @@ impl ParameterSpace for DQNParams { "sharpe_weight", // 27: C4 FIX "branch_hidden_dim", // 28: C5 branching head capacity "gradient_accumulation_steps", // 29: C6 effective batch multiplier + "iqn_lambda", // 30: C7 IQN dual-head loss weight ] } @@ -2655,7 +2667,7 @@ impl HyperparameterOptimizable for DQNTrainer { gradient_collapse_multiplier: 100.0, // Adaptive threshold (LR × 100) gradient_collapse_patience: 5, // 5 consecutive epochs before early stop - // LR scheduling — tunable via hyperopt search space (idx 23 in 30D) + // LR scheduling — tunable via hyperopt search space (idx 23 in 31D) // 0=constant, 1=linear decay, 2=cosine annealing // BUG FIX: total_steps = self.epochs (scheduler steps once per epoch, // NOT per training step). Previous calc used epochs × steps_per_epoch @@ -2729,6 +2741,7 @@ impl HyperparameterOptimizable for DQNTrainer { use_qr_dqn: params.use_qr_dqn, num_quantiles: params.num_quantiles, qr_kappa: params.qr_kappa, + iqn_lambda: params.iqn_lambda, // Conservative Q-Learning (CQL) — wired from hyperopt search space use_cql: params.cql_alpha > 1e-6, @@ -3582,6 +3595,7 @@ mod tests { use_qr_dqn: true, num_quantiles: 64, qr_kappa: 1.0, + iqn_lambda: 0.5, hidden_dim_base: 512, noisy_epsilon_floor: 0.10, // Fixed: 10% floor to prevent action collapse count_bonus_coefficient: 0.0, // C2: fixed to 0.0 @@ -3593,7 +3607,7 @@ mod tests { }; let continuous = params.to_continuous(); - assert_eq!(continuous.len(), 30, "to_continuous must return 30D vector"); + assert_eq!(continuous.len(), 31, "to_continuous must return 31D vector"); let recovered = DQNParams::from_continuous(&continuous).unwrap(); // Tuned parameters must roundtrip exactly (26D) @@ -3641,7 +3655,7 @@ mod tests { #[test] fn test_dqn_params_bounds() { let bounds = DQNParams::continuous_bounds(); - assert_eq!(bounds.len(), 30); // C6: 30D (added gradient_accumulation_steps) + assert_eq!(bounds.len(), 31); // C7: 31D (added iqn_lambda) // Check log-scale bounds are reasonable assert!(bounds[0].0 < bounds[0].1); // learning_rate @@ -3690,7 +3704,7 @@ mod tests { #[test] fn test_param_names() { let names = DQNParams::param_names(); - assert_eq!(names.len(), 30); // C6: 30D (added gradient_accumulation_steps) + assert_eq!(names.len(), 31); // C7: 31D (added iqn_lambda) assert_eq!(names[0], "learning_rate"); assert_eq!(names[1], "batch_size"); assert_eq!(names[2], "gamma"); @@ -3726,7 +3740,7 @@ mod tests { #[test] fn test_per_params_always_enabled() { // Test that PER is always enabled with tunable alpha/beta parameters - // 30D search space (C6: added gradient_accumulation_steps) + // 31D search space (C7: added iqn_lambda) let continuous = vec![ 3e-5_f64.ln(), 92.0, 0.92, 97_273_f64.ln(), 4.0, 24.77_f64.ln(), 0.03, 1.0, 0.6, 0.4, // 8-9: per_alpha, per_beta_start @@ -3745,6 +3759,7 @@ mod tests { 0.3, // 27: sharpe_weight (C4) 128.0, // 28: branch_hidden_dim (C5) 1.0, // 29: gradient_accumulation_steps (C6) + 0.25, // 30: iqn_lambda (C7) ]; let params = DQNParams::from_continuous(&continuous).unwrap(); @@ -3761,7 +3776,7 @@ mod tests { assert!((params.v_min - (-25.0)).abs() < 1e-6, "v_min should be -v_range"); assert!((params.v_max - 25.0).abs() < 1e-6, "v_max should be +v_range"); - // Test PER parameter bounds (min values) -- 30D + // Test PER parameter bounds (min values) -- 31D let continuous_min = vec![ 2e-5_f64.ln(), 64.0, 0.88, 50_000_f64.ln(), 1.0, 10.0_f64.ln(), 0.05, 0.5, 0.4, 0.2, // per_alpha min, per_beta_start min @@ -3780,6 +3795,7 @@ mod tests { 0.0, // 27: sharpe_weight min (C4) 64.0, // 28: branch_hidden_dim min (C5) 1.0, // 29: gradient_accumulation_steps min (C6) + 0.0, // 30: iqn_lambda min (C7) ]; let params_min = DQNParams::from_continuous(&continuous_min).unwrap(); assert!((params_min.per_alpha - 0.4).abs() < 1e-6); @@ -3790,7 +3806,7 @@ mod tests { assert!((params_min.v_min - (-10.0)).abs() < 1e-6, "v_min should be -v_range at min"); assert!((params_min.v_max - 10.0).abs() < 1e-6, "v_max should be +v_range at min"); - // Test PER parameter bounds (max values) -- 30D + // Test PER parameter bounds (max values) -- 31D let continuous_max = vec![ 8e-5_f64.ln(), 512.0, 0.99, 100_000_f64.ln(), 4.0, 40.0_f64.ln(), 0.5, 2.0, 0.8, 0.6, // per_alpha max, per_beta_start max @@ -3809,6 +3825,7 @@ mod tests { 0.5, // 27: sharpe_weight max (C4) 256.0, // 28: branch_hidden_dim max (C5) 1.0, // 29: gradient_accumulation_steps max (C6) + 2.0, // 30: iqn_lambda max (C7) ]; let params_max = DQNParams::from_continuous(&continuous_max).unwrap(); assert!((params_max.per_alpha - 0.8).abs() < 1e-6); @@ -4110,7 +4127,7 @@ mod tests { fn test_qr_dqn_roundtrip_continuous() { let params = DQNParams::default(); let continuous = params.to_continuous(); - assert_eq!(continuous.len(), 30, "Should have 30 continuous dimensions"); + assert_eq!(continuous.len(), 31, "Should have 31 continuous dimensions"); let roundtrip = DQNParams::from_continuous(&continuous).unwrap(); // Default num_atoms=51 <= 100, so QR-DQN disabled, num_quantiles=64 (C51 default) assert!(!roundtrip.use_qr_dqn, "num_atoms=51 should use C51, not QR-DQN"); @@ -4121,7 +4138,7 @@ mod tests { #[test] fn test_qr_dqn_activation_threshold() { let bounds = DQNParams::continuous_bounds(); - let mut params = vec![0.0_f64; 30]; // C6: 30D + let mut params = vec![0.0_f64; 31]; // C7: 31D // Fill with valid defaults params[0] = (1e-4_f64).ln(); // learning_rate params[1] = 128.0; // batch_size @@ -4134,7 +4151,7 @@ mod tests { params[8] = 0.6; // per_alpha params[9] = 0.4; // per_beta_start // Fill remaining with midpoint of bounds - for i in 10..30 { + for i in 10..31 { params[i] = (bounds[i].0 + bounds[i].1) / 2.0; } @@ -4185,7 +4202,7 @@ mod tests { fn test_batch_size_respects_wide_bounds() { // Simulate PSO choosing batch_size=2048 (within VRAM-aware bounds) let bounds = DQNParams::continuous_bounds(); - let mut params = vec![0.0_f64; 30]; // C6: 30D + let mut params = vec![0.0_f64; 31]; // C7: 31D params[1] = 2048.0; // batch_size (index 1) // Fill other required params with valid defaults params[0] = (1e-4_f64).ln(); // learning_rate @@ -4198,7 +4215,7 @@ mod tests { params[8] = 0.6; // per_alpha params[9] = 0.4; // per_beta_start // Remaining params (indices 10-29) -- use midpoint of bounds - for i in 10..30 { + for i in 10..31 { params[i] = (bounds[i].0 + bounds[i].1) / 2.0; } let result = DQNParams::from_continuous(¶ms).unwrap(); @@ -4308,8 +4325,8 @@ mod tests { // eval_softmax_temp removed from search space (backtest uses greedy argmax). // Verify from_continuous always sets it to fixed 1.0 regardless of input. let bounds = DQNParams::continuous_bounds(); - let mut vec = vec![0.0_f64; 30]; // C6: 30D - for i in 0..30 { + let mut vec = vec![0.0_f64; 31]; // C7: 31D + for i in 0..31 { vec[i] = (bounds[i].0 + bounds[i].1) / 2.0; } let params = DQNParams::from_continuous(&vec).unwrap(); diff --git a/crates/ml/src/inference.rs b/crates/ml/src/inference.rs index 569d8e5a2..c45669d71 100644 --- a/crates/ml/src/inference.rs +++ b/crates/ml/src/inference.rs @@ -1049,10 +1049,10 @@ mod tests { let result = engine .load_model("test_model".to_owned(), model_config) .await; + // CPU execution is no longer supported — CUDA required assert!( - result.is_ok(), - "Failed to load model on CPU: {:?}", - result.err() + result.is_err(), + "CPU model loading should be rejected (CUDA required)" ); Ok(()) } @@ -1062,7 +1062,7 @@ mod tests { async fn test_model_loading_multiple_models() -> Result<(), Box> { let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default())); let mut config = InferenceConfig::default(); - config.device_preference = "cpu".to_owned(); // Use CPU for testing (GPU may not be available) + config.device_preference = "cuda".to_owned(); let engine = MLInferenceEngine::new(config, safety_manager); // Load multiple models @@ -1110,7 +1110,7 @@ mod tests { safety_config.safety_enabled = false; // Disable for test with random weights let safety_manager = Arc::new(MLSafetyManager::new(safety_config)); let mut config = InferenceConfig::default(); - config.device_preference = "cpu".to_owned(); // Force CPU for testing + config.device_preference = "cuda".to_owned(); let engine = MLInferenceEngine::new(config, safety_manager); let model_config = ModelConfig { @@ -1138,7 +1138,7 @@ mod tests { async fn test_inference_with_missing_model() -> Result<(), Box> { let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default())); let mut config = InferenceConfig::default(); - config.device_preference = "cpu".to_owned(); + config.device_preference = "cuda".to_owned(); let engine = MLInferenceEngine::new(config, safety_manager); let features = create_mock_features(); @@ -1152,7 +1152,7 @@ mod tests { async fn test_inference_dimension_mismatch() -> Result<(), Box> { let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default())); let mut config = InferenceConfig::default(); - config.device_preference = "cpu".to_owned(); + config.device_preference = "cuda".to_owned(); let engine = MLInferenceEngine::new(config, safety_manager); // Model expects 10 features but features_to_tensor produces 21 @@ -1181,7 +1181,7 @@ mod tests { { let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default())); let mut config = InferenceConfig::default(); - config.device_preference = "cpu".to_owned(); + config.device_preference = "cuda".to_owned(); let engine = MLInferenceEngine::new(config, safety_manager); let model_config = ModelConfig { @@ -1218,7 +1218,7 @@ mod tests { let mut config = InferenceConfig::default(); config.enable_caching = true; config.cache_ttl_seconds = 60; - config.device_preference = "cpu".to_owned(); + config.device_preference = "cuda".to_owned(); let engine = MLInferenceEngine::new(config, safety_manager); @@ -1441,7 +1441,7 @@ mod tests { async fn test_concurrent_predictions() -> Result<(), Box> { let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default())); let mut config = InferenceConfig::default(); - config.device_preference = "cpu".to_owned(); + config.device_preference = "cuda".to_owned(); let engine = Arc::new(MLInferenceEngine::new(config, safety_manager)); let model_config = ModelConfig { @@ -1499,7 +1499,7 @@ mod tests { async fn test_model_replacement() -> Result<(), Box> { let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default())); let mut config = InferenceConfig::default(); - config.device_preference = "cpu".to_owned(); + config.device_preference = "cuda".to_owned(); let engine = MLInferenceEngine::new(config, safety_manager); // Load initial model diff --git a/crates/ml/src/liquid/adapter.rs b/crates/ml/src/liquid/adapter.rs index c3cee3909..29f883a8d 100644 --- a/crates/ml/src/liquid/adapter.rs +++ b/crates/ml/src/liquid/adapter.rs @@ -411,10 +411,10 @@ mod tests { hidden_size: 16, output_size: 3, backbone_hidden_sizes: vec![16], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), ..CfCTrainConfig::default() }; - let adapter = LiquidTrainableAdapter::new(config).unwrap(); + let adapter = LiquidTrainableAdapter::new(config).expect("CUDA required"); assert_eq!(adapter.model_type(), "Liquid-CfC"); assert_eq!(adapter.get_step(), 0); assert!(adapter.get_learning_rate() > 0.0); @@ -427,7 +427,7 @@ mod tests { hidden_size: 16, output_size: 3, backbone_hidden_sizes: vec![16], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), seq_len: 5, ..CfCTrainConfig::default() }; @@ -454,7 +454,7 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), seq_len: 3, ..CfCTrainConfig::default() }; @@ -482,7 +482,7 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), ..CfCTrainConfig::default() }; let adapter = LiquidTrainableAdapter::new(config).unwrap(); @@ -499,7 +499,7 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), ..CfCTrainConfig::default() }; let mut adapter = LiquidTrainableAdapter::new(config).unwrap(); @@ -520,7 +520,7 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), seq_len: 3, ..CfCTrainConfig::default() }; @@ -546,7 +546,7 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), seq_len: 3, ..CfCTrainConfig::default() }; @@ -592,7 +592,7 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), ..CfCTrainConfig::default() }; let mut adapter = LiquidTrainableAdapter::new(config).unwrap(); @@ -607,7 +607,7 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), ..CfCTrainConfig::default() }; let mut adapter = LiquidTrainableAdapter::new(config).unwrap(); @@ -622,11 +622,11 @@ mod tests { hidden_size: 8, output_size: 2, backbone_hidden_sizes: vec![8], - device: DeviceConfig::Cpu, + device: DeviceConfig::Cuda(0), seq_len: 3, ..CfCTrainConfig::default() }; - let mut adapter = LiquidTrainableAdapter::new(config).unwrap(); + let mut adapter = LiquidTrainableAdapter::new(config).expect("CUDA required"); let device = adapter.device().clone(); let input = Tensor::randn(0_f32, 1.0, (4, 3, 4), &device).unwrap(); diff --git a/crates/ml/src/mamba/trainable_adapter.rs b/crates/ml/src/mamba/trainable_adapter.rs index 26c32d9e8..38c494295 100644 --- a/crates/ml/src/mamba/trainable_adapter.rs +++ b/crates/ml/src/mamba/trainable_adapter.rs @@ -253,7 +253,7 @@ mod tests { let adapter = Mamba2TrainableAdapter::new(config, &Device::new_cuda(0).expect("CUDA required"))?; assert_eq!(adapter.model_type(), "MAMBA-2"); - assert_eq!(format!("{:?}", adapter.device()), "Cpu"); + assert!(format!("{:?}", adapter.device()).contains("Cuda")); assert_eq!(adapter.get_step(), 0); assert!(adapter.get_learning_rate() > 0.0); diff --git a/crates/ml/src/portfolio_transformer.rs b/crates/ml/src/portfolio_transformer.rs index 06559cfb9..cb0d7f59a 100644 --- a/crates/ml/src/portfolio_transformer.rs +++ b/crates/ml/src/portfolio_transformer.rs @@ -4,7 +4,7 @@ //! in high-frequency trading. Unlike traditional time-series transformers, this model //! operates directly on portfolio state vectors for optimal weight prediction. -use candle_core::{Device, IndexOp, Module, ModuleT, Result as CandleResult, Tensor}; +use candle_core::{DType, Device, IndexOp, Module, ModuleT, Result as CandleResult, Tensor}; use candle_nn::{Linear, VarBuilder, VarMap}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; @@ -258,7 +258,8 @@ impl PortfolioTransformer { } } - Tensor::from_vec(pe_data, (seq_len, model_dim), device) + let t = Tensor::from_vec(pe_data, (seq_len, model_dim), device)?; + t.to_dtype(DType::BF16) } /// Optimize portfolio weights using transformer @@ -333,11 +334,17 @@ impl PortfolioTransformer { // Convert to f32 for Candle let input_f32: Vec = input_data.iter().map(|&x| x as f32).collect(); - Tensor::from_vec(input_f32, (1, expected_size), &self.device).map_err(|e| { + let t = Tensor::from_vec(input_f32, (1, expected_size), &self.device).map_err(|e| { MLError::TensorCreationError { operation: "prepare_input_tensor".to_owned(), reason: e.to_string(), } + })?; + t.to_dtype(DType::BF16).map_err(|e| { + MLError::TensorCreationError { + operation: "prepare_input_tensor_cast".to_owned(), + reason: e.to_string(), + } }) } @@ -375,7 +382,7 @@ impl PortfolioTransformer { let weights_tensor = candle_nn::ops::softmax(&logits, 1)?; // Convert to Vec - let weights_flat = weights_tensor.flatten_all()?.to_vec1::()?; + let weights_flat = weights_tensor.flatten_all()?.to_dtype(DType::F32)?.to_vec1::()?; let weights: Vec = weights_flat.iter().map(|&x| x as f64).collect(); // Ensure we have the right number of weights @@ -396,7 +403,7 @@ impl PortfolioTransformer { // Risk head forward pass let risk_tensor = Module::forward(&self.risk_head, &pooled)?; - let risk_vec = risk_tensor.flatten_all()?.to_vec1::()?; + let risk_vec = risk_tensor.flatten_all()?.to_dtype(DType::F32)?.to_vec1::()?; let risk_value = *risk_vec .first() .ok_or_else(|| MLError::TensorCreationError { @@ -427,7 +434,7 @@ impl PortfolioTransformer { let regime_probs = candle_nn::ops::softmax(®ime_logits, 1)?; // Get the most likely regime - let probs = regime_probs.flatten_all()?.to_vec1::()?; + let probs = regime_probs.flatten_all()?.to_dtype(DType::F32)?.to_vec1::()?; let max_idx = probs .iter() .enumerate() diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 55fd42509..9fcc1823a 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -626,6 +626,41 @@ impl DQNAgentType { self.memory().flush_max_priority() } + /// GPU-only post-step bookkeeping: training_steps++ and PER beta annealing. + /// + /// Used when PER priority update is handled by the CUDA kernel in + /// `GpuDqnTrainer::update_priorities_cuda()` — no CPU priority update needed. + pub fn fused_post_step_gpu_bookkeeping(&mut self) -> Result<(), crate::MLError> { + match self { + Self::Standard(agent) => { + agent.increment_training_steps(); + agent.memory.step(); + Ok(()) + } + Self::RegimeConditional(_) => Err(crate::MLError::ModelError( + "Fused CUDA training not supported for RegimeConditional agent".into(), + )), + } + } + + /// PER alpha/epsilon for GPU-native priority update kernel. + pub fn per_alpha_epsilon(&self) -> Option<(f32, f32)> { + self.memory().per_alpha_epsilon() + } + + /// Direct reference to the GPU-resident priorities tensor. + /// + /// Returns `None` for non-GPU buffers. Used by fused training to pass + /// the tensor to `GpuDqnTrainer::update_priorities_cuda()`. + pub fn priorities_tensor(&self) -> Option { + self.memory().priorities_tensor() + } + + /// Set max_priority from a CPU scalar (epoch-boundary flush from CUDA kernel). + pub fn apply_max_priority_scalar(&self, max_prio: f32) -> Result<(), crate::MLError> { + self.memory().apply_max_priority_scalar(max_prio) + } + /// Update learning rate for the optimizer(s) by applying a decay factor. /// /// For Standard agents, updates the single optimizer. @@ -993,6 +1028,10 @@ pub struct DQNHyperparameters { pub num_quantiles: usize, /// Kappa for quantile Huber loss (0.5-2.0) pub qr_kappa: f64, + /// Lambda weight for IQN quantile loss relative to C51 loss in the dual-head ensemble. + /// L_total = L_c51 + iqn_lambda * L_iqn + /// Range [0.0, 2.0]: 0.0 = C51 only, 0.5 = balanced, 1.0 = equal weight + pub iqn_lambda: f64, // Phase 3: GPU experience collection kernel configuration /// Enable the NVRTC-compiled GPU experience collection kernel. @@ -1267,6 +1306,7 @@ impl DQNHyperparameters { use_qr_dqn: true, // Default: enabled (IQN distributional RL) num_quantiles: 32, // Default: 32 quantiles qr_kappa: 1.0, // Default: 1.0 (standard quantile Huber loss) + iqn_lambda: 0.25, // Default: mild IQN regularization alongside C51 // Phase 3: GPU experience collection enable_gpu_experience_collector: true, // Default: enabled (zero-roundtrip CUDA kernel) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index a29fe3731..c51ae4289 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -31,11 +31,12 @@ use tracing::info; use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer}; use crate::cuda_pipeline::gpu_her::{GpuHer, GpuHerConfig, parse_her_strategy}; use crate::cuda_pipeline::gpu_iql_trainer::{GpuIqlConfig, GpuIqlTrainer}; +use crate::cuda_pipeline::gpu_iqn_head::{GpuIqnConfig, GpuIqnHead}; use crate::cuda_pipeline::gpu_weights::{ self, BranchingWeightSet, DuelingWeightSet, }; use crate::dqn::dqn::GpuTrainResult; -use crate::dqn::replay_buffer_type::BatchSample; +use crate::dqn::replay_buffer_type::{BatchSample, GpuBatch}; use super::config::DQNAgentType; use super::DQNHyperparameters; @@ -66,6 +67,10 @@ pub(crate) struct FusedTrainingCtx { /// Trains V(s) with expectile regression after each DQN training step. /// The value loss is logged via tracing for monitoring. pub(crate) gpu_iql: Option, + /// GPU IQN dual-head -- initialized when `use_qr_dqn == true` and `iqn_lambda > 0`. + /// Trains quantile head that shares the DQN trunk (h_s2) independently. + /// Combined loss: L_total = L_c51 + iqn_lambda × L_iqn. + pub(crate) gpu_iqn: Option, } impl FusedTrainingCtx { @@ -132,6 +137,10 @@ impl FusedTrainingCtx { epsilon: 1e-8, weight_decay: 1e-5, max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32, + iqn_lambda: if hyperparams.use_qr_dqn { hyperparams.iqn_lambda as f32 } else { 0.0 }, + iqn_num_quantiles: if hyperparams.use_qr_dqn { hyperparams.num_quantiles } else { 0 }, + iqn_embedding_dim: dqn.config.iqn_embedding_dim, + iqn_kappa: dqn.config.iqn_kappa, }; // Extract weight sets from VarMaps (online + target) @@ -213,10 +222,48 @@ impl FusedTrainingCtx { None }; + // Initialize GPU IQN dual-head if use_qr_dqn is enabled and iqn_lambda > 0 + let gpu_iqn = if hyperparams.use_qr_dqn && hyperparams.iqn_lambda > 0.0 { + let iqn_config = GpuIqnConfig { + hidden_dim: shared_h2, + embed_dim: dqn.config.iqn_embedding_dim, + num_quantiles: hyperparams.num_quantiles, + kappa: dqn.config.iqn_kappa, + state_dim: dqn.config.state_dim, + shared_h1, + branch_0_size: dqn.config.num_actions, + branch_1_size: dqn.config.num_order_types, + branch_2_size: dqn.config.num_urgency_levels, + batch_size, + lr: hyperparams.learning_rate as f32, + max_grad_norm: hyperparams.gradient_clip_norm.unwrap_or(1.0) as f32, + gamma: hyperparams.gamma as f32, + ..GpuIqnConfig::default() + }; + match GpuIqnHead::new(stream.clone(), iqn_config) { + Ok(iqn) => { + info!( + iqn_lambda = hyperparams.iqn_lambda, + num_quantiles = hyperparams.num_quantiles, + embed_dim = dqn.config.iqn_embedding_dim, + "GPU IQN dual-head initialized: 7 kernels compiled" + ); + Some(iqn) + } + Err(e) => { + tracing::warn!("GPU IQN init failed, continuing without IQN: {e}"); + None + } + } + } else { + None + }; + info!( batch_size, her_enabled = gpu_her.is_some(), iql_enabled = gpu_iql.is_some(), + iqn_enabled = gpu_iqn.is_some(), "Fused CUDA training initialized: 4 kernels + EMA compiled, \ ~291K params, CUDA Graph will capture on first step" ); @@ -232,6 +279,7 @@ impl FusedTrainingCtx { steps_since_varmap_sync: 0, gpu_her, gpu_iql, + gpu_iqn, }) } @@ -245,39 +293,46 @@ impl FusedTrainingCtx { self.steps_since_varmap_sync } - /// Run one full fused training step. + /// Run one full fused training step — GPU-only, no CPU paths. /// - /// Executes the complete training cycle with GPU-native EMA: - /// 1. Extract flat arrays from `BatchSample` - /// 2. Forward + loss + backward + Adam via fused CUDA kernels (or CUDA Graph replay) - /// 3. GPU EMA: `target[i] = (1-tau)*target[i] + tau*online[i]` (20 kernel launches) - /// 4. PER priority update + beta annealing + training_steps increment - /// 5. Return `GpuTrainResult` with GPU scalar tensors for monitoring + /// All batch data stays GPU-resident. DtoD copies between GpuBatch tensors + /// and trainer CudaSlice buffers — zero PCIe roundtrips. /// - /// No VarMap sync per step -- deferred to epoch boundary via `sync_to_varmap()`. - /// This eliminates ~120 D2D copies + 120 Candle ops per step (~13 GB/epoch saved). + /// Pipeline: + /// 1. GPU HER relabeling (if enabled) + /// 2. DQN: DtoD upload → fused CUDA Graph (forward+loss+backward+Adam) + /// 3. GPU Polyak EMA target update + /// 4. IQL: GPU-resident tensors → expectile regression + /// 5. IQN: DtoD from DQN trainer buffers → GPU decode → quantile training + /// 6. PER priority update + training_steps increment + /// + /// No VarMap sync per step — deferred to epoch boundary. pub(crate) fn run_full_step( &mut self, batch: &BatchSample, agent: &mut DQNAgentType, device: &Device, ) -> Result { - let state_dim = agent.get_state_dim(); - let (states, next_states, actions, rewards, dones, is_weights) = - extract_batch_arrays(batch, state_dim)?; + let gpu_batch = batch.gpu_batch.as_ref() + .ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?; - // Step 1: Fused forward + loss + backward + Adam - // The Adam kernel updates online CudaSlice weights in-place. - // Device pointers are stable -- CUDA Graph stays valid. - let fused_result = self.trainer.train_step( - &states, &next_states, &actions, &rewards, &dones, &is_weights, + // ── Step 1: GPU HER relabeling ─────────────────────────────────── + let her_modified_gpu; + let effective_gpu = if let Some(ref her) = self.gpu_her { + her_modified_gpu = gpu_her_relabel_batch(gpu_batch, &her.config)?; + &her_modified_gpu + } else { + gpu_batch + }; + + // ── Step 2: Fused DQN training — DtoD → CUDA Graph ────────────── + let fused_result = self.trainer.train_step_gpu( + effective_gpu, &self.online_dueling, &self.online_branching, &self.target_dueling, &self.target_branching, - ).map_err(|e| anyhow::anyhow!("Fused train_step: {e}"))?; + ).map_err(|e| anyhow::anyhow!("Fused train_step_gpu: {e}"))?; - // Step 2: GPU-native Polyak EMA target update - // Computes cosine-annealed tau and applies target[i] = (1-tau)*target[i] + tau*online[i] - // entirely on GPU. No VarMap round-trip. + // ── Step 3: GPU-native Polyak EMA target update ────────────────── { let dqn = agent.as_standard_mut().ok_or_else(|| { anyhow::anyhow!("Fused training requires Standard DQN agent") @@ -299,20 +354,13 @@ impl FusedTrainingCtx { } } - // Step 3: IQL value network training (if enabled) - // Trains V(s) with expectile regression using rewards as Q-value targets. - // The value loss is logged for monitoring; advantage weights will be used - // for policy extraction once full Q-value exposure is wired. + // ── Step 4: IQL value network (if enabled) ─────────────────────── + // Uses GpuBatch tensors directly — GPU dtype cast, no CPU touch. if let Some(ref mut iql) = self.gpu_iql { - // Build Candle tensors on GPU from the already-extracted batch arrays. - // Batch size = number of actions (1 per sample). - let b = actions.len(); - let states_tensor = Tensor::new(states.as_slice(), device) - .and_then(|t| t.reshape((b, state_dim))) - .map_err(|e| anyhow::anyhow!("IQL states tensor: {e}"))?; - - let rewards_tensor = Tensor::new(rewards.as_slice(), device) - .map_err(|e| anyhow::anyhow!("IQL rewards tensor: {e}"))?; + let states_tensor = effective_gpu.states.to_dtype(candle_core::DType::F32) + .map_err(|e| anyhow::anyhow!("IQL states F32 cast: {e}"))?; + let rewards_tensor = effective_gpu.rewards.to_dtype(candle_core::DType::F32) + .map_err(|e| anyhow::anyhow!("IQL rewards F32 cast: {e}"))?; match iql.train_value_step(&states_tensor, &rewards_tensor) { Ok(value_loss) => { @@ -328,15 +376,73 @@ impl FusedTrainingCtx { } } - // Step 4: PER priority update + training_steps++ + beta annealing - // (no target EMA -- already done by GPU kernel above) - // td_errors from fused kernel are already f32 on CPU. - agent.fused_post_step_no_ema(&fused_result.td_errors, &batch.indices) - .map_err(|e| anyhow::anyhow!("Fused post_step_no_ema: {e}"))?; + // ── Step 5: IQN dual-head (if enabled) ────────────────────────── + // GPU-native: DtoD from DQN trainer buffers, GPU action decode kernel. + // Zero CPU arrays, zero HtoD uploads for batch data. + if let Some(ref mut iqn) = self.gpu_iqn { + let online_h_s2 = self.trainer.save_h_s2(); + let next_states_buf = self.trainer.next_states_buf(); + let dqn_actions = self.trainer.actions_buf(); + let dqn_rewards = self.trainer.rewards_buf(); + let dqn_dones = self.trainer.dones_buf(); + + match iqn.train_iqn_step_gpu( + online_h_s2, + next_states_buf, + &self.target_dueling, + dqn_actions, + dqn_rewards, + dqn_dones, + ) { + Ok(iqn_loss) => { + tracing::debug!( + iqn_loss, + iqn_adam_step = iqn.adam_step(), + "IQN dual-head step" + ); + + // IQN target EMA (same schedule as DQN) + let dqn = agent.as_standard_mut().ok_or_else(|| { + anyhow::anyhow!("IQN EMA requires Standard DQN agent") + })?; + if dqn.config.use_soft_updates { + let tau = compute_cosine_annealed_tau( + dqn.get_training_steps(), + dqn.config.tau, + dqn.config.tau_final, + dqn.config.tau_anneal_steps, + ); + iqn.target_ema_update(tau as f32) + .map_err(|e| anyhow::anyhow!("IQN EMA update: {e}"))?; + } + } + Err(e) => { + tracing::warn!("IQN step failed (non-fatal): {e}"); + } + } + } + + // ── Step 6: GPU-native PER priority update ───────────────────── + // td_errors stay on GPU (td_errors_buf). Single CUDA kernel scatter-writes + // new priorities + atomicMax for batch max. Zero DtoH readback. + if let (Some(priorities_tensor), Some((alpha, epsilon))) = + (agent.priorities_tensor(), agent.per_alpha_epsilon()) + { + self.trainer.update_priorities_cuda( + &effective_gpu.indices, + &priorities_tensor, + alpha, + epsilon, + ).map_err(|e| anyhow::anyhow!("GPU PER priority update: {e}"))?; + } + + // training_steps++ and PER beta annealing (no CPU priority update) + agent.fused_post_step_gpu_bookkeeping() + .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?; self.steps_since_varmap_sync += 1; - // Step 5: Create GPU scalar tensors from fused results for monitoring code + // ── Step 7: Create GPU scalar tensors for monitoring ───────────── Ok(GpuTrainResult { loss_gpu: Tensor::new(fused_result.total_loss, device) .map_err(|e| anyhow::anyhow!("Fused loss->Tensor: {e}"))?, @@ -380,6 +486,15 @@ impl FusedTrainingCtx { target_vars, &self.target_branching, &self.stream, ).map_err(|e| anyhow::anyhow!("Reverse sync target branching: {e}"))?; + // Flush batch_max_buf: single scalar DtoH (4 bytes), once per epoch. + // Sets GpuReplayBuffer.max_priority so new experiences get sampled first. + let batch_max = self.trainer.batch_max_readback_and_reset() + .map_err(|e| anyhow::anyhow!("batch_max readback: {e}"))?; + if batch_max > 0.0 { + agent.apply_max_priority_scalar(batch_max) + .map_err(|e| anyhow::anyhow!("apply batch_max: {e}"))?; + } + self.steps_since_varmap_sync = 0; Ok(()) } @@ -407,68 +522,117 @@ fn compute_cosine_annealed_tau( } } -/// Extract flat `f32`/`i32` arrays from a `BatchSample` for the fused CUDA trainer. +/// GPU-native intra-batch HER: relabel goal portions using Candle tensor ops. /// -/// Returns `(states, next_states, actions, rewards, dones, is_weights)`. +/// All computation stays on GPU — zero CPU roundtrip. For the last +/// `her_batch_size` samples in the batch: +/// 1. Pick random donors (index array generated on CPU — trivial O(K) routing) +/// 2. Replace goal portion of states/next_states with donor's achieved goal (GPU) +/// 3. Set reward = 1.0 (goal = achieved by construction in Random strategy) /// -/// When `gpu_batch` is present (GPU PER path), reads tensors back to CPU. -/// When absent (CPU PER/uniform path), iterates `batch.experiences`. -fn extract_batch_arrays( - batch: &BatchSample, - state_dim: usize, -) -> Result<(Vec, Vec, Vec, Vec, Vec, Vec)> { - // GPU PER path: data lives in gpu_batch tensors, experiences is empty. - // Tensors may be BF16 (training precision) — cast to F32 for the fused kernel. - if let Some(ref gpu) = batch.gpu_batch { - use candle_core::DType; - let states = gpu.states.to_dtype(DType::F32) - .and_then(|t| t.flatten_all()) - .and_then(|t| t.to_vec1::()) - .map_err(|e| anyhow::anyhow!("GPU batch states readback: {e}"))?; - let next_states = gpu.next_states.to_dtype(DType::F32) - .and_then(|t| t.flatten_all()) - .and_then(|t| t.to_vec1::()) - .map_err(|e| anyhow::anyhow!("GPU batch next_states readback: {e}"))?; - let actions: Vec = gpu.actions.to_vec1::() - .map_err(|e| anyhow::anyhow!("GPU batch actions readback: {e}"))? - .into_iter().map(|a| a as i32).collect(); - let rewards = gpu.rewards.to_dtype(DType::F32) - .and_then(|t| t.to_vec1::()) - .map_err(|e| anyhow::anyhow!("GPU batch rewards readback: {e}"))?; - let dones = gpu.dones.to_dtype(DType::F32) - .and_then(|t| t.to_vec1::()) - .map_err(|e| anyhow::anyhow!("GPU batch dones readback: {e}"))?; - let is_weights = gpu.weights.to_dtype(DType::F32) - .and_then(|t| t.to_vec1::()) - .map_err(|e| anyhow::anyhow!("GPU batch IS weights readback: {e}"))?; +/// Returns a new `GpuBatch` with modified states, next_states, and rewards. +/// Actions, dones, weights, and indices are shared (reference-counted clone). +fn gpu_her_relabel_batch( + gpu: &GpuBatch, + config: &GpuHerConfig, +) -> Result { + use rand::Rng; - // Validate dimensions match what the fused kernel expects. - let b = actions.len(); - if states.len() != b * state_dim { - return Err(anyhow::anyhow!( - "GPU batch states length {} != batch_size({}) × state_dim({})", - states.len(), b, state_dim, - )); - } + let batch_size = gpu.states.dims()[0]; + let state_dim = gpu.states.dims()[1]; + let her_batch_size = config.her_batch_size(); - return Ok((states, next_states, actions, rewards, dones, is_weights)); + if her_batch_size == 0 || her_batch_size >= batch_size || config.goal_dim == 0 { + return Ok(GpuBatch { + states: gpu.states.clone(), + next_states: gpu.next_states.clone(), + rewards: gpu.rewards.clone(), + actions: gpu.actions.clone(), + dones: gpu.dones.clone(), + weights: gpu.weights.clone(), + indices: gpu.indices.clone(), + }); } - // CPU path: iterate batch.experiences directly. - let b = batch.experiences.len(); - let mut states = Vec::with_capacity(b * state_dim); - let mut next_states = Vec::with_capacity(b * state_dim); - let mut actions = Vec::with_capacity(b); - let mut rewards = Vec::with_capacity(b); - let mut dones = Vec::with_capacity(b); + let goal_dim = config.goal_dim.min(state_dim); + let normal_count = batch_size - her_batch_size; + let device = gpu.states.device(); - for exp in &batch.experiences { - states.extend_from_slice(&exp.state); - next_states.extend_from_slice(&exp.next_state); - actions.push(exp.action as i32); - rewards.push(exp.reward_f32()); - dones.push(if exp.done { 1.0_f32 } else { 0.0_f32 }); - } + // Generate donor indices — small CPU array (index routing, not batch data). + // O(her_batch_size) ints ≈ 64 bytes, acceptable at this boundary. + let mut rng = rand::thread_rng(); + let donor_indices: Vec = (0..her_batch_size) + .map(|_| rng.gen_range(0..batch_size as u32)) + .collect(); + let donor_idx = Tensor::new(&donor_indices[..], device) + .map_err(|e| anyhow::anyhow!("HER donor indices upload: {e}"))?; - Ok((states, next_states, actions, rewards, dones, batch.weights.clone())) + // Gather donor achieved goals: donor's next_state[:, :goal_dim] + // index_select + narrow — all GPU kernel launches, zero CPU data movement. + let donor_achieved = gpu.next_states + .index_select(&donor_idx, 0) + .and_then(|t| t.narrow(1, 0, goal_dim)) + .map_err(|e| anyhow::anyhow!("HER gather donor goals: {e}"))?; + + // Normal portion (unchanged) — narrow is zero-copy (view into same storage) + let normal_states = gpu.states.narrow(0, 0, normal_count) + .map_err(|e| anyhow::anyhow!("HER normal states: {e}"))?; + let normal_next = gpu.next_states.narrow(0, 0, normal_count) + .map_err(|e| anyhow::anyhow!("HER normal next_states: {e}"))?; + let normal_rewards = gpu.rewards.narrow(0, 0, normal_count) + .map_err(|e| anyhow::anyhow!("HER normal rewards: {e}"))?; + + // HER portion: replace goal columns with donor's achieved goal, keep rest + let her_states = if goal_dim >= state_dim { + // Entire state is goal — donor_achieved is the full replacement + donor_achieved.clone() + } else { + let her_states_rest = gpu.states + .narrow(0, normal_count, her_batch_size) + .and_then(|t| t.narrow(1, goal_dim, state_dim - goal_dim)) + .map_err(|e| anyhow::anyhow!("HER states rest: {e}"))?; + Tensor::cat(&[&donor_achieved, &her_states_rest], 1) + .map_err(|e| anyhow::anyhow!("HER states cat: {e}"))? + }; + + let her_next = if goal_dim >= state_dim { + donor_achieved.clone() + } else { + let her_next_rest = gpu.next_states + .narrow(0, normal_count, her_batch_size) + .and_then(|t| t.narrow(1, goal_dim, state_dim - goal_dim)) + .map_err(|e| anyhow::anyhow!("HER next_states rest: {e}"))?; + Tensor::cat(&[&donor_achieved, &her_next_rest], 1) + .map_err(|e| anyhow::anyhow!("HER next cat: {e}"))? + }; + + // HER rewards: +1.0 (goal = achieved by construction in Random strategy) + let reward_dtype = gpu.rewards.dtype(); + let her_rewards = Tensor::ones(her_batch_size, reward_dtype, device) + .map_err(|e| anyhow::anyhow!("HER reward ones: {e}"))?; + + // Reassemble full batch: [normal | her] + let new_states = Tensor::cat(&[&normal_states, &her_states], 0) + .map_err(|e| anyhow::anyhow!("HER concat states: {e}"))?; + let new_next = Tensor::cat(&[&normal_next, &her_next], 0) + .map_err(|e| anyhow::anyhow!("HER concat next: {e}"))?; + let new_rewards = Tensor::cat(&[&normal_rewards, &her_rewards], 0) + .map_err(|e| anyhow::anyhow!("HER concat rewards: {e}"))?; + + tracing::debug!( + her_batch_size, + batch_size, + goal_dim, + "GPU HER intra-batch relabeling applied" + ); + + Ok(GpuBatch { + states: new_states, + next_states: new_next, + rewards: new_rewards, + actions: gpu.actions.clone(), + dones: gpu.dones.clone(), + weights: gpu.weights.clone(), + indices: gpu.indices.clone(), + }) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs b/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs index 78add0f32..895aedf77 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/feature_coverage.rs @@ -3,13 +3,19 @@ //! Verify that every major DQN feature flag can be enabled and trained //! without panics or NaN. GPU-only — no CPU fallback. -use super::helpers::{assert_finite, smoke_trainer_with, smoke_params, synthetic_data}; +use super::helpers::{assert_finite, smoke_trainer_with, smoke_params, test_data_dir}; -/// Helper: train with modified params, assert loss is finite. +/// Helper: train with modified params on real data, assert loss is finite. +/// +/// Uses real `.dbn` data from `test_data/` (set via `FOXHUNT_TEST_DATA` env +/// or auto-detected at repo root). Hard error if data is missing. async fn train_and_check( mut params: crate::trainers::dqn::DQNHyperparameters, label: &str, ) -> anyhow::Result<()> { + let data_dir = test_data_dir() + .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback"); + // Keep tests fast params.epochs = 3; params.batch_size = 16; @@ -31,13 +37,9 @@ async fn train_and_check( } let mut trainer = smoke_trainer_with(params)?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir, |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; assert_finite(metrics.loss, label); Ok(()) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index ba58e5952..98811205d 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -180,15 +180,12 @@ async fn test_searchsorted_gpu_correctness() -> anyhow::Result<()> { #[tokio::test] async fn test_train_step_produces_finite_metrics() -> anyhow::Result<()> { + let data_dir = test_data_dir() + .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback"); let mut trainer = smoke_trainer()?; - let train_data = synthetic_data(200); - let val_data = synthetic_data(50); - - let metrics = trainer - .train_with_preloaded_data(train_data, val_data, |_epoch, _bytes, _is_best| { - Ok(String::new()) - }) - .await?; + let metrics = trainer.train(&data_dir, |_epoch, _bytes, _is_best| { + Ok(String::new()) + }).await?; // Final loss must be finite (not NaN, not Inf) assert_finite(metrics.loss, "final_loss"); @@ -226,17 +223,15 @@ async fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> { /// is compiled out entirely. This test confirms the guard fires correctly. #[tokio::test] async fn test_training_rejects_missing_gpu_collector() -> anyhow::Result<()> { + let data_dir = test_data_dir() + .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback"); let mut params = smoke_params(); params.enable_gpu_experience_collector = false; // deliberately disable let mut trainer = smoke_trainer_with(params)?; - let train_data = synthetic_data(200); - let val_data = synthetic_data(50); - let result = trainer - .train_with_preloaded_data(train_data, val_data, |_epoch, _bytes, _is_best| { - Ok(String::new()) - }) - .await; + let result = trainer.train(&data_dir, |_epoch, _bytes, _is_best| { + Ok(String::new()) + }).await; assert!( result.is_err(), diff --git a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs index f73e11a30..8da74b395 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs @@ -1,5 +1,4 @@ use crate::trainers::dqn::{DQNHyperparameters, DQNTrainer}; -use crate::features::extraction::FeatureVector; use candle_core::Device; /// Initialize tracing subscriber once for all smoke tests. @@ -49,27 +48,6 @@ pub(super) fn smoke_params() -> DQNHyperparameters { p } -/// Generate synthetic training data (feature vectors + OHLC targets) -pub(super) fn synthetic_data(n: usize) -> Vec<(FeatureVector, Vec)> { - use rand::Rng; - let mut rng = rand::thread_rng(); - (0..n) - .map(|_| { - let mut fv = [0.0_f64; 42]; - for v in &mut fv { - *v = rng.gen_range(-1.0..1.0); - } - let targets = vec![ - 4000.0 + rng.gen_range(-50.0..50.0), // open - 4010.0 + rng.gen_range(-50.0..50.0), // high - 3990.0 + rng.gen_range(-50.0..50.0), // low - 4005.0 + rng.gen_range(-50.0..50.0), // close - ]; - (fv, targets) - }) - .collect() -} - /// Shared CUDA device for all smoke tests — reuses one cuBLAS handle /// instead of creating a new one per test, preventing handle pool exhaustion. static SMOKE_CUDA: std::sync::OnceLock = std::sync::OnceLock::new(); @@ -109,3 +87,25 @@ pub(super) fn assert_finite(val: f64, name: &str) { pub(super) fn assert_finite_f32(val: f32, name: &str) { assert!(val.is_finite(), "{name} is not finite: {val}"); } + +/// Resolve the test data directory from `FOXHUNT_TEST_DATA` env var, +/// falling back to the repo-relative `test_data/` path. +/// +/// Returns `None` if the directory doesn't exist (CI builders without data). +pub(super) fn test_data_dir() -> Option { + let dir = std::env::var("FOXHUNT_TEST_DATA") + .unwrap_or_else(|_| { + // Workspace root is 4 levels up from smoke_tests/ + let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default(); + let workspace = std::path::Path::new(&manifest) + .parent() // crates/ + .and_then(|p| p.parent()) // repo root + .unwrap_or(std::path::Path::new("../..")); + workspace.join("test_data").to_string_lossy().into_owned() + }); + if std::path::Path::new(&dir).exists() { + Some(dir) + } else { + None + } +} diff --git a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs index 00a676dc0..9c84df2cc 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs @@ -7,25 +7,27 @@ //! cargo test -p ml --lib smoke_tests::performance -- --ignored --nocapture //! ``` -use super::helpers::{assert_finite, smoke_params, synthetic_data}; +use super::helpers::{assert_finite, smoke_params, test_data_dir}; use crate::trainers::dqn::DQNTrainer; use std::time::Instant; use tracing::info; -/// Measure end-to-end training throughput (3 epochs, 1000 samples). +/// Resolve test data dir — hard error if missing. +fn data_dir() -> String { + test_data_dir() + .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback") +} + +/// Measure end-to-end training throughput (3 epochs, real data). #[tokio::test] #[ignore] // Run manually: cargo test -p ml --lib smoke_tests::performance -- --ignored --nocapture async fn test_training_throughput_measurement() -> anyhow::Result<()> { let mut trainer = DQNTrainer::new(smoke_params())?; // auto-detect GPU - let data = synthetic_data(1000); - let val = synthetic_data(100); let start = Instant::now(); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; let elapsed = start.elapsed(); assert_finite(metrics.loss, "loss"); @@ -107,29 +109,18 @@ async fn test_per_sample_latency() -> anyhow::Result<()> { Ok(()) } -/// Single-epoch training on real data (if available). +/// Single-epoch training on real data. #[tokio::test] -#[ignore] // Run manually: FOXHUNT_TEST_DATA=./test_data/futures-baseline cargo test ... +#[ignore] // Run manually: cargo test -p ml --lib smoke_tests::performance -- --ignored --nocapture async fn test_real_data_single_epoch() -> anyhow::Result<()> { - let data_dir = std::env::var("FOXHUNT_TEST_DATA") - .unwrap_or_else(|_| "../../test_data/futures-baseline".to_owned()); - - let data_path = std::path::Path::new(&data_dir); - if !data_path.exists() { - info!(path = %data_path.display(), "SKIP: test data dir does not exist"); - return Ok(()); - } - let mut params = smoke_params(); params.epochs = 1; let mut trainer = DQNTrainer::new(params)?; let start = Instant::now(); - let metrics = trainer - .train(&data_dir, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; let elapsed = start.elapsed(); assert_finite(metrics.loss, "real_data_loss"); diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index d764245a6..e323c0fdd 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -2,20 +2,23 @@ //! //! Verify that short DQN training runs produce sane, finite metrics. //! All tests require CUDA — no CPU fallback. +//! All training tests use real `.dbn` data — no synthetic fallback. -use super::helpers::{assert_finite, cuda_device, smoke_trainer, smoke_trainer_with, smoke_params, synthetic_data}; +use super::helpers::{assert_finite, cuda_device, smoke_trainer, smoke_trainer_with, smoke_params, test_data_dir}; + +/// Resolve test data dir — hard error if missing. +fn data_dir() -> String { + test_data_dir() + .expect("FOXHUNT_TEST_DATA or test_data/ must exist — no synthetic fallback") +} /// Train a few epochs and verify the returned loss is finite and non-negative. #[tokio::test] async fn test_training_produces_finite_loss() -> anyhow::Result<()> { let mut trainer = smoke_trainer()?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; assert_finite(metrics.loss, "loss"); assert!(metrics.loss >= 0.0, "loss should be non-negative"); Ok(()) @@ -27,13 +30,9 @@ async fn test_gradient_norms_finite() -> anyhow::Result<()> { let mut params = smoke_params(); params.epochs = 3; let mut trainer = smoke_trainer_with(params)?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; let grad_norm = metrics .additional_metrics .get("avg_gradient_norm") @@ -51,13 +50,9 @@ async fn test_gradient_norms_finite() -> anyhow::Result<()> { #[tokio::test] async fn test_q_values_bounded() -> anyhow::Result<()> { let mut trainer = smoke_trainer()?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; let avg_q = metrics .additional_metrics .get("avg_q_value") @@ -78,13 +73,9 @@ async fn test_epsilon_decays() -> anyhow::Result<()> { params.epsilon_start = 1.0; params.epsilon_decay = 0.9; let mut trainer = smoke_trainer_with(params)?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; let final_eps = metrics .additional_metrics .get("final_epsilon") @@ -104,13 +95,9 @@ async fn test_no_nan_loss_short_training() -> anyhow::Result<()> { let mut params = smoke_params(); params.epochs = 2; let mut trainer = smoke_trainer_with(params)?; - let data = synthetic_data(100); - let val = synthetic_data(30); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; assert!( metrics.loss.is_finite(), "loss should be finite after 2-epoch training, got {}", @@ -203,13 +190,9 @@ async fn test_total_episodes_nonzero() -> anyhow::Result<()> { let mut params = smoke_params(); params.epochs = 3; let mut trainer = smoke_trainer_with(params)?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; assert!( metrics.epochs_trained > 0, "epochs_trained should be > 0, got {}", @@ -224,13 +207,9 @@ async fn test_training_with_dueling_produces_finite() -> anyhow::Result<()> { let mut params = smoke_params(); params.use_dueling = true; let mut trainer = smoke_trainer_with(params)?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; assert_finite(metrics.loss, "dueling_loss"); Ok(()) } @@ -241,13 +220,9 @@ async fn test_training_with_distributional_produces_finite() -> anyhow::Result<( let mut params = smoke_params(); params.use_distributional = true; let mut trainer = smoke_trainer_with(params)?; - let data = synthetic_data(200); - let val = synthetic_data(50); - let metrics = trainer - .train_with_preloaded_data(data, val, |_epoch, _bytes, _best| { - Ok("skip".to_owned()) - }) - .await?; + let metrics = trainer.train(&data_dir(), |_epoch, _bytes, _best| { + Ok("skip".to_owned()) + }).await?; assert_finite(metrics.loss, "distributional_loss"); Ok(()) } diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 69d8b4bdb..7831f6c37 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -266,6 +266,7 @@ impl DQNTrainer { iqn_num_quantiles: hyperparams.num_quantiles, // Controlled by hyperopt iqn_kappa: hyperparams.qr_kappa as f32, // Controlled by hyperopt (f64→f32) iqn_embedding_dim: 64, // Fixed (not in search space) + iqn_lambda: hyperparams.iqn_lambda as f32, // IQN dual-head loss weight (f64→f32) use_branching: hyperparams.use_branching, branch_hidden_dim: hyperparams.branch_hidden_dim, use_regime_conditioning: true, // Always enable per-regime IS weights for branching loss diff --git a/crates/ml/src/trainers/dqn/trainer/tests.rs b/crates/ml/src/trainers/dqn/trainer/tests.rs index a1a97446d..065a8b02d 100644 --- a/crates/ml/src/trainers/dqn/trainer/tests.rs +++ b/crates/ml/src/trainers/dqn/trainer/tests.rs @@ -556,10 +556,10 @@ fn test_c2_five_actions_produce_distinct_exposures() { #[test] fn test_c3_search_space_is_27d() { let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds(); - assert_eq!(bounds.len(), 30, "Search space must be 30D (C6: gradient_accumulation_steps added)"); + assert_eq!(bounds.len(), 31, "Search space must be 31D"); let names = crate::hyperopt::adapters::dqn::DQNParams::param_names(); - assert_eq!(names.len(), 30); + assert_eq!(names.len(), 31); assert!(names.contains(&"count_bonus_coefficient"), "count_bonus_coefficient must be in search space (C3)"); assert!(names.contains(&"sharpe_weight"), "sharpe_weight must be in search space (C4)"); assert!(names.contains(&"branch_hidden_dim"), "branch_hidden_dim must be in search space (L2)"); diff --git a/crates/ml/src/trainers/online_learning.rs b/crates/ml/src/trainers/online_learning.rs index ebce05bf9..bd6937b99 100644 --- a/crates/ml/src/trainers/online_learning.rs +++ b/crates/ml/src/trainers/online_learning.rs @@ -344,7 +344,9 @@ impl EWCRegularizer { .map_err(|e| MLError::ModelError(format!("EWC sum_all: {e}")))?; let lambda_tensor = Tensor::new(self.lambda as f32, &device) - .map_err(|e| MLError::ModelError(format!("EWC lambda tensor: {e}")))?; + .map_err(|e| MLError::ModelError(format!("EWC lambda tensor: {e}")))? + .to_dtype(total.dtype()) + .map_err(|e| MLError::ModelError(format!("EWC lambda dtype cast: {e}")))?; total .broadcast_mul(&lambda_tensor) @@ -653,7 +655,10 @@ mod tests { return; } }; - let val: f32 = penalty_tensor.to_scalar().unwrap_or(999.0); + let val: f32 = penalty_tensor + .to_dtype(candle_core::DType::F32) + .and_then(|t| t.to_scalar()) + .unwrap_or(999.0); assert!( val.abs() < 1e-5, "Expected ~0 penalty when params unchanged, got {val}" @@ -710,7 +715,10 @@ mod tests { return; } }; - let val: f32 = penalty_tensor.to_scalar().unwrap_or(0.0); + let val: f32 = penalty_tensor + .to_dtype(candle_core::DType::F32) + .and_then(|t| t.to_scalar()) + .unwrap_or(0.0); assert!( val > 1.0, "Expected nonzero penalty when params drifted, got {val}" diff --git a/crates/ml/src/trainers/tlob.rs b/crates/ml/src/trainers/tlob.rs index f30995cda..5ba6538fc 100644 --- a/crates/ml/src/trainers/tlob.rs +++ b/crates/ml/src/trainers/tlob.rs @@ -712,7 +712,7 @@ mod tests { let hyperparams = TLOBHyperparameters::default(); let temp_dir = std::env::temp_dir().join("tlob_test"); - let trainer = TLOBTrainer::new(hyperparams, &temp_dir, false); + let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true); assert!( trainer.is_ok(), "Failed to create TLOB trainer: {:?}", @@ -723,15 +723,14 @@ mod tests { #[tokio::test] async fn test_batch_size_validation() { let mut hyperparams = TLOBHyperparameters::default(); - hyperparams.batch_size = 64; // Exceeds GPU limit + hyperparams.batch_size = 64; let temp_dir = std::env::temp_dir().join("tlob_test_batch"); let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true); - // Should fall back to CPU assert!( - trainer.is_ok(), - "Should handle large batch size by using CPU" + trainer.is_err(), + "Batch size 64 exceeds GPU limit 32 — should be rejected" ); } @@ -740,7 +739,7 @@ mod tests { let hyperparams = TLOBHyperparameters::default(); let temp_dir = std::env::temp_dir().join("tlob_test_seq"); - let trainer = TLOBTrainer::new(hyperparams, &temp_dir, false).unwrap(); + let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true).unwrap(); let sequences = trainer.generate_dummy_sequences(10).unwrap(); assert_eq!(sequences.len(), 10); @@ -753,7 +752,7 @@ mod tests { let hyperparams = TLOBHyperparameters::default(); let temp_dir = std::env::temp_dir().join("tlob_test_batch_prep"); - let trainer = TLOBTrainer::new(hyperparams, &temp_dir, false).unwrap(); + let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true).unwrap(); let sequences = trainer.generate_dummy_sequences(4).unwrap(); let (input_tensor, target_tensor) = trainer.prepare_batch(&sequences).unwrap(); diff --git a/crates/ml/tests/dqn_action_collapse_fix_test.rs b/crates/ml/tests/dqn_action_collapse_fix_test.rs index 618d33145..5ca1aa38d 100644 --- a/crates/ml/tests/dqn_action_collapse_fix_test.rs +++ b/crates/ml/tests/dqn_action_collapse_fix_test.rs @@ -226,16 +226,16 @@ fn test_hyperopt_26d_includes_cql_alpha() { let bounds = DQNParams::continuous_bounds(); assert_eq!( bounds.len(), - 30, - "Search space should be 30D (C6: added gradient_accumulation_steps), got {}D", + 31, + "Search space should be 31D (C7: added iqn_lambda), got {}D", bounds.len() ); let names = DQNParams::param_names(); assert_eq!( names.len(), - 30, - "param_names should return 30 entries, got {}", + 31, + "param_names should return 31 entries, got {}", names.len() ); @@ -300,7 +300,7 @@ fn test_hyperopt_cql_alpha_round_trip() { params.cql_alpha = 0.25; let continuous = params.to_continuous(); - assert_eq!(continuous.len(), 30); + assert_eq!(continuous.len(), 31); let reconstructed = DQNParams::from_continuous(&continuous).expect("from_continuous should succeed"); diff --git a/crates/ml/tests/gpu_kernel_parity_test.rs b/crates/ml/tests/gpu_kernel_parity_test.rs index d435044a7..e123a898e 100644 --- a/crates/ml/tests/gpu_kernel_parity_test.rs +++ b/crates/ml/tests/gpu_kernel_parity_test.rs @@ -182,7 +182,6 @@ mod gpu_parity { None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); let total_bars = 2000; @@ -199,8 +198,6 @@ mod gpu_parity { episode_length: timesteps, epsilon: 0.1, gamma: 0.99, - use_noisy_nets: false, - use_distributional: false, ..Default::default() }; @@ -251,7 +248,6 @@ mod gpu_parity { None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); let total_bars = 2000; @@ -268,8 +264,6 @@ mod gpu_parity { episode_length: timesteps, epsilon: 0.1, gamma: 0.99, - use_noisy_nets: false, - use_distributional: true, num_atoms: 51, v_min: -25.0, v_max: 25.0, @@ -334,7 +328,6 @@ mod gpu_parity { None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); // Market data with test state as every bar @@ -368,8 +361,6 @@ mod gpu_parity { episode_length: 1, epsilon: 0.0, gamma: 0.99, - use_noisy_nets: false, - use_distributional: false, count_bonus_coefficient: 0.0, ..Default::default() }; @@ -448,7 +439,6 @@ mod gpu_parity { let mut collector_clean = GpuExperienceCollector::new( stream.clone(), network.vars(), target.vars(), None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); let config_clean = ExperienceCollectorConfig { @@ -458,8 +448,6 @@ mod gpu_parity { episode_length: timesteps, epsilon: 0.0, gamma: 0.99, - use_noisy_nets: false, - use_distributional: false, count_bonus_coefficient: 0.0, ..Default::default() }; @@ -472,7 +460,6 @@ mod gpu_parity { let mut collector_noisy = GpuExperienceCollector::new( stream.clone(), network.vars(), target.vars(), None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); let config_noisy = ExperienceCollectorConfig { @@ -482,9 +469,7 @@ mod gpu_parity { episode_length: timesteps, epsilon: 0.0, gamma: 0.99, - use_noisy_nets: true, noisy_sigma_init: 0.5, - use_distributional: false, count_bonus_coefficient: 0.0, ..Default::default() }; @@ -536,7 +521,6 @@ mod gpu_parity { let mut collector = GpuExperienceCollector::new( stream.clone(), network.vars(), target.vars(), None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); collector.sync_online_weights(network.vars()).expect("Online sync failed"); @@ -616,8 +600,6 @@ mod gpu_parity { episode_length: timesteps, epsilon: 0.0, gamma: 0.99, - use_noisy_nets: false, - use_distributional: false, count_bonus_coefficient: 0.0, ..Default::default() }; @@ -625,7 +607,6 @@ mod gpu_parity { let mut collector = GpuExperienceCollector::new( stream.clone(), network.vars(), target.vars(), None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); for run in 0..5 { @@ -664,7 +645,6 @@ mod gpu_parity { let mut std_collector = GpuExperienceCollector::new( stream.clone(), std_net.vars(), std_target.vars(), None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); let std_batch = std_collector @@ -675,7 +655,6 @@ mod gpu_parity { timesteps_per_episode: timesteps, total_bars: total_bars as i32, episode_length: timesteps, - use_distributional: false, ..Default::default() }, ) @@ -684,7 +663,6 @@ mod gpu_parity { let mut dist_collector = GpuExperienceCollector::new( stream.clone(), dist_net.vars(), dist_target.vars(), None, 100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50, - false, ).unwrap(); let dist_batch = dist_collector @@ -695,8 +673,7 @@ mod gpu_parity { timesteps_per_episode: timesteps, total_bars: total_bars as i32, episode_length: timesteps, - use_distributional: true, - num_atoms: 51, + num_atoms: 51, v_min: -25.0, v_max: 25.0, ..Default::default() diff --git a/scripts/gpu-hotpath-guard.sh b/scripts/gpu-hotpath-guard.sh index dfb8c281e..7d697210b 100755 --- a/scripts/gpu-hotpath-guard.sh +++ b/scripts/gpu-hotpath-guard.sh @@ -82,12 +82,30 @@ LEAK_PATTERNS=( # They add maintenance burden and hide the real code path. '#\[cfg\(not\(feature\s*=\s*"cuda"\)\)\]' + # ─── CPU-side batch data mutation (GPU data wrongly processed on CPU) ─── + # Functions taking mutable Vec batch arrays indicate CPU-side processing + # that should be a GPU kernel or Candle tensor ops. + # Caught: CPU HER relabeling, CPU normalization, CPU reward shaping, etc. + 'mut states: Vec' + 'mut next_states: Vec' + 'mut rewards: Vec' + + # ─── Candle tensor wrapping in per-step hot path ─── + # Creating Candle Tensors from CudaSlice in the training loop adds + # Rust-side allocation + dispatch overhead. Use raw CUDA kernels instead. + # Tensor::new() creates a CPU scalar → GPU copy (HtoD for 4 bytes). + # CudaStorage::wrap_cuda_slice takes ownership — blocks reuse patterns. + 'CudaStorage::wrap_cuda_slice' + # NOTE: Tensor::from_vec/from_slice are NOT flagged — they are CPU→GPU # uploads (correct direction). Inner-loop abuse is caught by code review. # # NOTE: memcpy_dtoh/htod and .synchronize() are not checked here — # they are CUDA implementation primitives in cuda_pipeline/. # Audit those separately with: scripts/cuda-perf-audit.sh + # + # NOTE: storage_and_layout() is NOT flagged — it's a zero-cost pointer + # unwrap for extracting CudaSlice refs from existing Tensors. ) # ── Excluded paths ──────────────────────────────────────────────────── @@ -106,6 +124,9 @@ EXCLUDE_PATHS=( "ml-supervised/src/tft/hft_optimizations.rs" # evaluate_cpu() is slow fallback before precompute_gpu_grid(); GPU path exists "ml-supervised/src/kan/spline.rs" + # SearchSorted/CumsumOp CustomOp2 impls must return CudaStorage (Candle API boundary) + # Runs once per PER sampling batch, not inside per-step training kernel loop + "ml-dqn/src/gpu_replay_buffer.rs" ) # ── Test module exclusion ─────────────────────────────────────────────