docs: GPU PER sum-tree implementation plan (13 tasks)
TDD plan covering GpuReplayBuffer, proportional + rank-based GPU sampling, priority scatter updates, async loss readback, trainer integration, OOM fallback, and distribution correctness tests. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
615
docs/plans/2026-03-02-gpu-per-sumtree-implementation.md
Normal file
615
docs/plans/2026-03-02-gpu-per-sumtree-implementation.md
Normal file
@@ -0,0 +1,615 @@
|
||||
# GPU-Resident PER Sum-Tree Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Move Prioritized Experience Replay entirely to GPU — sampling, priority updates, experience storage — eliminating the last CPU bottleneck in the DQN training loop.
|
||||
|
||||
**Architecture:** A new `GpuReplayBuffer` stores experiences as contiguous GPU tensors in a ring buffer. Sampling uses CUDA parallel prefix-sum + binary search instead of a CPU segment tree. Priority updates happen GPU-side from TD errors (no `to_vec1()` needed). Loss readback uses async DMA to eliminate per-step GPU stalls.
|
||||
|
||||
**Tech Stack:** Rust, cudarc 0.17.3 (via candle-core/cuda), NVRTC for custom kernels, candle-core Tensors for integration.
|
||||
|
||||
**Design doc:** `docs/plans/2026-03-02-gpu-per-sumtree-design.md`
|
||||
|
||||
---
|
||||
|
||||
### Task 1: GpuBatch struct and BatchSample extension
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/dqn/replay_buffer_type.rs:14-29`
|
||||
|
||||
**Step 1: Add GpuBatch struct after BatchSample**
|
||||
|
||||
After `BatchSample` (line 29), add:
|
||||
|
||||
```rust
|
||||
/// Pre-built GPU tensors for a training batch.
|
||||
/// When present, compute_gradients() uses these directly — no CPU→GPU transfer.
|
||||
#[cfg(feature = "cuda")]
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GpuBatch {
|
||||
pub states: candle_core::Tensor, // [batch_size, state_dim] f32 on GPU
|
||||
pub actions: candle_core::Tensor, // [batch_size] u32 on GPU
|
||||
pub rewards: candle_core::Tensor, // [batch_size] f32 on GPU
|
||||
pub next_states: candle_core::Tensor, // [batch_size, state_dim] f32 on GPU
|
||||
pub dones: candle_core::Tensor, // [batch_size] f32 on GPU (0.0/1.0)
|
||||
pub weights: candle_core::Tensor, // [batch_size] f32 on GPU (IS weights)
|
||||
pub indices: candle_core::Tensor, // [batch_size] u32 on GPU (buffer indices)
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Add gpu_batch field to BatchSample**
|
||||
|
||||
Change `BatchSample` struct (lines 16-20) to:
|
||||
|
||||
```rust
|
||||
pub struct BatchSample {
|
||||
pub experiences: Vec<Experience>,
|
||||
pub weights: Vec<f32>,
|
||||
pub indices: Vec<usize>,
|
||||
#[cfg(feature = "cuda")]
|
||||
pub gpu_batch: Option<GpuBatch>,
|
||||
}
|
||||
```
|
||||
|
||||
**Step 3: Update BatchSample::uniform() constructor**
|
||||
|
||||
Update `uniform()` (line 25-28) to include `gpu_batch: None`.
|
||||
|
||||
**Step 4: Verify build**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | head -20`
|
||||
Expected: Fix any compilation errors from the new field in match arms across the codebase.
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/replay_buffer_type.rs
|
||||
git commit -m "feat(ml): add GpuBatch struct and gpu_batch field to BatchSample"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: GpuReplayBuffer core struct
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs`
|
||||
|
||||
**Step 1: Write the failing test**
|
||||
|
||||
At the bottom of the new file, add:
|
||||
|
||||
```rust
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_gpu_replay_buffer_creation() {
|
||||
// CPU fallback test — no CUDA required
|
||||
let config = GpuReplayBufferConfig {
|
||||
capacity: 1000,
|
||||
state_dim: 51,
|
||||
alpha: 0.6,
|
||||
beta_start: 0.4,
|
||||
beta_max: 1.0,
|
||||
beta_annealing_steps: 100_000,
|
||||
epsilon: 1e-6,
|
||||
};
|
||||
let buf = GpuReplayBuffer::new(config, &candle_core::Device::Cpu);
|
||||
assert!(buf.is_ok());
|
||||
let buf = buf.unwrap();
|
||||
assert_eq!(buf.len(), 0);
|
||||
assert_eq!(buf.capacity(), 1000);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Implement GpuReplayBuffer struct**
|
||||
|
||||
```rust
|
||||
//! GPU-Resident Replay Buffer with parallel prefix-sum PER sampling
|
||||
//!
|
||||
//! Stores all experience data as contiguous GPU tensors in a ring buffer.
|
||||
//! Sampling uses prefix-sum + binary search (proportional) or radix sort
|
||||
//! (rank-based) — all on GPU with zero CPU involvement.
|
||||
|
||||
use candle_core::{Device, Tensor, DType};
|
||||
use crate::MLError;
|
||||
|
||||
/// Configuration for GPU replay buffer
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GpuReplayBufferConfig {
|
||||
pub capacity: usize,
|
||||
pub state_dim: usize,
|
||||
pub alpha: f32,
|
||||
pub beta_start: f32,
|
||||
pub beta_max: f32,
|
||||
pub beta_annealing_steps: usize,
|
||||
pub epsilon: f32, // Small constant added to priorities
|
||||
}
|
||||
|
||||
/// GPU-resident ring buffer for experience replay
|
||||
pub struct GpuReplayBuffer {
|
||||
config: GpuReplayBufferConfig,
|
||||
device: Device,
|
||||
|
||||
// GPU tensor storage (pre-allocated, ring buffer)
|
||||
states: Tensor, // [capacity, state_dim] f32
|
||||
next_states: Tensor, // [capacity, state_dim] f32
|
||||
actions: Tensor, // [capacity] u32
|
||||
rewards: Tensor, // [capacity] f32
|
||||
dones: Tensor, // [capacity] f32
|
||||
priorities: Tensor, // [capacity] f32
|
||||
|
||||
// Ring buffer state (CPU-side, just indices)
|
||||
write_cursor: usize,
|
||||
size: usize,
|
||||
max_priority: f32,
|
||||
|
||||
// Beta annealing
|
||||
current_step: usize,
|
||||
}
|
||||
```
|
||||
|
||||
Implement `new()`, `len()`, `capacity()`, `is_empty()`, `can_sample()`.
|
||||
|
||||
The `new()` method pre-allocates zero tensors for all storage on the given device.
|
||||
|
||||
**Step 3: Run test**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --lib gpu_replay_buffer::tests::test_gpu_replay_buffer_creation`
|
||||
Expected: PASS
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs crates/ml/src/cuda_pipeline/mod.rs
|
||||
git commit -m "feat(ml): GpuReplayBuffer core struct with ring buffer storage"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Experience insertion (GPU→GPU ring buffer)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs`
|
||||
|
||||
**Step 1: Write the failing test**
|
||||
|
||||
```rust
|
||||
#[test]
|
||||
fn test_gpu_replay_buffer_insert_and_len() {
|
||||
let config = GpuReplayBufferConfig {
|
||||
capacity: 100, state_dim: 4, alpha: 0.6,
|
||||
beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6,
|
||||
};
|
||||
let device = Device::Cpu;
|
||||
let mut buf = GpuReplayBuffer::new(config, &device).unwrap();
|
||||
|
||||
// Insert a batch of 10 experiences
|
||||
let states = Tensor::zeros(&[10, 4], DType::F32, &device).unwrap();
|
||||
let next_states = Tensor::zeros(&[10, 4], DType::F32, &device).unwrap();
|
||||
let actions = Tensor::zeros(&[10], DType::U32, &device).unwrap();
|
||||
let rewards = Tensor::ones(&[10], DType::F32, &device).unwrap();
|
||||
let dones = Tensor::zeros(&[10], DType::F32, &device).unwrap();
|
||||
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap();
|
||||
assert_eq!(buf.len(), 10);
|
||||
|
||||
// Insert more, verify ring buffer wrapping
|
||||
for _ in 0..15 {
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap();
|
||||
}
|
||||
assert_eq!(buf.len(), 100); // Capped at capacity
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Implement insert_batch()**
|
||||
|
||||
Method copies incoming tensors into the ring buffer using `Tensor::narrow()` + `Tensor::copy_()` (or slice assignment). New experiences get `max_priority`. Write cursor wraps at capacity.
|
||||
|
||||
Key: When batch wraps around end of ring buffer, split into two copies (end segment + start segment).
|
||||
|
||||
**Step 3: Run test, verify pass**
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "feat(ml): GpuReplayBuffer::insert_batch with ring buffer wrapping"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Proportional sampling kernel (prefix-sum + binary search)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs`
|
||||
|
||||
**Step 1: Write the failing test**
|
||||
|
||||
```rust
|
||||
#[test]
|
||||
fn test_gpu_replay_buffer_proportional_sample() {
|
||||
let config = GpuReplayBufferConfig {
|
||||
capacity: 100, state_dim: 4, alpha: 0.6,
|
||||
beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6,
|
||||
};
|
||||
let device = Device::Cpu;
|
||||
let mut buf = GpuReplayBuffer::new(config, &device).unwrap();
|
||||
|
||||
// Fill with 50 experiences (varying rewards for priority diversity)
|
||||
let states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap();
|
||||
let next_states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap();
|
||||
let actions = Tensor::zeros(&[50], DType::U32, &device).unwrap();
|
||||
let rewards = Tensor::randn(0.0_f32, 1.0, &[50], &device).unwrap();
|
||||
let dones = Tensor::zeros(&[50], DType::F32, &device).unwrap();
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap();
|
||||
|
||||
// Sample batch of 16
|
||||
let batch = buf.sample_proportional(16).unwrap();
|
||||
assert_eq!(batch.states.dim(0).unwrap(), 16);
|
||||
assert_eq!(batch.states.dim(1).unwrap(), 4);
|
||||
assert_eq!(batch.weights.dim(0).unwrap(), 16);
|
||||
assert_eq!(batch.indices.dim(0).unwrap(), 16);
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Implement sample_proportional()**
|
||||
|
||||
For **CPU device** (test fallback): Use Candle tensor ops — `cumsum`, random uniform tensor, `searchsorted` equivalent via manual binary search in Rust.
|
||||
|
||||
For **CUDA device**: Compile NVRTC kernel that does:
|
||||
1. Generate random values via simple LCG (or use curand if available)
|
||||
2. Compute cumulative sum of priorities (or maintain incrementally)
|
||||
3. Binary search per thread: each of B threads finds its index
|
||||
4. Gather experiences into output tensors
|
||||
5. Compute IS weights: `w_i = (N * p_i / sum)^(-beta) / max_weight`
|
||||
|
||||
Return `GpuBatch` with all tensors on GPU.
|
||||
|
||||
**Step 3: Run test, verify pass**
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "feat(ml): proportional PER sampling on GPU with prefix-sum"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Rank-based sampling kernel
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs`
|
||||
|
||||
**Step 1: Write the failing test**
|
||||
|
||||
```rust
|
||||
#[test]
|
||||
fn test_gpu_replay_buffer_rank_sample() {
|
||||
// Same setup as proportional test
|
||||
let config = GpuReplayBufferConfig {
|
||||
capacity: 100, state_dim: 4, alpha: 0.6,
|
||||
beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6,
|
||||
};
|
||||
let device = Device::Cpu;
|
||||
let mut buf = GpuReplayBuffer::new(config, &device).unwrap();
|
||||
|
||||
let states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap();
|
||||
let next_states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap();
|
||||
let actions = Tensor::zeros(&[50], DType::U32, &device).unwrap();
|
||||
let rewards = Tensor::randn(0.0_f32, 1.0, &[50], &device).unwrap();
|
||||
let dones = Tensor::zeros(&[50], DType::F32, &device).unwrap();
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap();
|
||||
|
||||
let batch = buf.sample_rank_based(16).unwrap();
|
||||
assert_eq!(batch.states.dim(0).unwrap(), 16);
|
||||
assert_eq!(batch.weights.dim(0).unwrap(), 16);
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Implement sample_rank_based()**
|
||||
|
||||
For **CPU device**: Sort priorities descending, compute rank probabilities `1/rank^alpha`, normalize, cumulative sum, sample.
|
||||
|
||||
For **CUDA device**: Use `Tensor::sort()` (candle wraps cub sort) → rank probability kernel → sample.
|
||||
|
||||
**Step 3: Run test, verify pass**
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "feat(ml): rank-based PER sampling on GPU with radix sort"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: GPU-side priority update (no to_vec1)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs`
|
||||
|
||||
**Step 1: Write the failing test**
|
||||
|
||||
```rust
|
||||
#[test]
|
||||
fn test_gpu_replay_buffer_priority_update() {
|
||||
let config = GpuReplayBufferConfig {
|
||||
capacity: 100, state_dim: 4, alpha: 0.6,
|
||||
beta_start: 0.4, beta_max: 1.0, beta_annealing_steps: 1000, epsilon: 1e-6,
|
||||
};
|
||||
let device = Device::Cpu;
|
||||
let mut buf = GpuReplayBuffer::new(config, &device).unwrap();
|
||||
|
||||
let states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap();
|
||||
let next_states = Tensor::randn(0.0_f32, 1.0, &[50, 4], &device).unwrap();
|
||||
let actions = Tensor::zeros(&[50], DType::U32, &device).unwrap();
|
||||
let rewards = Tensor::randn(0.0_f32, 1.0, &[50], &device).unwrap();
|
||||
let dones = Tensor::zeros(&[50], DType::F32, &device).unwrap();
|
||||
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones).unwrap();
|
||||
|
||||
// Sample, then update priorities with fake TD errors
|
||||
let batch = buf.sample_proportional(16).unwrap();
|
||||
let td_errors = Tensor::randn(0.0_f32, 1.0, &[16], &device).unwrap();
|
||||
buf.update_priorities_gpu(&batch.indices, &td_errors).unwrap();
|
||||
|
||||
// Priorities should have changed — sample again and verify different weights
|
||||
let batch2 = buf.sample_proportional(16).unwrap();
|
||||
// Just verify it doesn't crash — statistical tests are in integration tests
|
||||
assert_eq!(batch2.weights.dim(0).unwrap(), 16);
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Implement update_priorities_gpu()**
|
||||
|
||||
Takes `indices: &Tensor` (u32 on GPU) and `td_errors: &Tensor` (f32 on GPU).
|
||||
Computes `new_priority = |td_error|^alpha + epsilon` via Candle ops (`.abs()`, `.powf()`, `.add()`).
|
||||
Scatter-writes into `self.priorities` at given indices.
|
||||
Updates `max_priority` (via `.max()` on the new priorities).
|
||||
|
||||
Key: This is ALL tensor ops — no `to_vec1()`, no CPU involvement.
|
||||
|
||||
**Step 3: Run test, verify pass**
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "feat(ml): GPU-side priority scatter update from TD errors"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: ReplayBufferType::GpuPrioritized variant
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/dqn/replay_buffer_type.rs:36-41`
|
||||
|
||||
**Step 1: Add the enum variant**
|
||||
|
||||
Add to `ReplayBufferType` enum (after line 40):
|
||||
|
||||
```rust
|
||||
#[cfg(feature = "cuda")]
|
||||
GpuPrioritized(Arc<parking_lot::Mutex<crate::cuda_pipeline::gpu_replay_buffer::GpuReplayBuffer>>),
|
||||
```
|
||||
|
||||
**Step 2: Implement sample/update/add dispatch for the new variant**
|
||||
|
||||
In `sample()` (line 84): Add `GpuPrioritized` match arm that calls `gpu_buf.lock().sample_proportional(batch_size)` and wraps in `BatchSample` with `gpu_batch: Some(...)`.
|
||||
|
||||
In `update_priorities()` (line 145): Add arm that calls `gpu_buf.lock().update_priorities_gpu()` — but needs `Tensor` inputs. Add a new method `update_priorities_gpu()` on `ReplayBufferType` that takes tensors directly.
|
||||
|
||||
In `add()` / `add_batch()`: For now, these are no-ops for GpuPrioritized — experiences are inserted via `insert_batch()` from the GPU experience collector directly.
|
||||
|
||||
**Step 3: Update Debug impl and all match arms across the file**
|
||||
|
||||
**Step 4: Verify build**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "feat(ml): ReplayBufferType::GpuPrioritized variant with dispatch"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: Wire GpuReplayBuffer into DQN trainer
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer.rs`
|
||||
|
||||
**Step 1: Buffer creation**
|
||||
|
||||
Find where `ReplayBufferType::Prioritized` or `ReplayBufferType::Uniform` is created during trainer init. Add `#[cfg(feature = "cuda")]` block that creates `GpuPrioritized` when device is CUDA and `use_per` is true.
|
||||
|
||||
**Step 2: Experience insertion**
|
||||
|
||||
Find where GPU experience collector inserts experiences (search for `store_experience` or `add_batch` near the collector code). Change to call `gpu_buffer.insert_batch()` with the GPU tensors directly when `GpuPrioritized`.
|
||||
|
||||
**Step 3: compute_gradients() — use GpuBatch**
|
||||
|
||||
In `crates/ml/src/dqn/dqn.rs:2039` (`compute_gradients`), which calls `compute_loss_internal()`:
|
||||
- Check if `batch.gpu_batch.is_some()`
|
||||
- If yes: use GPU tensors directly for states, actions, rewards, next_states, dones, weights
|
||||
- Skip the CPU→GPU tensor construction code (the `Tensor::from_vec` calls)
|
||||
- Skip `diff.detach().to_vec1()` (line 1961) — TD errors stay on GPU
|
||||
- Return `GradientResult` with `td_errors: Vec::new()` and add a new field `td_errors_gpu: Option<Tensor>`
|
||||
|
||||
**Step 4: Priority update wiring**
|
||||
|
||||
In `train_step_with_accumulation()` (trainer.rs ~line 3478), when `GpuPrioritized`:
|
||||
- Collect `td_errors_gpu` tensors from `GradientResult` instead of `Vec<f32>`
|
||||
- Call `buffer.update_priorities_gpu(&indices_tensor, &td_errors_tensor)` instead of CPU `update_priorities()`
|
||||
|
||||
**Step 5: Verify build**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --lib`
|
||||
|
||||
**Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "feat(ml): wire GpuReplayBuffer into DQN trainer hot path"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 9: Async loss readback
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/dqn/dqn.rs` (around line 1945-1948)
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer.rs` (epoch logging)
|
||||
|
||||
**Step 1: Defer loss.to_scalar()**
|
||||
|
||||
In `compute_loss_internal()` (dqn.rs line 1945-1948), instead of:
|
||||
```rust
|
||||
let loss_value = loss_tensor.to_scalar::<f32>()?;
|
||||
```
|
||||
|
||||
Return `loss_value: None` and keep `loss_tensor` alive. The caller accumulates loss tensors on GPU.
|
||||
|
||||
**Step 2: Batch loss readback at accumulation boundary**
|
||||
|
||||
In `train_step_with_accumulation()`, after the accumulation loop:
|
||||
- Sum all loss tensors on GPU: `loss_sum = accumulated_losses.sum()?`
|
||||
- Single `loss_sum.to_scalar::<f32>()` for logging
|
||||
- This is 1 sync per accumulation group instead of N
|
||||
|
||||
**Step 3: Test existing training tests still pass**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- dqn_trainer`
|
||||
Expected: All pass (CPU fallback path unchanged)
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "perf(ml): batch loss readback — 1 GPU sync per accumulation instead of N"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 10: Integration test — full GPU PER training loop
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml/tests/gpu_per_integration_test.rs`
|
||||
|
||||
**Step 1: Write integration test**
|
||||
|
||||
```rust
|
||||
//! Integration test: GPU-resident PER training loop
|
||||
//! Verifies that GpuReplayBuffer produces equivalent training dynamics
|
||||
//! to the CPU PrioritizedReplayBuffer.
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gpu_per_training_loop() {
|
||||
// Create DQN trainer with GpuReplayBuffer
|
||||
// Run 3 epochs on synthetic data
|
||||
// Verify: loss decreases, priorities update, no panics
|
||||
// Compare: GPU PER loss trajectory vs CPU PER (within 10% tolerance)
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Write KS distribution test**
|
||||
|
||||
```rust
|
||||
#[test]
|
||||
fn test_gpu_per_sampling_distribution_matches_cpu() {
|
||||
// Sample 10K indices from GPU PER and CPU PER with same priorities
|
||||
// KS test: distributions should not differ significantly (p > 0.01)
|
||||
}
|
||||
```
|
||||
|
||||
**Step 3: Run tests**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --test gpu_per_integration_test`
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "test(ml): GPU PER integration tests — training loop + distribution correctness"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 11: Fallback and OOM handling
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs`
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer.rs`
|
||||
|
||||
**Step 1: OOM fallback in GpuReplayBuffer::new()**
|
||||
|
||||
Wrap tensor allocations in a `match` that catches CUDA OOM. If allocation fails:
|
||||
- Log warning: "GPU replay buffer allocation failed, falling back to CPU PER"
|
||||
- Return `Err(MLError::GpuMemoryError(...))`
|
||||
|
||||
**Step 2: Trainer fallback**
|
||||
|
||||
In trainer init, wrap `GpuReplayBuffer::new()` in:
|
||||
```rust
|
||||
match GpuReplayBuffer::new(config, device) {
|
||||
Ok(buf) => ReplayBufferType::GpuPrioritized(Arc::new(Mutex::new(buf))),
|
||||
Err(e) => {
|
||||
warn!("GPU PER failed ({}), falling back to CPU PER", e);
|
||||
ReplayBufferType::Prioritized(...)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Step 3: Test fallback**
|
||||
|
||||
Write test that forces OOM (absurd capacity) and verifies CPU fallback.
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git commit -am "feat(ml): GPU PER graceful fallback to CPU on OOM"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 12: Clippy, warnings, final verification
|
||||
|
||||
**Files:** All modified files
|
||||
|
||||
**Step 1: Full workspace check**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check --workspace
|
||||
SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings
|
||||
```
|
||||
|
||||
**Step 2: Run all ml tests**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib
|
||||
```
|
||||
Expected: 2390+ tests pass, 0 warnings
|
||||
|
||||
**Step 3: Commit any fixes**
|
||||
|
||||
```bash
|
||||
git commit -am "fix(ml): clippy and warning cleanup for GPU PER"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 13: Update memory and backlog
|
||||
|
||||
**Files:**
|
||||
- Modify: `/home/jgrusewski/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/gpu-saturation-backlog.md`
|
||||
- Modify: `/home/jgrusewski/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md`
|
||||
|
||||
**Step 1:** Mark "DQN PER CPU sum-tree" as resolved in the backlog.
|
||||
|
||||
**Step 2:** Update MEMORY.md with GPU PER details.
|
||||
|
||||
**Step 3: Push and trigger pipeline**
|
||||
|
||||
```bash
|
||||
git push origin main
|
||||
```
|
||||
|
||||
Trigger `train-validate-rl` job to verify epoch time improvement (target: 31s → ~26s on L40S).
|
||||
Reference in New Issue
Block a user