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

fix: cross-stream race in PER update + pre-allocate scratch buffers

Browse Source 
The replay buffer used its own CUDA stream for seg_tree_update, but
the td_errors buffer was computed on the trainer's stream. Without
synchronization, the replay stream could read incomplete data — causing
a hang on H100 where streams execute truly concurrently.

Fix: update_priorities_gpu now accepts an optional ext_stream parameter.
The fused training path passes the trainer's main stream, ensuring all
GPU work (td_errors computation + seg_tree_update) runs on the same
stream with implicit ordering.

Also pre-allocate update_td_f32, update_batch_max, update_max_merge
scratch buffers at construction — eliminates 3 cuMemAlloc calls per
training step that could trigger implicit device synchronization.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 12:36:15 +02:00
parent 22f30fe37f
commit 907b5f3355
7 changed files with 71 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -183,6 +183,10 @@ pub struct GpuReplayBuffer {
sample_max_weight: CudaSlice<f32>,
sample_episode_ids: CudaSlice<i32>,
total_sum_buf: CudaSlice<f32>,
// Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step)
update_td_f32: CudaSlice<f32>, // [max_batch_size] bf16→f32 conversion
update_batch_max: CudaSlice<f32>, // [1] per-batch atomicMax accumulator
update_max_merge: CudaSlice<f32>, // [1] merge pending + batch max
rng_step: u32,
/// Size of the most recent `sample_proportional` call. Used by `sample_weights_ref` /
/// `sample_indices_ref` to return correctly-sized views (pre-allocated buffers are
@@ -244,6 +248,11 @@ impl GpuReplayBuffer {
let ep_ids = a32i(stream, cap, "episode_ids")?;
let s_ep = a32i(stream, mbs, "s_episode_ids")?;
// Pre-allocate PER update scratch buffers (zero cuMemAlloc per step)
let update_td_f32 = a32f(stream, mbs, "update_td_f32")?;
let update_batch_max = a32f(stream, 1, "update_batch_max")?;
let update_max_merge = a32f(stream, 1, "update_max_merge")?;
Ok(Self {
config, stream: Arc::clone(stream), kernels: k,
states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p,
@@ -258,6 +267,7 @@ impl GpuReplayBuffer {
sample_priorities: sp, sample_weights: sw,
sample_episode_ids: s_ep,
sample_max_weight: smw, total_sum_buf: tsb,
update_td_f32, update_batch_max, update_max_merge,
rng_step: 0,
last_batch_size: 0,
})
@@ -684,62 +694,73 @@ impl GpuReplayBuffer {
/// writes to priorities buffer AND tree leaves, propagates sums to root.
/// O(log n) per update instead of O(n) prefix sum rebuild.
/// Raw-pointer variant for tests and callers without CudaSlice access.
pub fn update_priorities_gpu_raw(&mut self, _indices_ptr: u64, td_errors: &CudaSlice<half::bf16>, bs: usize) -> Result<(), MLError> {
pub fn update_priorities_gpu_raw(&mut self, _indices_ptr: u64, td_errors: &CudaSlice<half::bf16>, bs: usize, ext_stream: Option<&Arc<CudaStream>>) -> Result<(), MLError> {
let idx_slice = unsafe { &*(&self.sample_indices_u32 as *const CudaSlice<u32>) };
self.update_priorities_gpu(idx_slice, td_errors, bs)
self.update_priorities_gpu(idx_slice, td_errors, bs, ext_stream)
}
/// Update PER priorities from TD errors using the segment tree kernel.
///
/// If `ext_stream` is provided, all GPU work runs on that stream instead of
/// the replay buffer's own stream. This avoids cross-stream synchronization
/// when the TD errors were computed on a different stream (e.g., the trainer's
/// main stream). All pre-allocated scratch buffers are used — zero per-step
/// GPU allocations.
pub fn update_priorities_gpu(
&mut self,
indices: &CudaSlice<u32>,
td_errors: &CudaSlice<half::bf16>,
bs: usize,
ext_stream: Option<&Arc<CudaStream>>,
) -> Result<(), MLError> {
let _nvtx = NvtxRange::new("per_update_priorities");
if bs == 0 { return Ok(()); }
// Convert bf16 td_errors to f32 for the CUDA kernel (which reads float*)
let kernels = get_cast_kernels(&self.stream)?;
let td_errors_f32 = {
let n = bs;
let ni = n as i32;
let out = self.stream.alloc_zeros::<f32>(n)
.map_err(|e| MLError::ModelError(format!("td bf16->f32 alloc: {e}")))?;
// SAFETY: out, td_errors are valid device allocations of at least n elements.
unsafe {
self.stream.launch_builder(&kernels.bf16_to_f32)
.arg(&out).arg(td_errors).arg(&ni)
.launch(lcfg(n))
.map_err(|e| MLError::ModelError(format!("td bf16_to_f32: {e}")))?;
}
out
};
let stream = ext_stream.unwrap_or(&self.stream);
// Convert bf16 td_errors to f32 using pre-allocated scratch buffer
let kernels = get_cast_kernels(stream)?;
let ni = bs as i32;
// SAFETY: update_td_f32, td_errors are valid device allocations of at least bs elements.
unsafe {
stream.launch_builder(&kernels.bf16_to_f32)
.arg(&self.update_td_f32).arg(td_errors).arg(&ni)
.launch(lcfg(bs))
.map_err(|e| MLError::ModelError(format!("td bf16_to_f32: {e}")))?;
}
let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32);
let cap_pow2_i = self.capacity_pow2 as i32;
let mut bm = a32f(&self.stream, 1, "bm")?;
self.stream.memset_zeros(&mut bm).map_err(|e| MLError::ModelError(format!("{e}")))?;
// seg_tree_update: writes priorities[], tree leaves, propagates, atomicMax.
// SAFETY: seg_tree, priorities, bm, indices, td_errors_f32 are valid device allocations.
stream.memset_zeros(&mut self.update_batch_max)
.map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?;
// seg_tree_update: writes priorities[], tree leaves, propagates via atomicAdd, atomicMax.
// SAFETY: seg_tree, priorities, update_batch_max, indices, update_td_f32 are valid device allocations.
unsafe {
self.stream.launch_builder(&self.kernels.seg_tree_update)
.arg(&self.seg_tree).arg(&self.priorities).arg(&mut bm)
.arg(indices).arg(&td_errors_f32)
stream.launch_builder(&self.kernels.seg_tree_update)
.arg(&self.seg_tree).arg(&self.priorities).arg(&self.update_batch_max)
.arg(indices).arg(&self.update_td_f32)
.arg(&al).arg(&ep).arg(&cap_pow2_i).arg(&bsi)
.launch(lcfg(bs)).map_err(|e| MLError::ModelError(format!("st update: {e}")))?;
}
self.pending_max_priority = Some(match self.pending_max_priority.take() {
// Merge batch max into pending max (GPU-only).
// Uses pre-allocated update_max_merge as output scratch.
match self.pending_max_priority.take() {
Some(prev) => {
let mut r = a32f(&self.stream, 1, "mm")?;
// SAFETY: r, prev, bm are valid 1-element device allocations on same context.
// SAFETY: update_max_merge, prev, update_batch_max are valid 1-element device allocations.
unsafe {
self.stream.launch_builder(&self.kernels.max_of_two_f32)
.arg(&mut r).arg(&prev).arg(&bm)
stream.launch_builder(&self.kernels.max_of_two_f32)
.arg(&self.update_max_merge).arg(&prev).arg(&self.update_batch_max)
.launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
.map_err(|e| MLError::ModelError(format!("max_of_two: {e}")))?;
}
r
// Rotate: update_max_merge becomes pending, prev becomes the new merge scratch.
self.pending_max_priority = Some(std::mem::replace(&mut self.update_max_merge, prev));
}
None => bm,
});
None => {
// First update this epoch: batch_max IS the pending.
// Rotate: update_batch_max becomes pending, allocate new batch_max scratch.
// This allocation happens once per epoch (after flush_max_priority), not per step.
let fresh = a32f(stream, 1, "bm_epoch")?;
self.pending_max_priority = Some(std::mem::replace(&mut self.update_batch_max, fresh));
}
}
Ok(())
}

10
crates/ml-dqn/src/replay_buffer_type.rs
View File

@@ -366,7 +366,7 @@ impl ReplayBufferType {
let mut buf = buffer.lock();
if bs == 0 { return Ok(()); }
// indices already CudaSlice<u32>, td_errors already CudaSlice<half::bf16>
buf.gpu.update_priorities_gpu(indices, td_errors, bs)
buf.gpu.update_priorities_gpu(indices, td_errors, bs, None)
}
Self::Uniform(_) | Self::Prioritized(_) => Ok(()), // No-op for non-GPU buffers
}
@@ -390,16 +390,22 @@ impl ReplayBufferType {
/// Uses the segment tree update kernel which writes priorities, updates
/// tree leaves, and propagates sums to root. The sample indices are stored
/// internally from the last `sample()` call.
/// Update PER priorities using td_errors and the indices from the last sample.
///
/// `ext_stream`: the stream on which `td_errors` was computed. All GPU work
/// runs on this stream to avoid cross-stream synchronization. Pass the
/// trainer's main stream here.
pub fn update_priorities_from_td(
&self,
td_errors: &cudarc::driver::CudaSlice<half::bf16>,
batch_size: usize,
ext_stream: &std::sync::Arc<cudarc::driver::CudaStream>,
) -> Result<(), MLError> {
match self {
Self::GpuPrioritized(buffer) => {
let mut buf = buffer.lock();
if batch_size == 0 { return Ok(()); }
buf.gpu.update_priorities_gpu_raw(0, td_errors, batch_size)
buf.gpu.update_priorities_gpu_raw(0, td_errors, batch_size, Some(ext_stream))
}
Self::Uniform(_) | Self::Prioritized(_) => Ok(()),
}

3
crates/ml/src/trainers/dqn/config.rs
View File

@@ -643,8 +643,9 @@ impl DQNAgentType {
&self,
td_errors: &cudarc::driver::CudaSlice<half::bf16>,
batch_size: usize,
stream: &std::sync::Arc<cudarc::driver::CudaStream>,
) -> Result<(), crate::MLError> {
self.memory().update_priorities_from_td(td_errors, batch_size)
self.memory().update_priorities_from_td(td_errors, batch_size, stream)
}
/// Set max_priority from a CPU scalar (epoch-boundary flush from CUDA kernel).

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

@@ -783,6 +783,7 @@ impl FusedTrainingCtx {
agent.update_priorities_from_td(
self.trainer.td_errors_buf(),
self.batch_size,
&self.stream,
).map_err(|e| anyhow::anyhow!("PER seg_tree_update: {e}"))?;
if step < 3 { eprintln!("H100_HANG4: step {step} — PER update done"); }

2
crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs
View File

@@ -190,7 +190,7 @@ fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> {
// Sample and get indices
let batch = buf.sample_proportional(16)?;
let td_errors = random_cuda_bf16(16, &stream)?;
let idx = batch.indices_ptr; buf.update_priorities_gpu_raw(idx, &td_errors, 16)?;
let idx = batch.indices_ptr; buf.update_priorities_gpu_raw(idx, &td_errors, 16, None)?;
// Verify sampling still works with updated priorities (proves update succeeded)
buf.sample_proportional(16)?;

2
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1558,7 +1558,7 @@ impl DQNTrainer {
let mut gpu_buf = agent.memory().as_gpu_buffer();
if let Some(ref mut buf) = gpu_buf {
if let Err(e) = buf.gpu.update_priorities_gpu(
&idx_slice, &td_slice, batch_size_refresh,
&idx_slice, &td_slice, batch_size_refresh, None,
) {
debug!("M2: GPU priority refresh failed (non-fatal): {}", e);
} else {

8
crates/ml/tests/gpu_per_integration_test.rs
View File

@@ -157,7 +157,7 @@ fn test_gpu_per_training_loop_cycle() {
let td_errors: CudaSlice<half::bf16> = stream.clone_htod(&td_host).unwrap();
let idx_ptr = batch.indices_ptr;
buf.update_priorities_gpu_raw(idx_ptr, &td_errors, batch_size).unwrap();
buf.update_priorities_gpu_raw(idx_ptr, &td_errors, batch_size, None).unwrap();
buf.step();
}
@@ -181,7 +181,7 @@ fn test_gpu_per_priorities_affect_sampling() {
// Set experience 0 to very high priority
let high_idx: CudaSlice<u32> = stream.clone_htod(&[0_u32]).unwrap();
let high_td: CudaSlice<half::bf16> = stream.clone_htod(&[half::bf16::from_f32(100.0)]).unwrap();
buf.update_priorities_gpu(&high_idx, &high_td, 1).unwrap();
buf.update_priorities_gpu(&high_idx, &high_td, 1, None).unwrap();
// Sample 1000 times in batches of 10, count how often index 0 appears
// Download indices to host, count matches
@@ -225,7 +225,7 @@ fn test_gpu_per_sampling_distribution_differs_from_uniform() {
for i in 0..50 {
let idx: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
let td: CudaSlice<half::bf16> = stream.clone_htod(&[half::bf16::from_f32(i as f32 + 1.0)]).unwrap();
buf.update_priorities_gpu(&idx, &td, 1).unwrap();
buf.update_priorities_gpu(&idx, &td, 1, None).unwrap();
}
// Sample 5000 indices via proportional sampling
@@ -267,7 +267,7 @@ fn test_gpu_per_is_weights_correct_range() {
for i in 0..200 {
let idx: CudaSlice<u32> = stream.clone_htod(&[i as u32]).unwrap();
let td: CudaSlice<half::bf16> = stream.clone_htod(&[half::bf16::from_f32((i % 10) as f32 + 1.0)]).unwrap();
buf.update_priorities_gpu(&idx, &td, 1).unwrap();
buf.update_priorities_gpu(&idx, &td, 1, None).unwrap();
}
let batch = buf.sample_proportional(64).unwrap();