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

fix: tiled IQL backward — 1.7GB→27MB per-sample gradient buffer

Browse Source 
At batch_size=16384, per-sample gradient buffer was B*P*4 = 1.7GB
per IQL trainer (3.4GB for dual tau). Caused OOM on H100 for attention.

Fix: tile backward+reduce into chunks of 256 samples. Same kernels,
launched multiple times with offset pointers. Per-sample buffer
shrinks to min(B,256)*P*4 = 27MB. 64x smaller.

Also fixed: weight_grad_reduce uses += (accumulate across tiles)
instead of = (overwrite). Backward kernel takes full_batch_size
param for correct 1/N gradient scaling.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 21:43:13 +02:00
parent e306f9f0e0
commit b11eec6ec5
2 changed files with 86 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -178,6 +178,7 @@ pub struct GpuIqlTrainer {
t_buf: CudaSlice<i32>,
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::<f32>();
// 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)

23
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -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 */
}
/* ------------------------------------------------------------------ */