diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index e7736f436..dda408b9a 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -69,6 +69,8 @@ struct ReplayKernels { per_sample: CudaFunction, /// C1/P1: Diversity-weighted priority kernel (health < 0.8 path). pow_alpha_diverse_f32: CudaFunction, + /// Graph-safe reduction: scaled mean of u32 actions → i32[1] on GPU. + actions_sum_scale_reduce_u32: CudaFunction, } impl ReplayKernels { @@ -111,6 +113,7 @@ impl ReplayKernels { per_prefix_scan: per_ld("per_prefix_scan")?, per_sample: per_ld("per_sample")?, pow_alpha_diverse_f32: ld("pow_alpha_diverse_f32")?, + actions_sum_scale_reduce_u32: ld("actions_sum_scale_reduce_u32")?, }) } } @@ -174,6 +177,14 @@ pub struct GpuReplayBuffer { update_batch_max: CudaSlice, // [1] per-batch atomicMax accumulator update_max_merge: CudaSlice, // [1] merge pending + batch max update_batch_spare: Option>, // [1] spare for None→Some swap (zero alloc) + /// C1/P1: Pinned device-mapped scalar for the scaled-mean action value. + /// Written by `actions_sum_scale_reduce_u32` (via `mean_action_dev_ptr`), + /// read by `pow_alpha_diverse_f32` in the same graph. Pinned (not pure + /// device) so the same buffer is available to host for zero-sync logging + /// without a dtoh. Pattern matches `size_dev_ptr` / `rng_step_dev_ptr` + /// below — allocated with `cuMemAllocHost_v2` + `cuMemHostGetDevicePointer_v2`. + _mean_action_pinned: usize, // host addr (i32 slot) — kept alive for device-mapped pointer + mean_action_dev_ptr: u64, // device-mapped addr // Pre-allocated insert scratch buffers (zero cuMemAlloc per insert) insert_prio_buf: CudaSlice, // [insert_scratch_cap] priority broadcast insert_idx_buf: CudaSlice, // [insert_scratch_cap] per_insert_pa indices @@ -300,6 +311,18 @@ impl GpuReplayBuffer { (host_ptr as usize, dev_ptr) // usize for Send/Sync }; + // Pinned device-mapped scaled-mean action slot (i32). GPU writes in + // `actions_sum_scale_reduce_u32` and reads back in `pow_alpha_diverse_f32` + // — both via the device pointer. No dtoh, graph-safe across captures. + let (mean_action_pinned, mean_action_dev_ptr) = unsafe { + let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut dev_ptr: u64 = 0; + cudarc::driver::sys::cuMemAllocHost_v2(&mut host_ptr, std::mem::size_of::()); + cudarc::driver::sys::cuMemHostGetDevicePointer_v2(&mut dev_ptr, host_ptr, 0); + *(host_ptr as *mut i32) = 0; + (host_ptr as usize, dev_ptr) + }; + Ok(Self { config, stream: Arc::clone(stream), kernels: k, states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p, @@ -316,6 +339,7 @@ impl GpuReplayBuffer { sample_episode_ids: s_ep, sample_max_weight: smw, total_sum_buf: tsb, update_batch_max, update_max_merge, update_batch_spare: Some(update_batch_spare), + _mean_action_pinned: mean_action_pinned, mean_action_dev_ptr, insert_prio_buf, insert_idx_buf, insert_ep_buf, insert_scratch_cap: isc, scratch_f32, _rng_step: 0, @@ -834,18 +858,20 @@ impl GpuReplayBuffer { if health < 0.8 { // C1/P1: during collapse, boost priorities of experiences whose action - // deviates from the batch mean. Requires a host-side mean readback to - // avoid a full reduction kernel — this path runs once per epoch at epoch boundary, - // so the DtoH sync cost is acceptable. + // deviates from the batch mean. Prior implementation did a stream sync + // + memcpy_dtoh to compute the mean on host — broken inside CUDA Graph + // capture (L40S smoke exposed the invalidation on Fold 1). Pipeline is + // now fully on-device and graph-safe: + // 1. gather_u32 → sample_actions[..gather_n] (u32) + // 2. actions_sum_scale_reduce_u32 → mean_action_dev_ptr (i32 scaled ×1000) + // 3. pow_alpha_diverse_f32 reads mean_action_dev_ptr // // Re-gather actions from the ring buffer using the current batch indices - // (sample_indices_i64 from the last sample_proportional call). This avoids - // depending on trainer_actions_ptr validity after graph replay. + // (sample_indices_i64 from the last sample_proportional call). This + // avoids depending on trainer_actions_ptr validity after graph replay. let mbs = self.config.max_batch_size.max(1); let gather_n = bs.min(mbs); let gather_ni = gather_n as i32; - // Reuse sample_actions (CudaSlice) as scratch. gather_u32 gathers - // the action values at the sampled indices from the ring buffer. let cap_i32 = self.config.capacity as i32; unsafe { stream.launch_builder(&self.kernels.gather_u32) @@ -857,15 +883,23 @@ impl GpuReplayBuffer { .launch(lcfg(gather_n)) .map_err(|e| MLError::ModelError(format!("gather_u32 for diversity: {e}")))?; } - // Synchronize to allow DtoH readback (epoch-boundary, not on hot path). + + // Single-block deterministic tree reduction → pinned i32 slot. + // 256 threads × 8 bytes (long long sdata) = 2 KB shared mem. + let mean_dev_ptr = self.mean_action_dev_ptr; + let reduce_cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256 * std::mem::size_of::() as u32, + }; unsafe { - cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); + stream.launch_builder(&self.kernels.actions_sum_scale_reduce_u32) + .arg(&self.sample_actions) + .arg(&mean_dev_ptr) + .arg(&gather_ni) + .launch(reduce_cfg) + .map_err(|e| MLError::ModelError(format!("actions_sum_scale_reduce_u32: {e}")))?; } - let mut actions_host = vec![0u32; gather_n]; - stream.memcpy_dtoh(&self.sample_actions.slice(..gather_n), &mut actions_host) - .map_err(|e| MLError::ModelError(format!("actions dtoh diversity: {e}")))?; - let sum_actions: i64 = actions_host.iter().map(|&a| a as i64).sum(); - let mean_action_scaled: i32 = ((sum_actions * 1000) / (gather_n as i64).max(1)) as i32; // Reinterpret sample_actions (u32) as i32 for the kernel argument. // The pointer type pun is safe: same device memory, same bit-width, kernel reads int*. @@ -878,7 +912,7 @@ impl GpuReplayBuffer { .arg(indices) .arg(td_errors) .arg(&actions_dev_ptr) - .arg(&mean_action_scaled) + .arg(&mean_dev_ptr) .arg(&al) .arg(&ep) .arg(&health) diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index 9de823983..2c3b8ecd3 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -471,6 +471,46 @@ void f32_idx_to_u32( out[i] = (unsigned int)input[i]; } +/* Graph-safe GPU reduction: scaled-mean of u32 actions → i32[1]. + * + * Writes (sum(actions) * 1000) / n into mean_scaled_out[0]. Launched as a + * single 256-thread block; deterministic tree reduction in shared memory. + * Replaces the prior dtoh-then-sum-on-host path, which was not graph-safe: + * memcpy_dtoh inside a CUDA Graph capture region invalidates the stream + * (L40S smoke surfaced this on Fold 1 re-capture). Keep n small — called + * with batch-sized windows (≤ max_batch_size), one block handles the lot. + */ +extern "C" __global__ void actions_sum_scale_reduce_u32( + const unsigned int* __restrict__ actions, + int* __restrict__ mean_scaled_out, /* i32[1]: (sum * 1000) / n */ + int n +) { + extern __shared__ long long sdata_ll[]; + int tid = threadIdx.x; + + long long local = 0; + for (int i = tid; i < n; i += blockDim.x) { + local += (long long)actions[i]; + } + sdata_ll[tid] = local; + __syncthreads(); + + /* Deterministic tree reduction — no atomics, no warp shuffles with + * data-dependent lane ordering. */ + for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) { + if (tid < offset) { + sdata_ll[tid] += sdata_ll[tid + offset]; + } + __syncthreads(); + } + + if (tid == 0) { + long long denom = n > 0 ? (long long)n : 1LL; + long long mean_scaled = (sdata_ll[0] * 1000LL) / denom; + mean_scaled_out[0] = (int)mean_scaled; + } +} + /* C1/P1: Health-weighted PER priority — boosts priorities of diverse-action * experiences during collapse (health < 0.8). * @@ -481,8 +521,9 @@ void f32_idx_to_u32( * priorities_pa[indices[i]] = new_prio^alpha * atomicMax(batch_max, new_prio) (IEEE 754 int trick) * - * `mean_action_scaled` is the batch mean of actions multiplied by 1000 (passed - * as int to avoid precision issues when reducing host-side). + * `mean_action_scaled_ptr` points to a device-side i32 set by the companion + * kernel `actions_sum_scale_reduce_u32` (same batch, scaled mean × 1000). + * Device-pointer passthrough keeps the update graph-safe — no host readback. * * Drop-in replacement for per_update_pa when health < 0.8. */ @@ -493,7 +534,7 @@ extern "C" __global__ void pow_alpha_diverse_f32( const unsigned int* __restrict__ indices, const float* __restrict__ td_errors, const int* __restrict__ actions, - int mean_action_scaled, + const int* __restrict__ mean_action_scaled_ptr, /* device i32[1] */ float alpha, float epsilon, float health, @@ -508,7 +549,13 @@ extern "C" __global__ void pow_alpha_diverse_f32( float base = powf(fabsf(td_errors[i]) + epsilon, alpha); - float mean_action = (float)mean_action_scaled * 0.001f; + /* Plain load (not __ldg): the backing buffer is device-mapped pinned host + * memory (cuMemHostAlloc + cuMemHostGetDevicePointer). __ldg routes through + * the read-only cache which is not guaranteed coherent for mapped host + * memory — stale reads manifested as a 13× drop in final q_gap on the + * multi-trial smoke. Plain load respects stream serialization between + * actions_sum_scale_reduce_u32 (writer) and this kernel (reader). */ + float mean_action = (float)(mean_action_scaled_ptr[0]) * 0.001f; float action_diff = fabsf((float)actions[i] - mean_action); float diversity_mult = 1.0f + 2.0f * (1.0f - health) * action_diff; diff --git a/infra/k8s/argo/train-template.yaml b/infra/k8s/argo/train-template.yaml index dc31434c7..0005f2c92 100644 --- a/infra/k8s/argo/train-template.yaml +++ b/infra/k8s/argo/train-template.yaml @@ -641,12 +641,15 @@ spec: esac echo "=== Evaluating: $MODEL ===" + # Args must match evaluate_baseline's Args struct — see crates/ml/examples/evaluate_baseline.rs + # `--models-dir` (not --checkpoint-dir), `--output` is a FILE path (not --output-dir). + mkdir -p /workspace/output/eval ${BINARY} \ --model "$MODEL" \ --symbol {{workflow.parameters.symbol}} \ --data-dir /data/futures-baseline \ - --checkpoint-dir /workspace/output \ - --output-dir /workspace/output/eval \ + --models-dir /workspace/output \ + --output /workspace/output/eval/evaluation_report.json \ --initial-capital {{workflow.parameters.initial-capital}} \ --tx-cost-bps {{workflow.parameters.tx-cost-bps}} \ --tick-size {{workflow.parameters.tick-size}} \