# Candle GpuVarStore Elimination 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:** Remove all GpuVarStore from the DQN training pipeline, replace with direct CudaSlice pointer buffers via DuelingWeightSet/BranchingWeightSet, and re-enable bottleneck_dim=16.

**Architecture:** The flat `params_buf` IS the single source of truth — zero copies. DuelingWeightSet/BranchingWeightSet become lightweight pointer structs (raw u64 device pointers + sizes) that reference OFFSETS into params_buf. No separate CudaSlice allocations, no D2D copies, no flatten/unflatten cycle. Checkpoint save/load reads from params_buf directly. Experience collector reads from the same pointer offsets.

**Tech Stack:** Rust, cudarc CudaSlice, safetensors, ml-dqn, ml crate

---

## File Structure

| Action | File | Purpose |
|--------|------|---------|
| Modify | `crates/ml/src/cuda_pipeline/gpu_weights.rs` | Allocate weight sets directly (no VarStore), checkpoint helpers |
| Modify | `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Remove flatten bottleneck skip, use direct weight allocation |
| Modify | `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Take DuelingWeightSet refs instead of GpuVarStore |
| Modify | `crates/ml/src/trainers/dqn/fused_training.rs` | Remove VarStore sync, pass weight sets to collector |
| Modify | `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Remove VarStore refs from weight sync calls |
| Modify | `crates/ml/src/trainers/dqn/config.rs` | Remove GpuVarStore imports, update checkpoint |
| Modify | `crates/ml-dqn/src/dqn.rs` | Remove vars field, load_from_safetensors works on weight sets |
| Modify | `crates/ml-dqn/src/branching.rs` | Remove vars, register_noisy_mu, VarStore polyak |
| Modify | `crates/ml-dqn/src/distributional_dueling.rs` | Remove vars |
| Modify | `crates/ml-dqn/src/target_update.rs` | Delete VarStore polyak functions |
| Modify | `crates/ml-core/src/checkpoint.rs` | Checkpoint from weight sets instead of VarStore |

---

### Task 1: Add Direct Weight Set Allocation (no VarStore)

Currently `extract_dueling_weights(vars: &GpuVarStore)` allocates CudaSlices by copying from VarStore. Create a NEW function that allocates weight sets directly with Xavier initialization — no VarStore involved.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_weights.rs`

- [ ] **Step 1: Add `allocate_dueling_weights_xavier` function**

This function creates a DuelingWeightSet with Xavier-initialized CudaSlices at the correct dimensions. When bottleneck_dim > 0, w_s1 uses the reduced s1_input_dim.

```rust
/// Allocate DuelingWeightSet with Xavier initialization — no GpuVarStore.
/// When bottleneck_dim > 0, w_s1 input dimension is (bottleneck_dim + portfolio_dim)
/// instead of state_dim.
pub fn allocate_dueling_weights_xavier(
    stream: &Arc<CudaStream>,
    state_dim: usize,
    market_dim: usize,
    bottleneck_dim: usize,
    shared_h1: usize,
    shared_h2: usize,
    value_h: usize,
    adv_h: usize,
    num_atoms: usize,
    branch_0_size: usize,
) -> Result<DuelingWeightSet, MLError> {
    let s1_input = if bottleneck_dim > 0 {
        bottleneck_dim + state_dim.saturating_sub(market_dim)
    } else {
        state_dim
    };

    Ok(DuelingWeightSet {
        w_s1: xavier_init_f32(stream, shared_h1, s1_input)?,
        b_s1: alloc_zeros_f32(stream, shared_h1)?,
        w_s2: xavier_init_f32(stream, shared_h2, shared_h1)?,
        b_s2: alloc_zeros_f32(stream, shared_h2)?,
        w_v1: xavier_init_f32(stream, value_h, shared_h2)?,
        b_v1: alloc_zeros_f32(stream, value_h)?,
        w_v2: xavier_init_f32(stream, num_atoms, value_h)?,
        b_v2: alloc_zeros_f32(stream, num_atoms)?,
        w_a1: xavier_init_f32(stream, adv_h, shared_h2)?,
        b_a1: alloc_zeros_f32(stream, adv_h)?,
        w_a2: xavier_init_f32(stream, branch_0_size * num_atoms, adv_h)?,
        b_a2: alloc_zeros_f32(stream, branch_0_size * num_atoms)?,
    })
}
```

- [ ] **Step 2: Add `allocate_branching_weights_xavier` function**

Same pattern for branches 1-3:

```rust
pub fn allocate_branching_weights_xavier(
    stream: &Arc<CudaStream>,
    shared_h2: usize,
    adv_h: usize,
    num_atoms: usize,
    branch_1_size: usize,
    branch_2_size: usize,
    branch_3_size: usize,
) -> Result<BranchingWeightSet, MLError> {
    Ok(BranchingWeightSet {
        w_bo1: xavier_init_f32(stream, adv_h, shared_h2)?,
        b_bo1: alloc_zeros_f32(stream, adv_h)?,
        w_bo2: xavier_init_f32(stream, branch_1_size * num_atoms, adv_h)?,
        b_bo2: alloc_zeros_f32(stream, branch_1_size * num_atoms)?,
        w_bu1: xavier_init_f32(stream, adv_h, shared_h2)?,
        b_bu1: alloc_zeros_f32(stream, adv_h)?,
        w_bu2: xavier_init_f32(stream, branch_2_size * num_atoms, adv_h)?,
        b_bu2: alloc_zeros_f32(stream, branch_2_size * num_atoms)?,
        w_bg1: xavier_init_f32(stream, adv_h, shared_h2)?,
        b_bg1: alloc_zeros_f32(stream, adv_h)?,
        w_bg2: xavier_init_f32(stream, branch_3_size * num_atoms, adv_h)?,
        b_bg2: alloc_zeros_f32(stream, branch_3_size * num_atoms)?,
    })
}
```

- [ ] **Step 3: Add xavier_init_f32 and alloc_zeros_f32 helpers**

```rust
fn xavier_init_f32(stream: &Arc<CudaStream>, out_dim: usize, in_dim: usize) -> Result<CudaSlice<f32>, MLError> {
    let n = out_dim * in_dim;
    let limit = (6.0_f32 / (in_dim + out_dim) as f32).sqrt();
    let mut rng = rand::thread_rng();
    let host: Vec<f32> = (0..n).map(|_| rng.gen_range(-limit..limit)).collect();
    stream.clone_htod(&host).map_err(|e| MLError::ModelError(format!("xavier HtoD: {e}")))
}

fn alloc_zeros_f32(stream: &Arc<CudaStream>, n: usize) -> Result<CudaSlice<f32>, MLError> {
    stream.alloc_zeros::<f32>(n).map_err(|e| MLError::ModelError(format!("zeros alloc: {e}")))
}
```

- [ ] **Step 4: Add checkpoint save/load on DuelingWeightSet**

```rust
/// Export DuelingWeightSet + BranchingWeightSet to host BTreeMap for safetensors.
pub fn weight_sets_to_host(
    online_d: &DuelingWeightSet,
    online_b: &BranchingWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<BTreeMap<String, (Vec<usize>, Vec<f32>)>, MLError> {
    let mut map = BTreeMap::new();
    let read = |slice: &CudaSlice<f32>, name: &str| -> Result<Vec<f32>, MLError> {
        let mut host = vec![0.0_f32; slice.len()];
        stream.memcpy_dtoh(slice, &mut host)
            .map_err(|e| MLError::ModelError(format!("{name} DtoH: {e}")))?;
        Ok(host)
    };
    // Shared trunk
    let s1_n = online_d.w_s1.len();
    let s1_out = online_d.b_s1.len();
    let s1_in = s1_n / s1_out;
    map.insert("shared_0.weight".into(), (vec![s1_out, s1_in], read(&online_d.w_s1, "w_s1")?));
    map.insert("shared_0.bias".into(), (vec![s1_out], read(&online_d.b_s1, "b_s1")?));
    // ... repeat for all 24 weight tensors with correct names ...
    Ok(map)
}

/// Import host BTreeMap into DuelingWeightSet + BranchingWeightSet.
pub fn host_to_weight_sets(
    map: &BTreeMap<String, (Vec<usize>, Vec<f32>)>,
    online_d: &mut DuelingWeightSet,
    online_b: &mut BranchingWeightSet,
    stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
    let write = |slice: &mut CudaSlice<f32>, data: &[f32], name: &str| -> Result<(), MLError> {
        if data.len() != slice.len() {
            return Err(MLError::ModelError(format!("{name}: expected {} got {}", slice.len(), data.len())));
        }
        stream.memcpy_htod(data, slice)
            .map_err(|e| MLError::ModelError(format!("{name} HtoD: {e}")))?;
        Ok(())
    };
    if let Some((_, data)) = map.get("shared_0.weight") { write(&mut online_d.w_s1, data, "w_s1")?; }
    if let Some((_, data)) = map.get("shared_0.bias") { write(&mut online_d.b_s1, data, "b_s1")?; }
    // ... repeat for all 24 ...
    Ok(())
}
```

- [ ] **Step 5: Compile check**

Run: `SQLX_OFFLINE=true cargo check -p ml --lib`

- [ ] **Step 6: Commit**

```bash
git commit -m "feat: direct weight set allocation + checkpoint without VarStore"
```

---

### Task 2: Wire Direct Allocation in FusedTrainingCtx

Replace VarStore weight extraction with direct allocation in the fused training context.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Replace weight extraction from VarStore**

Find where `extract_dueling_weights(vars, stream)` is called to create online/target DuelingWeightSets. Replace with `allocate_dueling_weights_xavier(stream, ...)`.

The key locations are in `FusedTrainingCtx::new()` where online_dueling and target_dueling are created. Currently they're extracted from the DQN model's VarStore. Change to direct allocation.

- [ ] **Step 2: Remove VarStore sync at epoch boundary**

Find `sync_gpu_weights` or `unflatten` calls that sync flat buffer → VarStore. These become no-ops since there's no VarStore to sync to. The flat buffer → DuelingWeightSet unflatten STAYS (it updates the individual CudaSlices that experience collector reads).

- [ ] **Step 3: Remove GpuVarStore imports**

Remove `use ml_core::cuda_autograd::GpuVarStore;` and any VarStore references.

- [ ] **Step 4: Compile check + commit**

---

### Task 3: Update Experience Collector to Take Weight Sets

Replace `&GpuVarStore` params with `&DuelingWeightSet` + `&BranchingWeightSet`.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_weights.rs`

- [ ] **Step 1: Change constructor signature**

Replace:
```rust
pub fn new(stream, online_vars: &GpuVarStore, target_vars: &GpuVarStore, ...)
```
With:
```rust
pub fn new(stream, online_d: &DuelingWeightSet, online_b: &BranchingWeightSet,
           target_d: &DuelingWeightSet, target_b: &BranchingWeightSet, ...)
```

The constructor currently calls `extract_dueling_weights_branching(online_vars)` to create internal weight copies. Change to D2D copy from the passed-in weight sets.

- [ ] **Step 2: Change sync methods**

Replace all `sync_*(&mut self, vars: &GpuVarStore)` with `sync_*(&mut self, d: &DuelingWeightSet, b: &BranchingWeightSet)`:
- `sync_online_weights`
- `sync_target_weights`
- `sync_online_branching`
- `sync_target_branching`

Each method copies from the weight set CudaSlices to internal buffers via D2D memcpy.

- [ ] **Step 3: Update all callers**

Search for every call to these sync methods and update to pass weight sets instead of VarStore refs. Key locations:
- `training_loop.rs` (epoch boundary sync)
- `fused_training.rs` (initialization)

- [ ] **Step 4: Delete extract_dueling_weights(vars: &GpuVarStore) and related functions**

These are now dead:
- `extract_dueling_weights`
- `extract_branching_weights`
- `extract_dueling_weights_branching_keys`
- `extract_one`
- All `reupload_*` functions

Keep `sync_dueling_weights_branching` if it's the D2D path, or delete if it goes through VarStore.

- [ ] **Step 5: Compile check + commit**

---

### Task 4: Remove GpuVarStore from DQN Network Structs

Delete `vars: GpuVarStore` from all network structs in ml-dqn.

**Files:**
- Modify: `crates/ml-dqn/src/dqn.rs`
- Modify: `crates/ml-dqn/src/branching.rs`
- Modify: `crates/ml-dqn/src/distributional_dueling.rs`
- Modify: `crates/ml-dqn/src/target_update.rs`

- [ ] **Step 1: Remove from Sequential (dqn.rs)**

Delete `vars: GpuVarStore` field (line 899). Delete `vars()`, `vars_mut()`. Remove `GpuVarStore::new()` from constructor. Remove `copy_weights_from()` VarStore path.

- [ ] **Step 2: Remove from DQN (dqn.rs)**

Delete `get_q_network_vars()`, `get_q_network_vars_mut()`. Update `load_from_safetensors` to work differently (Task 5 handles this).

- [ ] **Step 3: Remove from BranchingDuelingQNetwork (branching.rs)**

Delete `vars: GpuVarStore` field. Delete `vars()`, `vars_mut()`. Delete `register_noisy_mu_in_varstore()`. Delete VarStore-based `polyak_update_from()` and `copy_weights_from()`.

- [ ] **Step 4: Remove from DistributionalDuelingQNetwork (distributional_dueling.rs)**

Delete `vars: GpuVarStore` field. Delete `vars()`, `vars_mut()`. Remove `linear_xavier()` calls (weights are now allocated directly).

- [ ] **Step 5: Delete VarStore polyak functions (target_update.rs)**

Delete `polyak_update(online_vars: &GpuVarStore, target_vars: &mut GpuVarStore, ...)` and `polyak_update_cosine_annealed`. The GPU EMA kernel in `fused_training.rs` handles Polyak updates on the flat buffer.

- [ ] **Step 6: Compile check — expect many errors from callers**

Fix each caller. Key pattern: anywhere that called `agent.get_q_network_vars()` or `agent.vars()` now needs the DuelingWeightSet/BranchingWeightSet from the GpuDqnTrainer's context instead.

- [ ] **Step 7: Commit**

---

### Task 5: Update Checkpoint Save/Load

Replace VarStore-based checkpoint with direct weight set serialization.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/config.rs`
- Modify: `crates/ml-dqn/src/dqn.rs`
- Modify: `crates/ml-core/src/checkpoint.rs`

- [ ] **Step 1: Update save_checkpoint in config.rs**

Currently calls `vars.export_to_host()`. Replace with `weight_sets_to_host(online_d, online_b, stream)` from Task 1.

The checkpoint must be called from context that has access to the DuelingWeightSet — likely the training loop or FusedTrainingCtx.

- [ ] **Step 2: Update load_from_safetensors in dqn.rs**

Currently calls `get_q_network_vars_mut().import_from_host()`. Replace with:
1. Parse safetensors → host BTreeMap
2. Call `host_to_weight_sets(map, online_d, online_b, stream)` to write directly to weight sets
3. Call `flatten_online_weights()` to sync flat buffer

- [ ] **Step 3: Remove VarStore-based save/load from checkpoint.rs**

Delete or update `save_safetensors(store: &GpuVarStore, ...)` and `load_safetensors(store: &mut GpuVarStore, ...)` to take weight sets instead. Or keep the functions but mark as deprecated if other models (PPO, etc.) still use VarStore.

- [ ] **Step 4: Compile check + commit**

---

### Task 6: Remove GpuVarStore from Config + Clean Imports

Final cleanup — remove all remaining VarStore references.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/config.rs`
- Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs`
- Modify: `crates/ml/src/trainers/dqn/trainer/mod.rs`

- [ ] **Step 1: Remove GpuVarStore import from config.rs**

Delete `use ml_core::cuda_autograd::GpuVarStore;` (line 10).

- [ ] **Step 2: Update DQNAgentType methods that referenced VarStore**

Methods like `get_q_network_vars()` on DQNAgentType — these delegated to `agent.get_q_network_vars()`. Delete them or replace with weight set accessors.

- [ ] **Step 3: grep for any remaining VarStore references**

```bash
grep -rn 'GpuVarStore\|VarStore' crates/ml/src/trainers/dqn/ crates/ml-dqn/src/ --include='*.rs' | grep -v test | grep -v target
```

Fix each remaining reference.

- [ ] **Step 4: Compile check + commit**

---

### Task 7: Re-enable Bottleneck dim=16

Now that VarStore is gone, the dimension mismatch is resolved. w_s1 is allocated at the correct bottleneck-adjusted dimension.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/config.rs`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `config/training/dqn-*.toml`

- [ ] **Step 1: Set bottleneck_dim=16 in config defaults**

```rust
bottleneck_dim: 16,
```

- [ ] **Step 2: Set bottleneck_dim=16 in TOML configs**

All 3 TOML files: `dqn-production.toml`, `dqn-localdev.toml`, `dqn-smoketest.toml`.

- [ ] **Step 3: Remove the flatten bottleneck skip workaround**

In `gpu_dqn_trainer.rs`, find the `if self.config.bottleneck_dim > 0 && (i == 0 || i == 1) { continue; }` skip logic in `flatten_online_weights`. DELETE IT — with direct allocation, w_s1 is always the correct size.

- [ ] **Step 4: Compile check + smoke test**

```bash
SQLX_OFFLINE=true cargo check -p ml --lib
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::regression::test_no_hang_single_epoch --ignored --nocapture
```

The smoke test MUST pass with bottleneck_dim=16.

- [ ] **Step 5: Commit**

```bash
git commit -m "feat: re-enable bottleneck_dim=16 — VarStore eliminated, dimensions correct"
```

---

### Task 8: Full Verification + Deploy

- [ ] **Step 1: Run ALL test suites**

```bash
SQLX_OFFLINE=true cargo test -p ml-core --lib              # 302 tests
SQLX_OFFLINE=true cargo test -p ml-dqn --lib -- --test-threads=1  # 287 tests
SQLX_OFFLINE=true cargo test -p ml --lib                    # 903 tests
```

- [ ] **Step 2: Run smoke test with bottleneck=16**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
  cargo test -p ml --lib -- smoke_tests::regression::test_no_hang_single_epoch --ignored --nocapture
```

- [ ] **Step 3: Verify no VarStore references remain**

```bash
grep -rn 'GpuVarStore' crates/ml/src/ crates/ml-dqn/src/ --include='*.rs' | grep -v test | grep -v target | grep -v '//.*VarStore'
```
Expected: zero results.

- [ ] **Step 4: Push and deploy to H100**

```bash
git push origin main
./scripts/argo-train.sh dqn --baseline --epochs 5 --watch
```

### Success Criteria
1. Zero `GpuVarStore` references in DQN training pipeline
2. Smoke test passes with bottleneck_dim=16
3. All 1492 tests pass (302 + 287 + 903)
4. H100 training runs without errors
5. Magnitude diversity improved (target: Half >10%, Full >5%)
