From 9f5b7b58d4e832ed1c419d079ab8e7fa4a8f3894 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 6 Apr 2026 22:37:47 +0200 Subject: [PATCH] perf: eliminate 86ms/step PER double-sample + guard overhead MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three performance optimizations targeting the 15.6s/epoch training: 1. Single write lock for sample+step (was: read lock → sample → release → write lock → step → release = 43ms lock contention per step) 2. Vaccine batch sampled only when needed (every 10th step, not every step — saves 43ms on 90% of steps) 3. Training guard runs every 5th step instead of every step (12ms/step → 2.4ms/step amortized). NaN/loss checks persist across steps. 4. Causal intervention interval 10→50 (14 forward passes less often) 5. Removed all remaining H100_DIAG tracing::info in training_loop Expected impact: ~86ms → ~43ms sample, ~12ms → ~2.4ms guard = ~52ms/step savings = ~3.6s/epoch → target ~12s/epoch. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 2 +- .../src/trainers/dqn/trainer/training_loop.rs | 30 +++++++++++-------- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 55207f54f..c2daf89a0 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -265,7 +265,7 @@ impl Default for GpuDqnTrainConfig { pruning_fraction: 0.7, enable_causal_intervention: true, // always on causal_weight: 0.1, - causal_intervention_interval: 10, + causal_intervention_interval: 50, enable_gradient_vaccine: true, // always on bottleneck_dim: 2, market_dim: 42, // Default: 42 base features. Overridden to 50 when OFI (MBP-10) enabled. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index fbecdab13..5c700e3ab 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1004,13 +1004,11 @@ impl DQNTrainer { ..Default::default() }; - tracing::info!("H100_DIAG: starting collect_experiences_gpu"); let gpu_batch = collector.collect_experiences_gpu( features_buf, targets_buf, &episode_starts, &config, ).map_err(|e| anyhow::anyhow!( "GPU zero-roundtrip collection FAILED (no CPU fallback): {e}" ))?; - tracing::info!("H100_DIAG: collect_experiences_gpu done, inserting to PER"); let count = gpu_batch.n_episodes * gpu_batch.timesteps * 2; // counterfactual doubles experiences @@ -1067,7 +1065,6 @@ impl DQNTrainer { let agent = self.agent.read().await; let mut gpu_buf = agent.memory().as_gpu_buffer() .ok_or_else(|| anyhow::anyhow!("GPU PER buffer required"))?; - tracing::info!("H100_DIAG: PER insert_batch_bf16 starting (total={})", total); gpu_buf.gpu.insert_batch_bf16( &gpu_batch.states, &gpu_batch.next_states, @@ -1076,7 +1073,6 @@ impl DQNTrainer { &gpu_batch.dones, total, ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; - tracing::info!("H100_DIAG: PER insert_batch_bf16 done"); } match collector.read_epoch_dsr_state() { @@ -1175,23 +1171,28 @@ impl DQNTrainer { } for _step in 0..num_training_steps { + // Single write lock for both sample + fused step — eliminates + // read→write lock cycling overhead (~43ms per step from contention). + let mut agent = self.agent.write().await; + let sample_start = std::time::Instant::now(); - let (batch, vaccine_batch) = { - let agent = self.agent.read().await; + { let buffer = agent.memory(); if !buffer.can_sample(self.current_batch_size) { break; } - let b = buffer.sample(self.current_batch_size) - .map_err(|e| anyhow::anyhow!("PER sample: {e}"))?; - let vb = buffer.sample(self.current_batch_size).ok(); - (b, vb) + } + let batch = agent.memory().sample(self.current_batch_size) + .map_err(|e| anyhow::anyhow!("PER sample: {e}"))?; + // Vaccine batch: only sample when needed (every 10th step) + let vaccine_batch = if train_step_count % 10 == 0 { + agent.memory().sample(self.current_batch_size).ok() + } else { + None }; sample_total_us += sample_start.elapsed().as_micros() as u64; { - let mut agent = self.agent.write().await; - let fused_start = std::time::Instant::now(); let fused = self.fused_ctx.as_mut() .expect("Fused CUDA training is the only production path"); @@ -1204,6 +1205,10 @@ impl DQNTrainer { fused_total_us += fused_start.elapsed().as_micros() as u64; let guard_start = std::time::Instant::now(); + // Run guard every 5th step — NaN/loss checks add ~12ms/step overhead. + // Critical failures (NaN, loss explosion) persist across steps, + // so checking every 5th step catches them within <1s. + if train_step_count % 5 == 0 { if let Some(ref mut guard) = self.training_guard { let fused = self.fused_ctx.as_ref().unwrap(); let loss_raw = fused.loss_gpu_buf().raw_ptr(); @@ -1241,6 +1246,7 @@ impl DQNTrainer { })?; } } + } // end guard frequency check guard_total_us += guard_start.elapsed().as_micros() as u64; if train_step_count % 50 == 0 {