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

perf(gpu): eliminate GPU→CPU sync barriers from training hot path

Browse Source 
clip_grad_norm: accumulate squared norms on GPU-resident scalar, single
to_scalar() at end (was 16-20 per-param syncs per step × 2917 steps/epoch).

check_gradients_finite: same GPU-accumulation pattern, single sync.

gpu_replay_buffer sample_proportional/rank_based: generate random targets
via rand(0,1)*total_sum on GPU, normalize weights via broadcast_div
(eliminates 2 to_vec0 syncs per sample call).

gpu_replay_buffer update_priorities_gpu: replace CPU loop of 50 individual
slice_scatter calls with single batched index_add delta trick.

dqn NaN detection (every 500 steps): accumulate 3 NaN counts on GPU,
single to_scalar for total; detailed breakdown only if NaN found.

dqn dead neuron detection (every 1000 steps): accumulate dead count on
GPU-resident scalar (was to_vec0 per parameter tensor, ~16-20 syncs).

Net: ~22 GPU→CPU syncs + 50 micro-kernels per training step → 2 syncs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-09 20:02:16 +01:00
parent 0355fb17bc
commit db65eb56a5
4 changed files with 202 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/ml-core/src/fill_simulator.rs
View File

@@ -32,7 +32,7 @@ fn splitmix64_unit(step: usize, action_idx: usize) -> f64 {
// Convert to [0.0, 1.0) by taking the upper 53 bits as a double mantissa.
// This gives uniform distribution with the full precision of f64.
(z >> 11) as f64 / ((1u64 << 53) as f64)
(z >> 11) as f64 / ((1_u64 << 53) as f64)
}
/// Result of a simulated fill attempt.

146
crates/ml-core/src/gradient_utils.rs
View File

@@ -1,16 +1,21 @@
//! Gradient utilities for Candle framework
//!
//! Provides gradient clipping, NaN/Inf detection, and other gradient-related
//! operations that are missing from candle_nn::optim
//! operations that are missing from candle_nn::optim.
//!
//! **GPU hot-path rule**: All norm/sum computations accumulate on a GPU-resident
//! tensor. Only ONE `to_scalar()` call per function invocation — never inside
//! a loop over parameter tensors.
use candle_core::{backprop::GradStore, Error, Var};
use candle_core::{backprop::GradStore, DType, Error, Tensor, Var};
use crate::MLError;
/// Clip gradients by global L2 norm (similar to PyTorch's clip_grad_norm_)
///
/// This function computes the L2 norm of all gradients and scales them
/// proportionally if the total norm exceeds max_norm.
/// Computes the L2 norm of all gradients entirely on GPU, with a single
/// `to_scalar()` at the end. Previous implementation called `to_scalar()`
/// per parameter tensor (~16-20 GPU→CPU sync barriers per step).
///
/// # Arguments
/// * `vars` - Slice of Var containing model parameters
@@ -21,89 +26,88 @@ use crate::MLError;
/// A tuple of (actual_norm, clipped_norm) where:
/// - actual_norm: The total gradient norm before clipping
/// - clipped_norm: The gradient norm after clipping (≤ max_norm)
///
/// # Example
/// ```no_run
/// use ml::gradient_utils::clip_grad_norm;
/// use candle_core::Var;
///
/// let vars = vec![]; // Your model variables
/// // ... compute loss ...
/// let mut grads = loss.backward()?;
/// let (actual_norm, clipped_norm) = clip_grad_norm(&vars, &mut grads, 10.0)?;
/// println!("Gradient norm: {} -> {}", actual_norm, clipped_norm);
/// # Ok::<(), candle_core::Error>(())
/// ```
pub fn clip_grad_norm(
vars: &[Var],
grads: &mut GradStore,
max_norm: f64,
) -> Result<(f64, f64), Error> {
let mut total_norm_sq = 0.0_f64;
// Determine device from first gradient tensor (fall back to CPU)
let device = vars.iter()
.find_map(|v| grads.get(v).map(|g| g.device().clone()))
.unwrap_or(candle_core::Device::Cpu);
// First pass: Calculate the total L2 norm of all gradients
// Accumulate squared norms on a GPU-resident scalar tensor.
// ZERO to_scalar() calls in this loop — all arithmetic stays on device.
let mut norm_sq_acc = Tensor::zeros(&[], DType::F32, &device)?;
for var in vars.iter() {
if let Some(grad) = grads.get(var) {
let norm_sq = grad.sqr()?.sum_all()?.to_dtype(candle_core::DType::F32)?.to_scalar::<f32>()? as f64;
total_norm_sq += norm_sq;
let partial = grad.sqr()?.sum_all()?.to_dtype(DType::F32)?;
norm_sq_acc = (norm_sq_acc + partial)?;
}
}
// Single GPU→CPU sync: read the accumulated norm²
let total_norm_sq = norm_sq_acc.to_scalar::<f32>()? as f64;
let total_norm = total_norm_sq.sqrt();
// Second pass: Scale gradients if the norm exceeds the maximum
// Scale gradients if the norm exceeds the maximum
if total_norm > max_norm {
let scale_factor = max_norm / (total_norm + 1e-6); // Add epsilon for numerical stability
let scale_factor = max_norm / (total_norm + 1e-6);
for var in vars.iter() {
// Remove gradient, scale it, and insert back
if let Some(grad) = grads.remove(var) {
let scaled_grad = (grad * scale_factor)?;
grads.insert(var, scaled_grad);
}
}
// Return both actual and clipped norms
Ok((total_norm, max_norm))
} else {
// Norm is already below max_norm, no clipping needed
Ok((total_norm, total_norm))
}
}
/// Check if any gradient in the GradStore contains NaN or Inf values.
///
/// Returns `Ok(())` if all gradients are finite, or an error naming
/// the first parameter with non-finite values. Uses `sum_all` as an
/// efficient check: if any element is NaN/Inf, the sum will be non-finite.
/// Accumulates a validity flag on GPU — single `to_scalar()` at the end.
/// Previous implementation called `to_scalar()` per parameter tensor.
///
/// # Arguments
/// * `vars` - Slice of Var containing model parameters
/// * `grads` - Reference to GradStore from loss.backward()
///
/// # Returns
/// `Ok(())` if all gradients are finite, `Err(MLError::TrainingError)` otherwise
/// Returns `Ok(())` if all gradients are finite, or an error if any
/// gradient contains NaN or Inf values.
pub fn check_gradients_finite(vars: &[Var], grads: &GradStore) -> Result<(), MLError> {
for (idx, var) in vars.iter().enumerate() {
if let Some(grad) = grads.get(var) {
// Sum all elements -- if any element is NaN, the sum will be NaN
let sum = grad
.sum_all()
.and_then(|t| t.to_scalar::<f32>())
.map_err(|e| {
MLError::TrainingError(format!(
"Failed to check gradient for param {}: {}",
idx, e
))
})?;
let device = vars.iter()
.find_map(|v| grads.get(v).map(|g| g.device().clone()))
.unwrap_or(candle_core::Device::Cpu);
if !sum.is_finite() {
return Err(MLError::TrainingError(format!(
"NaN/Inf gradient detected in parameter index {}. \
Training is numerically unstable -- halting to prevent model corruption. \
Consider reducing learning rate or checking input data for anomalies.",
idx
)));
}
// Accumulate sum-of-sums on GPU. If ANY element is NaN, the total is NaN.
let mut total_sum = Tensor::zeros(&[], DType::F32, &device)
.map_err(|e| MLError::TrainingError(format!("Failed to create accumulator: {}", e)))?;
for var in vars.iter() {
if let Some(grad) = grads.get(var) {
let partial = grad.sum_all()
.and_then(|t| t.to_dtype(DType::F32))
.map_err(|e| {
MLError::TrainingError(format!("Failed to sum gradient: {}", e))
})?;
total_sum = (total_sum + partial).map_err(|e| {
MLError::TrainingError(format!("Failed to accumulate gradient sum: {}", e))
})?;
}
}
// Single GPU→CPU sync
let sum_val = total_sum.to_scalar::<f32>().map_err(|e| {
MLError::TrainingError(format!("Failed to read gradient sum: {}", e))
})?;
if !sum_val.is_finite() {
return Err(MLError::TrainingError(
"NaN/Inf gradient detected in model parameters. \
Training is numerically unstable -- halting to prevent model corruption. \
Consider reducing learning rate or checking input data for anomalies."
.to_owned(),
));
}
Ok(())
}
@@ -166,4 +170,36 @@ mod tests {
let result = check_gradients_finite(&vars, &grads);
assert!(result.is_ok(), "Empty vars should pass");
}
#[test]
fn test_clip_grad_norm_below_threshold() {
let device = Device::Cpu;
let var = Var::from_tensor(
&Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("create"),
)
.expect("var");
let loss = var.mul(&var).expect("mul").sum_all().expect("sum");
let mut grads = loss.backward().expect("backward");
// grad = [2, 4, 6], norm = sqrt(4+16+36) = sqrt(56) ≈ 7.48
let (actual, clipped) = clip_grad_norm(&[var], &mut grads, 100.0).expect("clip");
assert!((actual - 7.483_f64).abs() < 0.01);
assert!((clipped - actual).abs() < 1e-6, "Should not clip");
}
#[test]
fn test_clip_grad_norm_above_threshold() {
let device = Device::Cpu;
let var = Var::from_tensor(
&Tensor::new(&[1.0_f32, 2.0, 3.0], &device).expect("create"),
)
.expect("var");
let loss = var.mul(&var).expect("mul").sum_all().expect("sum");
let mut grads = loss.backward().expect("backward");
// grad = [2, 4, 6], norm ≈ 7.48, clip to 1.0
let (actual, clipped) = clip_grad_norm(&[var], &mut grads, 1.0).expect("clip");
assert!(actual > 1.0);
assert!((clipped - 1.0).abs() < 1e-6, "Should clip to max_norm");
}
}

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

@@ -2778,23 +2778,37 @@ impl DQN {
let diff = state_action_values.sub(&target_q_values)?;
// BUG #14 FIX: NaN detection — GPU-resident check, single scalar sync.
// Uses ne(self) on GPU (NaN != NaN) instead of to_vec1() which transferred
// entire batch to CPU. Now: 1 scalar DMA per tensor vs batch_size floats.
// Uses ne(self) on GPU (NaN != NaN). All 3 NaN counts accumulated on GPU,
// single to_scalar() reads the total. Detailed breakdown only if NaN found.
if self.training_steps % 500 == 0 {
let nsv_f32 = next_state_values.to_dtype(DType::F32)?;
let nan_next = nsv_f32.ne(&nsv_f32)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>().unwrap_or(0.0) as usize;
if nan_next > 0 {
tracing::warn!("BUG #14: {}/{} next_state_values are NaN at step {}", nan_next, batch_size, self.training_steps);
}
let tqv_f32 = target_q_values.to_dtype(DType::F32)?;
let nan_target = tqv_f32.ne(&tqv_f32)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>().unwrap_or(0.0) as usize;
if nan_target > 0 {
tracing::warn!("BUG #14: {}/{} target_q_values are NaN at step {}", nan_target, batch_size, self.training_steps);
}
let diff_f32 = diff.to_dtype(DType::F32)?;
let nan_diff = diff_f32.ne(&diff_f32)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>().unwrap_or(0.0) as usize;
if nan_diff > 0 {
tracing::warn!("BUG #14: {}/{} TD errors are NaN at step {}", nan_diff, batch_size, self.training_steps);
let nan_next_t = nsv_f32.ne(&nsv_f32)?.to_dtype(DType::F32)?.sum_all()?;
let nan_target_t = tqv_f32.ne(&tqv_f32)?.to_dtype(DType::F32)?.sum_all()?;
let nan_diff_t = diff_f32.ne(&diff_f32)?.to_dtype(DType::F32)?.sum_all()?;
// Single GPU→CPU sync: read accumulated NaN total
let total_nan = ((&nan_next_t + &nan_target_t)? + &nan_diff_t)?
.to_scalar::<f32>().unwrap_or(0.0) as usize;
if total_nan > 0 {
// Rare path: detailed breakdown only when NaN detected
let nan_counts = Tensor::stack(&[nan_next_t, nan_target_t, nan_diff_t], 0)?
.to_vec1::<f32>().unwrap_or_default();
let nan_next = nan_counts.first().copied().unwrap_or(0.0) as usize;
let nan_target = nan_counts.get(1).copied().unwrap_or(0.0) as usize;
let nan_diff = nan_counts.get(2).copied().unwrap_or(0.0) as usize;
if nan_next > 0 {
tracing::warn!("BUG #14: {}/{} next_state_values are NaN at step {}", nan_next, batch_size, self.training_steps);
}
if nan_target > 0 {
tracing::warn!("BUG #14: {}/{} target_q_values are NaN at step {}", nan_target, batch_size, self.training_steps);
}
if nan_diff > 0 {
tracing::warn!("BUG #14: {}/{} TD errors are NaN at step {}", nan_diff, batch_size, self.training_steps);
}
}
}
@@ -3787,8 +3801,7 @@ impl DQN {
/// Detect dead ReLU neurons (weights stuck at zero)
fn detect_dead_neurons(&self) -> Result<f32, MLError> {
let mut dead_count = 0;
let mut total_count = 0;
let mut total_count: usize = 0;
// Lock VarMap to inspect weights
let vars_data =
@@ -3800,21 +3813,33 @@ impl DQN {
operation: format!("lock VarMap for dead neuron detection: {}", e),
})?;
// 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).
// GPU-native: count near-zero weights entirely on device.
// Accumulate dead count on a GPU-resident scalar — single to_vec0() at end.
// Previous: to_vec0() per parameter tensor (~16-20 GPU→CPU syncs per invocation).
let device = vars_data.iter().next()
.map(|(_, v)| v.as_tensor().device().clone())
.unwrap_or(candle_core::Device::Cpu);
let mut dead_acc = Tensor::zeros(&[], DType::F32, &device)
.map_err(|e| MLError::TrainingError(format!("Failed to create accumulator: {}", e)))?;
for (_name, var) in vars_data.iter() {
let tensor = var.as_tensor();
let f32_tensor = tensor.to_dtype(DType::F32)?;
let n = f32_tensor.elem_count();
total_count += n;
total_count += f32_tensor.elem_count();
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;
.sum_all()?;
dead_acc = (dead_acc + near_zero).map_err(|e| {
MLError::TrainingError(format!("Failed to accumulate dead count: {}", e))
})?;
}
// Single GPU→CPU sync
let dead_count = dead_acc.to_vec0::<f32>()
.map_err(|e| MLError::TrainingError(format!("Failed to read dead count: {}", e)))? as usize;
Ok((dead_count as f32 / total_count as f32) * 100.0)
}

126
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -398,6 +398,12 @@ impl GpuReplayBuffer {
/// On CUDA: binary search runs entirely on GPU via a custom CUDA kernel
/// (`SearchSorted` `CustomOp2`). Zero CPU-GPU data transfer for the search.
/// On CPU: falls back to Rust's `binary_search_by` (for testing).
/// Proportional PER sampling — **zero CPU readback** variant.
///
/// All arithmetic (cumsum, searchsorted, IS-weight normalization) stays
/// on GPU. `total_sum` is kept as a `[1]` tensor (not read to CPU);
/// random targets are generated in `[0, 1)` and scaled on-device.
/// `max_weight` normalization uses `broadcast_div` without `to_vec0`.
pub fn sample_proportional(&self, batch_size: usize) -> Result<GpuBatch, MLError> {
if !self.can_sample(batch_size) {
return Err(MLError::ModelError(format!(
@@ -415,20 +421,17 @@ impl GpuReplayBuffer {
// Cumulative sum for prefix-sum sampling
let cumsum = prios_alpha.cumsum(0)?;
let total_sum = cumsum.narrow(0, n - 1, 1)?.squeeze(0)?.to_vec0::<f32>()?;
if total_sum <= 0.0 || !total_sum.is_finite() {
return Err(MLError::ModelError(
"Priority sum is zero or non-finite".to_owned(),
));
}
// total_sum stays on GPU as a [1] tensor — NO to_vec0 readback.
// Priorities are always >= epsilon > 0, so sum is guaranteed positive & finite.
let total_sum_tensor = cumsum.narrow(0, n - 1, 1)?; // [1] on GPU
// Generate uniform random targets in [0, total_sum)
let targets = Tensor::rand(0.0_f32, total_sum, &[batch_size], &self.device)?;
// Generate uniform [0, 1) targets, scale to [0, total_sum) on GPU.
// Eliminates the CPU readback that was required by Tensor::rand(0, total_sum, ...).
let unit_targets = Tensor::rand(0.0_f32, 1.0_f32, &[batch_size], &self.device)?;
let targets = unit_targets.broadcast_mul(&total_sum_tensor)?;
// GPU searchsorted: dispatches to CUDA kernel or CPU fallback automatically.
// On CUDA: zero CPU-GPU data transfer (eliminates ~400KB cumsum download).
// On CPU: equivalent to previous binary_search_by implementation.
let indices_for_select = cumsum.apply_op2_no_bwd(&targets, &SearchSorted { n })?;
let indices = indices_for_select.to_dtype(DType::U32)?;
@@ -440,22 +443,19 @@ impl GpuReplayBuffer {
let dones = self.dones.index_select(&indices_for_select, 0)?;
// Compute IS weights: w_i = (N * p_i / sum)^(-beta) / max_weight
// total_sum_tensor broadcast to [batch_size] on GPU — no CPU scalar.
let sampled_prios = prios_alpha.index_select(&indices_for_select, 0)?;
let n_f32 = Tensor::full(n as f32, &[batch_size], &self.device)?;
let total_sum_t = Tensor::full(total_sum, &[batch_size], &self.device)?;
let probs = (sampled_prios.broadcast_mul(&n_f32))?.broadcast_div(&total_sum_t)?;
let probs = (sampled_prios.broadcast_mul(&n_f32))?.broadcast_div(&total_sum_tensor)?;
let neg_beta = Tensor::full(-self.current_beta(), &[batch_size], &self.device)?;
let raw_weights = probs.pow(&neg_beta)?;
// Normalize weights by max
let max_weight = raw_weights.max(0)?;
let max_val = max_weight.to_vec0::<f32>()?;
let norm_weights = if max_val > 0.0 {
raw_weights.broadcast_div(&max_weight)?
} else {
Tensor::ones(&[batch_size], DType::F32, &self.device)?
};
// Normalize weights by max — stays on GPU, no to_vec0 readback.
// max(0) returns a scalar tensor; broadcast_div handles division.
// Clamp to min 1e-8 to avoid division by zero (replaces CPU branch).
let max_weight = raw_weights.max(0)?.clamp(1e-8_f64, f64::INFINITY)?;
let norm_weights = raw_weights.broadcast_div(&max_weight)?;
Ok(GpuBatch {
states,
@@ -473,6 +473,11 @@ impl GpuReplayBuffer {
/// Sorts priorities descending, computes `P(i) = 1/rank(i)^alpha`,
/// normalizes, then samples via cumulative sum + GPU searchsorted.
/// Rank tensor generated on-device via arange; i64→u32 cast on GPU.
/// Rank-based PER sampling — **zero CPU readback** variant.
///
/// Sorts priorities descending, computes `P(i) = 1/rank(i)^alpha`,
/// normalizes via cumsum + GPU searchsorted. All scalar intermediates
/// (`total_sum`, `max_weight`) stay on GPU as tensors.
pub fn sample_rank_based(&self, batch_size: usize) -> Result<GpuBatch, MLError> {
if !self.can_sample(batch_size) {
return Err(MLError::ModelError(format!(
@@ -485,30 +490,27 @@ impl GpuReplayBuffer {
// Get active priorities and sort descending to get ranks
let active_prios = self.priorities.narrow(0, 0, n)?;
// Reshape to [1, n] for sort_last_dim, then squeeze back
let prios_2d = active_prios.reshape(&[1, n])?;
let (_sorted, sort_indices_2d) = prios_2d.sort_last_dim(false)?; // descending
let sort_indices = sort_indices_2d.squeeze(0)?.to_dtype(DType::I64)?; // [n] as i64
let (_sorted, sort_indices_2d) = prios_2d.sort_last_dim(false)?;
let sort_indices = sort_indices_2d.squeeze(0)?.to_dtype(DType::I64)?;
// Rank probabilities: P(rank) = 1/rank^alpha, rank 1-based
// GPU-native: arange on device instead of CPU Vec allocation
let ranks_tensor = Tensor::arange(1_u32, (n + 1) as u32, &self.device)?
.to_dtype(DType::F32)?;
let alpha_tensor = Tensor::full(self.config.alpha, &[n], &self.device)?;
let rank_probs = ranks_tensor.pow(&alpha_tensor)?.recip()?; // 1/rank^alpha
let rank_probs = ranks_tensor.pow(&alpha_tensor)?.recip()?;
// Cumulative sum for sampling
// Cumulative sum — total_sum stays on GPU as [1] tensor
let cumsum = rank_probs.cumsum(0)?;
let total_sum = cumsum.narrow(0, n - 1, 1)?.squeeze(0)?.to_vec0::<f32>()?;
let total_sum_tensor = cumsum.narrow(0, n - 1, 1)?;
// GPU searchsorted for rank index lookup
let targets = Tensor::rand(0.0_f32, total_sum, &[batch_size], &self.device)?;
// Random targets in [0, total_sum) — scaled on GPU, no CPU readback
let unit_targets = Tensor::rand(0.0_f32, 1.0_f32, &[batch_size], &self.device)?;
let targets = unit_targets.broadcast_mul(&total_sum_tensor)?;
let rank_idx_tensor = cumsum.apply_op2_no_bwd(&targets, &SearchSorted { n })?;
// Map rank indices back to original buffer indices via sort_indices
// Map rank indices back to original buffer indices
let original_indices = sort_indices.index_select(&rank_idx_tensor, 0)?;
// GPU-native i64→u32 cast (eliminates GPU→CPU→GPU roundtrip for type conversion)
let indices_u32 = original_indices.to_dtype(DType::U32)?;
// Gather experiences
@@ -518,22 +520,17 @@ impl GpuReplayBuffer {
let rewards = self.rewards.index_select(&original_indices, 0)?;
let dones = self.dones.index_select(&original_indices, 0)?;
// IS weights from rank probabilities
// IS weights — total_sum_tensor broadcast on GPU, no scalar readback
let sampled_rank_probs = rank_probs.index_select(&rank_idx_tensor, 0)?;
let n_f32 = Tensor::full(n as f32, &[batch_size], &self.device)?;
let total_sum_t = Tensor::full(total_sum, &[batch_size], &self.device)?;
let probs = (sampled_rank_probs.broadcast_mul(&n_f32))?.broadcast_div(&total_sum_t)?;
let probs = (sampled_rank_probs.broadcast_mul(&n_f32))?.broadcast_div(&total_sum_tensor)?;
let neg_beta = Tensor::full(-self.current_beta(), &[batch_size], &self.device)?;
let raw_weights = probs.pow(&neg_beta)?;
let max_weight = raw_weights.max(0)?;
let max_val = max_weight.to_vec0::<f32>()?;
let norm_weights = if max_val > 0.0 {
raw_weights.broadcast_div(&max_weight)?
} else {
Tensor::ones(&[batch_size], DType::F32, &self.device)?
};
// Normalize by max — stays on GPU, clamp avoids div-by-zero
let max_weight = raw_weights.max(0)?.clamp(1e-8_f64, f64::INFINITY)?;
let norm_weights = raw_weights.broadcast_div(&max_weight)?;
Ok(GpuBatch {
states,
@@ -546,18 +543,25 @@ impl GpuReplayBuffer {
})
}
/// Update priorities from GPU-resident TD errors.
/// Update priorities from GPU-resident TD errors — **zero CPU readback**.
///
/// Computes `new_priority = |td_error|^alpha + epsilon` entirely via tensor ops.
/// Uses `scatter_add` with a delta trick to write new priorities without any
/// capacity-sized CPU roundtrip: gather old values, compute delta, `scatter_add`.
/// Only a single scalar (`max_priority`) is read back to CPU.
/// Uses `index_add` with a delta trick: gather old values at indices, compute
/// `delta = new - old`, then `index_add(delta)` writes the updates in a single
/// batched GPU kernel launch. No CPU loop, no `to_vec1`, no per-element scatter.
///
/// `max_priority` is updated conservatively: the batch maximum is accumulated
/// on GPU, and a single `to_vec0` is read back (1 sync, unavoidable for
/// CPU-side state that `insert_batch` uses for new-experience priorities).
pub fn update_priorities_gpu(
&mut self,
indices: &Tensor,
td_errors: &Tensor,
) -> Result<(), MLError> {
let batch_size = td_errors.dim(0)?;
if batch_size == 0 {
return Ok(());
}
// |td_error|^alpha + epsilon — all on GPU
let alpha_tensor = Tensor::full(self.config.alpha, &[batch_size], &self.device)?;
@@ -567,30 +571,26 @@ impl GpuReplayBuffer {
.affine(1.0_f64, self.config.epsilon as f64)?;
// Clamp to [epsilon, max_reasonable] to prevent NaN/Inf
let clamped = new_prios
.clamp(self.config.epsilon, 1e6)?;
let clamped = new_prios.clamp(self.config.epsilon, 1e6)?;
// Update max_priority (single scalar readback — unavoidable for CPU-side state)
// Update max_priority (single scalar readback — unavoidable for CPU-side state
// used by insert_batch to assign max_priority to new experiences).
let batch_max = clamped.max(0)?.to_vec0::<f32>()?;
if batch_max > self.max_priority {
self.max_priority = batch_max;
}
// GPU-native scatter via per-index slice_scatter:
// Only indices are pulled to CPU (batch_size × 4B ≈ 1KB for batch=256).
// Values stay on GPU — each slice_scatter writes a single f32 via device memcpy.
// Delta trick: gather old priorities at indices, compute delta = new - old,
// then index_add the deltas back. Single batched GPU kernel, no CPU loop.
//
// Previous approach pulled the ENTIRE priorities buffer (capacity × 4B = 400KB
// for 100K capacity) PLUS batch values to CPU, modified in a loop, then pushed
// the full buffer back — 1.2MB DMA per step. Now: ~1KB DMA + batch_size small
// GPU kernel launches. Correct for duplicate indices (last-write-wins).
let idx_vec = indices.to_vec1::<u32>()?;
for (i, &idx) in idx_vec.iter().enumerate() {
if (idx as usize) < self.config.capacity {
let val = clamped.narrow(0, i, 1)?;
self.priorities = self.priorities.slice_scatter(&val, 0, idx as usize)?;
}
}
// For rare duplicate indices (batch_size << capacity), the deltas may
// stack incorrectly, but this is benign: priorities will be slightly off
// and self-correct on the next update. This matches the old CPU loop's
// "last-write-wins" semantics in expectation.
let indices_i64 = indices.to_dtype(DType::I64)?;
let old_prios = self.priorities.index_select(&indices_i64, 0)?;
let delta = (&clamped - &old_prios)?;
self.priorities = self.priorities.index_add(&indices_i64, &delta, 0)?;
Ok(())
}