feat(generalization): #32 vaccine batch sampling from replay buffer
Wire the gradient vaccine's validation batch through the training loop: - Pre-sample a SECOND batch from PER alongside each training batch - Pass via FusedTrainingCtx::pending_vaccine_batch (consumed per step) - Separate PER indices ensure train and val batches are independent The vaccine's non-graph forward+backward runs on this held-out batch to produce g_val. Gradient projection (dot + SAXPY) ensures g_train only moves in directions where both batches agree. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1356,7 +1356,7 @@ impl DQNTrainer {
|
||||
let chunk_end = (chunk_start + PREFETCH_K).min(num_training_steps);
|
||||
|
||||
let sample_start = std::time::Instant::now();
|
||||
let batches = {
|
||||
let (batches, vaccine_batches) = {
|
||||
let agent = self.agent.read().await;
|
||||
let buffer = agent.memory();
|
||||
let can = buffer.can_sample(self.current_batch_size);
|
||||
@@ -1369,14 +1369,24 @@ impl DQNTrainer {
|
||||
"Pre-sampling replay buffer"
|
||||
);
|
||||
let mut b = Vec::with_capacity(chunk_end - chunk_start);
|
||||
// #32 Gradient Vaccine: sample a second batch for validation
|
||||
let mut vaccine_batches = Vec::with_capacity(chunk_end - chunk_start);
|
||||
let vaccine_enabled = self.hyperparams.enable_gradient_vaccine;
|
||||
for _ in chunk_start..chunk_end {
|
||||
b.push(buffer.can_sample(self.current_batch_size).then(|| {
|
||||
buffer
|
||||
.sample(self.current_batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Pre-sample: {e}"))
|
||||
}).transpose()?);
|
||||
// Sample separate vaccine batch (different indices from PER)
|
||||
if vaccine_enabled && buffer.can_sample(self.current_batch_size) {
|
||||
vaccine_batches.push(buffer.sample(self.current_batch_size)
|
||||
.map_err(|e| anyhow::anyhow!("Vaccine sample: {e}")).ok());
|
||||
} else {
|
||||
vaccine_batches.push(None);
|
||||
}
|
||||
}
|
||||
b
|
||||
(b, vaccine_batches)
|
||||
};
|
||||
sample_total_us += sample_start.elapsed().as_micros() as u64;
|
||||
|
||||
@@ -1386,12 +1396,17 @@ impl DQNTrainer {
|
||||
let accum_steps = self.hyperparams.gradient_accumulation_steps;
|
||||
|
||||
if accum_steps <= 1 {
|
||||
let mut vaccine_iter = vaccine_batches.into_iter();
|
||||
for explicit_batch in batches {
|
||||
let fused_start = std::time::Instant::now();
|
||||
let _gpu_result = if let Some(ref mut fused) = self.fused_ctx {
|
||||
let batch_data = explicit_batch.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("No batch data for fused training step")
|
||||
})?;
|
||||
// #32 Set vaccine batch for this step (consumed by run_full_step)
|
||||
if let Some(Some(vb)) = vaccine_iter.next() {
|
||||
fused.pending_vaccine_batch = vb.gpu_batch;
|
||||
}
|
||||
let result = fused.run_full_step(batch_data, &mut *agent, &self.device)
|
||||
.map_err(|e| { eprintln!("!!! FUSED STEP ERROR: {:#}", e); e })
|
||||
.context("Fused CUDA training step failed")?;
|
||||
|
||||
Reference in New Issue
Block a user