diff --git a/config/training/dqn-production.toml b/config/training/dqn-production.toml index 2b3aaec08..be51392af 100644 --- a/config/training/dqn-production.toml +++ b/config/training/dqn-production.toml @@ -86,7 +86,7 @@ c51_alpha_max = 0.5 her_ratio = 0.2 cql_alpha = 1.0 curiosity_weight = 0.1 -iqn_lambda = 0.0 # DIAG: disabled to isolate NaN source +iqn_lambda = 0.25 spectral_norm_sigma_max = 1.5 spectral_decoupling_lambda = 0.01 gradient_clip_norm = 1.0 diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index e3274fa20..b9cd644ce 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -853,20 +853,6 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?; self.steps_since_varmap_sync += 1; - - // Per-step gradient diagnostic — zero cost (uses existing readback, no sync). - // Log every 50 steps to see gradient trajectory during collapse. - if self.steps_since_varmap_sync % 50 == 1 { - let alpha = self.trainer.c51_alpha(); - tracing::warn!( - step = self.steps_since_varmap_sync, - grad_norm = fused_result.grad_norm, - loss = fused_result.total_loss, - c51_alpha = alpha, - batch_size = self.batch_size, - "STEP_DIAG: per-step gradient readback" - ); - } self.last_combined_norm = fused_result.grad_norm; // Capture mega-graph at step 2 (all sub-trainers initialized). diff --git a/docs/superpowers/plans/2026-04-09-h100-nan-fix.md b/docs/superpowers/plans/2026-04-09-h100-nan-fix.md new file mode 100644 index 000000000..1415add3e --- /dev/null +++ b/docs/superpowers/plans/2026-04-09-h100-nan-fix.md @@ -0,0 +1,194 @@ +# H100 NaN Fix 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:** Fix the three remaining H100 training instabilities — IQN type mismatch NaN, IQN gradient amplification, and C51 premature convergence. + +**Architecture:** Three independent fixes applied to the CUDA training pipeline. Fix 1 (type mismatch) is the NaN root cause. Fix 2 (gradient amplification) prevents IQN from destabilizing Adam. Fix 3 (alpha cap) keeps MSE gradient alive after C51 converges. Each fix is independently testable. + +**Tech Stack:** Rust + CUDA kernels (`.cu`), cuBLAS backward pass, fused DQN training pipeline. + +--- + +### Task 1: Fix f32-as-bf16 type mismatch in apply_iqn_trunk_gradient + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:1082-1135` + +The IQN's `d_h_s2_buf` is allocated as `CudaSlice` (gpu_iqn_head.rs:235) and the backward kernel writes f32 via `atomicAdd`. But `apply_iqn_trunk_gradient` at line 1114-1133 calls `bf16_to_f32_kernel` on this f32 pointer, reinterpreting f32 bits as bf16 — producing garbage/NaN. + +- [ ] **Step 1: Remove the bf16→f32 cast and use direct f32 DtoD copy** + +In `apply_iqn_trunk_gradient()`, replace the bf16_to_f32_kernel call (lines 1114-1134) with a direct f32→f32 DtoD copy. The IQN d_h_s2 is already f32 — no cast needed. + +Replace: +```rust + // ── 2. Cast IQN d_h_s2 (bf16) → bw_d_h_s2 (f32) ───────────────── + // IQN head produces bf16 gradient. Cast to f32 for the backward scratch. + { + let src = iqn_d_h_s2_ptr; + let dst = self.ptrs.bw_d_h_s2; + let n_elems = (b * sh2) as i32; + let blocks = ((b * sh2 + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.bf16_to_f32_kernel) + .arg(&src) + .arg(&dst) + .arg(&n_elems) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQN d_h_s2 bf16→f32: {e}")))?; + } + } +``` + +With: +```rust + // ── 2. DtoD copy IQN d_h_s2 (f32) → bw_d_h_s2 (f32) ──────────── + // IQN d_h_s2_buf is f32 (atomicAdd in backward kernel). No cast needed. + { + let n_bytes = b * sh2 * std::mem::size_of::(); + dtod_copy(self.ptrs.bw_d_h_s2, iqn_d_h_s2_ptr, n_bytes, &self.stream, 0, "iqn_d_h_s2_copy") + .map_err(|e| MLError::ModelError(format!("IQN d_h_s2 f32 DtoD: {e}")))?; + } +``` + +- [ ] **Step 2: Build and verify** + +Run: `SQLX_OFFLINE=true cargo check -p ml` +Expected: compiles clean. + +- [ ] **Step 3: Run smoke test** + +Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::regression --ignored` +Expected: 4 passed, 0 failed. + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "fix(critical): IQN d_h_s2 is f32, not bf16 — remove bogus cast" +``` + +--- + +### Task 2: Fix IQN gradient clipping amplification + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu:728-731` + +With mean-reduced IQN gradients (~1e-6 magnitude), the IQN Adam's `clip_scale = max_grad_norm / grad_norm = 1.0 / 2.7e-4 ≈ 3700×`. This AMPLIFIES gradients instead of clipping them. The fix: clamp clip_scale to max 1.0 so it only ever reduces, never amplifies. + +- [ ] **Step 1: Clamp clip_scale to ≤ 1.0** + +In `iqn_adam_kernel`, change line 729-731 from: +```cuda + float clip_scale = (grad_norm_val > max_grad_norm && grad_norm_val > 0.0f) + ? max_grad_norm / grad_norm_val : 1.0f; +``` + +To: +```cuda + /* Clamp clip_scale to [0, 1]: only reduce, never amplify. + * With mean-reduced gradients (inv_batch), grad_norm ≈ 1e-4 which is + * below max_grad_norm — clip_scale stays 1.0 (no clipping). */ + float clip_scale = (grad_norm_val > max_grad_norm) + ? max_grad_norm / grad_norm_val : 1.0f; +``` + +Note: the existing code already does this correctly (clip_scale is 1.0 when norm ≤ max). The issue is that `max_grad_norm=1.0` is too large relative to the mean-reduced gradient (~2.7e-4). The clipping never fires, and Adam amplifies the tiny gradient via the m_hat/sqrt(v_hat) ratio. The real fix is to not clip at all for IQN (set max_grad_norm=Inf) and let Adam's adaptive LR handle scaling. OR reduce IQN's LR. + +Actually, the simplest fix: IQN should use the same LR as the main DQN. Check that `config.lr` is set correctly from the TOML learning_rate. + +- [ ] **Step 2: Verify IQN LR matches production config** + +In `crates/ml/src/trainers/dqn/fused_training.rs` line 437: +```rust +lr: hyperparams.learning_rate as f32, +``` +This should be `1e-5` from the production TOML. Confirm by searching for the value in the constructor. + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +git commit -m "fix: document IQN gradient clipping behavior with mean-reduced gradients" +``` + +--- + +### Task 3: Set c51_alpha_max=0.5 for production H100 + +**Files:** +- Modify: `config/training/dqn-production.toml` +- Modify: `crates/ml/src/trainers/dqn/config.rs` + +The IQN-disabled H100 run showed fold 2 collapsing to grad_norm=0 (C51 converged). With IQN fixed, this may not be the binding constraint, but c51_alpha_max=0.5 is a safety net. + +- [ ] **Step 1: Set default c51_alpha_max to 0.5** + +In config.rs default constructor, change `c51_alpha_max: 1.0` to `c51_alpha_max: 0.5`. +In dqn-production.toml, change `c51_alpha_max = 1.0` to `c51_alpha_max = 0.5`. + +- [ ] **Step 2: Run full smoke test suite** + +Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability --ignored` +Expected: 6/6 pass. + +- [ ] **Step 3: Commit** + +```bash +git add config/training/dqn-production.toml crates/ml/src/trainers/dqn/config.rs +git commit -m "fix: set c51_alpha_max=0.5 default — prevent C51 gradient convergence" +``` + +--- + +### Task 4: Re-enable IQN, remove diagnostics, deploy baseline + +**Files:** +- Modify: `config/training/dqn-production.toml` — restore `iqn_lambda = 0.25` +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` — remove STEP_DIAG logging +- Modify: `infra/k8s/argo/compile-and-train-template.yaml` — remove FOXHUNT_GRAD_DIAG env + +- [ ] **Step 1: Restore iqn_lambda, remove diagnostics** + +dqn-production.toml: `iqn_lambda = 0.25` +fused_training.rs: remove the `STEP_DIAG` warn! block (~14 lines) +compile-and-train-template.yaml: remove `FOXHUNT_GRAD_DIAG` env var + +- [ ] **Step 2: Full test suite** + +Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability --ignored` +Expected: 6/6 pass. + +- [ ] **Step 3: Commit and push** + +```bash +git add -A +git commit -m "fix: re-enable IQN + remove diagnostics for H100 baseline" +git push +``` + +- [ ] **Step 4: Deploy H100 baseline (20 epochs)** + +```bash +./scripts/argo-train.sh dqn --baseline --epochs 20 --watch +``` + +Expected: grad_norm stable (0.5-1.0) through all 20 epochs, both folds. No NaN. + +--- + +### Task 5: Update memory and documentation + +**Files:** +- Modify: `/home/jgrusewski/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/followup_gradient_collapse_h100.md` + +- [ ] **Step 1: Update memory with final root causes** + +Add the IQN f32-as-bf16 type mismatch as root cause #4. Update the verification section with the successful H100 baseline results. diff --git a/infra/k8s/argo/compile-and-train-template.yaml b/infra/k8s/argo/compile-and-train-template.yaml index 2b2131985..ef4f17923 100644 --- a/infra/k8s/argo/compile-and-train-template.yaml +++ b/infra/k8s/argo/compile-and-train-template.yaml @@ -526,8 +526,6 @@ spec: value: ":4096:8" - name: FOXHUNT_FEATURE_CACHE_DIR value: "{{workflow.parameters.feature-cache-dir}}" - - name: FOXHUNT_GRAD_DIAG - value: "1" args: - | set -e