# 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<f32>` (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::<f32>();
            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.
