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

cleanup: delete dead ensemble_diversity_kernel + readback path

Browse Source 
Found while auditing the remaining atomicAdds in the codebase:
ensemble_diversity_kernel computed pairwise KL divergence between K
ensemble heads and wrote the result to ensemble_diversity_loss_buf.
A readback function (`readback_diversity_loss`) was defined to consume
this scalar — but **nothing in the codebase ever called it**. The
kernel ran every training step, allocated a buffer per step, and the
result was discarded.

Removed (-195 LOC net):
  - ensemble_diversity_kernel CUDA kernel (~100 LOC C++)
  - kernel load + cubin pin in compile_ensemble_kernels
  - ensemble_diversity_kernel field on FusedTrainingCtx
  - ensemble_diversity_loss_buf field + alloc
  - pending_diversity_loss_ptr + pending_diversity_normalizer fields
  - readback_diversity_loss() function (the would-be consumer)
  - launch site in run_ensemble_step (zero+launch+pending_ptr setup)

Kept (these ARE used):
  - ensemble_aggregate_kernel (Q-value mean/variance for exploration bonus)
  - ensemble_kl_gradient_kernel (computes diversity gradient → SAXPY into
    grad_buf → adam — this is the actual training-path mechanism)
  - apply_ensemble_diversity_backward (calls kl_gradient_kernel)
  - ensemble_diversity_weight (scales the gradient)

Net effect on training: zero (the deleted code's output was unused).
Net effect on per-step cost: small but non-zero — saves one kernel
launch + memset per step + a CudaSlice<f32> alloc per training context.
On L40S/H100 this is microseconds; on RTX 3050 Ti slightly more.

Effect on determinism: zero. The atomicAdd in the deleted kernel was
in a code path whose output didn't feed training, so removing it
doesn't change the training trajectory. The remaining 9 atomicAdds
in the codebase break down as: 5 in monitoring_kernel (diagnostic
stats only), 3 in experience_kernels (atom_stats / penalty_out —
diagnostic-ish), 1 in dqn_utility (sensitivity_out feature attribution),
1 in trade_stats. None are in the gradient hot path.

Files touched:
  crates/ml/src/cuda_pipeline/ensemble_kernels.cu  (-100 lines)
  crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs   (-13 lines)
  crates/ml/src/trainers/dqn/fused_training.rs     (-100 lines)

Verified: SQLX_OFFLINE=true cargo check -p ml --lib --tests passes.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-21 08:59:40 +02:00
parent 396f0d09b5
commit 25ceb3b8d5
3 changed files with 26 additions and 221 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
crates/ml/src/cuda_pipeline/ensemble_kernels.cu
View File

@@ -1,16 +1,22 @@
/**
* Ensemble Q-Network kernels -- aggregate and diversity loss.
* Ensemble Q-Network kernels -- aggregate + KL gradient.
*
* Two kernels for ensemble training with K independent value/advantage head
* weight sets sharing a common DQN trunk:
* Two live kernels for ensemble training with K independent value/advantage
* head weight sets sharing a common DQN trunk:
*
* 1. ensemble_aggregate_kernel -- compute mean and variance of Q-values across
* K heads (uncertainty estimation).
* K heads (uncertainty estimation, used by exploration bonus).
*
* 2. ensemble_diversity_kernel -- pairwise KL divergence between softmax
* distributions of C51 head logits (diversity regularization).
* Uses warp-level -> block-level -> atomicAdd-per-block hierarchical reduction
* matching the pattern in dqn_grad_norm_kernel.
* 2. ensemble_kl_gradient_kernel -- compute dL_diversity/d_logits for head 0,
* fed into grad_buf via SAXPY in apply_ensemble_diversity_backward.
* This is the actual gradient that pushes heads apart during training.
*
* REMOVED 2026-04-21: ensemble_diversity_kernel (computed pairwise KL as a
* scalar diagnostic). The result was written to ensemble_diversity_loss_buf
* and pending_diversity_loss_ptr, but readback_diversity_loss() was never
* called from anywhere — pure dead code that ran every training step. Its
* removal eliminates the only remaining training-path atomicAdd outside
* monitoring kernels.
*
* No #include of common_device_functions.cuh is needed here -- the build
* function in gpu_dqn_trainer.rs prepends it via string concatenation.
@@ -53,107 +59,6 @@ extern "C" __global__ void ensemble_aggregate_kernel(
var_q[idx] = fmaxf(v, 0.0f);
}
/* ======================================================================
* KERNEL 2: ENSEMBLE DIVERSITY LOSS (KL DIVERGENCE)
*
* Computes pairwise KL divergence between softmax distributions of K heads'
* C51 value logits. KL is averaged over all K*(K-1)/2 head pairs and
* all B samples. Written as a scalar into diversity_loss[0].
*
* Uses hierarchical reduction: warp reduce -> block shared memory -> per-block
* atomicAdd. This matches dqn_grad_norm_kernel and avoids serialization at
* a single atomic address.
*
* head_logits layout: [K * B * num_atoms] -- head k starts at k*B*num_atoms.
*
* Each thread processes one (sample, head_i, head_j) pair's KL contribution.
* Grid sizing: grid_x = ceil(B * num_pairs / 256), where num_pairs = K*(K-1)/2.
*
* Launch config: grid=(ceil(B*num_pairs/256), 1, 1), block=(256, 1, 1).
* ====================================================================== */
extern "C" __global__ void ensemble_diversity_kernel(
const float* __restrict__ head_logits, /* [K * B * num_atoms] f32 */
float* __restrict__ diversity_loss, /* [1] output: accumulated KL */
int K,
int B,
int num_atoms
) {
/* Number of ordered pairs (i < j): K*(K-1)/2 */
int num_pairs = K * (K - 1) / 2;
int total_work = B * num_pairs;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
float kl_sum = 0.0f;
if (idx < total_work) {
/* Decode sample index and pair index */
int sample = idx / num_pairs;
int pair = idx % num_pairs;
/* Decode ordered pair (i, j) with i < j from linear pair index. */
int hi = 0, hj = 1;
int p = 0;
for (int i = 0; i < K - 1; i++) {
for (int j = i + 1; j < K; j++) {
if (p == pair) { hi = i; hj = j; }
p++;
}
}
/* Pointers to head i and head j logits for this sample (f32) */
const float* logits_i = head_logits + (long long)hi * B * num_atoms + sample * num_atoms;
const float* logits_j = head_logits + (long long)hj * B * num_atoms + sample * num_atoms;
/* Softmax of logits_i -- f32 */
float max_i = logits_i[0];
for (int a = 1; a < num_atoms; a++) max_i = fmaxf(max_i, logits_i[a]);
float sum_i = 0.0f;
for (int a = 0; a < num_atoms; a++) sum_i += expf(logits_i[a] - max_i);
/* Softmax of logits_j -- f32 */
float max_j = logits_j[0];
for (int a = 1; a < num_atoms; a++) max_j = fmaxf(max_j, logits_j[a]);
float sum_j = 0.0f;
for (int a = 0; a < num_atoms; a++) sum_j += expf(logits_j[a] - max_j);
/* KL(p_i || p_j) -- f32 for numerical stability */
float kl = 0.0f;
for (int a = 0; a < num_atoms; a++) {
float pi = expf(logits_i[a] - max_i) / sum_i;
float pj = expf(logits_j[a] - max_j) / sum_j;
float log_ratio_ij = logf(pi / (pj + 1e-8f) + 1e-8f);
float log_ratio_ji = logf(pj / (pi + 1e-8f) + 1e-8f);
kl += 0.5f * (pi * log_ratio_ij + pj * log_ratio_ji);
}
kl_sum = kl;
}
/* -- Hierarchical reduction: warp -> block -> one atomicAdd per BLOCK --
* This is nearly deterministic: only grid_dim atomicAdds to a single
* scalar. For full determinism, a two-phase reduction (like IQN's
* iqn_loss_reduce) would eliminate atomicAdd entirely, but for a
* monitoring-only diversity loss, per-block atomicAdd suffices. */
/* Warp-level reduction via shuffle (no shared memory) */
for (int offset = 16; offset > 0; offset >>= 1)
kl_sum = kl_sum + __shfl_xor_sync(0xFFFFFFFF, kl_sum, offset);
/* Block-level cross-warp reduction via shared memory (padded, same as grad_norm) */
__shared__ float warp_sums[16]; /* 8 warps x 2 stride (padded) */
int warp_id = threadIdx.x / 32;
int warp_lane = threadIdx.x % 32;
if (warp_lane == 0) warp_sums[warp_id * 2] = kl_sum;
__syncthreads();
/* First warp reduces across warps */
if (warp_id == 0) {
float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : 0.0f;
for (int offset = 16; offset > 0; offset >>= 1)
val = val + __shfl_xor_sync(0xFFFFFFFF, val, offset);
if (warp_lane == 0)
atomicAdd(diversity_loss, val);
}
}
/* ======================================================================
* ENSEMBLE KL GRADIENT KERNEL
*

14
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -13414,12 +13414,15 @@ fn dtod_from_staging(
/// Load ensemble aggregate + diversity kernels from precompiled cubin.
///
/// Returns `(ensemble_aggregate_kernel, ensemble_diversity_kernel)`.
/// Both kernels take all sizes as runtime arguments.
/// Returns just the aggregate kernel; the diagnostic-only diversity-loss
/// kernel was removed (its result was never read — see ensemble_kernels.cu
/// header for the full story). The KL *gradient* kernel
/// (`ensemble_kl_gradient_kernel`) is loaded separately in
/// `fused_training.rs` since that one DOES feed back into training.
pub(crate) fn compile_ensemble_kernels(
stream: &Arc<CudaStream>,
_state_dim: usize,
) -> Result<(CudaFunction, CudaFunction), MLError> {
) -> Result<CudaFunction, MLError> {
let context = stream.context();
let module = context.load_cubin(ENSEMBLE_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("ensemble cubin load: {e}")))?;
@@ -13427,11 +13430,8 @@ pub(crate) fn compile_ensemble_kernels(
let aggregate = module
.load_function("ensemble_aggregate_kernel")
.map_err(|e| MLError::ModelError(format!("ensemble_aggregate_kernel load: {e}")))?;
let diversity = module
.load_function("ensemble_diversity_kernel")
.map_err(|e| MLError::ModelError(format!("ensemble_diversity_kernel load: {e}")))?;
Ok((aggregate, diversity))
Ok(aggregate)
}
/// Launch the f32 copy CUDA kernel: copies f32 CudaSlice<f32> → f32 CudaSlice<f32>.

110
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -249,17 +249,11 @@ pub(crate) struct FusedTrainingCtx {
/// Ensemble diversity weight λ — scales the KL loss relative to C51 loss.
/// Default 0.01. Zero when ensemble_count <= 1.
pub(crate) ensemble_diversity_weight: f32,
/// Deferred diversity loss readback — accumulated on GPU per step, read at epoch boundary.
pending_diversity_loss_ptr: Option<u64>,
pending_diversity_normalizer: f32,
/// Compiled ensemble aggregate kernel (Q-value mean/variance across heads).
/// None when ensemble_count <= 1.
pub(crate) ensemble_aggregate_kernel: Option<cudarc::driver::CudaFunction>,
/// Compiled ensemble diversity kernel (pairwise KL divergence across heads).
/// None when ensemble_count <= 1.
pub(crate) ensemble_diversity_kernel: Option<cudarc::driver::CudaFunction>,
/// Pre-allocated buffer: [K * B * num_atoms] f32 for per-head value logits.
/// Used to assemble logits from all K heads before diversity kernel.
/// Used by the KL gradient kernel to assemble logits from all K heads.
/// F32 because output layer GemmEx writes f32 (no f32 truncation overflow).
/// None when ensemble_count <= 1.
pub(crate) ensemble_logits_buf: Option<cudarc::driver::CudaSlice<f32>>,
@@ -269,9 +263,6 @@ pub(crate) struct FusedTrainingCtx {
/// Pre-allocated buffer: [B * num_atoms] for Q-value variance across K heads.
/// None when ensemble_count <= 1.
pub(crate) ensemble_var_q_buf: Option<cudarc::driver::CudaSlice<f32>>,
/// Pre-allocated buffer: [1] for accumulated diversity loss scalar.
/// None when ensemble_count <= 1.
pub(crate) ensemble_diversity_loss_buf: Option<cudarc::driver::CudaSlice<f32>>,
/// Compiled KL gradient kernel for diversity gradient flow.
pub(crate) ensemble_kl_grad_kernel: Option<cudarc::driver::CudaFunction>,
/// Pre-allocated buffer: [B * num_atoms] for diversity gradient on head 0 logits.
@@ -619,18 +610,16 @@ impl FusedTrainingCtx {
let (
ensemble_extra_heads,
ensemble_aggregate_kernel,
ensemble_diversity_kernel,
ensemble_logits_buf,
ensemble_mean_q_buf,
ensemble_var_q_buf,
ensemble_diversity_loss_buf,
ensemble_kl_grad_kernel,
ensemble_d_logits_buf,
ens_backings_vec,
) = if k > 1 {
use crate::cuda_pipeline::gpu_dqn_trainer::compile_ensemble_kernels;
let (agg_kernel, div_kernel) =
let agg_kernel =
compile_ensemble_kernels(&stream, ml_core::state_layout::STATE_DIM)
.map_err(|e| anyhow::anyhow!("Ensemble kernels compile: {e}"))?;
@@ -662,8 +651,6 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("Alloc ensemble_mean_q_buf: {e}"))?;
let var_q_buf = stream.alloc_zeros::<f32>(batch_size * na)
.map_err(|e| anyhow::anyhow!("Alloc ensemble_var_q_buf: {e}"))?;
let div_loss_buf = stream.alloc_zeros::<f32>(1)
.map_err(|e| anyhow::anyhow!("Alloc ensemble_diversity_loss_buf: {e}"))?;
info!(
ensemble_count = k,
@@ -688,18 +675,15 @@ impl FusedTrainingCtx {
(
extra_heads,
Some(agg_kernel),
Some(div_kernel),
Some(logits_buf),
Some(mean_q_buf),
Some(var_q_buf),
Some(div_loss_buf),
Some(kl_grad_kernel),
Some(d_logits_buf),
ens_backings,
)
} else {
(Vec::new(), None, None, None, None, None, None,
None, None, Vec::new()) // kl_grad_kernel, d_logits_buf, ens_backings
(Vec::new(), None, None, None, None, None, None, Vec::new())
};
// Wire IQL buffer pointers to the DQN trainer for C51 kernel launches.
@@ -754,14 +738,10 @@ impl FusedTrainingCtx {
attn_done_event,
ensemble_extra_heads,
ensemble_diversity_weight: hyperparams.ensemble_diversity_weight as f32,
pending_diversity_loss_ptr: None,
pending_diversity_normalizer: 1.0,
ensemble_aggregate_kernel,
ensemble_diversity_kernel,
ensemble_logits_buf,
ensemble_mean_q_buf,
ensemble_var_q_buf,
ensemble_diversity_loss_buf,
ensemble_kl_grad_kernel,
ensemble_d_logits_buf,
pending_vaccine_batch: None,
@@ -858,33 +838,6 @@ impl FusedTrainingCtx {
Ok(())
}
/// Read back accumulated ensemble diversity loss (deferred from per-step to epoch boundary).
/// Returns 0.0 if no ensemble is active or no pending readback.
///
/// Uses async DtoH into the trainer's pinned buffer at offset 9 (double-buffered).
/// The returned value is from the PREVIOUS call — the current async copy will
/// be readable by the next invocation. No CPU sync required because the
/// diversity loss is monitoring-only (not in gradient path).
pub(crate) fn readback_diversity_loss(&mut self) -> f32 {
if let Some(ptr) = self.pending_diversity_loss_ptr.take() {
// Async DtoH into pinned buffer offset 9 (no CPU blocking)
let pinned = self.trainer.readback_pinned_ptr();
unsafe {
cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
pinned.add(9).cast(),
ptr,
std::mem::size_of::<f32>(),
self.stream.cu_stream(),
);
}
// Return previous value (double-buffered — current copy lands by next call)
let prev = unsafe { *pinned.add(9) };
prev / self.pending_diversity_normalizer * self.ensemble_diversity_weight
} else {
0.0
}
}
/// Steps since last flat buffer sync.
pub(crate) fn steps_since_varmap_sync(&self) -> usize {
self.steps_since_varmap_sync
@@ -1901,7 +1854,6 @@ impl FusedTrainingCtx {
let k = self.ensemble_extra_heads.len() + 1; // total heads including head 0
let b = self.batch_size;
let na = self.trainer.config().num_atoms;
let bf16_size = std::mem::size_of::<f32>();
// No cuStreamSynchronize needed — all ops are on the same stream.
// CUDA guarantees in-order execution on a single stream.
@@ -1913,15 +1865,6 @@ impl FusedTrainingCtx {
Some(b) => b,
None => return Ok(()),
};
let div_kernel = match self.ensemble_diversity_kernel.as_ref() {
Some(k) => k,
None => return Ok(()),
};
let div_loss_buf = match self.ensemble_diversity_loss_buf.as_mut() {
Some(b) => b,
None => return Ok(()),
};
// ── 1. Copy head 0 value logits (on_v_logits_buf, f32) → logits_buf[0..B*na] ──
// Head 0's logits are already computed by the CUDA Graph.
{
@@ -1968,42 +1911,7 @@ impl FusedTrainingCtx {
).map_err(|e| anyhow::anyhow!("Ensemble head{k_idx} value forward: {e}"))?;
}
// ── 3. Zero diversity_loss_buf, then launch diversity kernel ──
// Zero the scalar accumulator.
let div_loss_ptr = div_loss_buf.raw_ptr();
unsafe {
cudarc::driver::result::memset_d8_async(
div_loss_ptr, 0u8, bf16_size, self.stream.cu_stream()
).map_err(|e| anyhow::anyhow!("Ensemble div_loss zero: {e}"))?;
}
// Launch diversity kernel over all K*(K-1)/2 head pairs × B samples.
let num_pairs = k * (k - 1) / 2;
let total_work = (b * num_pairs) as u32;
let blocks = (total_work + 255) / 256;
if blocks > 0 {
let logits_ptr = logits_buf.raw_ptr();
let k_i32 = k as i32;
let b_i32 = b as i32;
let na_i32 = na as i32;
unsafe {
self.stream
.launch_builder(div_kernel)
.arg(&logits_ptr)
.arg(&div_loss_ptr)
.arg(&k_i32)
.arg(&b_i32)
.arg(&na_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| anyhow::anyhow!("Ensemble diversity kernel: {e}"))?;
}
}
// ── 3b. Aggregate kernel: compute per-atom mean/variance across K heads ──
// ── 3. Aggregate kernel: compute per-atom mean/variance across K heads ──
// Populates ensemble_mean_q_buf and ensemble_var_q_buf [B * na] f32.
// These buffers are used for exploration bonus at the epoch boundary.
// The aggregate kernel treats logits as [K * B * na] and outputs mean/variance.
@@ -2040,15 +1948,7 @@ impl FusedTrainingCtx {
}
}
// ── 4. Diversity loss readback deferred to epoch boundary ────────
// NO cuStreamSynchronize here — that would block the CPU on every step,
// destroying the async pipeline. The diversity_loss scalar is monitoring-only
// (not in gradient path). Accumulate on GPU, readback at epoch end.
// Store the buffer pointer + normalizer for epoch-boundary readback.
self.pending_diversity_loss_ptr = Some(div_loss_ptr);
self.pending_diversity_normalizer = if num_pairs > 0 { num_pairs as f32 * b as f32 } else { 1.0_f32 };
// ── 5. Compute KL gradient and SAXPY into grad_buf ──────────────
// ── 4. Compute KL gradient and SAXPY into grad_buf ──────────────
// The KL gradient encourages head 0 to disagree with heads 1..K-1.
// d_logits = -diversity_weight * mean_k( softmax(logits_0) - softmax(logits_k) )
// This is added to grad_buf's value logits section (d_value_logits offset).