From 5191fa022aa3d57588892efc651e9977c3bfd5db Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 8 Mar 2026 23:22:39 +0100 Subject: [PATCH] perf(ml): replace remaining CPU tensor allocations with GPU-native ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eliminate all remaining from_vec(vec![const; N]) patterns and CPU tensor roundtrips in the training hot path: - gamma discount: from_vec → Tensor::full (DQN + distributional C51) - C51 atoms: from_vec(computed Vec) → arange+affine, 4 instances across forward pass and C51 loss paths (online/standard/double-DQN) - IQN cosine embedding indices: from_vec → arange+reshape - IQN uniform quantiles: from_vec → arange+affine (τ_i = (i+0.5)/N) - PPO clip epsilon: from_vec+ones+sub/add+tensor_clamp → scalar clamp (eliminates 6 intermediate tensor allocations, 2 call sites) - Dead neuron detection: flatten+to_vec1+CPU loop → abs().le().sum_all() (single scalar per parameter tensor instead of full weight transfer) 548 tests pass (350 ml-dqn + 198 ml-ppo), ml crate compiles clean. Co-Authored-By: Claude Opus 4.6 --- crates/ml-dqn/src/distributional.rs | 7 +-- crates/ml-dqn/src/dqn.rs | 70 +++++++++++------------- crates/ml-dqn/src/quantile_regression.rs | 23 +++----- crates/ml-ppo/src/ppo.rs | 56 ++++--------------- 4 files changed, 54 insertions(+), 102 deletions(-) diff --git a/crates/ml-dqn/src/distributional.rs b/crates/ml-dqn/src/distributional.rs index a3e3a2b27..557490e7b 100644 --- a/crates/ml-dqn/src/distributional.rs +++ b/crates/ml-dqn/src/distributional.rs @@ -287,11 +287,8 @@ impl CategoricalDistribution { // Compute T_z = r + γ × z_j × (1 - done) // Shape: [batch, num_atoms] let rewards_broadcast = rewards_f32.unsqueeze(1)?.broadcast_as((batch_size, num_atoms))?; - let gamma_tensor = Tensor::from_vec( - vec![gamma; batch_size * num_atoms], - (batch_size, num_atoms), - rewards_f32.device(), - )?; + // GPU-native: Tensor::full creates constant tensor on device (no CPU Vec) + let gamma_tensor = Tensor::full(gamma, (batch_size, num_atoms), rewards_f32.device())?; // (1 - done) mask let dones_broadcast = dones_f32.unsqueeze(1)?.broadcast_as((batch_size, num_atoms))?; diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 3dfb5644c..4159a8cbd 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -1367,17 +1367,15 @@ impl DQN { let v_max = self.config.v_max; let delta_z = (v_max - v_min) / (num_atoms as f32 - 1.0); - // Create atom values tensor: [v_min, v_min + delta_z, ..., v_max] - let atoms: Vec = (0..num_atoms) - .map(|i| v_min + i as f32 * delta_z) - .collect(); - let batch_size = state.dim(0)?; let num_actions = self.config.num_actions; - // Broadcast atoms to [batch, num_actions, num_atoms], matching training dtype + // GPU-native: arange + affine creates atoms on device (no CPU Vec) + // atoms[i] = v_min + i * delta_z let dtype = training_dtype(&self.device); - let atoms_tensor = Tensor::from_vec(atoms, num_atoms, &self.device)? + let atoms_tensor = Tensor::arange(0u32, num_atoms as u32, &self.device)? + .to_dtype(DType::F32)? + .affine(delta_z as f64, v_min as f64)? .to_dtype(dtype)? .unsqueeze(0)? .unsqueeze(0)? @@ -2304,16 +2302,13 @@ impl DQN { let v_max = self.config.v_max; let delta_z = (v_max - v_min) / (num_atoms as f32 - 1.0); - // Create atom values tensor - let atoms: Vec = (0..num_atoms) - .map(|i| v_min + i as f32 * delta_z) - .collect(); - let num_actions = self.config.num_actions; - // Broadcast atoms to [batch, num_actions, num_atoms] + // GPU-native: arange + affine creates atoms on device (no CPU Vec) // BF16 FIX: atoms stay F32 to match z_probs (network output is F32) - let atoms_tensor = Tensor::from_vec(atoms, num_atoms, device)? + let atoms_tensor = Tensor::arange(0u32, num_atoms as u32, device)? + .to_dtype(DType::F32)? + .affine(delta_z as f64, v_min as f64)? .unsqueeze(0)? .unsqueeze(0)? .broadcast_as((batch_size, num_actions, num_atoms))?; @@ -2365,12 +2360,11 @@ impl DQN { // R_t + gamma*R_{t+1} + ... + gamma^{n-1}*R_{t+n-1} in the reward. // The bootstrap discount must therefore be gamma^n (not gamma^1). let gamma_n = self.config.gamma.powi(self.config.n_steps as i32); - let gamma_tensor = - Tensor::from_vec(vec![gamma_n; batch_size], batch_size, device).map_err( - |e| MLError::TrainingError(format!("Failed to create gamma tensor: {}", e)), - )?.to_dtype(dtype).map_err( - |e| MLError::TrainingError(format!("Failed to cast gamma tensor: {}", e)), - )?; + // GPU-native: Tensor::full creates constant tensor on device (no CPU Vec) + let gamma_tensor = Tensor::full(gamma_n, &[batch_size], device) + .map_err(|e| MLError::TrainingError(format!("Failed to create gamma tensor: {}", e)))? + .to_dtype(dtype) + .map_err(|e| MLError::TrainingError(format!("Failed to cast gamma tensor: {}", e)))?; let not_done = (Tensor::ones(&[batch_size], dtype, device)? - &dones_tensor)?; let gamma_next = (&gamma_tensor * &next_state_values) @@ -2550,10 +2544,10 @@ impl DQN { let v_min = self.config.v_min; let v_max = self.config.v_max; let delta_z = (v_max - v_min) / (num_atoms as f32 - 1.0); - let atoms: Vec = (0..num_atoms) - .map(|i| v_min + i as f32 * delta_z) - .collect(); - let atoms_tensor = Tensor::from_vec(atoms, num_atoms, device)? + // GPU-native: arange + affine (no CPU Vec) + let atoms_tensor = Tensor::arange(0u32, num_atoms as u32, device)? + .to_dtype(DType::F32)? + .affine(delta_z as f64, v_min as f64)? .to_dtype(dtype)? .unsqueeze(0)? .unsqueeze(0)? @@ -2570,10 +2564,10 @@ impl DQN { let v_min = self.config.v_min; let v_max = self.config.v_max; let delta_z = (v_max - v_min) / (num_atoms as f32 - 1.0); - let atoms: Vec = (0..num_atoms) - .map(|i| v_min + i as f32 * delta_z) - .collect(); - let atoms_tensor = Tensor::from_vec(atoms, num_atoms, device)? + // GPU-native: arange + affine (no CPU Vec) + let atoms_tensor = Tensor::arange(0u32, num_atoms as u32, device)? + .to_dtype(DType::F32)? + .affine(delta_z as f64, v_min as f64)? .to_dtype(dtype)? .unsqueeze(0)? .unsqueeze(0)? @@ -3320,17 +3314,19 @@ impl DQN { operation: format!("lock VarMap for dead neuron detection: {}", e), })?; - // Check each layer's weights + // GPU-native: count near-zero weights without transferring all params to CPU. + // abs().le(1e-6) creates boolean mask on device, sum_all counts matches (1 scalar readback). for (_name, var) in vars_data.iter() { let tensor = var.as_tensor(); - let values = tensor.flatten_all()?.to_dtype(DType::F32)?.to_vec1::()?; - - for &val in values.iter() { - total_count += 1; - if val.abs() < 1e-6 { - dead_count += 1; - } - } + let f32_tensor = tensor.to_dtype(DType::F32)?; + let n = f32_tensor.elem_count(); + total_count += n; + let near_zero = f32_tensor.abs()? + .le(1e-6_f32)? + .to_dtype(DType::F32)? + .sum_all()? + .to_vec0::()? as usize; + dead_count += near_zero; } Ok((dead_count as f32 / total_count as f32) * 100.0) diff --git a/crates/ml-dqn/src/quantile_regression.rs b/crates/ml-dqn/src/quantile_regression.rs index 1e1762b0c..17d19b32d 100644 --- a/crates/ml-dqn/src/quantile_regression.rs +++ b/crates/ml-dqn/src/quantile_regression.rs @@ -207,13 +207,10 @@ impl QuantileNetwork { let num_quantiles = taus.dim(1)?; let embed_dim = self.config.quantile_embedding_dim; - // Create indices [1, 2, 3, ..., embedding_dim] in same dtype as taus - let indices: Vec = (1..=embed_dim).map(|i| i as f32).collect(); - let indices_tensor = Tensor::from_vec( - indices, - (1, 1, embed_dim), - device, - )?.to_dtype(target_dtype)?; + // GPU-native: arange creates indices on device (no CPU Vec) + let indices_tensor = Tensor::arange(1u32, (embed_dim + 1) as u32, device)? + .to_dtype(target_dtype)? + .reshape((1, 1, embed_dim))?; // Broadcast taus to [batch, num_quantiles, 1] let taus_broadcast = taus.unsqueeze(2)?; @@ -244,13 +241,11 @@ impl QuantileNetwork { pub fn sample_uniform_quantiles(&self, batch_size: usize, device: &Device) -> CandleResult { let num_quantiles = self.config.num_quantiles; - // τ_i = (i + 0.5) / N - let taus: Vec = (0..num_quantiles) - .map(|i| (i as f32 + 0.5) / num_quantiles as f32) - .collect(); - - // Create [num_quantiles] tensor - let taus_tensor = Tensor::from_vec(taus, (num_quantiles,), device)?; + // GPU-native: arange + affine creates τ_i = (i + 0.5) / N on device + // affine(1/N, 0.5/N) maps i → (i + 0.5) / N + let taus_tensor = Tensor::arange(0u32, num_quantiles as u32, device)? + .to_dtype(DType::F32)? + .affine(1.0 / num_quantiles as f64, 0.5 / num_quantiles as f64)?; // Broadcast to [batch, num_quantiles] taus_tensor.unsqueeze(0)?.broadcast_as((batch_size, num_quantiles)) diff --git a/crates/ml-ppo/src/ppo.rs b/crates/ml-ppo/src/ppo.rs index 8841b2be2..fe12569c5 100644 --- a/crates/ml-ppo/src/ppo.rs +++ b/crates/ml-ppo/src/ppo.rs @@ -1569,27 +1569,11 @@ impl PPO { let log_ratio = (&seq_new_log_probs - &seq_old_log_probs)?; let ratio = log_ratio.exp()?; - let clip_epsilon_tensor = Tensor::from_vec( - vec![self.config.clip_epsilon; seq_len], - (seq_len,), - device, - )?; - - // BF16 FIX: Use F32 for loss-path arithmetic (same as MLP path) - let one_tensor = Tensor::ones((seq_len,), DType::F32, device)?; - let clip_min = (&one_tensor - &clip_epsilon_tensor)?; - let clip_max = if let Some(eps_high) = self.config.clip_epsilon_high { - let clip_high_tensor = Tensor::from_vec( - vec![eps_high; seq_len], - (seq_len,), - device, - ).map_err(|e| MLError::TrainingError(format!("Failed to create clip_high tensor: {}", e)))?; - (&one_tensor + &clip_high_tensor)? - } else { - (&one_tensor + &clip_epsilon_tensor)? - }; - - let clipped_ratio = ratio.clamp(&clip_min, &clip_max)?; + // GPU-native: scalar clamp eliminates 3 tensor allocations + let clip_min_val = 1.0 - self.config.clip_epsilon as f64; + let clip_max_val = 1.0 + self.config.clip_epsilon_high + .unwrap_or(self.config.clip_epsilon) as f64; + let clipped_ratio = ratio.clamp(clip_min_val, clip_max_val)?; let surr1 = (&ratio * &seq_advantages)?; let surr2 = (&clipped_ratio * &seq_advantages)?; @@ -1776,31 +1760,11 @@ impl PPO { let log_ratio = (&new_log_probs - &batch.log_probs)?; let ratio = log_ratio.exp()?; - // Clipped surrogate objective - let clip_epsilon_tensor = Tensor::from_vec( - vec![self.config.clip_epsilon; batch.advantages.dims()[0]], - batch.advantages.dims(), - self.actor.device(), - ) - .map_err(|e| MLError::TrainingError(format!("Failed to create clip tensor: {}", e)))?; - - // BF16 FIX: Use F32 for loss-path arithmetic — network forward outputs are F32 - // (for autograd), so all loss constants must match to avoid Candle dtype mismatch. - let one_tensor = Tensor::ones(batch.advantages.dims(), DType::F32, self.actor.device())?; - let clip_min = (&one_tensor - &clip_epsilon_tensor)?; - let clip_max = if let Some(eps_high) = self.config.clip_epsilon_high { - let clip_high_tensor = Tensor::from_vec( - vec![eps_high; batch.advantages.dims()[0]], - batch.advantages.dims(), - self.actor.device(), - ).map_err(|e| MLError::TrainingError(format!("Failed to create clip_high tensor: {}", e)))?; - (&one_tensor + &clip_high_tensor)? - } else { - (&one_tensor + &clip_epsilon_tensor)? - }; - - // Clamp ratio to [1-ε, 1+ε] (or [1-ε, 1+ε_high] with clip-higher) - let clipped_ratio = ratio.clamp(&clip_min, &clip_max)?; + // GPU-native: scalar clamp eliminates 3 tensor allocations + let clip_min_val = 1.0 - self.config.clip_epsilon as f64; + let clip_max_val = 1.0 + self.config.clip_epsilon_high + .unwrap_or(self.config.clip_epsilon) as f64; + let clipped_ratio = ratio.clamp(clip_min_val, clip_max_val)?; // PPO objective: min(ratio * advantage, clipped_ratio * advantage) let surr1 = (&ratio * &batch.advantages)?;