diff --git a/docs/superpowers/specs/2026-04-11-candle-varstore-elimination.md b/docs/superpowers/specs/2026-04-11-candle-varstore-elimination.md new file mode 100644 index 000000000..228fbfe98 --- /dev/null +++ b/docs/superpowers/specs/2026-04-11-candle-varstore-elimination.md @@ -0,0 +1,112 @@ +# Candle GpuVarStore Elimination — Design Spec + +## Problem + +The Candle `GpuVarStore` is a legacy weight management layer that: + +1. **Blocks bottleneck activation**: bottleneck_dim>0 causes dimension mismatch between VarStore (full state_dim) and flat params_buf (reduced dimension). EMA target sync crashes. +2. **Is redundant**: The flat `params_buf` / `target_params_buf` in GpuDqnTrainer is the real source of truth. Adam optimizer, CUDA graph forward/backward, and EMA all operate on the flat buffer. +3. **Adds complexity**: flatten/unflatten cycles, sync at epoch boundaries, dual weight sources that can drift. +4. **Wastes VRAM**: Duplicate weight storage (VarStore CudaSlices + flat buffer CudaSlices). + +## Scope + +40+ call sites across 15 files. Three subsystems depend on VarStore: + +### 1. Checkpoint save/load +- `save_safetensors()` reads from `GpuVarStore.export_to_host()` +- `load_safetensors()` writes to `GpuVarStore.import_from_host()` +- **Replacement**: Serialize flat params_buf directly as raw f32 bytes + metadata JSON + +### 2. Weight extraction for experience collector +- `extract_dueling_weights(vars: &GpuVarStore)` looks up named parameters +- 7 `sync_*()` methods re-extract after training +- **Replacement**: Extract from DuelingWeightSet/BranchingWeightSet CudaSlices directly (already used by flatten/unflatten) + +### 3. Polyak target updates (Candle path) +- `polyak_update(online_vars, target_vars, tau)` operates on VarStore params +- **Replacement**: Already replaced by GPU EMA kernel on flat buffer (fused_training.rs). Delete the VarStore polyak path. + +### 4. Network weight initialization +- GpuLinear layers registered via `vars.linear_xavier()` +- NoisyLinear mu params registered via `register_noisy_mu_in_varstore()` +- **Replacement**: Xavier init directly on CudaSlice buffers (already done in flat buffer path) + +## What to Delete + +### Struct fields +- `Sequential.vars: GpuVarStore` +- `DistributionalDuelingQNetwork.vars: GpuVarStore` +- `BranchingDuelingQNetwork.vars: GpuVarStore` + +### Methods +- `vars()` / `vars_mut()` on all 3 network structs +- `DQN.get_q_network_vars()` / `get_q_network_vars_mut()` +- `register_noisy_mu_in_varstore()` +- `flatten_online_weights()` / `flatten_target_weights()` (replace with direct buffer init) +- `unflatten_online_weights()` / `unflatten_target_weights()` (delete — flat buffer IS the source) +- All `polyak_update` functions in `target_update.rs` (GPU EMA kernel replaces) + +### Functions in gpu_weights.rs +- `extract_weight()`, `extract_dueling_weights()`, `extract_branching_weights()` +- All `reupload_*()` functions +- **Replace with**: functions that take `&DuelingWeightSet` / `&BranchingWeightSet` instead of `&GpuVarStore` + +### Experience collector methods +- All 7 `sync_*()` methods that take `&GpuVarStore` +- **Replace with**: methods that take weight set references directly + +## What to Keep + +- `params_buf` / `target_params_buf` flat buffers (the real source of truth) +- `DuelingWeightSet` / `BranchingWeightSet` (individual CudaSlice references into params_buf) +- GPU EMA kernel for Polyak updates +- `compute_param_sizes()` (defines the flat buffer layout) + +## Checkpoint Format Change + +### Current (safetensors via VarStore) +``` +model.safetensors: { "shared_1.weight": [64, 80], "shared_1.bias": [64], ... } +``` + +### New (raw flat buffer + metadata) +``` +model.fxparams: raw f32 bytes (params_buf dumped to disk) +model.fxparams.json: { "version": 2, "total_params": 336548, "param_layout": [...] } +``` + +Or keep safetensors format by constructing the BTreeMap> from the flat buffer using compute_param_sizes() to determine offsets + names. This preserves backward compatibility. + +## Implementation Order + +1. Replace checkpoint save/load to use flat buffer +2. Replace weight extraction to use DuelingWeightSet directly +3. Replace experience collector sync methods +4. Delete polyak_update VarStore path (GPU EMA already handles this) +5. Delete GpuVarStore fields from network structs +6. Delete GpuVarStore imports +7. Re-enable bottleneck_dim=16 +8. Full test suite verification + +## Success Criteria + +1. All 903 + 287 + 302 tests pass +2. Smoke test passes with bottleneck_dim=16 +3. No `GpuVarStore` references remain in DQN training pipeline +4. Checkpoint save/load roundtrip verified +5. H100 deployment with bottleneck=16 shows improved magnitude diversity + +## Key Files + +| File | Changes | +|------|---------| +| `crates/ml-dqn/src/dqn.rs` | Delete vars field, vars(), get_q_network_vars() | +| `crates/ml-dqn/src/branching.rs` | Delete vars, register_noisy_mu, polyak_update_from VarStore path | +| `crates/ml-dqn/src/distributional_dueling.rs` | Delete vars, GpuLinear registration | +| `crates/ml-dqn/src/target_update.rs` | Delete polyak_update VarStore functions | +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Delete flatten/unflatten, update checkpoint | +| `crates/ml/src/cuda_pipeline/gpu_weights.rs` | Replace extract/reupload to use weight sets | +| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Replace sync methods | +| `crates/ml/src/trainers/dqn/config.rs` | Remove GpuVarStore imports, update checkpoint | +| `crates/ml-core/src/checkpoint.rs` | New flat buffer serialization |