From a662e54e86b5d2c781be25671b38808fbe06cfdc Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 11 Apr 2026 13:35:15 +0200 Subject: [PATCH] =?UTF-8?q?plan:=20VarStore=20elimination=20=E2=80=94=20ze?= =?UTF-8?q?ro-copy=20pointer=20views=20into=20flat=20params=5Fbuf=20(8=20t?= =?UTF-8?q?asks)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DuelingWeightSet becomes raw u64 pointers into params_buf, not separate CudaSlice allocations. Eliminates ALL D2D copies between weight representations. flatten/unflatten cycle deleted entirely. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../plans/2026-04-11-varstore-elimination.md | 439 ++++++++++++++++++ 1 file changed, 439 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-11-varstore-elimination.md diff --git a/docs/superpowers/plans/2026-04-11-varstore-elimination.md b/docs/superpowers/plans/2026-04-11-varstore-elimination.md new file mode 100644 index 000000000..c0e485859 --- /dev/null +++ b/docs/superpowers/plans/2026-04-11-varstore-elimination.md @@ -0,0 +1,439 @@ +# 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, + 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 { + 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, + shared_h2: usize, + adv_h: usize, + num_atoms: usize, + branch_1_size: usize, + branch_2_size: usize, + branch_3_size: usize, +) -> Result { + 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, out_dim: usize, in_dim: usize) -> Result, 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 = (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, n: usize) -> Result, MLError> { + stream.alloc_zeros::(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, +) -> Result, Vec)>, MLError> { + let mut map = BTreeMap::new(); + let read = |slice: &CudaSlice, name: &str| -> Result, 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, Vec)>, + online_d: &mut DuelingWeightSet, + online_b: &mut BranchingWeightSet, + stream: &Arc, +) -> Result<(), MLError> { + let write = |slice: &mut CudaSlice, 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%)