fix: scope event tracking disable to CUDA Graph capture only

disable_event_tracking() was called globally in GpuDqnTrainer::new(),
polluting the CUDA context for ALL subsequent operations including
other models (TFT, PPO). This caused test_tft_adapter_deterministic
to fail when run after DQN tests.

Fix: disable/enable only inside capture_training_graph(), not at
construction. The capture region is the only place where CudaEvents
are disallowed.

Result: 855/855 ml lib tests pass (was 854/855 with 1 flaky failure).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-20 12:19:56 +01:00
parent 2752758b25
commit 667e61734c

View File

@@ -403,9 +403,9 @@ impl GpuDqnTrainer {
let b = config.batch_size;
let total_params = compute_total_params(&config);
// Disable event tracking — required for CUDA Graph capture (no CudaEvents
// inside captured regions). Safe: single stream, explicit ordering.
unsafe { stream.context().disable_event_tracking(); }
// Event tracking: kept ENABLED during normal ops (buffer allocation, DtoD).
// DISABLED only during CUDA Graph capture (in capture_training_graph).
// Global disable pollutes the context for other tests/models.
// Stack size for the training kernel: ~46KB/thread with NUM_ATOMS=51,
// ~12KB/thread with NUM_ATOMS=11. DIST_SIZE(SCRATCH1_DIM=256) * 4 = 11KB