gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: separate Adam optimizer for exposure aux loss — breaks i%3 permanently

Browse Source 
The exposure aux loss now has its OWN Adam state (m, v, t) and updates
the exposure output weights INDEPENDENTLY from the main C51 optimizer.
No shared gradient buffer, no gradient clipping competition.

This is multi-task learning done right — each objective has its own
optimizer that doesn't interfere with the other.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-08 10:28:19 +02:00
parent 58ff8db406
commit e3393205f4
2 changed files with 194 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

197
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -892,6 +892,33 @@ pub struct GpuDqnTrainer {
/// v8: GPU-resident aux weight scalar [1] — enables graph capture of aux GEMM.
pub(crate) exposure_aux_weight_gpu: CudaSlice<f32>,
// ── Separate Adam optimizer for exposure aux loss (breaks i%3 degeneracy) ──
// The aux loss gets its OWN Adam state and updates exposure output weights
// INDEPENDENTLY from the main C51 optimizer. No shared gradient buffer,
// no gradient clipping competition.
/// First moment [aux_param_count] — w_b0out + b_b0out.
aux_adam_m: CudaSlice<f32>,
/// Second moment [aux_param_count].
aux_adam_v: CudaSlice<f32>,
/// Step counter [1].
aux_adam_t: CudaSlice<i32>,
/// Separate gradient accumulator [aux_param_count].
aux_grad_buf: CudaSlice<f32>,
/// Grad norm f32 accumulator [1].
aux_grad_norm_f32: CudaSlice<f32>,
/// Grad norm bf16 output [1].
aux_grad_norm_bf16: CudaSlice<half::bf16>,
/// Per-block partial sums for grad norm reduction.
aux_grad_norm_partials: CudaSlice<f32>,
/// Number of exposure output parameters (w_b0out + b_b0out).
aux_param_count: usize,
/// Host-side step counter for aux Adam.
aux_adam_step: i32,
/// Standalone Adam kernel for aux (not captured in any CUDA graph).
aux_adam_update_kernel: CudaFunction,
/// Standalone grad_norm_finalize for aux (not captured in any CUDA graph).
aux_grad_norm_finalize_kernel: CudaFunction,
/// v8: PopArt reward normalization kernel (running mean/variance).
popart_normalize_kernel: CudaFunction,
/// v8: PopArt running mean [1].
@@ -926,6 +953,15 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("reset iqn t: {e}")))?;
self.iqn_trunk_adam_step = 0;
// Reset exposure aux Adam state (separate optimizer for exposure output weights).
self.stream.memset_zeros(&mut self.aux_adam_m)
.map_err(|e| MLError::ModelError(format!("reset aux_adam_m: {e}")))?;
self.stream.memset_zeros(&mut self.aux_adam_v)
.map_err(|e| MLError::ModelError(format!("reset aux_adam_v: {e}")))?;
self.stream.memset_zeros(&mut self.aux_adam_t)
.map_err(|e| MLError::ModelError(format!("reset aux_adam_t: {e}")))?;
self.aux_adam_step = 0;
tracing::info!("Adam optimizer state reset for new fold");
Ok(())
}
@@ -2001,12 +2037,13 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("f32_to_bf16 exposure_aux: {e}")))?;
}
// Phase 3: Backward FC layer -- accumulate weight + bias gradients
let sizes = compute_param_sizes(&self.config);
let goff_w_b0out: usize = sizes[..10].iter().sum::<usize>();
let goff_b_b0out: usize = sizes[..11].iter().sum::<usize>();
// Phase 3: Zero aux_grad_buf before backward (backward_fc_layer beta=1.0 accumulates)
self.stream.memset_zeros(&mut self.aux_grad_buf)
.map_err(|e| MLError::ModelError(format!("zero aux_grad_buf: {e}")))?;
let grad_base = self.ptrs.grad_buf;
// Phase 4: Backward FC layer -> aux_grad_buf (NOT shared grad_buf)
let sizes = compute_param_sizes(&self.config);
let aux_grad_ptr = self.aux_grad_buf.raw_ptr();
let h_b0_ptr = self.save_h_b0.raw_ptr();
self.cublas_backward.backward_fc_layer(
@@ -2014,14 +2051,117 @@ impl GpuDqnTrainer {
staging_ptr,
h_b0_ptr,
0u64,
grad_base + (goff_w_b0out as u64) * 4,
grad_base + (goff_b_b0out as u64) * 4,
aux_grad_ptr, // dW at offset 0
aux_grad_ptr + (sizes[10] as u64) * 4, // db at offset w_b0out_size
0u64,
b0 * na,
ah,
b,
)?;
// Phase 5: Compute grad norm of aux_grad_buf
let aux_pc = self.aux_param_count;
let aux_blocks = ((aux_pc + 255) / 256) as u32;
let aux_partials_ptr = self.aux_grad_norm_partials.raw_ptr();
let aux_norm_f32_ptr = self.aux_grad_norm_f32.raw_ptr();
let aux_norm_bf16_ptr = self.aux_grad_norm_bf16.raw_ptr();
let aux_n_i32 = aux_pc as i32;
let aux_nb = aux_blocks as i32;
// Phase 1: per-block partials (reuse grad_norm_standalone — not captured in any graph)
unsafe {
self.stream
.launch_builder(&self.grad_norm_standalone)
.arg(&aux_grad_ptr)
.arg(&aux_partials_ptr)
.arg(&aux_n_i32)
.launch(LaunchConfig {
grid_dim: (aux_blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("aux grad_norm phase1: {e}")))?;
}
// Phase 2: reduce partials → f32 sum-of-squares + bf16 L2 norm
unsafe {
self.stream
.launch_builder(&self.aux_grad_norm_finalize_kernel)
.arg(&aux_partials_ptr)
.arg(&aux_norm_f32_ptr)
.arg(&aux_norm_bf16_ptr)
.arg(&aux_nb)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("aux grad_norm phase2: {e}")))?;
}
// Phase 6: Separate Adam update on exposure output weights only
self.aux_adam_step += 1;
unsafe {
cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
self.aux_adam_t.raw_ptr(),
(&self.aux_adam_step as *const i32).cast(),
std::mem::size_of::<i32>(),
self.stream.cu_stream(),
);
}
let goff_w_b0out: usize = sizes[..10].iter().sum::<usize>();
let aux_params_ptr = self.ptrs.params_f32_ptr + (goff_w_b0out as u64) * 4;
let aux_lr = 0.001_f32; // higher LR for aux (independent from main)
let aux_beta1 = 0.9_f32;
let aux_beta2 = 0.999_f32;
let aux_eps = 1e-8_f32;
let aux_wd = 0.0_f32; // no weight decay for aux
let aux_max_gn = 1.0_f32; // already small gradients
let aux_m_ptr = self.aux_adam_m.raw_ptr();
let aux_v_ptr = self.aux_adam_v.raw_ptr();
let aux_t_ptr = self.aux_adam_t.raw_ptr();
unsafe {
self.stream
.launch_builder(&self.aux_adam_update_kernel)
.arg(&aux_params_ptr)
.arg(&aux_grad_ptr)
.arg(&aux_m_ptr)
.arg(&aux_v_ptr)
.arg(&aux_norm_f32_ptr)
.arg(&aux_lr)
.arg(&aux_beta1)
.arg(&aux_beta2)
.arg(&aux_eps)
.arg(&aux_wd)
.arg(&aux_max_gn)
.arg(&aux_t_ptr)
.arg(&aux_n_i32)
.launch(LaunchConfig {
grid_dim: (aux_blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("aux Adam update: {e}")))?;
}
// Phase 7: Cast updated f32 weights → bf16 shadow
let bf16_w_ptr = self.ptrs.params_buf + (goff_w_b0out as u64) * 2;
let cast_count = aux_pc as i32;
let cast_blk = aux_blocks;
unsafe {
self.stream
.launch_builder(&self.f32_to_bf16_kernel)
.arg(&aux_params_ptr)
.arg(&bf16_w_ptr)
.arg(&cast_count)
.launch(LaunchConfig {
grid_dim: (cast_blk, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("aux f32→bf16 cast: {e}")))?;
}
Ok(())
}
@@ -2966,6 +3106,38 @@ impl GpuDqnTrainer {
let exposure_aux_weight_gpu = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc exposure_aux_weight_gpu: {e}")))?;
// Separate Adam optimizer for exposure aux loss (breaks i%3 degeneracy)
let sizes_for_aux = compute_param_sizes(&config);
let aux_param_count = sizes_for_aux[10] + sizes_for_aux[11]; // w_b0out + b_b0out
let aux_adam_m = stream.alloc_zeros::<f32>(aux_param_count)
.map_err(|e| MLError::ModelError(format!("alloc aux_adam_m: {e}")))?;
let aux_adam_v = stream.alloc_zeros::<f32>(aux_param_count)
.map_err(|e| MLError::ModelError(format!("alloc aux_adam_v: {e}")))?;
let aux_adam_t = stream.alloc_zeros::<i32>(1)
.map_err(|e| MLError::ModelError(format!("alloc aux_adam_t: {e}")))?;
let aux_grad_buf = stream.alloc_zeros::<f32>(aux_param_count)
.map_err(|e| MLError::ModelError(format!("alloc aux_grad_buf: {e}")))?;
let aux_grad_norm_f32 = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc aux_grad_norm_f32: {e}")))?;
let aux_grad_norm_bf16 = stream.alloc_zeros::<half::bf16>(1)
.map_err(|e| MLError::ModelError(format!("alloc aux_grad_norm_bf16: {e}")))?;
let aux_gn_blocks = (aux_param_count + 255) / 256;
let aux_grad_norm_partials = stream.alloc_zeros::<f32>(aux_gn_blocks)
.map_err(|e| MLError::ModelError(format!("alloc aux_grad_norm_partials: {e}")))?;
// Standalone kernel handles for aux Adam (not captured in any CUDA graph)
let aux_adam_update_kernel = {
let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("aux_adam cubin: {e}")))?;
module.load_function("dqn_adam_update_kernel")
.map_err(|e| MLError::ModelError(format!("aux_adam load: {e}")))?
};
let aux_grad_norm_finalize_kernel = {
let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("aux_gnf cubin: {e}")))?;
module.load_function("dqn_grad_norm_finalize")
.map_err(|e| MLError::ModelError(format!("aux_gnf load: {e}")))?
};
// v8: PopArt running statistics buffers
let popart_mean = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc popart_mean: {e}")))?;
@@ -3214,6 +3386,17 @@ impl GpuDqnTrainer {
exposure_aux_grad_kernel,
exposure_aux_scratch,
exposure_aux_weight_gpu,
aux_adam_m,
aux_adam_v,
aux_adam_t,
aux_grad_buf,
aux_grad_norm_f32,
aux_grad_norm_bf16,
aux_grad_norm_partials,
aux_param_count,
aux_adam_step: 0,
aux_adam_update_kernel,
aux_grad_norm_finalize_kernel,
popart_normalize_kernel,
popart_mean,
popart_var,

10
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -1052,8 +1052,8 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;
// v7.1: Exposure auxiliary gradient (breaks i%3 degeneracy).
// Moved here from run_full_step so it's captured inside graph_aux / graph_mega.
// The kernel reads aux_weight from GPU-resident scalar (set_exposure_aux_weight).
// Now uses a SEPARATE Adam optimizer — aux gradients go into aux_grad_buf,
// not the shared grad_buf. No gradient clipping competition with C51.
if self.exposure_aux_weight > 1e-6 {
if let Err(e) = self.trainer.launch_exposure_aux_grad(
&self.exposure_targets,
@@ -1061,10 +1061,8 @@ impl FusedTrainingCtx {
) {
tracing::warn!("Exposure aux grad failed (non-fatal): {e}");
}
// Clip combined gradient to prevent aux GEMM backward from dominating.
// Same pattern as CQL gradient budget — keeps total norm within max_grad_norm.
self.trainer.clip_grad_buf_inplace(self.trainer.config().max_grad_norm)
.map_err(|e| anyhow::anyhow!("Aux gradient clip: {e}"))?;
// No clip_grad_buf_inplace needed — aux gradients are in a separate
// buffer with their own Adam optimizer. They never enter grad_buf.
}
Ok(())