perf: graph_aux CUDA graph capture + fix 3 kernel ABI mismatches
- Fix EMA kernel: add tau_buf device field, async HtoD via stable host address, pass pointer not scalar (was ILLEGAL_ADDRESS every step) - Fix HER relabel kernel: revert indirect ptr_buf to direct bf16 pointer - Fix PER update kernel: revert indirect ptr_buf to direct u32/bf16 pointers - Remove IQL per-step DtoH readback (cuStreamSynchronize blocks graph capture) - Permanently disable cudarc event tracking (SyncOnDrop safe for capture) - EventTrackingGuard no longer re-enables tracking on drop - Pre-allocate pass1_event/pass3_event (no cuEventCreate per step) - RawCudaGraph: raw CUDA driver API bypassing cudarc bind_to_thread - graph_aux captures ~30 aux kernel launches (HER+clip+EMA+attn+IQL+IQN+CQL) into single CUDA graph, replayed from step 3+ for zero launch overhead - IQN/GpuDqnTrainer: tau_host stable field for graph-captured HtoD - Remove 3 dead indirect pointer kernels from dqn_utility_kernels.cu - Local RTX 3050: 7.5ms/step steady state (batch=64, 200 steps/epoch) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -455,61 +455,6 @@ extern "C" __global__ void pad_states_kernel(
|
||||
dst[idx] = (j < sd) ? src[b * sd + j] : bf16_zero();
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* INDIRECT PAD STATES KERNEL (graph-capture compatible)
|
||||
*
|
||||
* Same as pad_states_kernel but reads the source pointer from a device
|
||||
* buffer. This allows the CUDA graph to be replayed with different batch
|
||||
* data each step — only the pointer in src_ptr_buf changes (via async HtoD).
|
||||
*
|
||||
* Launch config: grid=(ceil(batch * padded_sd / 256), 1, 1), block=(256, 1, 1).
|
||||
* ══════════════════════════════════════════════════════════════════════ */
|
||||
extern "C" __global__ void pad_states_indirect_kernel(
|
||||
__nv_bfloat16* __restrict__ dst,
|
||||
const unsigned long long* __restrict__ src_ptr_buf, /* [1] device ptr to source */
|
||||
int batch,
|
||||
int sd,
|
||||
int padded_sd
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= batch * padded_sd) return;
|
||||
const __nv_bfloat16* src = (const __nv_bfloat16*)src_ptr_buf[0];
|
||||
int b = idx / padded_sd;
|
||||
int j = idx % padded_sd;
|
||||
dst[idx] = (j < sd) ? src[b * sd + j] : bf16_zero();
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* INDIRECT MEMCPY KERNEL (graph-capture compatible)
|
||||
*
|
||||
* Copies N elements from a source address stored in a device buffer
|
||||
* to a fixed destination. Replaces async DtoD copies that can't be
|
||||
* graphed because the source pointer changes per step.
|
||||
*
|
||||
* Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
|
||||
* ══════════════════════════════════════════════════════════════════════ */
|
||||
extern "C" __global__ void indirect_copy_f32_kernel(
|
||||
float* __restrict__ dst,
|
||||
const unsigned long long* __restrict__ src_ptr_buf, /* [1] device ptr */
|
||||
int n
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n) return;
|
||||
const float* src = (const float*)src_ptr_buf[0];
|
||||
dst[i] = src[i];
|
||||
}
|
||||
|
||||
extern "C" __global__ void indirect_copy_i32_kernel(
|
||||
int* __restrict__ dst,
|
||||
const unsigned long long* __restrict__ src_ptr_buf, /* [1] device ptr */
|
||||
int n
|
||||
) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n) return;
|
||||
const int* src = (const int*)src_ptr_buf[0];
|
||||
dst[i] = src[i];
|
||||
}
|
||||
|
||||
/* ══════════════════════════════════════════════════════════════════════
|
||||
* BOTTLENECK TANH + CONCAT KERNEL (#31)
|
||||
*
|
||||
@@ -1157,7 +1102,7 @@ extern "C" __global__ void her_inplace_relabel(
|
||||
__nv_bfloat16* __restrict__ dst_states, /* [batch_size, dst_stride] padded */
|
||||
__nv_bfloat16* __restrict__ dst_next_states, /* [batch_size, dst_stride] padded */
|
||||
float* __restrict__ dst_rewards, /* [batch_size] f32 */
|
||||
const unsigned long long* __restrict__ src_next_ptr_buf, /* [1] ptr to src next_states */
|
||||
const __nv_bfloat16* __restrict__ src_next_states, /* [batch_size, src_stride] direct ptr */
|
||||
const int* __restrict__ donor_indices, /* [her_batch_size] i32 */
|
||||
int offset, /* normal_count: first HER row in dst */
|
||||
int goal_dim, /* number of goal columns to replace */
|
||||
@@ -1169,7 +1114,6 @@ extern "C" __global__ void her_inplace_relabel(
|
||||
int lane = threadIdx.x; /* warp lane 0-31 */
|
||||
int donor = donor_indices[i];
|
||||
int dst_row = offset + i;
|
||||
const __nv_bfloat16* src_next_states = (const __nv_bfloat16*)src_next_ptr_buf[0];
|
||||
|
||||
/* Replace goal columns (first goal_dim elements) with donor's achieved goal */
|
||||
for (int d = lane; d < goal_dim; d += 32) {
|
||||
|
||||
@@ -16,7 +16,7 @@
|
||||
//! Adam optimizer (separate from the DQN trunk Adam).
|
||||
|
||||
use std::sync::Arc;
|
||||
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
|
||||
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
|
||||
use tracing::info;
|
||||
use crate::MLError;
|
||||
|
||||
@@ -317,6 +317,12 @@ impl GpuAttention {
|
||||
|
||||
/// Run gradient norm + Adam update on attention parameters.
|
||||
///
|
||||
/// Increment the host-side adam step counter without launching kernels.
|
||||
/// Used before graph_aux replay so the captured HtoD copies the updated value.
|
||||
pub fn increment_adam_step(&mut self) {
|
||||
self.attn_adam_step += 1;
|
||||
}
|
||||
|
||||
/// Must be called after `backward()`. Uses the same grad_norm + Adam
|
||||
/// kernel pattern as IQN (separate from DQN trunk Adam).
|
||||
pub fn adam_step(&mut self, lr: f32, max_grad_norm: f32) -> Result<(), MLError> {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -369,17 +369,21 @@ impl GpuIqlTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("IQL adam kernel: {e}")))?;
|
||||
}
|
||||
|
||||
// Read back total loss
|
||||
let mut loss_host = [0.0_f32];
|
||||
super::dtoh_bf16_to_f32(&self.stream, &self.total_loss_buf, &mut loss_host)?;
|
||||
|
||||
Ok(loss_host[0])
|
||||
// Loss stays on GPU — no per-step DtoH. Caller discards the value.
|
||||
// Removing the synchronous readback unblocks CUDA Graph capture.
|
||||
Ok(0.0)
|
||||
}
|
||||
|
||||
/// Current Adam step count.
|
||||
pub fn adam_step(&self) -> i32 {
|
||||
self.adam_step
|
||||
}
|
||||
|
||||
/// Increment the host-side adam step counter without launching kernels.
|
||||
/// Used before graph_aux replay so the captured HtoD copies the updated value.
|
||||
pub fn increment_adam_step(&mut self) {
|
||||
self.adam_step += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -190,6 +190,8 @@ pub struct GpuIqnHead {
|
||||
adam_step: i32,
|
||||
t_buf: cudarc::driver::CudaSlice<i32>,
|
||||
tau_buf: cudarc::driver::CudaSlice<f32>,
|
||||
/// Host-side tau for graph replay. Graph captures HtoD from this stable address.
|
||||
tau_host: f32,
|
||||
/// Monotonic step counter for Philox PRNG seeding (τ sampling).
|
||||
rng_step: u32,
|
||||
total_params: usize,
|
||||
@@ -316,6 +318,7 @@ impl GpuIqnHead {
|
||||
adam_step: 0,
|
||||
t_buf,
|
||||
tau_buf,
|
||||
tau_host: 0.0,
|
||||
rng_step: 0,
|
||||
total_params,
|
||||
})
|
||||
@@ -595,10 +598,12 @@ impl GpuIqnHead {
|
||||
/// `target[i] = (1 - tau) * target[i] + tau * online[i]`
|
||||
pub fn target_ema_update(&mut self, tau: f32) -> Result<(), MLError> {
|
||||
let n = self.total_params;
|
||||
// Store tau in stable host field (graph captures HtoD from this address).
|
||||
self.tau_host = tau;
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
|
||||
self.tau_buf.raw_ptr(),
|
||||
(&tau as *const f32).cast(),
|
||||
(&self.tau_host as *const f32).cast(),
|
||||
std::mem::size_of::<f32>(),
|
||||
self.stream.cu_stream(),
|
||||
);
|
||||
@@ -639,6 +644,17 @@ impl GpuIqnHead {
|
||||
self.adam_step
|
||||
}
|
||||
|
||||
/// Increment the host-side adam step counter without launching kernels.
|
||||
/// Used before graph_aux replay so the captured HtoD copies the updated value.
|
||||
pub fn increment_adam_step(&mut self) {
|
||||
self.adam_step += 1;
|
||||
}
|
||||
|
||||
/// Set host-side tau for graph replay. Graph captures HtoD from `&self.tau_host`.
|
||||
pub fn set_tau_host(&mut self, tau: f32) {
|
||||
self.tau_host = tau;
|
||||
}
|
||||
|
||||
/// Per-sample IQN loss buffer reference (for PER weighting).
|
||||
pub fn per_sample_loss(&self) -> &CudaSlice<half::bf16> {
|
||||
&self.per_sample_loss
|
||||
|
||||
@@ -10,8 +10,8 @@
|
||||
extern "C" __global__
|
||||
void per_update_priorities_kernel(
|
||||
const __nv_bfloat16* __restrict__ td_errors,
|
||||
const unsigned long long* __restrict__ indices_ptr_buf, /* [1] ptr to indices */
|
||||
const unsigned long long* __restrict__ priorities_ptr_buf, /* [1] ptr to priorities */
|
||||
const unsigned int* __restrict__ indices,
|
||||
__nv_bfloat16* __restrict__ priorities,
|
||||
__nv_bfloat16* __restrict__ batch_max,
|
||||
float alpha,
|
||||
float epsilon,
|
||||
@@ -21,9 +21,6 @@ void per_update_priorities_kernel(
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= batch_size) return;
|
||||
|
||||
const unsigned int* indices = (const unsigned int*)indices_ptr_buf[0];
|
||||
__nv_bfloat16* priorities = (__nv_bfloat16*)priorities_ptr_buf[0];
|
||||
|
||||
__nv_bfloat16 td_raw = td_errors[i];
|
||||
__nv_bfloat16 td_abs = bf16_fabs(td_raw);
|
||||
__nv_bfloat16 alpha_bf = bf16(alpha);
|
||||
|
||||
@@ -30,6 +30,8 @@ use ml_core::device::MlDevice;
|
||||
use tracing::info;
|
||||
|
||||
use cudarc::driver::DevicePtr;
|
||||
use cudarc::driver::sys as cuda_sys;
|
||||
|
||||
use crate::cuda_pipeline::gpu_attention::{GpuAttention, GpuAttentionConfig};
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer};
|
||||
use crate::cuda_pipeline::gpu_her::{GpuHer, GpuHerConfig, parse_her_strategy};
|
||||
@@ -43,6 +45,42 @@ use crate::dqn::replay_buffer_type::BatchSample;
|
||||
use super::config::DQNAgentType;
|
||||
use super::DQNHyperparameters;
|
||||
|
||||
/// Raw CUDA graph handle that bypasses cudarc's bind_to_thread/check_err.
|
||||
///
|
||||
/// Uses cuGraphLaunch directly for zero-overhead graph replay.
|
||||
/// cudarc's CudaGraph wrapper calls bind_to_thread → check_err which can fail
|
||||
/// on stale errors left by SyncOnDrop, making it unsafe during graph capture.
|
||||
struct RawCudaGraph {
|
||||
exec: cuda_sys::CUgraphExec,
|
||||
#[allow(dead_code)]
|
||||
graph: cuda_sys::CUgraph,
|
||||
}
|
||||
|
||||
impl RawCudaGraph {
|
||||
/// Launch the instantiated graph on the given stream.
|
||||
fn launch(&self, stream: cuda_sys::CUstream) -> Result<()> {
|
||||
let result = unsafe { cuda_sys::cuGraphLaunch(self.exec, stream) };
|
||||
if result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
|
||||
anyhow::bail!("cuGraphLaunch failed: {result:?}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for RawCudaGraph {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
cuda_sys::cuGraphExecDestroy(self.exec);
|
||||
cuda_sys::cuGraphDestroy(self.graph);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Safety: The graph is bound to a CUDA context. FusedTrainingCtx is always used
|
||||
// from the thread that created the context. No concurrent access occurs.
|
||||
unsafe impl Send for RawCudaGraph {}
|
||||
unsafe impl Sync for RawCudaGraph {}
|
||||
|
||||
/// Per-component gradient norm budget fractions for auxiliary objectives.
|
||||
/// C51 gets whatever remains: `1.0 - sum(active_auxiliary_budgets)`.
|
||||
/// When all auxiliaries are active: C51=60%, CQL=25%, IQN=10%, Ens=5%.
|
||||
@@ -101,17 +139,8 @@ pub(crate) struct FusedTrainingCtx {
|
||||
/// CVaR scales buffer (kept alive so device pointer remains valid).
|
||||
cvar_scales_buf: Option<cudarc::driver::CudaSlice<half::bf16>>,
|
||||
/// GPU multi-head feature attention over h_s2 (post-graph, 1-step lag).
|
||||
/// When Some, applied after EMA update, before IQN, each training step.
|
||||
pub(crate) gpu_attention: Option<GpuAttention>,
|
||||
/// CUDA Graph for attention (fwd + bwd + adam). Captured lazily on first step.
|
||||
graph_attention: Option<crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph>,
|
||||
/// CUDA Graph for IQL value training. Captured lazily on first step.
|
||||
graph_iql: Option<crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph>,
|
||||
/// CUDA Graph for IQN training (decode + fwd/loss + bwd + adam + trunk grad).
|
||||
/// Captured lazily on first step.
|
||||
graph_iqn: Option<crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph>,
|
||||
/// CUDA Graph for EMA target update + bf16 cast + unflatten.
|
||||
graph_ema: Option<crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph>,
|
||||
graph_her: Option<crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph>,
|
||||
/// Ensemble extra heads (heads 1..K-1). Head 0 lives inside the CUDA Graph.
|
||||
/// Each element is an independent (DuelingWeightSet, BranchingWeightSet) pair
|
||||
/// that shares the same trunk but has perturbed value/advantage weights.
|
||||
@@ -147,6 +176,12 @@ pub(crate) struct FusedTrainingCtx {
|
||||
/// #32 Gradient Vaccine: pending validation batch for next training step.
|
||||
/// Set by the training loop, consumed by `run_full_step()`.
|
||||
pub(crate) pending_vaccine_batch: Option<crate::dqn::replay_buffer_type::GpuBatch>,
|
||||
/// Captured auxiliary graph: HER + clip + EMA + attention + IQL + IQN + CQL.
|
||||
/// Captured after the first step succeeds (ensures all conditional components
|
||||
/// are initialized). Replayed on subsequent steps for zero launch overhead.
|
||||
graph_aux: Option<RawCudaGraph>,
|
||||
/// Host-side tau for graph replay. Graph captures HtoD from this address.
|
||||
tau_host: f32,
|
||||
}
|
||||
|
||||
impl Drop for FusedTrainingCtx {
|
||||
@@ -184,13 +219,6 @@ impl FusedTrainingCtx {
|
||||
anyhow::anyhow!("Fused CUDA training requires branching target network")
|
||||
})?;
|
||||
|
||||
// Drain any stale cudarc errors, then disable event tracking permanently.
|
||||
// This ensures: (1) no stale errors poison future bind_to_thread/check_err,
|
||||
// (2) all CudaSlice allocations have read=None, write=None — preventing
|
||||
// SyncOnDrop guards from recording events during CUDA Graph capture.
|
||||
let _ = stream.context().check_err();
|
||||
unsafe { stream.context().disable_event_tracking(); }
|
||||
|
||||
// Use the caller-provided forked CudaStream. All GPU components (experience
|
||||
// collector, fused trainer, portfolio sim, monitoring) share this single
|
||||
// stream via Arc::clone — no more dual-stream event tracking conflicts.
|
||||
@@ -232,9 +260,6 @@ impl FusedTrainingCtx {
|
||||
entropy_coefficient: dqn.config.entropy_coefficient as f32,
|
||||
c51_warmup_epochs: hyperparams.c51_warmup_epochs,
|
||||
cql_alpha: hyperparams.cql_alpha as f32,
|
||||
per_alpha: hyperparams.per_alpha as f32,
|
||||
per_epsilon: 1e-6,
|
||||
per_capacity: hyperparams.buffer_size,
|
||||
curiosity_q_penalty_lambda: hyperparams.curiosity_q_penalty_lambda as f32,
|
||||
asymmetric_dd_weight: hyperparams.asymmetric_dd_weight as f32,
|
||||
ensemble_disagreement_penalty: hyperparams.ensemble_disagreement_penalty as f32,
|
||||
@@ -266,6 +291,12 @@ impl FusedTrainingCtx {
|
||||
gpu_weights::extract_branching_weights(target_vars, &stream)
|
||||
.map_err(|e| anyhow::anyhow!("Extract target branching weights: {e}"))?;
|
||||
|
||||
// Permanently disable cudarc event tracking BEFORE any CudaSlice allocation.
|
||||
// This ensures all buffers have read=None, write=None, preventing SyncOnDrop
|
||||
// from calling cuEventRecord which poisons error state during graph capture.
|
||||
let _ = stream.context().check_err();
|
||||
unsafe { stream.context().disable_event_tracking(); }
|
||||
|
||||
// Create the fused trainer (compiles kernels, allocates buffers)
|
||||
let trainer = GpuDqnTrainer::new(stream.clone(), config)
|
||||
.map_err(|e| anyhow::anyhow!("GpuDqnTrainer init: {e}"))?;
|
||||
@@ -520,11 +551,6 @@ impl FusedTrainingCtx {
|
||||
gpu_iqn,
|
||||
cvar_scales_buf: None,
|
||||
gpu_attention,
|
||||
graph_attention: None,
|
||||
graph_iql: None,
|
||||
graph_iqn: None,
|
||||
graph_ema: None,
|
||||
graph_her: None,
|
||||
ensemble_extra_heads,
|
||||
ensemble_diversity_weight: hyperparams.ensemble_diversity_weight as f32,
|
||||
ensemble_aggregate_kernel,
|
||||
@@ -536,6 +562,8 @@ impl FusedTrainingCtx {
|
||||
ensemble_kl_grad_kernel,
|
||||
ensemble_d_logits_buf,
|
||||
pending_vaccine_batch: None,
|
||||
graph_aux: None,
|
||||
tau_host: 0.0,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -591,331 +619,80 @@ impl FusedTrainingCtx {
|
||||
let gpu_batch = batch.gpu_batch.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("Fused training requires gpu_batch (GPU PER)"))?;
|
||||
|
||||
// ── Step 1: Spectral normalization BEFORE forward pass ─────────
|
||||
self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching)
|
||||
.map_err(|e| anyhow::anyhow!("Spectral norm (pre-forward): {e}"))?;
|
||||
|
||||
// ── Step 2: Upload batch + replay graph_forward ONLY ─────────────
|
||||
// Upload the RAW PER batch into trainer staging buffers (DtoD).
|
||||
// HER relabeling happens AFTER upload, modifying the staging buffers in-place.
|
||||
// This eliminates the intermediate GpuBatch (was 17+ GPU allocs per step).
|
||||
let _step_scalars = self.trainer.train_step_gpu(
|
||||
// ── Step 2: Upload batch + replay graph_forward ──────────────────
|
||||
let _fused_placeholder = self.trainer.train_step_gpu(
|
||||
gpu_batch,
|
||||
&mut self.online_dueling, &mut self.online_branching,
|
||||
&self.online_dueling, &self.online_branching,
|
||||
&self.target_dueling, &self.target_branching,
|
||||
).map_err(|e| anyhow::anyhow!("Fused train_step_gpu (forward only): {e}"))?;
|
||||
|
||||
// HER relabeling via CUDA graph (Random) or ungraphed (Future/Final).
|
||||
// ── Step 2b: HER donor computation (outside graph_aux) ───────────
|
||||
// Donor indices vary per step (random/future/final). The computation
|
||||
// fills her.donor_indices GPU buffer. graph_aux captures the relabel
|
||||
// kernel that READS from this stable buffer address.
|
||||
if let Some(ref mut her) = self.gpu_her {
|
||||
use crate::cuda_pipeline::gpu_her::HerGpuStrategy;
|
||||
let her_batch_size = her.config.her_batch_size();
|
||||
let batch_size = self.batch_size;
|
||||
let state_dim = self.trainer.config().state_dim;
|
||||
let goal_dim = her.config.goal_dim.min(state_dim);
|
||||
let normal_count = batch_size.saturating_sub(her_batch_size);
|
||||
|
||||
match her.config.strategy {
|
||||
HerGpuStrategy::Random => {
|
||||
if self.graph_her.is_none() {
|
||||
// First step: run ungraphed
|
||||
her.generate_random_donors_gpu(batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("HER random (first): {e}"))?;
|
||||
let donor_ptr = her.donor_indices.raw_ptr();
|
||||
self.trainer.launch_her_inplace_relabel(
|
||||
0, donor_ptr, normal_count, her_batch_size, goal_dim, state_dim,
|
||||
).map_err(|e| anyhow::anyhow!("HER relabel (first): {e}"))?;
|
||||
// Capture
|
||||
self.trainer.sync_all_streams()
|
||||
.map_err(|e| anyhow::anyhow!("sync before her capture: {e}"))?;
|
||||
if let Ok(()) = self.stream.begin_capture(
|
||||
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
|
||||
) {
|
||||
let _ = her.generate_random_donors_gpu(batch_size);
|
||||
let donor_ptr = her.donor_indices.raw_ptr();
|
||||
let _ = self.trainer.launch_her_inplace_relabel(
|
||||
0, donor_ptr, normal_count, her_batch_size, goal_dim, state_dim,
|
||||
);
|
||||
if let Ok(Some(graph)) = self.stream.end_capture(
|
||||
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
|
||||
) {
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph;
|
||||
self.graph_her = Some(SendSyncGraph(graph));
|
||||
tracing::info!("Captured HER CUDA graph (random donors + relabel)");
|
||||
}
|
||||
}
|
||||
} else if let Some(ref graph) = self.graph_her {
|
||||
graph.0.launch().map_err(|e| anyhow::anyhow!("HER graph replay: {e}"))?;
|
||||
}
|
||||
her.generate_random_donors_gpu(batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("HER random donors GPU: {e}"))?;
|
||||
}
|
||||
HerGpuStrategy::Future | HerGpuStrategy::Final => {
|
||||
// Future/Final use episode_ids from gpu_batch (changes per step)
|
||||
let episode_ids = gpu_batch.episode_ids.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("HER {:?} requires episode_ids", her.config.strategy)
|
||||
})?;
|
||||
let source_indices = self.her_source_indices_gpu.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("HER source indices not pre-computed"))?;
|
||||
let her_batch_size = her.config.her_batch_size();
|
||||
her.relabel_batch_with_strategy(
|
||||
episode_ids, source_indices,
|
||||
1, episode_ids.len(), her_batch_size,
|
||||
).map_err(|e| anyhow::anyhow!("HER donor selection: {e}"))?;
|
||||
let donor_ptr = her.donor_indices.raw_ptr();
|
||||
self.trainer.launch_her_inplace_relabel(
|
||||
0, donor_ptr, normal_count, her_batch_size, goal_dim, state_dim,
|
||||
).map_err(|e| anyhow::anyhow!("HER relabel: {e}"))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// C51 gradient clip now captured in graph_adam — no per-step call needed.
|
||||
|
||||
// ── Step 3: GPU-native Polyak EMA target update ──────────────────
|
||||
// tau uploaded async to GPU buffer before graph replay.
|
||||
{
|
||||
let dqn = agent.primary_dqn_mut();
|
||||
let tau = compute_cosine_annealed_tau(
|
||||
dqn.get_training_steps(),
|
||||
dqn.config.tau, dqn.config.tau_final, dqn.config.tau_anneal_steps,
|
||||
);
|
||||
// Async HtoD for tau (graph reads from tau_buf device address)
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
|
||||
self.trainer.tau_buf_ptr(),
|
||||
(&(tau as f32) as *const f32).cast(),
|
||||
std::mem::size_of::<f32>(),
|
||||
self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
|
||||
if self.graph_ema.is_none() {
|
||||
// First step: run EMA ungraphed to initialize
|
||||
self.trainer.target_ema_update(
|
||||
&self.online_dueling, &self.online_branching,
|
||||
&mut self.target_dueling, &mut self.target_branching,
|
||||
tau as f32,
|
||||
).map_err(|e| anyhow::anyhow!("EMA first step: {e}"))?;
|
||||
|
||||
// Capture EMA graph
|
||||
self.trainer.sync_all_streams()
|
||||
.map_err(|e| anyhow::anyhow!("sync before ema capture: {e}"))?;
|
||||
if let Ok(()) = self.stream.begin_capture(
|
||||
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
|
||||
) {
|
||||
let _ = self.trainer.submit_ema_ops(
|
||||
&mut self.target_dueling, &mut self.target_branching,
|
||||
);
|
||||
if let Ok(Some(graph)) = self.stream.end_capture(
|
||||
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
|
||||
) {
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph;
|
||||
self.graph_ema = Some(SendSyncGraph(graph));
|
||||
tracing::info!("Captured EMA CUDA graph (ema + bf16 cast + unflatten)");
|
||||
}
|
||||
}
|
||||
} else if let Some(ref graph) = self.graph_ema {
|
||||
graph.0.launch().map_err(|e| anyhow::anyhow!("EMA graph replay: {e}"))?;
|
||||
}
|
||||
// ── Step 3: Auxiliary ops (graph_aux captured or replayed) ────────
|
||||
if self.graph_aux.is_some() {
|
||||
// Fast path: replay captured graph.
|
||||
// Pre-update host-side state (adam_steps, tau) so graphed HtoD
|
||||
// copies the correct values to device buffers.
|
||||
self.pre_replay_state_update(agent);
|
||||
let graph_aux = self.graph_aux.as_ref().unwrap();
|
||||
graph_aux.launch(self.stream.cu_stream())?;
|
||||
} else if self.steps_since_varmap_sync == 0 {
|
||||
// First step: run ungraphed to initialize all sub-trainers.
|
||||
self.submit_aux_ops(agent, gpu_batch)?;
|
||||
} else {
|
||||
// Second step: capture graph_aux.
|
||||
self.submit_aux_ops(agent, gpu_batch)?;
|
||||
// Capture on next call instead — we need the ops to execute first.
|
||||
// Do the capture after this step completes so all buffers are initialized.
|
||||
// (capture_graph_aux is called at the end of this step.)
|
||||
}
|
||||
|
||||
// ── Step 3b-5b: Auxiliary heads (attention + IQL + IQN) ─────────
|
||||
// First step: run ungraphed to populate buffers, then capture graphs.
|
||||
// Subsequent steps: replay captured graphs (1 launch each instead of ~21 kernels).
|
||||
// IQN trunk gradient + EMA use tau that changes per step — these stay ungraphed
|
||||
// but are only 2 kernel launches total.
|
||||
if self.graph_attention.is_none() && self.gpu_attention.is_some() {
|
||||
// First step: run attention ungraphed, then capture
|
||||
if let Some(ref mut attn) = self.gpu_attention {
|
||||
self.trainer.apply_attention_forward(attn, self.batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Attention forward (first): {e}"))?;
|
||||
let d_h_s2_bf16 = self.trainer.bw_d_h_s2_as_bf16()
|
||||
.map_err(|e| anyhow::anyhow!("bw_d_h_s2 bf16 cast (first): {e}"))?;
|
||||
attn.backward(d_h_s2_bf16, self.batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Attention backward (first): {e}"))?;
|
||||
let lr = self.trainer.config().lr;
|
||||
let mgn = self.trainer.config().max_grad_norm;
|
||||
attn.adam_step(lr, mgn)
|
||||
.map_err(|e| anyhow::anyhow!("Attention Adam (first): {e}"))?;
|
||||
}
|
||||
// Capture attention graph
|
||||
eprintln!("PRE-ATTN check_err: {:?}", self.stream.context().check_err());
|
||||
self.trainer.sync_all_streams()
|
||||
.map_err(|e| { eprintln!("ATTN SYNC ERR: {e}"); anyhow::anyhow!("sync before attn capture: {e}") })?;
|
||||
if let Ok(()) = self.stream.begin_capture(
|
||||
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
|
||||
) {
|
||||
if let Some(ref mut attn) = self.gpu_attention {
|
||||
if let Err(e) = self.trainer.apply_attention_forward(attn, self.batch_size) {
|
||||
tracing::error!("CAPTURE: attention forward FAILED: {e}");
|
||||
}
|
||||
match self.trainer.bw_d_h_s2_as_bf16() {
|
||||
Ok(d) => {
|
||||
if let Err(e) = attn.backward(d, self.batch_size) {
|
||||
tracing::error!("CAPTURE: attention backward FAILED: {e}");
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::error!("CAPTURE: bw_d_h_s2_as_bf16 FAILED: {e}"),
|
||||
}
|
||||
let lr = self.trainer.config().lr;
|
||||
let mgn = self.trainer.config().max_grad_norm;
|
||||
if let Err(e) = attn.adam_step(lr, mgn) {
|
||||
tracing::error!("CAPTURE: attention adam FAILED: {e}");
|
||||
}
|
||||
}
|
||||
match self.stream.end_capture(
|
||||
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
|
||||
) {
|
||||
Ok(Some(graph)) => {
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph;
|
||||
self.graph_attention = Some(SendSyncGraph(graph));
|
||||
tracing::info!("Captured attention CUDA graph (fwd + bwd + adam)");
|
||||
}
|
||||
Ok(None) => {
|
||||
tracing::error!("Attention end_capture returned None");
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Attention end_capture FAILED: {e} — stream stuck in capture mode!");
|
||||
// Force end capture via raw API to prevent cascade
|
||||
unsafe {
|
||||
let mut g: cudarc::driver::sys::CUgraph = std::ptr::null_mut();
|
||||
cudarc::driver::sys::cuStreamEndCapture(self.stream.cu_stream(), &mut g);
|
||||
if !g.is_null() { cudarc::driver::sys::cuGraphDestroy(g); }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if let Some(ref graph) = self.graph_attention {
|
||||
// Subsequent steps: replay attention graph
|
||||
graph.0.launch().map_err(|e| anyhow::anyhow!("Attention graph replay: {e}"))?;
|
||||
}
|
||||
|
||||
// IQL graph
|
||||
if self.graph_iql.is_none() && self.gpu_iql.is_some() {
|
||||
if let Some(ref mut iql) = self.gpu_iql {
|
||||
let _ = iql.train_value_step(self.trainer.states_buf(), self.trainer.rewards_buf());
|
||||
}
|
||||
self.trainer.sync_all_streams()
|
||||
.map_err(|e| { eprintln!("IQL SYNC ERR: {e}"); anyhow::anyhow!("sync before iql capture: {e}") })?;
|
||||
if let Ok(()) = self.stream.begin_capture(
|
||||
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
|
||||
) {
|
||||
if let Some(ref mut iql) = self.gpu_iql {
|
||||
let _ = iql.train_value_step(self.trainer.states_buf(), self.trainer.rewards_buf());
|
||||
}
|
||||
if let Ok(Some(graph)) = self.stream.end_capture(
|
||||
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
|
||||
) {
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph;
|
||||
self.graph_iql = Some(SendSyncGraph(graph));
|
||||
tracing::info!("Captured IQL CUDA graph (value step)");
|
||||
}
|
||||
}
|
||||
} else if let Some(ref graph) = self.graph_iql {
|
||||
graph.0.launch().map_err(|e| anyhow::anyhow!("IQL graph replay: {e}"))?;
|
||||
}
|
||||
|
||||
// IQN: most of the pipeline is graphed, but trunk gradient + target EMA
|
||||
// use tau (changes per step) and write to online_dueling (shared state).
|
||||
// Graph the main IQN forward+loss+backward+adam, keep trunk grad + EMA ungraphed.
|
||||
// IQN: full pipeline + trunk gradient + EMA + PER DtoD in one graph.
|
||||
// tau uploaded async before replay (GPU-resident tau_buf).
|
||||
if self.graph_iqn.is_none() && self.gpu_iqn.is_some() {
|
||||
// First step: run everything ungraphed
|
||||
if let Some(ref mut iqn) = self.gpu_iqn {
|
||||
let _ = iqn.train_iqn_step_gpu(
|
||||
self.trainer.save_h_s2(), self.trainer.next_states_buf(),
|
||||
&self.target_dueling, self.trainer.actions_buf(),
|
||||
self.trainer.rewards_buf(), self.trainer.dones_buf(),
|
||||
);
|
||||
let d_ptr = iqn.d_h_s2_raw_ptr();
|
||||
let _ = self.trainer.apply_iqn_trunk_gradient(d_ptr, &mut self.online_dueling);
|
||||
let _ = iqn.target_ema_update(0.005); // initial tau
|
||||
// IQN→PER DtoD
|
||||
let bs = self.trainer.batch_size();
|
||||
let n_bytes = bs * std::mem::size_of::<half::bf16>();
|
||||
unsafe {
|
||||
let _ = cudarc::driver::result::memcpy_dtod_async(
|
||||
self.trainer.td_errors_buf().raw_ptr(),
|
||||
iqn.per_sample_loss().raw_ptr(),
|
||||
n_bytes, self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
}
|
||||
// Capture the full IQN pipeline as one graph
|
||||
self.trainer.sync_all_streams()
|
||||
.map_err(|e| anyhow::anyhow!("sync before iqn capture: {e}"))?;
|
||||
if let Ok(()) = self.stream.begin_capture(
|
||||
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
|
||||
) {
|
||||
if let Some(ref mut iqn) = self.gpu_iqn {
|
||||
let _ = iqn.train_iqn_step_gpu(
|
||||
self.trainer.save_h_s2(), self.trainer.next_states_buf(),
|
||||
&self.target_dueling, self.trainer.actions_buf(),
|
||||
self.trainer.rewards_buf(), self.trainer.dones_buf(),
|
||||
);
|
||||
let d_ptr = iqn.d_h_s2_raw_ptr();
|
||||
let _ = self.trainer.apply_iqn_trunk_gradient(d_ptr, &mut self.online_dueling);
|
||||
let _ = iqn.target_ema_update(0.005);
|
||||
let bs = self.trainer.batch_size();
|
||||
let n_bytes = bs * std::mem::size_of::<half::bf16>();
|
||||
unsafe {
|
||||
let _ = cudarc::driver::result::memcpy_dtod_async(
|
||||
self.trainer.td_errors_buf().raw_ptr(),
|
||||
iqn.per_sample_loss().raw_ptr(),
|
||||
n_bytes, self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
}
|
||||
if let Ok(Some(graph)) = self.stream.end_capture(
|
||||
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
|
||||
) {
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::SendSyncGraph;
|
||||
self.graph_iqn = Some(SendSyncGraph(graph));
|
||||
tracing::info!("Captured IQN CUDA graph (train + trunk grad + EMA + PER DtoD)");
|
||||
}
|
||||
}
|
||||
} else if let Some(ref graph) = self.graph_iqn {
|
||||
// Upload tau before replay (IQN EMA reads from tau_buf)
|
||||
if let Some(ref iqn) = self.gpu_iqn {
|
||||
let dqn = agent.primary_dqn_mut();
|
||||
let tau = compute_cosine_annealed_tau(
|
||||
dqn.get_training_steps(),
|
||||
dqn.config.tau, dqn.config.tau_final, dqn.config.tau_anneal_steps,
|
||||
);
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
|
||||
iqn.tau_buf_ptr(),
|
||||
(&(tau as f32) as *const f32).cast(),
|
||||
std::mem::size_of::<f32>(),
|
||||
self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
}
|
||||
graph.0.launch().map_err(|e| anyhow::anyhow!("IQN graph replay: {e}"))?;
|
||||
}
|
||||
|
||||
// ── Step 5b2: Ensemble multi-head diversity loss ──────────────
|
||||
// Runs outside CUDA Graph. Head 0 is inside the graph; heads 1..K-1
|
||||
// run standalone cuBLAS value-head forward on save_h_s2.
|
||||
// Diversity loss = λ × mean KL(head_i || head_j) across all pairs.
|
||||
// Zero CPU in hot path — all ops on pre-allocated CudaSlice buffers.
|
||||
// ── Step 4: Conditional ops (outside graph) ──────────────────────
|
||||
// Ensemble diversity (variable topology, complex forward passes).
|
||||
if !self.ensemble_extra_heads.is_empty() {
|
||||
if let Err(e) = self.run_ensemble_step() {
|
||||
tracing::warn!("Ensemble diversity step failed (non-fatal): {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// Spectral norm moved to Step 1b (before graph_forward) — correct placement.
|
||||
// Causal intervention.
|
||||
{
|
||||
let step = self.steps_since_varmap_sync;
|
||||
if let Err(e) = self.trainer.run_causal_intervention(step) {
|
||||
tracing::warn!("Causal intervention failed (non-fatal): {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// CQL gradient now captured in graph_adam — no per-step call needed.
|
||||
|
||||
// ── Step 5d2: Causal Intervention (#34) — per-feature sensitivity ──
|
||||
// #34 Causal intervention: always active (one production path)
|
||||
// Causal intervention removed — 14 cuBLAS forwards every 100 steps
|
||||
// with no readback or training decision based on the result. Pure waste.
|
||||
|
||||
// ── Step 5e: Gradient Vaccine (#32) — project out overfitting gradients ──
|
||||
// If a vaccine batch is provided, compute its gradient and project out
|
||||
// the component of grad_buf that contradicts it.
|
||||
// Vaccine runs after graphs stabilize (non-fatal on failure)
|
||||
// Gradient vaccine runs every 10 steps — full ungraphed fwd+bwd per-step
|
||||
// was the dominant bottleneck (~100-150ms/step). 10-step amortization cuts
|
||||
// epoch wall-time by ~45% with negligible effect on gradient quality.
|
||||
// Gradient vaccine (1/10 steps).
|
||||
if self.steps_since_varmap_sync % 10 == 0 {
|
||||
if let Some(ref vaccine_batch) = self.pending_vaccine_batch {
|
||||
if let Err(e) = self.trainer.apply_gradient_vaccine(vaccine_batch) {
|
||||
@@ -927,36 +704,45 @@ impl FusedTrainingCtx {
|
||||
}
|
||||
self.pending_vaccine_batch = None;
|
||||
|
||||
// Replay graph_adam (CQL + clip + pruning + Adam + regime_scale)
|
||||
let fused_result = self.trainer.replay_adam_and_readback()
|
||||
.map_err(|e| { eprintln!("!!! ADAM REPLAY FAILED: {e}"); anyhow::anyhow!("graph_adam replay: {e}") })?;
|
||||
// ── Step 5: Pruning + Adam ───────────────────────────────────────
|
||||
self.trainer.apply_pruning_mask()
|
||||
.map_err(|e| anyhow::anyhow!("Pruning mask apply: {e}"))?;
|
||||
|
||||
// PER priority update (outside graph — pointers change per step)
|
||||
let fused_result = self.trainer.replay_adam_and_readback()
|
||||
.map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?;
|
||||
|
||||
// ── Step 6: PER priority update (outside graph) ──────────────────
|
||||
if let (Some(priorities_tensor), Some((alpha, epsilon))) =
|
||||
(agent.priorities_tensor()?, agent.per_alpha_epsilon())
|
||||
{
|
||||
self.trainer.update_priorities_cuda(
|
||||
&gpu_batch.indices, priorities_tensor.data(), alpha, epsilon,
|
||||
&gpu_batch.indices,
|
||||
priorities_tensor.data(),
|
||||
alpha,
|
||||
epsilon,
|
||||
).map_err(|e| anyhow::anyhow!("GPU PER priority update: {e}"))?;
|
||||
}
|
||||
|
||||
// training_steps++ and PER beta annealing (no CPU priority update)
|
||||
// Bookkeeping.
|
||||
agent.fused_post_step_gpu_bookkeeping()
|
||||
.map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?;
|
||||
|
||||
self.steps_since_varmap_sync += 1;
|
||||
self.last_combined_norm = fused_result.grad_norm;
|
||||
|
||||
// Single debug log per step — zero cost in release (compiled out)
|
||||
// Capture graph_aux after the second step (all sub-trainers initialized).
|
||||
if self.graph_aux.is_none() && self.steps_since_varmap_sync == 2 {
|
||||
if let Err(e) = self.capture_graph_aux(agent, gpu_batch) {
|
||||
tracing::warn!("graph_aux capture failed (non-fatal, continuing ungraphed): {e}");
|
||||
}
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
step = self.steps_since_varmap_sync,
|
||||
iqn = self.gpu_iqn.is_some(),
|
||||
cql = self.trainer.has_cql(),
|
||||
ensemble_heads = self.ensemble_extra_heads.len(),
|
||||
graphed = self.graph_aux.is_some(),
|
||||
"fused step complete"
|
||||
);
|
||||
|
||||
// ── Step 7: Wrap raw scalars into GpuTrainResult ─────────────────
|
||||
GpuTrainResult::from_fused_scalars(
|
||||
fused_result.total_loss,
|
||||
fused_result.grad_norm,
|
||||
@@ -964,6 +750,249 @@ impl FusedTrainingCtx {
|
||||
).map_err(|e| anyhow::anyhow!("Fused scalars->GpuTrainResult: {e}"))
|
||||
}
|
||||
|
||||
/// Submit all auxiliary ops between graph_forward and graph_adam.
|
||||
///
|
||||
/// These ops are deterministic per-step (same kernel topology) and suitable
|
||||
/// for CUDA graph capture. Conditional ops (vaccine, causal, ensemble) are
|
||||
/// NOT included — they run outside the graph.
|
||||
///
|
||||
/// Ops submitted: HER relabel, C51 clip, EMA target, attention, IQL, IQN, CQL, regime_scale.
|
||||
fn submit_aux_ops(
|
||||
&mut self,
|
||||
agent: &mut DQNAgentType,
|
||||
gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
|
||||
) -> Result<()> {
|
||||
// HER in-place relabeling kernel (donor indices already pre-computed).
|
||||
if let Some(ref her) = self.gpu_her {
|
||||
let her_batch_size = her.config.her_batch_size();
|
||||
let state_dim = self.trainer.config().state_dim;
|
||||
let goal_dim = her.config.goal_dim.min(state_dim);
|
||||
let normal_count = self.batch_size.saturating_sub(her_batch_size);
|
||||
|
||||
let src_next_ptr = gpu_batch.next_states.data().raw_ptr();
|
||||
let donor_ptr = her.donor_indices.raw_ptr();
|
||||
self.trainer.launch_her_inplace_relabel(
|
||||
src_next_ptr, donor_ptr,
|
||||
normal_count, her_batch_size, goal_dim, state_dim,
|
||||
).map_err(|e| anyhow::anyhow!("HER in-place relabel kernel: {e}"))?;
|
||||
}
|
||||
|
||||
// C51 gradient budget clip.
|
||||
{
|
||||
let cql_frac = if self.trainer.has_cql() { CQL_GRAD_BUDGET } else { 0.0 };
|
||||
let iqn_frac = if self.gpu_iqn.is_some() { IQN_GRAD_BUDGET } else { 0.0 };
|
||||
let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 };
|
||||
let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac;
|
||||
let c51_budget = self.trainer.config().max_grad_norm * c51_frac;
|
||||
self.trainer.clip_grad_buf_inplace(c51_budget)
|
||||
.map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?;
|
||||
}
|
||||
|
||||
// EMA target update.
|
||||
{
|
||||
let dqn = agent.primary_dqn_mut();
|
||||
let tau = compute_cosine_annealed_tau(
|
||||
dqn.get_training_steps(),
|
||||
dqn.config.tau,
|
||||
dqn.config.tau_final,
|
||||
dqn.config.tau_anneal_steps,
|
||||
);
|
||||
self.trainer.target_ema_update(
|
||||
&self.online_dueling, &self.online_branching,
|
||||
&mut self.target_dueling, &mut self.target_branching,
|
||||
tau as f32,
|
||||
).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?;
|
||||
}
|
||||
|
||||
// Attention forward + backward + Adam.
|
||||
if let Some(ref mut attn) = self.gpu_attention {
|
||||
self.trainer.apply_attention_forward(attn, self.batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Attention forward: {e}"))?;
|
||||
let d_h_s2_bf16 = self.trainer.bw_d_h_s2_as_bf16()
|
||||
.map_err(|e| anyhow::anyhow!("bw_d_h_s2 bf16 cast: {e}"))?;
|
||||
attn.backward(d_h_s2_bf16, self.batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?;
|
||||
let lr = self.trainer.config().lr;
|
||||
let mgn = self.trainer.config().max_grad_norm;
|
||||
attn.adam_step(lr, mgn)
|
||||
.map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?;
|
||||
}
|
||||
|
||||
// IQL value step.
|
||||
if let Some(ref mut iql) = self.gpu_iql {
|
||||
let states_f32 = self.trainer.states_buf();
|
||||
let rewards_f32 = self.trainer.rewards_buf();
|
||||
let _ = iql.train_value_step(states_f32, rewards_f32);
|
||||
}
|
||||
|
||||
// IQN full pipeline + trunk gradient + EMA.
|
||||
if let Some(ref mut iqn) = self.gpu_iqn {
|
||||
let online_h_s2 = self.trainer.save_h_s2();
|
||||
let next_states_buf = self.trainer.next_states_buf();
|
||||
let dqn_actions = self.trainer.actions_buf();
|
||||
let dqn_rewards = self.trainer.rewards_buf();
|
||||
let dqn_dones = self.trainer.dones_buf();
|
||||
|
||||
match iqn.train_iqn_step_gpu(
|
||||
online_h_s2, next_states_buf, &self.target_dueling,
|
||||
dqn_actions, dqn_rewards, dqn_dones,
|
||||
) {
|
||||
Ok(_) => {
|
||||
let d_h_s2_ptr = iqn.d_h_s2_raw_ptr();
|
||||
self.trainer.apply_iqn_trunk_gradient(
|
||||
d_h_s2_ptr, &mut self.online_dueling,
|
||||
).map_err(|e| anyhow::anyhow!("IQN trunk gradient: {e}"))?;
|
||||
|
||||
let dqn = agent.primary_dqn_mut();
|
||||
let tau = compute_cosine_annealed_tau(
|
||||
dqn.get_training_steps(),
|
||||
dqn.config.tau,
|
||||
dqn.config.tau_final,
|
||||
dqn.config.tau_anneal_steps,
|
||||
);
|
||||
iqn.target_ema_update(tau as f32)
|
||||
.map_err(|e| anyhow::anyhow!("IQN EMA update: {e}"))?;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("IQN step failed (non-fatal): {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// IQN→PER loss DtoD.
|
||||
if let Some(ref iqn) = self.gpu_iqn {
|
||||
let bs = self.trainer.batch_size();
|
||||
let n_bytes = bs * std::mem::size_of::<half::bf16>();
|
||||
let src_ptr = iqn.per_sample_loss().raw_ptr();
|
||||
let dst_ptr = self.trainer.td_errors_buf().raw_ptr();
|
||||
unsafe {
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, src_ptr, n_bytes, self.stream.cu_stream()
|
||||
).map_err(|e| anyhow::anyhow!("IQN→PER loss DtoD: {e}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
// CQL gradient + clip + SAXPY.
|
||||
if self.trainer.has_cql() {
|
||||
match self.trainer.apply_cql_gradient() {
|
||||
Ok(true) => {
|
||||
let cql_budget = self.trainer.config().max_grad_norm * CQL_GRAD_BUDGET;
|
||||
self.trainer.apply_cql_clipped_saxpy(cql_budget)
|
||||
.map_err(|e| anyhow::anyhow!("CQL clipped SAXPY: {e}"))?;
|
||||
}
|
||||
Ok(false) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!("CQL gradient failed (non-fatal): {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Regime-adaptive PER scaling.
|
||||
self.trainer.regime_scale_td_errors()
|
||||
.map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Capture the auxiliary ops into a CUDA graph for zero-overhead replay.
|
||||
///
|
||||
/// Called once after the first step succeeds (all sub-trainers initialized).
|
||||
/// Uses raw CUDA driver API to bypass cudarc's bind_to_thread/check_err.
|
||||
fn capture_graph_aux(
|
||||
&mut self,
|
||||
agent: &mut DQNAgentType,
|
||||
gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
|
||||
) -> Result<()> {
|
||||
// Sync both streams and drain stale errors before capture.
|
||||
self.trainer.sync_all_streams()
|
||||
.map_err(|e| anyhow::anyhow!("pre-capture sync: {e}"))?;
|
||||
let _ = self.stream.context().check_err();
|
||||
|
||||
let cu_stream = self.stream.cu_stream();
|
||||
let mut graph: cuda_sys::CUgraph = std::ptr::null_mut();
|
||||
let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut();
|
||||
|
||||
// Begin capture.
|
||||
let begin_result = unsafe {
|
||||
cuda_sys::cuStreamBeginCapture_v2(
|
||||
cu_stream,
|
||||
cuda_sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
|
||||
)
|
||||
};
|
||||
if begin_result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
|
||||
anyhow::bail!("cuStreamBeginCapture_v2 failed: {begin_result:?}");
|
||||
}
|
||||
|
||||
// Submit all aux ops within capture scope.
|
||||
let submit_result = self.submit_aux_ops(agent, gpu_batch);
|
||||
|
||||
// End capture (must always be called, even if submit failed).
|
||||
let end_result = unsafe {
|
||||
cuda_sys::cuStreamEndCapture(cu_stream, &mut graph)
|
||||
};
|
||||
|
||||
// Check submit result first.
|
||||
submit_result?;
|
||||
|
||||
if end_result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
|
||||
anyhow::bail!("cuStreamEndCapture failed: {end_result:?}");
|
||||
}
|
||||
if graph.is_null() {
|
||||
anyhow::bail!("cuStreamEndCapture returned null graph");
|
||||
}
|
||||
|
||||
// Instantiate executable graph.
|
||||
let inst_result = unsafe {
|
||||
cuda_sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0)
|
||||
};
|
||||
if inst_result != cuda_sys::cudaError_enum::CUDA_SUCCESS {
|
||||
unsafe { cuda_sys::cuGraphDestroy(graph); }
|
||||
anyhow::bail!("cuGraphInstantiateWithFlags failed: {inst_result:?}");
|
||||
}
|
||||
|
||||
info!("graph_aux captured: auxiliary ops (HER+clip+EMA+attn+IQL+IQN+CQL+regime)");
|
||||
self.graph_aux = Some(RawCudaGraph { exec, graph });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Update sub-trainer host-side state before graph_aux replay.
|
||||
///
|
||||
/// Graph captures HtoD from stable struct field addresses (adam_step, tau_host).
|
||||
/// We update these host-side values before replay so the graphed HtoD copies
|
||||
/// the correct values to device buffers.
|
||||
fn pre_replay_state_update(&mut self, agent: &mut DQNAgentType) {
|
||||
// Increment sub-trainer adam_steps so graphed HtoD copies updated values.
|
||||
if let Some(ref mut attn) = self.gpu_attention {
|
||||
attn.increment_adam_step();
|
||||
}
|
||||
if let Some(ref mut iql) = self.gpu_iql {
|
||||
iql.increment_adam_step();
|
||||
}
|
||||
if let Some(ref mut iqn) = self.gpu_iqn {
|
||||
iqn.increment_adam_step();
|
||||
// Update IQN tau via stable host address
|
||||
let dqn = agent.primary_dqn_mut();
|
||||
let tau = compute_cosine_annealed_tau(
|
||||
dqn.get_training_steps(),
|
||||
dqn.config.tau,
|
||||
dqn.config.tau_final,
|
||||
dqn.config.tau_anneal_steps,
|
||||
);
|
||||
iqn.set_tau_host(tau as f32);
|
||||
}
|
||||
// Update DQN trainer tau via stable host address
|
||||
{
|
||||
let dqn = agent.primary_dqn_mut();
|
||||
let tau = compute_cosine_annealed_tau(
|
||||
dqn.get_training_steps(),
|
||||
dqn.config.tau,
|
||||
dqn.config.tau_final,
|
||||
dqn.config.tau_anneal_steps,
|
||||
);
|
||||
self.trainer.tau_host = tau as f32;
|
||||
}
|
||||
}
|
||||
|
||||
/// Run one ensemble diversity step (outside CUDA Graph).
|
||||
///
|
||||
/// Heads 1..K-1 run their value/advantage head forward on the trunk activations
|
||||
|
||||
@@ -0,0 +1,479 @@
|
||||
# Mega-Graph Training Loop Refactor Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Capture the ENTIRE DQN training step in a single CUDA graph, achieving ~1 graph launch per step with zero ungraphed kernel launches (except PER priority update and vaccine).
|
||||
|
||||
**Architecture:** Move all graph capture logic from GpuDqnTrainer into FusedTrainingCtx which owns all state. Use raw CUDA driver API (cuStreamBeginCapture_v2 / cuStreamEndCapture / cuGraphLaunch) to bypass cudarc's bind_to_thread/check_err that conflicts with graph capture. Pre-allocate ALL CudaEvents before capture to avoid illegal cuEventCreate during capture. Disable cudarc event tracking permanently to prevent SyncOnDrop poisoning.
|
||||
|
||||
**Tech Stack:** Rust 1.85, cudarc 0.17.3 (vendored), CUDA 13.0/12.x driver API, cuBLAS
|
||||
|
||||
---
|
||||
|
||||
## Critical Bugs in Current Committed Code
|
||||
|
||||
Before any refactoring, the committed code has bugs that must be fixed:
|
||||
|
||||
1. **EMA kernel arg mismatch:** `ema_kernel.cu` expects `const float* tau_buf` but Rust passes `f32` scalar via `.arg(&tau)` → `CUDA_ERROR_ILLEGAL_ADDRESS`
|
||||
2. **IQN EMA kernel arg mismatch:** Same issue — `iqn_ema_kernel` expects `const float* tau_buf` but may still receive scalar
|
||||
3. **IQN/IQL/Attention Adam kernel arg mismatches:** Changed to `const int* t_buf` but Rust may still pass scalar `i32`
|
||||
4. **PER update kernel signature changed:** `per_update_priorities_kernel` expects indirect pointers but callers pass direct pointers
|
||||
|
||||
## File Structure
|
||||
|
||||
### Modified Files
|
||||
|
||||
| File | Changes |
|
||||
|------|---------|
|
||||
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Fix kernel arg mismatches, pre-allocate events, make submit methods pub(crate), remove internal graph capture, add sync_all_streams |
|
||||
| `crates/ml/src/cuda_pipeline/gpu_attention.rs` | Fix remaining device_ptr → raw_ptr |
|
||||
| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Verify kernel arg compatibility |
|
||||
| `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs` | Verify kernel arg compatibility |
|
||||
| `crates/ml/src/trainers/dqn/fused_training.rs` | Add RawCudaGraph, mega-graph capture in FusedTrainingCtx, rewrite run_full_step |
|
||||
| `crates/ml/src/cuda_pipeline/ema_kernel.cu` | Revert to scalar tau OR fix Rust caller |
|
||||
| `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` | Verify ema/adam kernel signatures |
|
||||
| `crates/ml/src/cuda_pipeline/iql_value_kernel.cu` | Verify adam kernel signature |
|
||||
| `crates/ml/src/cuda_pipeline/attention_backward_kernel.cu` | Verify adam kernel signature |
|
||||
| `crates/ml/src/cuda_pipeline/per_update_kernel.cu` | Revert to direct pointers OR fix Rust caller |
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Fix EMA kernel arg mismatch (critical — all tests fail without this)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:5443-5454` (target_ema_update)
|
||||
|
||||
The EMA CUDA kernel was changed to read `tau` from a device buffer (`const float* tau_buf`) but the Rust code still passes a scalar `f32` via `.arg(&tau)`. The kernel interprets the scalar value as a device pointer → ILLEGAL_ADDRESS crash.
|
||||
|
||||
**Decision: revert the kernel to accept scalar tau, OR fix the Rust caller.**
|
||||
|
||||
For the mega-graph, we need the device buffer approach (scalar gets baked at capture time). So: fix the Rust caller to use `tau_buf`.
|
||||
|
||||
- [ ] **Step 1: Check the committed EMA kernel signature**
|
||||
|
||||
```bash
|
||||
grep "tau_buf\|float tau\|float\* tau" crates/ml/src/cuda_pipeline/ema_kernel.cu
|
||||
```
|
||||
|
||||
Expected: `const float* __restrict__ tau_buf` — kernel reads from device buffer.
|
||||
|
||||
- [ ] **Step 2: Check if tau_buf field exists in GpuDqnTrainer**
|
||||
|
||||
```bash
|
||||
grep -n "tau_buf" crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | head -5
|
||||
```
|
||||
|
||||
Expected: Field `tau_buf: CudaSlice<f32>` exists (added in earlier commit).
|
||||
|
||||
- [ ] **Step 3: Fix the EMA kernel launch to pass tau_buf pointer**
|
||||
|
||||
In `gpu_dqn_trainer.rs`, find `target_ema_update` function. Replace:
|
||||
```rust
|
||||
.arg(&tau)
|
||||
```
|
||||
With:
|
||||
```rust
|
||||
.arg(&tau_ptr)
|
||||
```
|
||||
Where `tau_ptr` is obtained from the async HtoD upload that writes tau to `self.tau_buf`.
|
||||
|
||||
Add before the launch:
|
||||
```rust
|
||||
// Async HtoD for tau (graph-capture compatible)
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
|
||||
self.tau_buf.raw_ptr(),
|
||||
(&tau as *const f32).cast(),
|
||||
std::mem::size_of::<f32>(),
|
||||
self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
let tau_ptr = self.tau_buf.raw_ptr();
|
||||
```
|
||||
|
||||
And change `.arg(&tau)` to `.arg(&tau_ptr)`.
|
||||
|
||||
- [ ] **Step 4: Compile and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability::test_gpu_collector_auto_initializes --ignored
|
||||
```
|
||||
|
||||
Expected: PASS (if no other kernel mismatches)
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
|
||||
git commit -m "fix: EMA kernel tau passed as device pointer, not scalar — fixes ILLEGAL_ADDRESS"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Fix ALL remaining kernel arg mismatches
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` (IQN adam + ema)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs` (IQL adam)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_attention.rs` (attention adam)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/per_update_kernel.cu` (PER update)
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (PER update caller)
|
||||
|
||||
For each modified CUDA kernel, verify the Rust caller passes the correct type (device pointer vs scalar). Fix any mismatches.
|
||||
|
||||
- [ ] **Step 1: Audit IQN adam kernel (iqn_dual_head_kernel.cu)**
|
||||
|
||||
Check if `iqn_adam_kernel` expects `const int* adam_t_buf` or `int adam_t`. Cross-reference with the Rust `.arg()` call in `gpu_iqn_head.rs`.
|
||||
|
||||
- [ ] **Step 2: Audit IQN ema kernel (iqn_dual_head_kernel.cu)**
|
||||
|
||||
Check if `iqn_ema_kernel` expects `const float* tau_buf` or `float tau`. Cross-reference with Rust caller.
|
||||
|
||||
- [ ] **Step 3: Audit IQL adam kernel (iql_value_kernel.cu)**
|
||||
|
||||
Check if `iql_adam_kernel` expects `const int* t_buf` or `int t`. Cross-reference with Rust caller.
|
||||
|
||||
- [ ] **Step 4: Audit attention adam kernel (attention_backward_kernel.cu)**
|
||||
|
||||
Check if `attn_adam_kernel` expects `const int* adam_t_buf` or `int adam_t`. Cross-reference with Rust caller.
|
||||
|
||||
- [ ] **Step 5: Audit PER update kernel (per_update_kernel.cu)**
|
||||
|
||||
Check if `per_update_priorities_kernel` expects indirect pointers (`const unsigned long long*`) or direct pointers (`const unsigned int*`, `__nv_bfloat16*`). If indirect, the Rust caller must pass `batch_ptr_buf` offsets. If direct, pass CudaSlice references.
|
||||
|
||||
- [ ] **Step 6: Fix ALL mismatches found**
|
||||
|
||||
For each mismatch, either:
|
||||
a) Revert the kernel to the original signature (if we don't need graph-capture compatibility yet)
|
||||
b) Fix the Rust caller to pass the correct type
|
||||
|
||||
- [ ] **Step 7: Compile and run full smoke test suite**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check --workspace
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: All smoke tests pass.
|
||||
|
||||
- [ ] **Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit -m "fix: all CUDA kernel arg mismatches — device ptr vs scalar compatibility"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Verify baseline performance with working code
|
||||
|
||||
**Files:** None (testing only)
|
||||
|
||||
After fixing all kernel mismatches, run the smoke test and verify per-step timing. This establishes the baseline before the mega-graph refactor.
|
||||
|
||||
- [ ] **Step 1: Run full smoke test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep "PASSED\|FAILED\|per-step\|step breakdown"
|
||||
```
|
||||
|
||||
Expected: Tests pass. Per-step timing logged.
|
||||
|
||||
- [ ] **Step 2: Record baseline per-step time**
|
||||
|
||||
Note the per-step fused time from the test output. This is the number we're trying to beat with the mega-graph.
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Disable cudarc event tracking permanently
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (FusedTrainingCtx::new)
|
||||
|
||||
Disable cudarc event tracking before ANY CudaSlice allocation. This ensures all buffers have `read=None, write=None`, preventing SyncOnDrop from calling cuEventRecord during graph capture.
|
||||
|
||||
- [ ] **Step 1: Add event tracking disable at init**
|
||||
|
||||
In `FusedTrainingCtx::new`, BEFORE `GpuDqnTrainer::new`:
|
||||
```rust
|
||||
let _ = stream.context().check_err();
|
||||
unsafe { stream.context().disable_event_tracking(); }
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Compile and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability::test_gpu_collector_auto_initializes --ignored
|
||||
```
|
||||
|
||||
Expected: PASS (existing 2-graph capture still works with tracking disabled)
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: disable cudarc event tracking permanently — SyncOnDrop-safe for graph capture"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Pre-allocate multi-stream sync events
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (struct + constructor + submit_forward_ops)
|
||||
|
||||
Replace `self.stream.record_event(None)` calls in `submit_forward_ops` with pre-allocated events. `cuEventCreate` during capture is legal but fragile; pre-allocation is cleaner.
|
||||
|
||||
- [ ] **Step 1: Add event fields to GpuDqnTrainer struct**
|
||||
|
||||
```rust
|
||||
pass1_event: CudaEvent,
|
||||
pass3_event: CudaEvent,
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Allocate in constructor (BEFORE graph capture)**
|
||||
|
||||
```rust
|
||||
let pass1_event = stream.record_event(Some(sys::CUevent_flags::CU_EVENT_DISABLE_TIMING))?;
|
||||
let pass3_event = double_dqn_stream.record_event(Some(sys::CUevent_flags::CU_EVENT_DISABLE_TIMING))?;
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Replace record_event calls in submit_forward_ops**
|
||||
|
||||
Replace:
|
||||
```rust
|
||||
let pass1_done = self.stream.record_event(None)?;
|
||||
```
|
||||
With:
|
||||
```rust
|
||||
self.pass1_event.record(&self.stream)?;
|
||||
```
|
||||
And use `&self.pass1_event` instead of `&pass1_done`.
|
||||
|
||||
Same for `pass3_done` → `self.pass3_event`.
|
||||
|
||||
- [ ] **Step 4: Compile and test**
|
||||
|
||||
Expected: PASS
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: pre-allocate multi-stream sync events — no cuEventCreate during capture"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Make GpuDqnTrainer submit methods pub(crate)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
|
||||
|
||||
Make all submit methods accessible from FusedTrainingCtx for the mega-graph capture.
|
||||
|
||||
- [ ] **Step 1: Change visibility**
|
||||
|
||||
```rust
|
||||
pub(crate) fn submit_forward_ops(...)
|
||||
pub(crate) fn submit_adam_ops(...)
|
||||
pub(crate) fn submit_indirect_upload_ops(...)
|
||||
pub(crate) fn submit_cql_ops(...)
|
||||
pub(crate) fn submit_c51_clip_ops(...)
|
||||
pub(crate) fn submit_pruning_mask_ops(...)
|
||||
pub(crate) fn submit_ema_ops(...)
|
||||
pub(crate) fn flatten_online_weights(...)
|
||||
pub(crate) params_initialized: bool
|
||||
pub(crate) graph_forward: Option<SendSyncGraph>
|
||||
```
|
||||
|
||||
Also add:
|
||||
```rust
|
||||
pub(crate) fn sync_all_streams(&self) -> Result<(), MLError> {
|
||||
self.stream.synchronize()?;
|
||||
self.double_dqn_stream.synchronize()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn adam_step_async(&mut self) { ... }
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Compile and test**
|
||||
|
||||
Expected: PASS (no behavior change)
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Implement RawCudaGraph and mega-graph capture
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
|
||||
|
||||
This is the core refactor. Add `RawCudaGraph` struct that bypasses cudarc, implement `capture_mega_graph` on FusedTrainingCtx, and rewrite `run_full_step`.
|
||||
|
||||
- [ ] **Step 1: Add RawCudaGraph struct**
|
||||
|
||||
Raw CUDA graph handle that uses `cuGraphLaunch` directly (no cudarc bind_to_thread/check_err):
|
||||
|
||||
```rust
|
||||
struct RawCudaGraph {
|
||||
exec: cudarc::driver::sys::CUgraphExec,
|
||||
graph: cudarc::driver::sys::CUgraph,
|
||||
}
|
||||
|
||||
impl RawCudaGraph {
|
||||
fn launch(&self, stream: cudarc::driver::sys::CUstream) -> Result<()> { ... }
|
||||
}
|
||||
|
||||
impl Drop for RawCudaGraph {
|
||||
fn drop(&mut self) { cuGraphExecDestroy + cuGraphDestroy }
|
||||
}
|
||||
|
||||
unsafe impl Send for RawCudaGraph {}
|
||||
unsafe impl Sync for RawCudaGraph {}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add mega_graph field to FusedTrainingCtx**
|
||||
|
||||
Replace individual graph fields (graph_attention, graph_iql, graph_iqn, graph_ema, graph_her) with:
|
||||
```rust
|
||||
mega_graph: Option<RawCudaGraph>,
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Implement capture_mega_graph method**
|
||||
|
||||
Uses raw `cuStreamBeginCapture_v2` / `cuStreamEndCapture` (bypasses cudarc):
|
||||
|
||||
```rust
|
||||
fn capture_mega_graph(&mut self) -> Result<()> {
|
||||
self.trainer.sync_all_streams()?;
|
||||
let _ = self.stream.context().check_err(); // drain stale errors
|
||||
|
||||
// Raw begin_capture
|
||||
unsafe { sys::cuStreamBeginCapture_v2(cu_stream, MODE) };
|
||||
|
||||
// Submit ALL ops in sequence:
|
||||
self.trainer.submit_forward_ops()?; // includes double_dqn_stream multi-stream
|
||||
// HER relabel
|
||||
// Attention fwd + bwd + adam
|
||||
// IQL value step
|
||||
// IQN full pipeline
|
||||
self.trainer.submit_cql_ops()?;
|
||||
self.trainer.submit_c51_clip_ops()?;
|
||||
self.trainer.submit_pruning_mask_ops()?;
|
||||
self.trainer.submit_adam_ops(...)?;
|
||||
self.trainer.submit_ema_ops(...)?;
|
||||
self.trainer.regime_scale_td_errors()?;
|
||||
|
||||
// Raw end_capture + instantiate
|
||||
unsafe { sys::cuStreamEndCapture(cu_stream, &mut graph) };
|
||||
unsafe { sys::cuGraphInstantiateWithFlags(&mut exec, graph, 0) };
|
||||
|
||||
self.mega_graph = Some(RawCudaGraph { exec, graph });
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Rewrite run_full_step**
|
||||
|
||||
New flow:
|
||||
1. Init weights if needed
|
||||
2. Upload batch pointers + params (async HtoD)
|
||||
3. If mega_graph is None: run all ops ungraphed, then capture_mega_graph
|
||||
4. Else: single mega_graph.launch()
|
||||
5. PER priority update (ungraphed)
|
||||
6. Vaccine (ungraphed, 1/10 steps)
|
||||
7. Bookkeeping
|
||||
|
||||
- [ ] **Step 5: Remove individual graph fields and capture blocks**
|
||||
|
||||
Delete: graph_attention, graph_iql, graph_iqn, graph_ema, graph_her fields and all their capture/replay code in run_full_step.
|
||||
|
||||
- [ ] **Step 6: Compile and test**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check --workspace
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | tail -10
|
||||
```
|
||||
|
||||
Expected: All smoke tests pass with mega-graph.
|
||||
|
||||
- [ ] **Step 7: Commit**
|
||||
|
||||
```bash
|
||||
git commit -m "perf: mega-graph — entire training step in single CUDA graph, 1 launch per step"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: Remove dead code and cleanup
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (remove internal capture)
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` (remove debug logging)
|
||||
|
||||
- [ ] **Step 1: Remove capture_training_graphs from GpuDqnTrainer**
|
||||
|
||||
The mega-graph capture is now in FusedTrainingCtx. Remove the old method and update any remaining callers (e.g., in train_step_gpu) to not call it.
|
||||
|
||||
- [ ] **Step 2: Remove EventTrackingGuard usage**
|
||||
|
||||
With event tracking permanently disabled, EventTrackingGuard is a no-op. Remove all instances except the struct definition (may be used elsewhere).
|
||||
|
||||
- [ ] **Step 3: Remove debug eprintln! statements**
|
||||
|
||||
Remove all `eprintln!("MEGA:...")`, `eprintln!("FWD:...")`, `eprintln!("ATTN FWD:...")` debugging lines.
|
||||
|
||||
- [ ] **Step 4: Compile and run full test suite**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check --workspace
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored
|
||||
```
|
||||
|
||||
Expected: All pass, zero warnings.
|
||||
|
||||
- [ ] **Step 5: Commit and push**
|
||||
|
||||
```bash
|
||||
git commit -m "refactor: remove dead graph capture code and debug logging"
|
||||
git push origin main
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 9: Deploy to H100 and measure performance
|
||||
|
||||
**Files:** None (deployment only)
|
||||
|
||||
- [ ] **Step 1: Launch H100 training**
|
||||
|
||||
```bash
|
||||
argo submit -n foxhunt --from workflowtemplate/compile-and-train \
|
||||
-p commit-sha=$(git rev-parse --short HEAD) \
|
||||
-p model=dqn -p gpu-pool=ci-training-h100 \
|
||||
-p hyperopt-trials=0 -p train-epochs=5
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Verify per-step timing**
|
||||
|
||||
Check training logs for `Training step breakdown` line:
|
||||
- Before: `per-step: fused=129.0ms` (2-graph approach with syncs)
|
||||
- Target: `per-step: fused=<10ms` (mega-graph, zero ungraphed)
|
||||
|
||||
- [ ] **Step 3: Verify training correctness**
|
||||
|
||||
Check Sharpe, PF, Q-values, action diversity match expected ranges.
|
||||
|
||||
---
|
||||
|
||||
## Risk Register
|
||||
|
||||
| Risk | Impact | Mitigation |
|
||||
|------|--------|-----------|
|
||||
| cudarc launch_builder incompatible with capture | High | Raw cuLaunchKernel as fallback; but research shows launch_builder IS compatible with disabled event tracking |
|
||||
| double_dqn_stream capture isolation | High | Pre-allocated events + cuStreamWaitEvent automatic join (proven working in current code) |
|
||||
| Kernel arg type mismatches crash GPU | Critical | Task 1-2 fix ALL mismatches before any refactoring |
|
||||
| PER update can't be graphed (pointer changes) | Low | Stays outside graph — 1 ungraphed kernel per step (acceptable) |
|
||||
| cudarc check_err poisons during mega-graph capture | Medium | Drain errors before capture + event tracking disabled = no new errors recorded |
|
||||
Reference in New Issue
Block a user