Files

jgrusewski b4e1ce30e1 docs: add gradient accumulation implementation plan

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

2026-02-20 21:59:14 +01:00

39 KiB

Raw Blame History

True Gradient Accumulation — Implementation Plan

For Claude: REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

Goal: Replace fake gradient accumulation (N independent optimizer steps) with true gradient accumulation (N forward+backward passes, 1 optimizer step).

Architecture: Add backward_and_clip() + apply_grads() to MonitoredOptimizer (lib.rs). Extract compute_loss() from DQN's train_step() to share forward-pass logic. Add compute_gradients() + apply_accumulated_gradients() to DQN. Rewrite trainer's train_step_with_accumulation() to accumulate GradStore objects and step once. Add convergence test.

Tech Stack: Rust, candle-core (GradStore, Tensor, Var), candle-optimisers (Adam)

Context

MonitoredOptimizer: ml/src/lib.rs:87 — wrapper around candle's Adam. Has backward_step_with_monitoring() which does backward + clip + step in one call.
DQN::train_step(): ml/src/dqn/dqn.rs:1275-1989 — ~700 lines covering batch prep, optimizer init, forward pass, 3 loss paths (IQN, C51, standard), backward+step, and bookkeeping.
Trainer: ml/src/trainers/dqn/trainer.rs:2808 — train_step_with_accumulation() runs N sequential agent.train_step() calls (each does its own optimizer step — not true accumulation).
GradStore: candle_core::backprop::GradStore — wraps HashMap<TensorId, Tensor>. Public API: get(tensor), remove(tensor), insert(tensor, grad), get_ids(). Constructor new() is private — first batch's GradStore becomes the accumulator.
DQNAgentType: ml/src/trainers/dqn/config.rs:57 — dispatches to Standard(DQN) or RegimeConditional(RegimeConditionalDQN). Only Standard supports accumulation (RegimeConditional has per-regime heads with separate optimizers).
Clippy rules: Denies unwrap_used, expect_used, panic, indexing_slicing. Use get() + error handling.
Build command: SQLX_OFFLINE=true cargo test -p ml ...

Task 1: Add `backward_and_clip()` + `apply_grads()` to MonitoredOptimizer

Files:

Modify: ml/src/lib.rs:189-235 (add two methods after backward_step_with_monitoring)

Step 1: Add backward_and_clip() method

Add this method to the impl Adam block in ml/src/lib.rs, right after backward_step_with_monitoring() (after line 235):

    /// Compute gradients and clip them without stepping the optimizer.
    ///
    /// Used for gradient accumulation: call this N times to accumulate
    /// gradients, then call `apply_grads()` once to update weights.
    ///
    /// Returns (clipped_gradients, gradient_norm).
    pub fn backward_and_clip(
        &self,
        loss: &Tensor,
        max_norm: f64,
    ) -> Result<(GradStore, f64), MLError> {
        let mut grads = loss
            .backward()
            .map_err(|e| MLError::TrainingError(format!("Backward pass failed: {}", e)))?;

        let (_actual_grad_norm, clipped_grad_norm) =
            crate::gradient_utils::clip_grad_norm(&self.vars, &mut grads, max_norm)
                .map_err(|e| MLError::TrainingError(format!("Gradient clipping failed: {}", e)))?;

        Ok((grads, clipped_grad_norm))
    }

Step 2: Add apply_grads() method

Add right after backward_and_clip():

    /// Apply pre-computed gradients to update weights.
    ///
    /// This is the "step" half of gradient accumulation. Call after
    /// accumulating gradients from multiple mini-batches.
    pub fn apply_grads(&mut self, grads: &GradStore) -> Result<(), MLError> {
        Optimizer::step(&mut self.optimizer, grads)
            .map_err(|e| MLError::TrainingError(format!("Optimizer step failed: {}", e)))?;
        Ok(())
    }

Step 3: Add import for GradStore to the file header

At the top of ml/src/lib.rs, add to the existing imports:

use candle_core::backprop::GradStore;

Check if this import already exists. If GradStore is already referenced via full path (candle_core::backprop::GradStore) in the compute_gradient_norm method, add the use and simplify the references.

Step 4: Verify compilation

Run: SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20 Expected: Compiles with no new errors.

Step 5: Commit

git add ml/src/lib.rs
git commit -m "feat(ml): add backward_and_clip and apply_grads to Adam optimizer"

Task 2: Add gradient accumulation utility functions

Files:

Create: ml/src/gradient_accumulation.rs
Modify: ml/src/lib.rs (add pub mod gradient_accumulation;)

Step 1: Create the utility module

Create ml/src/gradient_accumulation.rs:

//! Gradient accumulation utilities for true multi-step gradient accumulation.
//!
//! These functions accumulate `GradStore` objects across mini-batches,
//! then scale the result so a single optimizer step trains on an
//! effective batch of `batch_size * accumulation_steps`.

use candle_core::{backprop::GradStore, Var};
use crate::MLError;

/// Accumulate gradients from `source` into `target`.
///
/// On the first call (`target` is `None`), `source` becomes the accumulator.
/// On subsequent calls, each gradient tensor in `source` is added element-wise
/// to the corresponding tensor in `target`.
///
/// Only accumulates gradients for variables in `vars` (optimizer parameters).
pub fn accumulate_grads(
    target: &mut Option<GradStore>,
    source: GradStore,
    vars: &[Var],
) -> Result<(), MLError> {
    match target {
        None => {
            // First mini-batch: use its GradStore directly as the accumulator
            *target = Some(source);
            Ok(())
        }
        Some(ref mut acc) => {
            // Subsequent mini-batches: add gradients element-wise
            for var in vars {
                if let Some(new_grad) = source.get(var) {
                    if let Some(existing_grad) = acc.remove(var) {
                        let summed = (&existing_grad + new_grad).map_err(|e| {
                            MLError::TrainingError(format!(
                                "Failed to accumulate gradients: {}",
                                e
                            ))
                        })?;
                        acc.insert(var, summed);
                    } else {
                        // Variable had no gradient in accumulator — insert the new one
                        acc.insert(var, new_grad.clone());
                    }
                }
            }
            Ok(())
        }
    }
}

/// Scale all gradients in the store by a constant factor.
///
/// Typically called with `1.0 / accumulation_steps` to average gradients
/// before the optimizer step.
pub fn scale_grads(
    grads: &mut GradStore,
    vars: &[Var],
    scale: f64,
) -> Result<(), MLError> {
    for var in vars {
        if let Some(grad) = grads.remove(var) {
            let scaled = (grad * scale).map_err(|e| {
                MLError::TrainingError(format!("Failed to scale gradient: {}", e))
            })?;
            grads.insert(var, scaled);
        }
    }
    Ok(())
}

Step 2: Register the module

In ml/src/lib.rs, add at an appropriate location among the other module declarations:

pub mod gradient_accumulation;

Search for the existing pub mod declarations (e.g. pub mod gradient_utils;) and add it nearby.

Step 3: Verify compilation

Run: SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20 Expected: Compiles with no new errors.

Step 4: Commit

git add ml/src/gradient_accumulation.rs ml/src/lib.rs
git commit -m "feat(ml): add gradient accumulation utility functions"

Task 3: Add `GradientResult` struct and extract `compute_loss_internal()` from DQN

This is the largest task. We extract the forward-pass + loss computation from train_step() into a shared compute_loss_internal() method, keeping train_step() behavior identical.

Files:

Modify: ml/src/dqn/dqn.rs:1275-1862 (extract method, refactor train_step)

Step 1: Add GradientResult struct

At the top of ml/src/dqn/dqn.rs, near other struct definitions, add:

use candle_core::backprop::GradStore;

/// Result of computing gradients for a mini-batch without stepping the optimizer.
///
/// Used for true gradient accumulation: collect multiple `GradientResult`s,
/// accumulate their `grads`, then apply once.
pub struct GradientResult {
    /// Scalar loss value for this mini-batch
    pub loss: f32,
    /// Gradient norm after clipping
    pub grad_norm: f32,
    /// Clipped gradients (not yet applied to weights)
    pub grads: GradStore,
    /// TD errors for PER priority updates
    pub td_errors: Vec<f32>,
    /// Replay buffer indices for PER priority updates
    pub indices: Vec<usize>,
}

Step 2: Define ComputeLossResult internal struct

Below GradientResult, add a struct for the compute_loss return type:

/// Internal result from forward pass + loss computation (no backward pass).
struct ComputeLossResult {
    /// Loss tensor (for backward pass)
    loss_tensor: Tensor,
    /// Scalar loss value
    loss_value: f32,
    /// TD errors for PER
    td_errors: Vec<f32>,
    /// PER indices
    indices: Vec<usize>,
    /// States tensor (needed for post-step diagnostics like log_q_values)
    states_tensor: Tensor,
}

Step 3: Extract compute_loss_internal()

This method contains lines 1275-1862 of the current train_step(), up to but NOT including the loss extraction at line 1863. The method signature:

    /// Shared forward pass + loss computation used by both `train_step()` and `compute_gradients()`.
    ///
    /// Handles: warmup check, batch sampling, optimizer init, tensor creation,
    /// forward pass (IQN/C51/standard), TD error computation, and loss (Huber/MSE + entropy + CQL).
    fn compute_loss_internal(
        &mut self,
        batch: Option<Vec<Experience>>,
    ) -> Result<ComputeLossResult, MLError> {

The body is the existing train_step() code from line 1275 through line 1862, with these changes:

Remove the warmup check (lines 1276-1279) — move it to callers (train_step and compute_gradients). Both callers will check warmup before calling this method.
At the end (after the loss is computed on line 1861), extract the scalar loss value and return:

        // Extract loss value
        let loss_value = loss
            .to_scalar::<f32>()
            .map_err(|e| MLError::TrainingError(format!("Failed to extract loss: {}", e)))?;

        Ok(ComputeLossResult {
            loss_tensor: loss,
            loss_value,
            td_errors: td_errors_vec,
            indices,
            states_tensor,
        })

Note: td_errors_vec is computed at line 1542 (let td_errors_vec: Vec<f32> = diff.detach().to_vec1()?;), indices comes from the batch sampling at line 1282-1295, and states_tensor is created at line 1362.

Step 4: Refactor train_step() to use compute_loss_internal()

Replace the body of train_step() with:

    pub fn train_step(&mut self, batch: Option<Vec<Experience>>) -> Result<(f32, f32), MLError> {
        // Skip gradient updates during warmup period
        if self.total_steps < self.config.warmup_steps as u64 {
            return Ok((0.0, 0.0));
        }

        let result = self.compute_loss_internal(batch)?;

        // Backward pass with gradient monitoring
        let grad_norm = if let Some(ref mut optimizer) = self.optimizer {
            let norm = optimizer
                .backward_step_with_monitoring(&result.loss_tensor, self.gradient_clip_norm)
                .map_err(|e| {
                    MLError::TrainingError(format!("Backward step with monitoring failed: {}", e))
                })?;
            tracing::debug!("Gradient norm: {:.4}", norm);
            norm as f32
        } else {
            return Err(MLError::TrainingError(
                "Optimizer not initialized".to_string(),
            ));
        };

        // --- Bookkeeping (identical to original) ---
        self.training_steps += 1;

        // Update PER priorities
        if !result.indices.is_empty() {
            self.memory.update_priorities(&result.indices, &result.td_errors)?;
        }
        self.memory.step();

        // Diagnostics
        if self.training_steps % 10 == 0 {
            self.log_q_values(&result.states_tensor)?;
        }
        if self.training_steps % 100 == 0 {
            self.log_diagnostics(grad_norm)?;
        }

        // Target network update (copy the existing code from lines 1905-1987 verbatim)
        // ... (the entire soft/hard update block) ...

        Ok((result.loss_value, grad_norm))
    }

The target network update block (lines 1905-1987) must be copied into the refactored train_step() verbatim. It handles soft EMA updates (cosine-annealed), hard updates, and IQN target updates.

Step 5: Verify compilation and test

Run: SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20 Expected: Compiles.

Run: SQLX_OFFLINE=true cargo test -p ml --test dqn_training_smoke_test -- --nocapture 2>&1 | tail -30 Expected: Smoke test still passes (refactoring preserved behavior).

If the smoke test data isn't available, at minimum verify: Run: SQLX_OFFLINE=true cargo test -p ml dqn 2>&1 | tail -20 Expected: All DQN unit tests pass.

Step 6: Commit

git add ml/src/dqn/dqn.rs
git commit -m "refactor(ml): extract compute_loss_internal from DQN train_step"

Task 4: Add `compute_gradients()` and `apply_accumulated_gradients()` to DQN

Files:

Modify: ml/src/dqn/dqn.rs (add two new public methods)

Step 1: Add compute_gradients()

Add this method to the impl DQN block, after train_step():

    /// Compute forward pass, loss, and gradients WITHOUT stepping the optimizer.
    ///
    /// Used for true gradient accumulation: call N times to accumulate
    /// gradients across mini-batches, then call `apply_accumulated_gradients()` once.
    pub fn compute_gradients(
        &mut self,
        batch: Option<Vec<Experience>>,
    ) -> Result<GradientResult, MLError> {
        // Skip during warmup
        if self.total_steps < self.config.warmup_steps as u64 {
            return Err(MLError::TrainingError(
                "Cannot compute gradients during warmup period".to_string(),
            ));
        }

        let result = self.compute_loss_internal(batch)?;

        // Backward + clip WITHOUT optimizer step
        let (grads, grad_norm) = if let Some(ref optimizer) = self.optimizer {
            optimizer.backward_and_clip(&result.loss_tensor, self.gradient_clip_norm)?
        } else {
            return Err(MLError::TrainingError(
                "Optimizer not initialized".to_string(),
            ));
        };

        Ok(GradientResult {
            loss: result.loss_value,
            grad_norm: grad_norm as f32,
            grads,
            td_errors: result.td_errors,
            indices: result.indices,
        })
    }

Step 2: Add apply_accumulated_gradients()

    /// Apply pre-accumulated gradients to update weights and perform bookkeeping.
    ///
    /// Call this once after accumulating gradients from N mini-batches.
    /// Handles: optimizer step, training_steps increment, target network update.
    ///
    /// PER priority updates and diagnostics are the caller's responsibility
    /// since they depend on accumulated TD errors from multiple batches.
    pub fn apply_accumulated_gradients(
        &mut self,
        grads: &GradStore,
    ) -> Result<(), MLError> {
        // Apply accumulated gradients
        if let Some(ref mut optimizer) = self.optimizer {
            optimizer.apply_grads(grads)?;
        } else {
            return Err(MLError::TrainingError(
                "Optimizer not initialized".to_string(),
            ));
        }

        // Increment training steps (once per effective batch, not per mini-batch)
        self.training_steps += 1;

        // Target network update (same as train_step)
        if self.config.use_soft_updates {
            let current_tau = if self.config.tau_anneal_steps > 0 {
                let progress = (self.training_steps as f64
                    / self.config.tau_anneal_steps as f64)
                    .min(1.0);
                let cosine_factor = (std::f64::consts::PI * progress).cos();
                self.config.tau_final
                    - (self.config.tau_final - self.config.tau) * (cosine_factor + 1.0) / 2.0
            } else {
                self.config.tau
            };

            // Soft update all network variants
            if let (Some(ref dist_dueling_net), Some(ref mut dist_dueling_target)) =
                (&self.dist_dueling_q_network, &mut self.dist_dueling_target_network)
            {
                polyak_update(dist_dueling_net.vars(), dist_dueling_target.vars(), current_tau)
                    .map_err(|e| MLError::TrainingError(format!("Hybrid EMA update failed: {}", e)))?;
            } else if let (Some(ref dueling_net), Some(ref mut dueling_target)) =
                (&self.dueling_q_network, &mut self.dueling_target_network)
            {
                polyak_update(dueling_net.vars(), dueling_target.vars(), current_tau)
                    .map_err(|e| MLError::TrainingError(format!("Dueling EMA update failed: {}", e)))?;
            } else {
                polyak_update(self.q_network.vars(), self.target_network.vars(), current_tau)
                    .map_err(|e| MLError::TrainingError(format!("EMA update failed: {}", e)))?;
            }
            if let (Some(ref iqn_net), Some(ref iqn_target)) =
                (&self.iqn_network, &self.iqn_target_network)
            {
                polyak_update(iqn_net.vars(), iqn_target.vars(), current_tau)
                    .map_err(|e| MLError::TrainingError(format!("IQN EMA update failed: {}", e)))?;
            }
        } else if self.training_steps % self.config.target_update_freq as u64 == 0 {
            // Hard update path
            if let (Some(ref dist_dueling_net), Some(ref mut dist_dueling_target)) =
                (&self.dist_dueling_q_network, &mut self.dist_dueling_target_network)
            {
                hard_update(dist_dueling_net.vars(), dist_dueling_target.vars())
                    .map_err(|e| MLError::TrainingError(format!("Hybrid hard update failed: {}", e)))?;
            } else if let (Some(ref dueling_net), Some(ref mut dueling_target)) =
                (&self.dueling_q_network, &mut self.dueling_target_network)
            {
                hard_update(dueling_net.vars(), dueling_target.vars())
                    .map_err(|e| MLError::TrainingError(format!("Dueling hard update failed: {}", e)))?;
            } else {
                hard_update(self.q_network.vars(), self.target_network.vars())
                    .map_err(|e| MLError::TrainingError(format!("Hard update failed: {}", e)))?;
            }
            if let (Some(ref iqn_net), Some(ref mut iqn_target)) =
                (&self.iqn_network, &mut self.iqn_target_network)
            {
                iqn_target.copy_weights_from(iqn_net)
                    .map_err(|e| MLError::TrainingError(format!("IQN hard update failed: {}", e)))?;
            }
        }

        Ok(())
    }

Step 3: Add method to expose optimizer vars for accumulation

The accumulation utility needs access to the optimizer's variable list. Add:

    /// Get the optimizer's tracked variables (needed for gradient accumulation).
    pub fn optimizer_vars(&self) -> Result<&[Var], MLError> {
        if let Some(ref optimizer) = self.optimizer {
            Ok(optimizer.vars())
        } else {
            Err(MLError::TrainingError(
                "Optimizer not initialized — call train_step or compute_gradients first".to_string(),
            ))
        }
    }

Step 4: Verify compilation

Run: SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20 Expected: Compiles.

Step 5: Commit

git add ml/src/dqn/dqn.rs
git commit -m "feat(ml): add compute_gradients and apply_accumulated_gradients to DQN"

Task 5: Add dispatch to DQNAgentType

Files:

Modify: ml/src/trainers/dqn/config.rs

Step 1: Add compute_gradients dispatch

Add to the impl DQNAgentType block (after train_step):

    /// Compute gradients without stepping the optimizer (for true gradient accumulation).
    /// Only supported for Standard DQN. RegimeConditional returns an error.
    pub fn compute_gradients(
        &mut self,
        batch: Option<Vec<crate::dqn::experience::Experience>>,
    ) -> Result<crate::dqn::dqn::GradientResult, crate::MLError> {
        match self {
            Self::Standard(agent) => agent.compute_gradients(batch),
            Self::RegimeConditional(_) => Err(crate::MLError::TrainingError(
                "Gradient accumulation not supported for RegimeConditionalDQN".to_string(),
            )),
        }
    }

Step 2: Add apply_accumulated_gradients dispatch

    /// Apply accumulated gradients (for true gradient accumulation).
    /// Only supported for Standard DQN.
    pub fn apply_accumulated_gradients(
        &mut self,
        grads: &candle_core::backprop::GradStore,
    ) -> Result<(), crate::MLError> {
        match self {
            Self::Standard(agent) => agent.apply_accumulated_gradients(grads),
            Self::RegimeConditional(_) => Err(crate::MLError::TrainingError(
                "Gradient accumulation not supported for RegimeConditionalDQN".to_string(),
            )),
        }
    }

Step 3: Add optimizer_vars dispatch

    /// Get optimizer variables for gradient accumulation.
    pub fn optimizer_vars(&self) -> Result<&[candle_core::Var], crate::MLError> {
        match self {
            Self::Standard(agent) => agent.optimizer_vars(),
            Self::RegimeConditional(_) => Err(crate::MLError::TrainingError(
                "Gradient accumulation not supported for RegimeConditionalDQN".to_string(),
            )),
        }
    }

Step 4: Add PER update dispatch

The trainer needs to update PER priorities after accumulation. Check if a update_priorities method already exists on DQNAgentType. If not, add:

    /// Update PER priorities with accumulated TD errors.
    pub fn update_priorities(&self, indices: &[usize], td_errors: &[f32]) -> Result<(), crate::MLError> {
        let buffer = self.memory();
        buffer.update_priorities(indices, td_errors)
    }

    /// Step the replay buffer beta annealing.
    pub fn step_replay_buffer(&self) {
        self.memory().step();
    }

Step 5: Verify compilation

Run: SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20 Expected: Compiles.

Step 6: Commit

git add ml/src/trainers/dqn/config.rs
git commit -m "feat(ml): add gradient accumulation dispatch to DQNAgentType"

Task 6: Rewrite `train_step_with_accumulation()` in the trainer

Files:

Modify: ml/src/trainers/dqn/trainer.rs:2808-2890

Step 1: Rewrite the method

Replace the entire train_step_with_accumulation() method body (lines 2808-end of method) with:

    async fn train_step_with_accumulation(&mut self) -> Result<(f64, f64, f64)> {
        let accumulation_steps = self.hyperparams.gradient_accumulation_steps;

        debug!(
            "Starting true gradient accumulation with {} steps (effective batch: {})",
            accumulation_steps,
            self.current_batch_size * accumulation_steps
        );

        let mut agent = self.agent.write().await;

        // === Phase 1: Accumulate gradients across N mini-batches ===
        let mut accumulated_grads: Option<candle_core::backprop::GradStore> = None;
        let mut total_loss = 0.0_f64;
        let mut all_td_errors = Vec::new();
        let mut all_indices = Vec::new();
        let mut final_grad_norm = 0.0_f32;

        for step in 0..accumulation_steps {
            // Sample mini-batch (respecting OOM-reduced batch size)
            let explicit_batch = if self.current_batch_size < self.hyperparams.batch_size {
                let buffer = agent.memory();
                if buffer.can_sample(self.current_batch_size) {
                    let batch_sample = buffer
                        .sample(self.current_batch_size)
                        .map_err(|e| {
                            anyhow::anyhow!(
                                "Failed to sample reduced batch (accum step {}): {}",
                                step,
                                e
                            )
                        })?;
                    Some(batch_sample.experiences)
                } else {
                    None
                }
            } else {
                None
            };

            // Compute forward pass + backward WITHOUT optimizer step
            let result = agent
                .compute_gradients(explicit_batch)
                .map_err(|e| anyhow::anyhow!("Gradient computation step {} failed: {}", step, e))?;

            // Accumulate gradients
            let vars = agent
                .optimizer_vars()
                .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?;
            crate::gradient_accumulation::accumulate_grads(
                &mut accumulated_grads,
                result.grads,
                vars,
            )
            .map_err(|e| anyhow::anyhow!("Gradient accumulation step {} failed: {}", step, e))?;

            total_loss += result.loss as f64;
            final_grad_norm = result.grad_norm;
            all_td_errors.extend(result.td_errors);
            all_indices.extend(result.indices);

            // Emergency brake
            if result.loss > 1e6 {
                warn!(
                    "Loss spike in accumulation step {}/{}: {:.2e}",
                    step + 1,
                    accumulation_steps,
                    result.loss
                );
            }
        }

        // === Phase 2: Average and apply gradients (single optimizer step) ===
        if let Some(ref mut grads) = accumulated_grads {
            let vars = agent
                .optimizer_vars()
                .map_err(|e| anyhow::anyhow!("Failed to get optimizer vars: {}", e))?;
            crate::gradient_accumulation::scale_grads(
                grads,
                vars,
                1.0 / accumulation_steps as f64,
            )
            .map_err(|e| anyhow::anyhow!("Gradient scaling failed: {}", e))?;

            agent
                .apply_accumulated_gradients(grads)
                .map_err(|e| anyhow::anyhow!("Apply accumulated gradients failed: {}", e))?;
        }

        // === Phase 3: Bookkeeping ===
        // Update PER priorities with all TD errors from all mini-batches
        if !all_indices.is_empty() {
            agent
                .update_priorities(&all_indices, &all_td_errors)
                .map_err(|e| anyhow::anyhow!("PER priority update failed: {}", e))?;
        }
        agent.step_replay_buffer();

        let avg_loss = total_loss / accumulation_steps as f64;

        // Gradient collapse detection
        agent.log_diagnostics(final_grad_norm)
            .map_err(|e| {
                tracing::info!("Early stopping triggered (gradient collapse): {}", e);
                anyhow::anyhow!("Early stopping: {}", e)
            })?;

        // Clip averaged loss
        let loss_clipped = if avg_loss > 1e6 {
            warn!(
                "Averaged loss clipped from {:.2e} to 1.0e6",
                avg_loss,
            );
            1e6
        } else {
            avg_loss
        };

        // Get Q-values for monitoring
        let avg_q_value = self.estimate_avg_q_value_with_early_stopping(&mut agent).await?;

        Ok((loss_clipped, final_grad_norm as f64, avg_q_value))
    }

Important: The update_priorities and step_replay_buffer calls need to be accessible through DQNAgentType. If Task 5 added them, use them. If they're already accessible via agent.memory(), use that pattern instead (check how the existing code calls them).

Step 2: Verify compilation

Run: SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20 Expected: Compiles.

Step 3: Commit

git add ml/src/trainers/dqn/trainer.rs
git commit -m "feat(ml): rewrite train_step_with_accumulation with true gradient accumulation"

Task 7: Unit test — single optimizer step per accumulation call

Files:

Modify: ml/tests/dqn_training_pipeline_test.rs (or create ml/tests/dqn_gradient_accumulation_test.rs)

Step 1: Write the test

Create ml/tests/dqn_gradient_accumulation_test.rs:

//! Tests for true gradient accumulation in DQN training.
//!
//! Verifies that `train_step_with_accumulation()` performs exactly 1 optimizer
//! step per call (not N steps), proving true accumulation works correctly.

#![allow(unused_crate_dependencies)]

use anyhow::{Context, Result};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use std::path::PathBuf;

fn get_6e_fut_data_dir() -> Result<String> {
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .context("Failed to get workspace root")?
        .to_path_buf();
    let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");
    if !data_dir.exists() {
        anyhow::bail!("6E.FUT data not found: {}", data_dir.display());
    }
    Ok(data_dir.to_string_lossy().to_string())
}

/// Verify that gradient accumulation steps=4 produces 1 optimizer step per epoch,
/// not 4 separate steps. We check this by comparing training_steps counters.
#[tokio::test]
async fn test_accumulation_single_optimizer_step() -> Result<()> {
    let data_dir = match get_6e_fut_data_dir() {
        Ok(dir) => dir,
        Err(e) => {
            eprintln!("Skipping: {}", e);
            return Ok(());
        }
    };

    let checkpoint_dir = tempfile::tempdir()?;

    // Train with accumulation_steps=4, 5 epochs
    let mut hyperparams = DQNHyperparameters::conservative();
    hyperparams.epochs = 5;
    hyperparams.batch_size = 32;
    hyperparams.learning_rate = 0.0001;
    hyperparams.gradient_accumulation_steps = 4;
    hyperparams.early_stopping_enabled = false;
    hyperparams.checkpoint_frequency = 100; // Don't checkpoint

    let mut trainer = DQNTrainer::new(hyperparams)?;
    let metrics = trainer
        .train(&data_dir, |_epoch, checkpoint_data, _is_best| {
            let path = checkpoint_dir.path().join("accum_test.safetensors");
            std::fs::write(&path, &checkpoint_data)?;
            Ok(path.to_string_lossy().to_string())
        })
        .await?;

    // With 5 epochs and accumulation_steps=4, each epoch does 1 optimizer step
    // (4 forward+backward, 1 step). So total optimizer steps should equal epochs_trained,
    // NOT epochs_trained * accumulation_steps.
    assert_eq!(
        metrics.epochs_trained, 5,
        "Should complete all 5 epochs"
    );

    // Loss should still decrease (training still works)
    let loss_history = trainer.loss_history();
    assert!(
        loss_history.len() >= 2,
        "Need at least 2 epochs of loss history"
    );
    let initial_loss = loss_history.first().copied().unwrap_or(0.0);
    let final_loss = loss_history.last().copied().unwrap_or(0.0);
    assert!(
        initial_loss.is_finite() && final_loss.is_finite(),
        "Losses must be finite: initial={}, final={}",
        initial_loss,
        final_loss
    );

    println!("\n=== GRADIENT ACCUMULATION TEST ===");
    println!("  Epochs:              {}", metrics.epochs_trained);
    println!("  Accumulation steps:  4");
    println!("  Initial loss:        {:.6}", initial_loss);
    println!("  Final loss:          {:.6}", final_loss);
    println!("  Effective batch:     {} (32 * 4)", 32 * 4);
    println!("  Training time:       {:.1}s", metrics.training_time_seconds);
    println!("=================================");

    Ok(())
}

Step 2: Verify it compiles

Run: SQLX_OFFLINE=true cargo test -p ml --test dqn_gradient_accumulation_test --no-run 2>&1 | tail -10 Expected: Compiles.

Step 3: Run the test (if data available)

Run: SQLX_OFFLINE=true cargo test -p ml --test dqn_gradient_accumulation_test -- --nocapture 2>&1 | tail -30 Expected: PASS.

Step 4: Commit

git add ml/tests/dqn_gradient_accumulation_test.rs
git commit -m "test(ml): add gradient accumulation single-step verification test"

Task 8: Convergence test — accumulated vs non-accumulated training

Files:

Create: ml/tests/dqn_accumulation_convergence_test.rs

Step 1: Write the convergence test

//! Convergence comparison test for gradient accumulation.
//!
//! Compares loss trajectories between:
//! - Training with batch_size=16, accumulation_steps=4 (effective batch=64)
//! - Training with batch_size=64, accumulation_steps=1 (direct batch=64)
//!
//! Both should produce similar loss curves since the effective batch size is identical.
//! We allow 30% tolerance due to sampling order differences and floating-point effects.

#![allow(unused_crate_dependencies)]

use anyhow::{Context, Result};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use std::path::PathBuf;

fn get_6e_fut_data_dir() -> Result<String> {
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .context("Failed to get workspace root")?
        .to_path_buf();
    let data_dir = workspace_root.join("test_data/real/databento/ml_training_small");
    if !data_dir.exists() {
        anyhow::bail!("6E.FUT data not found: {}", data_dir.display());
    }
    Ok(data_dir.to_string_lossy().to_string())
}

/// Compare accumulated training (batch=16, accum=4) vs direct training (batch=64).
/// Both have effective batch size 64. Loss trajectories should be within 30%.
#[tokio::test]
#[ignore] // Requires real data + runs two training sessions
async fn test_accumulation_convergence_similar_to_direct() -> Result<()> {
    let data_dir = get_6e_fut_data_dir()?;
    let epochs = 10;

    // --- Run 1: Accumulated training (batch=16, accum=4) ---
    let mut hp_accum = DQNHyperparameters::conservative();
    hp_accum.epochs = epochs;
    hp_accum.batch_size = 16;
    hp_accum.gradient_accumulation_steps = 4;
    hp_accum.learning_rate = 0.0001;
    hp_accum.early_stopping_enabled = false;
    hp_accum.checkpoint_frequency = 100;

    let checkpoint_dir_1 = tempfile::tempdir()?;
    let mut trainer_accum = DQNTrainer::new(hp_accum)?;
    let _metrics_accum = trainer_accum
        .train(&data_dir, |_epoch, data, _is_best| {
            let p = checkpoint_dir_1.path().join("accum.safetensors");
            std::fs::write(&p, &data)?;
            Ok(p.to_string_lossy().to_string())
        })
        .await?;
    let loss_accum = trainer_accum.loss_history().to_vec();

    // --- Run 2: Direct training (batch=64, accum=1) ---
    let mut hp_direct = DQNHyperparameters::conservative();
    hp_direct.epochs = epochs;
    hp_direct.batch_size = 64;
    hp_direct.gradient_accumulation_steps = 1;
    hp_direct.learning_rate = 0.0001;
    hp_direct.early_stopping_enabled = false;
    hp_direct.checkpoint_frequency = 100;

    let checkpoint_dir_2 = tempfile::tempdir()?;
    let mut trainer_direct = DQNTrainer::new(hp_direct)?;
    let _metrics_direct = trainer_direct
        .train(&data_dir, |_epoch, data, _is_best| {
            let p = checkpoint_dir_2.path().join("direct.safetensors");
            std::fs::write(&p, &data)?;
            Ok(p.to_string_lossy().to_string())
        })
        .await?;
    let loss_direct = trainer_direct.loss_history().to_vec();

    // --- Compare loss trajectories ---
    let min_len = loss_accum.len().min(loss_direct.len());
    assert!(min_len >= 5, "Both runs should complete at least 5 epochs");

    // Compare final losses — should be within 30% of each other
    let final_accum = loss_accum.get(min_len - 1).copied().unwrap_or(f64::NAN);
    let final_direct = loss_direct.get(min_len - 1).copied().unwrap_or(f64::NAN);

    assert!(
        final_accum.is_finite() && final_direct.is_finite(),
        "Final losses must be finite: accum={:.6}, direct={:.6}",
        final_accum,
        final_direct
    );

    // Relative difference
    let avg = (final_accum + final_direct) / 2.0;
    let rel_diff = (final_accum - final_direct).abs() / avg;

    // Both should show loss decrease
    let accum_decrease = 1.0 - final_accum / loss_accum.first().copied().unwrap_or(1.0);
    let direct_decrease = 1.0 - final_direct / loss_direct.first().copied().unwrap_or(1.0);

    println!("\n=== ACCUMULATION CONVERGENCE TEST ===");
    println!("  Accumulated (16x4):  initial={:.4} final={:.4} decrease={:.1}%",
        loss_accum.first().copied().unwrap_or(0.0), final_accum, accum_decrease * 100.0);
    println!("  Direct (64x1):       initial={:.4} final={:.4} decrease={:.1}%",
        loss_direct.first().copied().unwrap_or(0.0), final_direct, direct_decrease * 100.0);
    println!("  Relative difference: {:.1}% (threshold: 30%)", rel_diff * 100.0);
    println!("=====================================");

    assert!(
        rel_diff < 0.30,
        "Loss trajectories diverged too much: accum={:.6}, direct={:.6}, rel_diff={:.1}%",
        final_accum,
        final_direct,
        rel_diff * 100.0
    );

    // Both should show meaningful decrease (>2%)
    assert!(
        accum_decrease > 0.02,
        "Accumulated training should decrease loss >2% (got {:.1}%)",
        accum_decrease * 100.0
    );
    assert!(
        direct_decrease > 0.02,
        "Direct training should decrease loss >2% (got {:.1}%)",
        direct_decrease * 100.0
    );

    Ok(())
}

Step 2: Verify compilation

Run: SQLX_OFFLINE=true cargo test -p ml --test dqn_accumulation_convergence_test --no-run 2>&1 | tail -10 Expected: Compiles.

Step 3: Run if desired (slow, marked #[ignore])

Run: SQLX_OFFLINE=true cargo test -p ml --test dqn_accumulation_convergence_test -- --ignored --nocapture 2>&1 | tail -30 Expected: PASS.

Step 4: Commit

git add ml/tests/dqn_accumulation_convergence_test.rs
git commit -m "test(ml): add gradient accumulation convergence comparison test"

Task 9: Final verification — all existing tests pass

Files: None (verification only)

Step 1: Run all DQN tests

Run: SQLX_OFFLINE=true cargo test -p ml dqn 2>&1 | tail -30 Expected: All non-ignored tests pass. No regressions.

Step 2: Run the smoke test (if data available)

Run: SQLX_OFFLINE=true cargo test -p ml --test dqn_training_smoke_test -- --nocapture 2>&1 | tail -30 Expected: All 7 assertions pass (refactored train_step() has identical behavior).

Step 3: Run the new accumulation test (if data available)

Run: SQLX_OFFLINE=true cargo test -p ml --test dqn_gradient_accumulation_test -- --nocapture 2>&1 | tail -30 Expected: PASS.

Step 4: Run clippy

Run: SQLX_OFFLINE=true cargo clippy -p ml -- -D warnings 2>&1 | tail -20 Expected: No new warnings in our changed files.

Step 5: Commit if any fixes needed, otherwise done

If clippy or tests reveal issues, fix and commit with: fix(ml): address clippy/test issues in gradient accumulation

39 KiB Raw Blame History

True Gradient Accumulation — Implementation Plan

Context

Task 1: Add backward_and_clip() + apply_grads() to MonitoredOptimizer

Task 2: Add gradient accumulation utility functions

Task 3: Add GradientResult struct and extract compute_loss_internal() from DQN

Task 4: Add compute_gradients() and apply_accumulated_gradients() to DQN

Task 5: Add dispatch to DQNAgentType

Task 6: Rewrite train_step_with_accumulation() in the trainer

Task 7: Unit test — single optimizer step per accumulation call

Task 8: Convergence test — accumulated vs non-accumulated training

Task 9: Final verification — all existing tests pass

39 KiB

Raw Blame History

Task 1: Add `backward_and_clip()` + `apply_grads()` to MonitoredOptimizer

Task 3: Add `GradientResult` struct and extract `compute_loss_internal()` from DQN

Task 4: Add `compute_gradients()` and `apply_accumulated_gradients()` to DQN

Task 6: Rewrite `train_step_with_accumulation()` in the trainer