diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 42df6b49c..175bd41ac 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -178,6 +178,7 @@ pub struct GpuIqlTrainer { t_buf: CudaSlice, total_params: usize, grad_norm_blocks: usize, + grad_tile_size: usize, // per-sample grad tile (min(B, 256)) } impl GpuIqlTrainer { @@ -203,7 +204,10 @@ impl GpuIqlTrainer { let m_buf = alloc_f32(&stream, total_params, "iql_m")?; let v_buf = alloc_f32(&stream, total_params, "iql_v")?; let grad_buf = alloc_f32(&stream, total_params, "iql_grad")?; - let grads_per_sample = alloc_f32(&stream, b * total_params, "iql_grads_per_sample")?; + // Tiled per-sample gradients: process TILE samples at a time instead of all B. + // At B=16384, P=27009: full buffer = 1.7GB. Tiled at 256: 27MB. 64x smaller. + let grad_tile_size = b.min(256); + let grads_per_sample = alloc_f32(&stream, grad_tile_size * total_params, "iql_grads_tile")?; // Grad norm buffers (two-phase) let grad_norm_blocks = (total_params + 255) / 256; @@ -316,6 +320,7 @@ impl GpuIqlTrainer { t_buf, total_params, grad_norm_blocks, + grad_tile_size, }) } @@ -363,46 +368,75 @@ impl GpuIqlTrainer { .map_err(|e| MLError::ModelError(format!("IQL forward+loss: {e}")))?; } - // 2. Backward per-sample (256 threads per sample, no atomicAdd) - unsafe { - self.stream - .launch_builder(&self.backward_per_sample_kernel) - .arg(states_f32) - .arg(&self.q_taken_buf) - .arg(&self.v_out_buf) - .arg(&self.params_buf) - .arg(&self.save_pre1) - .arg(&self.save_pre2) - .arg(&self.save_h1) - .arg(&self.save_h2) - .arg(&mut self.grads_per_sample) - .arg(&batch_size_i32) - .arg(&state_dim_i32) - .arg(&total_params_i32) - .arg(&expectile_tau) - .launch(LaunchConfig { - grid_dim: (b as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }) - .map_err(|e| MLError::ModelError(format!("IQL backward_per_sample: {e}")))?; - } - - // 3. Weight grad reduce (deterministic sum across samples) + // 2+3. Tiled backward + reduce: process TILE samples at a time. + // backward_per_sample writes to grads_tile [TILE * P] + // weight_grad_reduce sums TILE samples → accumulates into grad_buf [P] + // All launches are async on the same stream — no host sync between tiles. + let tile = self.grad_tile_size; let reduce_blocks = (self.total_params + 255) / 256; - unsafe { - self.stream - .launch_builder(&self.weight_grad_reduce_kernel) - .arg(&self.grads_per_sample) - .arg(&mut self.grad_buf) - .arg(&batch_size_i32) - .arg(&total_params_i32) - .launch(LaunchConfig { - grid_dim: (reduce_blocks as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }) - .map_err(|e| MLError::ModelError(format!("IQL weight_grad_reduce: {e}")))?; + let grads_tile_ptr = self.grads_per_sample.raw_ptr(); + let grad_buf_ptr = self.grad_buf.raw_ptr(); + + // Zero grad_buf before tiled accumulation + self.stream.memset_zeros(&mut self.grad_buf) + .map_err(|e| MLError::ModelError(format!("IQL zero grad_buf: {e}")))?; + + for tile_start in (0..b).step_by(tile) { + let tile_b = (b - tile_start).min(tile); + let tile_b_i32 = tile_b as i32; + let f32_sz = std::mem::size_of::(); + + // Backward: process tile_b samples starting at tile_start + // The kernel reads from states/q_taken/v_out/saves at [tile_start..] offsets. + // We pass offset pointers so the kernel sees sample indices 0..tile_b. + let states_off = states_f32.raw_ptr() + (tile_start * self.config.state_dim * f32_sz) as u64; + let qt_off = self.q_taken_buf.raw_ptr() + (tile_start * f32_sz) as u64; + let vo_off = self.v_out_buf.raw_ptr() + (tile_start * f32_sz) as u64; + let sp1_off = self.save_pre1.raw_ptr() + (tile_start * self.config.value_hidden_dim * f32_sz) as u64; + let sp2_off = self.save_pre2.raw_ptr() + (tile_start * self.config.value_hidden_dim * f32_sz) as u64; + let sh1_off = self.save_h1.raw_ptr() + (tile_start * self.config.value_hidden_dim * f32_sz) as u64; + let sh2_off = self.save_h2.raw_ptr() + (tile_start * self.config.value_hidden_dim * f32_sz) as u64; + + unsafe { + self.stream + .launch_builder(&self.backward_per_sample_kernel) + .arg(&states_off) + .arg(&qt_off) + .arg(&vo_off) + .arg(&self.params_buf) + .arg(&sp1_off) + .arg(&sp2_off) + .arg(&sh1_off) + .arg(&sh2_off) + .arg(&grads_tile_ptr) + .arg(&tile_b_i32) // grid size (samples in this tile) + .arg(&state_dim_i32) + .arg(&total_params_i32) + .arg(&expectile_tau) + .arg(&batch_size_i32) // FULL batch_size for 1/N mean reduction + .launch(LaunchConfig { + grid_dim: (tile_b as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL backward tile {tile_start}: {e}")))?; + } + + // Reduce tile_b samples into grad_buf (accumulates via +=) + unsafe { + self.stream + .launch_builder(&self.weight_grad_reduce_kernel) + .arg(&grads_tile_ptr) + .arg(&grad_buf_ptr) + .arg(&tile_b_i32) + .arg(&total_params_i32) + .launch(LaunchConfig { + grid_dim: (reduce_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("IQL reduce tile {tile_start}: {e}")))?; + } } // 4. Loss reduce (deterministic sequential sum) diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index b5a0ba9a5..24e560ec5 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -212,15 +212,16 @@ void iql_backward_per_sample( const float* __restrict__ save_pre2, const float* __restrict__ save_h1, const float* __restrict__ save_h2, - float* __restrict__ grads_per_sample, /* [B, total_params] */ - int batch_size, + float* __restrict__ grads_per_sample, /* [TILE, total_params] */ + int tile_size, /* samples in this tile (grid dim) */ int state_dim, int total_params, - float expectile_tau + float expectile_tau, + int full_batch_size /* total B for 1/N mean reduction */ ) { int sample = blockIdx.x; - if (sample >= batch_size) return; + if (sample >= tile_size) return; int tid = threadIdx.x; /* 0..255 */ @@ -230,7 +231,7 @@ void iql_backward_per_sample( const float* w2 = params + off_w2; const float* w3 = params + off_w3; - /* Per-sample gradient slice — no overlap with other samples */ + /* Per-sample gradient slice within tile buffer */ float* g = grads_per_sample + sample * total_params; float* gw1 = g + off_w1; float* gb1 = g + off_b1; @@ -248,10 +249,10 @@ void iql_backward_per_sample( float v_val = v_out[sample]; float q_val = q_values[sample]; - /* dL/dV = -2 * weight * (Q - V) / batch_size */ + /* dL/dV = -2 * weight * (Q - V) / full_batch_size */ float u = q_val - v_val; float weight = (u >= 0.0f) ? expectile_tau : (1.0f - expectile_tau); - float dldv = -2.0f * weight * u / (float)batch_size; + float dldv = -2.0f * weight * u / (float)full_batch_size; /* ---- Output layer gradient: dL/dw3, dL/db3 ---- */ for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { @@ -305,9 +306,9 @@ void iql_backward_per_sample( */ extern "C" __global__ void iql_weight_grad_reduce( - const float* __restrict__ grads_per_sample, /* [B, total_params] */ - float* __restrict__ grads, /* [total_params] */ - int batch_size, + const float* __restrict__ grads_per_sample, /* [TILE, total_params] */ + float* __restrict__ grads, /* [total_params] — accumulates (+=) */ + int batch_size, /* tile size (not full B) */ int total_params ) { @@ -318,7 +319,7 @@ void iql_weight_grad_reduce( for (int b = 0; b < batch_size; b++) { sum += grads_per_sample[b * total_params + i]; } - grads[i] = sum; + grads[i] += sum; /* accumulate across tiles — caller zeros grad_buf before first tile */ } /* ------------------------------------------------------------------ */