diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 785ea13ea..de71f2ef4 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -560,6 +560,7 @@ fn train_dqn_fold( // to leave room for training kernels + replay buffer. gpu_n_episodes: if detect_vram_mb() < 8192 { 16 } else { 256 }, gpu_timesteps_per_episode: if detect_vram_mb() < 8192 { 100 } else { 500 }, + max_training_steps_per_epoch: args.max_steps_per_epoch, ..DQNHyperparameters::default() }; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 4769a9a51..d3fed2187 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -947,9 +947,13 @@ impl DQNTrainer { // Batch pre-sampling: sample K batches under one READ lock const PREFETCH_K: usize = 32; + let mut sample_total_us = 0_u64; + let mut fused_total_us = 0_u64; + let mut guard_total_us = 0_u64; for chunk_start in (0..num_training_steps).step_by(PREFETCH_K) { let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps); + let sample_start = std::time::Instant::now(); let batches = { let agent = self.agent.read().await; let buffer = agent.memory(); @@ -972,6 +976,7 @@ impl DQNTrainer { } b }; + sample_total_us += sample_start.elapsed().as_micros() as u64; // GPU train steps (single WRITE lock) { @@ -980,8 +985,7 @@ impl DQNTrainer { if accum_steps <= 1 { for explicit_batch in batches { - // Fused CUDA path: 3 kernels + CUDA Graph, no Candle dispatch. - // Fused training failures are hard errors — no silent fallback. + let fused_start = std::time::Instant::now(); let _gpu_result = if let Some(ref mut fused) = self.fused_ctx { let batch_data = explicit_batch.as_ref().ok_or_else(|| { anyhow::anyhow!("No batch data for fused training step") @@ -993,6 +997,9 @@ impl DQNTrainer { .map_err(|e| anyhow::anyhow!("Train step FAILED: {e}"))? }; + fused_total_us += fused_start.elapsed().as_micros() as u64; + + let guard_start = std::time::Instant::now(); if let Some(ref mut guard) = self.training_guard { let loss_slice = _gpu_result.loss_cuda_slice() .map_err(|e| anyhow::anyhow!("guard loss CudaSlice: {e}"))?; @@ -1018,6 +1025,7 @@ impl DQNTrainer { })?; } } + guard_total_us += guard_start.elapsed().as_micros() as u64; train_step_count += 1; self.gradient_logging_step += 1; } @@ -1153,6 +1161,18 @@ impl DQNTrainer { } } + if train_step_count > 0 { + info!( + "Training step breakdown ({} steps): sample={:.0}ms fused={:.0}ms guard={:.0}ms (per-step: sample={:.1}ms fused={:.1}ms guard={:.1}ms)", + train_step_count, + sample_total_us as f64 / 1000.0, + fused_total_us as f64 / 1000.0, + guard_total_us as f64 / 1000.0, + sample_total_us as f64 / 1000.0 / train_step_count as f64, + fused_total_us as f64 / 1000.0 / train_step_count as f64, + guard_total_us as f64 / 1000.0 / train_step_count as f64, + ); + } Ok(train_step_count) }