fix(bf16): EventTrackingGuard for post-graph kernel launches

Root cause: cudarc event tracking left stale state after CUDA graph
replay, causing INVALID_VALUE on subsequent kernel launches.
Fix: add EventTrackingGuard to clip_grad_buf_inplace (same pattern
as read_grad_norm_sync, apply_spectral_norm, etc.).

Smoke test now runs 3 epochs / 600 steps successfully.
Q-values diverge (bf16 precision) — needs hyperparameter tuning
for bf16 (lower lr, smaller Huber delta, tighter gradient clip).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-28 15:50:38 +01:00
parent de968acf81
commit 03ad33fda4

View File

@@ -1521,7 +1521,9 @@ impl GpuDqnTrainer {
/// from overwhelming the subsequent auxiliary gradient additions.
/// All operations are async on the same stream — zero CPU sync.
pub fn clip_grad_buf_inplace(&mut self, max_norm: f32) -> Result<(), MLError> {
// clip_grad_buf_inplace: standalone kernel launch (not captured in CUDA graph)
// Disable cudarc event tracking — graph replay leaves stale events
// that cause INVALID_VALUE on post-graph kernel launches.
let _evt_guard = EventTrackingGuard::new(self.stream.context());
// Zero the grad_norm accumulator
self.stream
.memset_zeros(&mut self.grad_norm_buf)
@@ -3571,7 +3573,6 @@ impl GpuDqnTrainer {
fn launch_grad_norm(&self) -> Result<(), MLError> {
let tp = self.total_params as i32;
let blocks = ((self.total_params + 255) / 256) as u32;
tp, blocks, self.ptrs.grad_buf, self.ptrs.grad_norm_buf);
let launch_cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),