gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(ml): replace remaining CPU tensor allocations with GPU-native ops

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-08 23:22:39 +01:00
parent 801781a472
commit 5191fa022a
4 changed files with 54 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml-dqn/src/distributional.rs
View File

@@ -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))?;

70
crates/ml-dqn/src/dqn.rs
View File

@@ -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<f32> = (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<f32> = (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<f32> = (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<f32> = (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::<f32>()?;
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::<f32>()? as usize;
dead_count += near_zero;
}
Ok((dead_count as f32 / total_count as f32) * 100.0)

23
crates/ml-dqn/src/quantile_regression.rs
View File

@@ -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<f32> = (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<Tensor> {
let num_quantiles = self.config.num_quantiles;
// τ_i = (i + 0.5) / N
let taus: Vec<f32> = (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))

56
crates/ml-ppo/src/ppo.rs
View File

@@ -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)?;