diff --git a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu index 75aa60f67..6a6661511 100644 --- a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu +++ b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu @@ -10,7 +10,7 @@ * * Requires common_device_functions.cuh to be prepended for CUR_INPUT/CUR_HIDDEN/CUR_OUTPUT. * - * Native float everywhere. + * Native float everywhere. Fully deterministic (no atomicAdd). */ /* Total trainable parameters */ @@ -49,28 +49,27 @@ extern "C" __global__ void curiosity_zero_grads(float* grads, int n) { } /* ------------------------------------------------------------------ */ -/* Kernel 3: Forward + backward pass, accumulate gradients */ +/* Kernel 3: Forward + backward per-block (deterministic) */ /* ------------------------------------------------------------------ */ -/* Warp-level sum reduction via butterfly shuffle (all lanes get result) */ -__device__ __forceinline__ float warp_sum_cur(float val) { - for (int offset = 16; offset > 0; offset >>= 1) - val = val + __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - /** * One thread per sample. Computes forward pass, MSE loss, and backpropagates - * gradients via atomicAdd with warp-level pre-reduction (32x fewer atomics). - * Threads within a warp process different samples but accumulate gradients - * to the same weight indices, so warp reduction before atomicAdd is valid. + * gradients. Each thread's gradient contributions are reduced within the block + * using shared memory, then written to a per-block partial gradient buffer. + * No atomicAdd -- fully deterministic. * - * DETERMINISM NOTE: Residual per-block atomicAdd remains (one per warp per weight). - * Full elimination would require a separate per-weight reduction kernel. - * Impact: curiosity is a small auxiliary model (~11K params), not on the - * primary DQN gradient path, so residual non-determinism is negligible. + * Output: partial_grads[blockIdx.x * CUR_TOTAL_PARAMS + param_idx] contains + * the sum of gradient contributions from all threads in that block. + * + * The partial_grads layout is a flat vector of CUR_TOTAL_PARAMS per block: + * [0 .. w1_len) = grad_w1 partials + * [w1_len .. w1_len+b1_len) = grad_b1 partials + * [w1_len+b1_len .. w1_len+b1_len+w2_len) = grad_w2 partials + * [w1_len+b1_len+w2_len .. total) = grad_b2 partials + * + * Launch: grid=(ceil(N/BLOCK_SIZE)), block=(BLOCK_SIZE), shared_mem=BLOCK_SIZE*sizeof(float) */ -extern "C" __global__ void curiosity_forward_backward( +extern "C" __global__ void curiosity_fwd_bwd_per_block( const float* __restrict__ states, /* [N, state_dim] */ const int* __restrict__ actions, /* [N] */ const float* __restrict__ next_states, /* [N, state_dim] */ @@ -78,110 +77,216 @@ extern "C" __global__ void curiosity_forward_backward( const float* __restrict__ b1, /* [CUR_HIDDEN] */ const float* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ const float* __restrict__ b2, /* [CUR_OUTPUT] */ - float* __restrict__ grad_w1, /* [CUR_HIDDEN, CUR_INPUT] */ - float* __restrict__ grad_b1, /* [CUR_HIDDEN] */ - float* __restrict__ grad_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - float* __restrict__ grad_b2, /* [CUR_OUTPUT] */ + float* __restrict__ partial_grads, /* [num_blocks, CUR_TOTAL_PARAMS] */ int N, int state_dim ) { int tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= N) return; + int lane = threadIdx.x; + int block_size = blockDim.x; - const float* state_bf = states + tid * state_dim; - int action_idx = actions[tid]; - const float* next_state_bf = next_states + tid * state_dim; + /* Shared memory for intra-block reduction: one float per thread */ + extern __shared__ float shmem[]; - /* ---- Forward pass ---- */ + /* Base offset for this block's partial gradient output */ + float* block_out = partial_grads + blockIdx.x * CUR_TOTAL_PARAMS; - /* Build input: first MARKET_DIM state features + 3-class action one-hot */ + /* ---- Forward + backward per sample ---- */ + + /* We'll compute each thread's gradient contribution into local arrays, + * then reduce across the block for each parameter. Since CUR_TOTAL_PARAMS + * is ~11K, we can't store the full gradient vector per thread. Instead, + * we compute forward+backward, and then iterate over each parameter index, + * contributing to the block-level reduction one parameter at a time. + * + * But that would require re-running forward+backward for each parameter -- + * too expensive. Instead, each thread stores its intermediate values + * (input, hidden, pre_act, d_pred, d_hidden) on the stack, then we + * iterate over parameters and reduce. */ + + /* Forward pass -- only if this thread has a valid sample */ float input[CUR_INPUT]; - for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf[i]; } - input[MARKET_DIM + 0] = 0.0f; - input[MARKET_DIM + 1] = 0.0f; - input[MARKET_DIM + 2] = 0.0f; - - /* Action to category one-hot -- decode exposure from factored action. - * Factored action = exposure_idx * (b1*b2) + order_idx * b2 + urgency_idx. - * With b0=9, b1=3, b2=3: exposure_idx = action / 9. - * Categories: Short (idx 0-3), Flat (idx 4), Long (idx 5-8). */ - int exposure_idx = action_idx / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS); - int category; - if (exposure_idx < DQN_NUM_ACTIONS / 2) category = 0; /* Short */ - else if (exposure_idx == DQN_NUM_ACTIONS / 2) category = 1; /* Flat */ - else category = 2; /* Long */ - input[MARKET_DIM + category] = 1.0f; - - /* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, alpha=0.01) */ float pre_act[CUR_HIDDEN]; float hidden[CUR_HIDDEN]; - for (int h = 0; h < CUR_HIDDEN; h++) { - float sum = b1[h]; - for (int i = 0; i < CUR_INPUT; i++) { - sum += w1[h * CUR_INPUT + i] * input[i]; - } - pre_act[h] = sum; - hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum; - } - - /* Layer 2: pred = w2 * hidden + b2 (no activation) */ - float pred[CUR_OUTPUT]; - for (int o = 0; o < CUR_OUTPUT; o++) { - float sum = b2[o]; - for (int h = 0; h < CUR_HIDDEN; h++) { - sum += w2[o * CUR_HIDDEN + h] * hidden[h]; - } - pred[o] = sum; - } - - /* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */ - /* d_loss/d_pred[o] = 2 * (pred[o] - next_state_bf[o]) / CUR_OUTPUT */ float d_pred[CUR_OUTPUT]; - float inv_out = (2.0f / (float)CUR_OUTPUT); - for (int o = 0; o < CUR_OUTPUT; o++) { - d_pred[o] = (pred[o] - next_state_bf[o]) * inv_out; - } - - /* ---- Backward: Layer 2 (warp-reduced atomics) ---- */ float d_hidden[CUR_HIDDEN]; - for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f; + int valid = (tid < N) ? 1 : 0; - for (int o = 0; o < CUR_OUTPUT; o++) { - { - float val = warp_sum_cur(d_pred[o]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b2[o], val); - } + if (valid) { + const float* state_bf = states + tid * state_dim; + int action_idx = actions[tid]; + const float* next_state_bf = next_states + tid * state_dim; + + /* Build input: first MARKET_DIM state features + 3-class action one-hot */ + for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf[i]; } + input[MARKET_DIM + 0] = 0.0f; + input[MARKET_DIM + 1] = 0.0f; + input[MARKET_DIM + 2] = 0.0f; + + /* Action to category one-hot -- decode exposure from factored action. + * Factored action = exposure_idx * (b1*b2) + order_idx * b2 + urgency_idx. + * With b0=9, b1=3, b2=3: exposure_idx = action / 9. + * Categories: Short (idx 0-3), Flat (idx 4), Long (idx 5-8). */ + int exposure_idx = action_idx / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS); + int category; + if (exposure_idx < DQN_NUM_ACTIONS / 2) category = 0; /* Short */ + else if (exposure_idx == DQN_NUM_ACTIONS / 2) category = 1; /* Flat */ + else category = 2; /* Long */ + input[MARKET_DIM + category] = 1.0f; + + /* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, alpha=0.01) */ for (int h = 0; h < CUR_HIDDEN; h++) { - { - float val = warp_sum_cur(d_pred[o] * hidden[h]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w2[o * CUR_HIDDEN + h], val); + float sum = b1[h]; + for (int i = 0; i < CUR_INPUT; i++) { + sum += w1[h * CUR_INPUT + i] * input[i]; } - d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o]; + pre_act[h] = sum; + hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum; + } + + /* Layer 2: pred = w2 * hidden + b2 (no activation) */ + float pred[CUR_OUTPUT]; + for (int o = 0; o < CUR_OUTPUT; o++) { + float sum = b2[o]; + for (int h = 0; h < CUR_HIDDEN; h++) { + sum += w2[o * CUR_HIDDEN + h] * hidden[h]; + } + pred[o] = sum; + } + + /* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */ + /* d_loss/d_pred[o] = 2 * (pred[o] - next_state_bf[o]) / CUR_OUTPUT */ + float inv_out = (2.0f / (float)CUR_OUTPUT); + for (int o = 0; o < CUR_OUTPUT; o++) { + d_pred[o] = (pred[o] - next_state_bf[o]) * inv_out; + } + + /* ---- Backward: compute d_hidden from Layer 2 ---- */ + for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f; + for (int o = 0; o < CUR_OUTPUT; o++) { + for (int h = 0; h < CUR_HIDDEN; h++) { + d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o]; + } + } + + /* ---- Backward: LeakyReLU ---- */ + for (int h = 0; h < CUR_HIDDEN; h++) { + if (pre_act[h] <= 0.0f) d_hidden[h] = d_hidden[h] * 0.01f; } } - /* ---- Backward: LeakyReLU ---- */ - for (int h = 0; h < CUR_HIDDEN; h++) { - if (pre_act[h] <= 0.0f) d_hidden[h] = d_hidden[h] * 0.01f; - } + /* ---- Block-level reduction for each parameter ---- */ + /* Iterate over all CUR_TOTAL_PARAMS parameters. For each parameter, + * each thread computes its gradient contribution, then we do a + * block-level tree reduction in shared memory. Thread 0 writes the + * result to the per-block output buffer. + * + * Layout: [grad_w1 | grad_b1 | grad_w2 | grad_b2] + * grad_w1[h, i] = d_hidden[h] * input[i] (for Layer 1) + * grad_b1[h] = d_hidden[h] (for Layer 1) + * grad_w2[o, h] = d_pred[o] * hidden[h] (for Layer 2) + * grad_b2[o] = d_pred[o] (for Layer 2) + */ - /* ---- Backward: Layer 1 (warp-reduced atomics) ---- */ - for (int h = 0; h < CUR_HIDDEN; h++) { - { - float val = warp_sum_cur(d_hidden[h]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b1[h], val); - } - for (int i = 0; i < CUR_INPUT; i++) { - { - float val = warp_sum_cur(d_hidden[h] * input[i]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w1[h * CUR_INPUT + i], val); + int w1_len = CUR_HIDDEN * CUR_INPUT; + int b1_len = CUR_HIDDEN; + int w2_len = CUR_OUTPUT * CUR_HIDDEN; + /* b2_len = CUR_OUTPUT */ + + for (int p = 0; p < CUR_TOTAL_PARAMS; p++) { + /* Compute this thread's gradient contribution for parameter p */ + float my_grad = 0.0f; + if (valid) { + if (p < w1_len) { + /* grad_w1[h][i]: h = p / CUR_INPUT, i = p % CUR_INPUT */ + int h = p / CUR_INPUT; + int i = p % CUR_INPUT; + my_grad = d_hidden[h] * input[i]; + } else if (p < w1_len + b1_len) { + /* grad_b1[h] */ + int h = p - w1_len; + my_grad = d_hidden[h]; + } else if (p < w1_len + b1_len + w2_len) { + /* grad_w2[o][h]: idx = p - w1_len - b1_len, o = idx / CUR_HIDDEN, h = idx % CUR_HIDDEN */ + int idx = p - w1_len - b1_len; + int o = idx / CUR_HIDDEN; + int h = idx % CUR_HIDDEN; + my_grad = d_pred[o] * hidden[h]; + } else { + /* grad_b2[o] */ + int o = p - w1_len - b1_len - w2_len; + my_grad = d_pred[o]; } } + + /* Block-level tree reduction in shared memory */ + shmem[lane] = my_grad; + __syncthreads(); + + for (int stride = block_size / 2; stride > 0; stride >>= 1) { + if (lane < stride) { + shmem[lane] += shmem[lane + stride]; + } + __syncthreads(); + } + + /* Thread 0 writes the block's partial sum */ + if (lane == 0) { + block_out[p] = shmem[0]; + } + __syncthreads(); } } /* ------------------------------------------------------------------ */ -/* Kernel 4: Adam optimizer step (single param group, legacy) */ +/* Kernel 3b: Reduce per-block partial gradients (deterministic) */ +/* ------------------------------------------------------------------ */ + +/** + * One thread per parameter. Loops over all block partials and sums them + * into the final gradient buffer. Fully deterministic -- fixed summation + * order (block 0, 1, 2, ...). + * + * The gradient is written into 4 separate buffers (grad_w1, grad_b1, + * grad_w2, grad_b2) matching the Adam optimizer's per-group layout. + * + * Launch: grid=(ceil(CUR_TOTAL_PARAMS / 256)), block=(256) + */ +extern "C" __global__ void curiosity_grad_reduce( + const float* __restrict__ partial_grads, /* [num_blocks, CUR_TOTAL_PARAMS] */ + float* __restrict__ grad_w1, /* [CUR_HIDDEN * CUR_INPUT] */ + float* __restrict__ grad_b1, /* [CUR_HIDDEN] */ + float* __restrict__ grad_w2, /* [CUR_OUTPUT * CUR_HIDDEN] */ + float* __restrict__ grad_b2, /* [CUR_OUTPUT] */ + int num_blocks +) { + int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p >= CUR_TOTAL_PARAMS) return; + + int w1_len = CUR_HIDDEN * CUR_INPUT; + int b1_len = CUR_HIDDEN; + int w2_len = CUR_OUTPUT * CUR_HIDDEN; + + /* Sum across all blocks for this parameter */ + float sum = 0.0f; + for (int blk = 0; blk < num_blocks; blk++) { + sum += partial_grads[blk * CUR_TOTAL_PARAMS + p]; + } + + /* Write to the correct gradient buffer */ + if (p < w1_len) { + grad_w1[p] = sum; + } else if (p < w1_len + b1_len) { + grad_b1[p - w1_len] = sum; + } else if (p < w1_len + b1_len + w2_len) { + grad_w2[p - w1_len - b1_len] = sum; + } else { + grad_b2[p - w1_len - b1_len - w2_len] = sum; + } +} + +/* ------------------------------------------------------------------ */ +/* Kernel 4: Adam optimizer step (single param group) */ /* ------------------------------------------------------------------ */ /** @@ -222,273 +327,3 @@ extern "C" __global__ void curiosity_adam_step( float v_hat = v[i] / (1.0f - powf(beta2, (float)step)); params[i] = params[i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf); } - -/* ------------------------------------------------------------------ */ -/* Kernel 4b: Fused Adam step -- all 4 param groups in one launch */ -/* ------------------------------------------------------------------ */ - -/** - * Processes w1, b1, w2, b2 in a single kernel launch. Each thread - * determines which parameter group it belongs to via offset comparison. - * Eliminates 3 kernel launch overheads (~150 us/step on H100). - */ -extern "C" __global__ void curiosity_adam_step_fused( - float* __restrict__ p0, const float* __restrict__ g0, float* __restrict__ m0, float* __restrict__ v0, int n0, - float* __restrict__ p1, const float* __restrict__ g1, float* __restrict__ m1, float* __restrict__ v1, int n1, - float* __restrict__ p2, const float* __restrict__ g2, float* __restrict__ m2, float* __restrict__ v2, int n2, - float* __restrict__ p3, const float* __restrict__ g3, float* __restrict__ m3, float* __restrict__ v3, int n3, - int batch_size, - float lr, float beta1, float beta2, float eps, int step -) { - int total = n0 + n1 + n2 + n3; - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) return; - - /* Determine which param group and local offset */ - float* p; const float* g; float* mm; float* vv; - int local_i; - if (idx < n0) { - p = p0; g = g0; mm = m0; vv = v0; local_i = idx; - } else if (idx < n0 + n1) { - p = p1; g = g1; mm = m1; vv = v1; local_i = idx - n0; - } else if (idx < n0 + n1 + n2) { - p = p2; g = g2; mm = m2; vv = v2; local_i = idx - n0 - n1; - } else { - p = p3; g = g3; mm = m3; vv = v3; local_i = idx - n0 - n1 - n2; - } - - float lr_bf = lr; - float beta1_bf = beta1; - float beta2_bf = beta2; - float eps_bf = eps; - float one_bf = 1.0f; - - float grad = g[local_i] / (float)batch_size; - - /* Skip NaN/Inf gradients */ - if (!isfinite((float)grad)) return; - - mm[local_i] = beta1_bf * mm[local_i] + (one_bf - beta1_bf) * grad; - vv[local_i] = beta2_bf * vv[local_i] + (one_bf - beta2_bf) * grad * grad; - float m_hat = mm[local_i] / (1.0f - powf(beta1, (float)step)); - float v_hat = vv[local_i] / (1.0f - powf(beta2, (float)step)); - p[local_i] = p[local_i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf); -} - -/* ------------------------------------------------------------------ */ -/* Kernel 5: Fully fused zero + forward/backward + Adam */ -/* ------------------------------------------------------------------ */ - -/** - * Fuses gradient zeroing, forward+backward pass, and Adam optimizer - * update into a single kernel launch. Uses an atomic block-arrival - * counter for grid-wide synchronization between the fwd/bwd phase - * (sample-parallel) and the Adam phase (parameter-parallel). - * - * Phase 1 (all threads): Forward + backward pass (one sample per thread), - * accumulate gradients via warp-reduced atomicAdd. - * Phase 2 (last-arriving block only): After all blocks complete phase 1, - * the last block to arrive (detected via atomic counter) applies - * Adam update across all parameters in a grid-stride loop. - * - * Saves 4 memset dispatches + 1 kernel launch = 5 fewer GPU dispatches - * per training step (~30-50 us on H100 at high training frequency). - * - * DETERMINISM NOTE: Same warp-reduced atomicAdd pattern as curiosity_forward_backward. - * Residual per-warp atomicAdd contention is negligible for this auxiliary model. - */ -extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam( - const float* __restrict__ states, /* [N, state_dim] */ - const int* __restrict__ actions, /* [N] */ - const float* __restrict__ next_states, /* [N, state_dim] */ - /* Weights (updated in-place by Adam in phase 3) */ - float* __restrict__ w1, /* [CUR_HIDDEN, CUR_INPUT] */ - float* __restrict__ b1, /* [CUR_HIDDEN] */ - float* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - float* __restrict__ b2, /* [CUR_OUTPUT] */ - /* Gradient accumulators (zeroed in phase 1, accumulated in phase 2) */ - float* __restrict__ grad_w1, /* [CUR_HIDDEN, CUR_INPUT] */ - float* __restrict__ grad_b1, /* [CUR_HIDDEN] */ - float* __restrict__ grad_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ - float* __restrict__ grad_b2, /* [CUR_OUTPUT] */ - /* Adam first moment */ - float* __restrict__ adam_m_w1, /* [CUR_HIDDEN * CUR_INPUT] */ - float* __restrict__ adam_m_b1, /* [CUR_HIDDEN] */ - float* __restrict__ adam_m_w2, /* [CUR_OUTPUT * CUR_HIDDEN] */ - float* __restrict__ adam_m_b2, /* [CUR_OUTPUT] */ - /* Adam second moment */ - float* __restrict__ adam_v_w1, /* [CUR_HIDDEN * CUR_INPUT] */ - float* __restrict__ adam_v_b1, /* [CUR_HIDDEN] */ - float* __restrict__ adam_v_w2, /* [CUR_OUTPUT * CUR_HIDDEN] */ - float* __restrict__ adam_v_b2, /* [CUR_OUTPUT] */ - /* Atomic block-arrival counter (must be zeroed before launch) */ - int* __restrict__ block_counter, - /* Scalar parameters */ - int N, /* number of training samples */ - int state_dim, - int batch_size, /* == N, for gradient averaging */ - float lr, float beta1, float beta2, float eps, - int adam_step /* 1-based step counter */ -) { - int tid = blockIdx.x * blockDim.x + threadIdx.x; - - /* Gradient buffers are pre-zeroed by host-side memset_zeros on the same - * CUDA stream before this kernel launch. Stream ordering guarantees the - * memsets complete before this kernel starts. No in-kernel zeroing needed. - * - * The previous in-kernel Phase 1 zeroing had an inter-block race: - * __syncthreads() only syncs within a block, so fast blocks could start - * atomicAdd-ing gradients while slow blocks were still zeroing. */ - - int w1_len = CUR_HIDDEN * CUR_INPUT; - int b1_len = CUR_HIDDEN; - int w2_len = CUR_OUTPUT * CUR_HIDDEN; - int total_grad_elems = CUR_TOTAL_PARAMS; - - /* ================================================================ */ - /* PHASE 1: Forward + backward pass (one sample per thread) */ - /* ================================================================ */ - if (tid < N) { - const float* state_bf2 = states + tid * state_dim; - int action_idx = actions[tid]; - const float* next_state_bf2 = next_states + tid * state_dim; - - /* ---- Forward pass ---- */ - - /* Build input: first MARKET_DIM state features + 3-class action one-hot */ - float input[CUR_INPUT]; - for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf2[i]; } - input[MARKET_DIM + 0] = 0.0f; - input[MARKET_DIM + 1] = 0.0f; - input[MARKET_DIM + 2] = 0.0f; - - /* Action to category one-hot */ - int category; - if (action_idx <= 1) category = 0; /* Short100/Short50 */ - else if (action_idx == 2) category = 1; /* Flat */ - else category = 2; /* Long50/Long100 */ - input[MARKET_DIM + category] = 1.0f; - - /* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, 0.01) */ - float pre_act[CUR_HIDDEN]; - float hidden[CUR_HIDDEN]; - for (int h = 0; h < CUR_HIDDEN; h++) { - float sum = b1[h]; - for (int i = 0; i < CUR_INPUT; i++) { - sum += w1[h * CUR_INPUT + i] * input[i]; - } - pre_act[h] = sum; - hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum; - } - - /* Layer 2: pred = w2 * hidden + b2 (no activation) */ - float pred[CUR_OUTPUT]; - for (int o = 0; o < CUR_OUTPUT; o++) { - float sum = b2[o]; - for (int h = 0; h < CUR_HIDDEN; h++) { - sum += w2[o * CUR_HIDDEN + h] * hidden[h]; - } - pred[o] = sum; - } - - /* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */ - float d_pred[CUR_OUTPUT]; - float inv_out = (2.0f / (float)CUR_OUTPUT); - for (int o = 0; o < CUR_OUTPUT; o++) { - d_pred[o] = (pred[o] - next_state_bf2[o]) * inv_out; - } - - /* ---- Backward: Layer 2 (warp-reduced atomics) ---- */ - float d_hidden[CUR_HIDDEN]; - for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f; - - for (int o = 0; o < CUR_OUTPUT; o++) { - { - float val = warp_sum_cur(d_pred[o]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b2[o], val); - } - for (int h = 0; h < CUR_HIDDEN; h++) { - { - float val = warp_sum_cur(d_pred[o] * hidden[h]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w2[o * CUR_HIDDEN + h], val); - } - d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o]; - } - } - - /* ---- Backward: LeakyReLU ---- */ - for (int h = 0; h < CUR_HIDDEN; h++) { - if (pre_act[h] <= 0.0f) d_hidden[h] = d_hidden[h] * 0.01f; - } - - /* ---- Backward: Layer 1 (warp-reduced atomics) ---- */ - for (int h = 0; h < CUR_HIDDEN; h++) { - { - float val = warp_sum_cur(d_hidden[h]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b1[h], val); - } - for (int i = 0; i < CUR_INPUT; i++) { - { - float val = warp_sum_cur(d_hidden[h] * input[i]); - if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w1[h * CUR_INPUT + i], val); - } - } - } - } - - /* ================================================================ */ - /* Grid-wide barrier via atomic block-arrival counter */ - /* ================================================================ */ - /* Ensure all global memory writes (gradient atomicAdds) from this - * block are visible to all other blocks before signalling arrival. */ - __threadfence(); - __syncthreads(); - - /* Thread 0 of each block increments the arrival counter. - * The last block to arrive (counter == gridDim.x - 1) proceeds - * to phase 3. All other blocks exit. */ - __shared__ int is_last_block; - if (threadIdx.x == 0) { - int arrived = atomicAdd(block_counter, 1); - is_last_block = (arrived == (int)gridDim.x - 1) ? 1 : 0; - } - __syncthreads(); - if (!is_last_block) return; - - /* ================================================================ */ - /* PHASE 3: Adam optimizer update (last block, grid-stride loop) */ - /* ================================================================ */ - /* The last-arriving block processes all CUR_TOTAL_PARAMS parameters - * using blockDim.x threads in a stride loop. This avoids a second - * kernel launch entirely. */ - int block_tid = threadIdx.x; - int block_size = blockDim.x; - - float lr_bf = lr; - float beta1_bf = beta1; - float beta2_bf = beta2; - float eps_bf = eps; - float one_bf = 1.0f; - - for (int i = block_tid; i < total_grad_elems; i += block_size) { - /* Determine which param group and local offset */ - float* p; float* g; float* mmv; float* vvv; - int local_i; - if (i < w1_len) { - p = w1; g = grad_w1; mmv = adam_m_w1; vvv = adam_v_w1; local_i = i; - } else if (i < w1_len + b1_len) { - p = b1; g = grad_b1; mmv = adam_m_b1; vvv = adam_v_b1; local_i = i - w1_len; - } else if (i < w1_len + b1_len + w2_len) { - p = w2; g = grad_w2; mmv = adam_m_w2; vvv = adam_v_w2; local_i = i - w1_len - b1_len; - } else { - p = b2; g = grad_b2; mmv = adam_m_b2; vvv = adam_v_b2; local_i = i - w1_len - b1_len - w2_len; - } - - float grad = g[local_i] / (float)batch_size; - mmv[local_i] = beta1_bf * mmv[local_i] + (one_bf - beta1_bf) * grad; - vvv[local_i] = beta2_bf * vvv[local_i] + (one_bf - beta2_bf) * grad * grad; - float m_hat = mmv[local_i] / (1.0f - powf(beta1, (float)adam_step)); - float v_hat = vvv[local_i] / (1.0f - powf(beta2, (float)adam_step)); - p[local_i] = p[local_i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf); - } -} diff --git a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs index f740859aa..57084ce56 100644 --- a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs @@ -9,13 +9,11 @@ //! //! Architecture: `[MARKET_DIM+3] -> [128] LeakyReLU -> [MARKET_DIM]` (MARKET_DIM=42: 11_954 params) //! -//! Kernels (fused path -- 2 launches per step): +//! Kernels (deterministic path -- 4 launches per step): //! - `curiosity_shift_states`: builds shifted next_states from states buffer -//! - `curiosity_fused_zero_fwd_bwd_adam`: zeros grads + forward/backward + Adam in one launch -//! -//! Legacy kernels (kept for fallback, not used in hot path): -//! - `curiosity_forward_backward`: standalone forward + backward pass -//! - `curiosity_adam_step` / `curiosity_adam_step_fused`: standalone Adam optimizers +//! - `curiosity_fwd_bwd_per_block`: forward + backward, block-level reduce to partial gradients +//! - `curiosity_grad_reduce`: deterministic reduction of per-block partials into final gradients +//! - `curiosity_adam_step`: per-param-group Adam optimizer update use std::sync::Arc; @@ -42,18 +40,17 @@ const CUR_B1_LEN: usize = CUR_HIDDEN; // [128] const CUR_W2_LEN: usize = CUR_OUTPUT * CUR_HIDDEN; // [42, 128] = 5376 const CUR_B2_LEN: usize = CUR_OUTPUT; // [42] +const CUR_TOTAL_PARAMS: usize = CUR_W1_LEN + CUR_B1_LEN + CUR_W2_LEN + CUR_B2_LEN; // 11306 + +/// Block size for the forward+backward kernel. +const FWD_BWD_BLOCK_SIZE: u32 = 256; + /// Adam optimizer hyperparameters. const ADAM_LR: f32 = 0.001; const ADAM_BETA1: f32 = 0.9; const ADAM_BETA2: f32 = 0.999; const ADAM_EPS: f32 = 1e-8; -// --------------------------------------------------------------------------- -// PTX cache -// --------------------------------------------------------------------------- - -// PTX cache removed -- using precompiled cubin via CURIOSITY_TRAINING_CUBIN - // --------------------------------------------------------------------------- // GpuCuriosityTrainer // --------------------------------------------------------------------------- @@ -63,17 +60,18 @@ const ADAM_EPS: f32 = 1e-8; /// Trains the curiosity MLP entirely on GPU using experience data that is /// already device-resident. Maintains gradient buffers and Adam optimizer /// state. Modifies [`CuriosityWeightSet`] in-place -- zero CPU traffic. +/// +/// Gradient accumulation is fully deterministic: per-block shared-memory +/// reduction followed by a sequential sum over block partials. No atomicAdd. #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct GpuCuriosityTrainer { stream: Arc, // Kernel functions shift_func: CudaFunction, - fwd_bwd_func: CudaFunction, + fwd_bwd_per_block_func: CudaFunction, + grad_reduce_func: CudaFunction, adam_func: CudaFunction, - adam_fused_func: CudaFunction, - /// Fully fused kernel: zero grads + fwd/bwd + Adam in one launch. - fused_zero_fwd_bwd_adam_func: CudaFunction, // Gradient buffers grad_w1: CudaSlice, // [CUR_W1_LEN] @@ -81,6 +79,10 @@ pub struct GpuCuriosityTrainer { grad_w2: CudaSlice, // [CUR_W2_LEN] grad_b2: CudaSlice, // [CUR_B2_LEN] + // Per-block partial gradient buffer for deterministic reduction + partial_grads: CudaSlice, // [max_blocks * CUR_TOTAL_PARAMS] + max_blocks: usize, + // Adam first moment (per-param-group) adam_m_w1: CudaSlice, // [CUR_W1_LEN] adam_m_b1: CudaSlice, // [CUR_B1_LEN] @@ -96,9 +98,6 @@ pub struct GpuCuriosityTrainer { // Shifted next_states buffer next_states_buf: CudaSlice, - /// Atomic block-arrival counter for fused kernel grid-wide sync (single i32 on GPU). - block_counter: CudaSlice, - // Adam step counter (1-based) step: i32, @@ -156,7 +155,7 @@ fn launch_adam_step( impl GpuCuriosityTrainer { /// Create a new GPU curiosity trainer. /// - /// Compiles the CUDA training kernel via NVRTC, allocates gradient and + /// Loads the precompiled CUDA training cubin, allocates gradient and /// Adam optimizer state buffers as zeros on GPU. /// /// # Arguments @@ -168,14 +167,6 @@ impl GpuCuriosityTrainer { state_dim: usize, max_samples: usize, ) -> Result { - // The fused curiosity kernel allocates ~3KB/thread on stack (6 arrays of - // 42-128 floats: input[45], pre_act[128], hidden[128], pred[42], d_pred[42], - // Stack sizing: curiosity kernel needs ~3KB/thread (513 floats for - // input/hidden/pred arrays). The fused training kernel (GpuDqnTrainer) - // sets 64KB stack when it inits — that covers all kernels including - // curiosity. Don't set stack here: cuCtxSetLimit reserves stack_bytes × - // max_threads of VRAM upfront, causing OOM on 4GB GPUs. - // ---- Load precompiled cubin ---- let context = stream.context(); let module = context.load_cubin(CURIOSITY_TRAINING_CUBIN.to_vec()).map_err(|e| { @@ -185,20 +176,15 @@ impl GpuCuriosityTrainer { let shift_func = module.load_function("curiosity_shift_states").map_err(|e| { MLError::ModelError(format!("curiosity_shift_states load: {e}")) })?; - let fwd_bwd_func = module.load_function("curiosity_forward_backward").map_err(|e| { - MLError::ModelError(format!("curiosity_forward_backward load: {e}")) + let fwd_bwd_per_block_func = module.load_function("curiosity_fwd_bwd_per_block").map_err(|e| { + MLError::ModelError(format!("curiosity_fwd_bwd_per_block load: {e}")) + })?; + let grad_reduce_func = module.load_function("curiosity_grad_reduce").map_err(|e| { + MLError::ModelError(format!("curiosity_grad_reduce load: {e}")) })?; let adam_func = module.load_function("curiosity_adam_step").map_err(|e| { MLError::ModelError(format!("curiosity_adam_step load: {e}")) })?; - let adam_fused_func = module.load_function("curiosity_adam_step_fused").map_err(|e| { - MLError::ModelError(format!("curiosity_adam_step_fused load: {e}")) - })?; - let fused_zero_fwd_bwd_adam_func = module - .load_function("curiosity_fused_zero_fwd_bwd_adam") - .map_err(|e| { - MLError::ModelError(format!("curiosity_fused_zero_fwd_bwd_adam load: {e}")) - })?; // ---- Allocate gradient buffers ---- let grad_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { @@ -214,6 +200,16 @@ impl GpuCuriosityTrainer { MLError::ModelError(format!("alloc grad_b2: {e}")) })?; + // ---- Allocate per-block partial gradient buffer ---- + // With FWD_BWD_BLOCK_SIZE=256 threads/block and max_samples samples: + // max_blocks = ceil(max_samples / 256). Buffer = max_blocks * CUR_TOTAL_PARAMS floats. + let max_blocks = (max_samples + FWD_BWD_BLOCK_SIZE as usize - 1) / FWD_BWD_BLOCK_SIZE as usize; + let partial_grads = stream + .alloc_zeros::(max_blocks * CUR_TOTAL_PARAMS) + .map_err(|e| { + MLError::ModelError(format!("alloc partial_grads: {e}")) + })?; + // ---- Allocate Adam first moment buffers ---- let adam_m_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { MLError::ModelError(format!("alloc adam_m_w1: {e}")) @@ -249,29 +245,27 @@ impl GpuCuriosityTrainer { MLError::ModelError(format!("alloc next_states_buf: {e}")) })?; - // ---- Allocate atomic block-arrival counter for fused kernel ---- - let block_counter = stream.alloc_zeros::(1).map_err(|e| { - MLError::ModelError(format!("alloc block_counter: {e}")) - })?; - debug!( state_dim, max_samples, - total_params = CUR_W1_LEN + CUR_B1_LEN + CUR_W2_LEN + CUR_B2_LEN, - "GPU curiosity trainer initialized (Adam optimizer)" + max_blocks, + partial_grads_bytes = max_blocks * CUR_TOTAL_PARAMS * 4, + total_params = CUR_TOTAL_PARAMS, + "GPU curiosity trainer initialized (deterministic per-block reduce)" ); Ok(Self { stream, shift_func, - fwd_bwd_func, + fwd_bwd_per_block_func, + grad_reduce_func, adam_func, - adam_fused_func, - fused_zero_fwd_bwd_adam_func, grad_w1, grad_b1, grad_w2, grad_b2, + partial_grads, + max_blocks, adam_m_w1, adam_m_b1, adam_m_w2, @@ -281,7 +275,6 @@ impl GpuCuriosityTrainer { adam_v_w2, adam_v_b2, next_states_buf, - block_counter, step: 0, state_dim, buf_capacity: max_samples, @@ -295,6 +288,10 @@ impl GpuCuriosityTrainer { /// next_states -- episode boundary noise is negligible for this tiny /// auxiliary model. /// + /// Gradient accumulation is fully deterministic: + /// 1. `curiosity_fwd_bwd_per_block` -- block-level shared-memory reduce + /// 2. `curiosity_grad_reduce` -- sequential sum over block partials + /// /// # Arguments /// * `weights` - Curiosity model weights to update in-place on GPU /// * `states` - State observations `[n_samples * state_dim]` on GPU @@ -333,9 +330,7 @@ impl GpuCuriosityTrainer { self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("curiosity PRE-SHIFT sync FAILED: {e}")))?; - // ---- Launch 1/2: Build shifted next_states buffer ---- - // next_states[i] = states[i + state_dim] (shift by one timestep) - // Separate launch because grid dimensions differ from the fwd/bwd grid. + // ---- Launch 1/4: Build shifted next_states buffer ---- let shift_total = n_train * sd; let shift_cfg = LaunchConfig { grid_dim: (((shift_total as u32) + 255) / 256, 1, 1), @@ -355,39 +350,27 @@ impl GpuCuriosityTrainer { })?; } - // Sync after shift kernel to catch crashes self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("curiosity shift kernel CRASHED: {e}")))?; - // ---- Launch 2/2: Fused zero + forward/backward + Adam ---- - self.step += 1; - let step = self.step; - let _bs_i32 = n_train as i32; + // ---- Launch 2/4: Forward + backward per-block ---- + let num_blocks = ((n_train as u32) + FWD_BWD_BLOCK_SIZE - 1) / FWD_BWD_BLOCK_SIZE; - // Zero gradient buffers + block-arrival counter before fused kernel launch. - // Gradients are accumulated via atomicAdd in the kernel — they MUST start at zero. - // Previously this was done inside the kernel (Phase 1) but had an inter-block race. - // memset_zeros is GPU-side cuMemsetD8Async, ordered on the same stream. - self.stream.memset_zeros(&mut self.grad_w1) - .map_err(|e| MLError::ModelError(format!("memset grad_w1: {e}")))?; - self.stream.memset_zeros(&mut self.grad_b1) - .map_err(|e| MLError::ModelError(format!("memset grad_b1: {e}")))?; - self.stream.memset_zeros(&mut self.grad_w2) - .map_err(|e| MLError::ModelError(format!("memset grad_w2: {e}")))?; - self.stream.memset_zeros(&mut self.grad_b2) - .map_err(|e| MLError::ModelError(format!("memset grad_b2: {e}")))?; - self.stream.memset_zeros(&mut self.block_counter) - .map_err(|e| MLError::ModelError(format!("memset block_counter: {e}")))?; + if (num_blocks as usize) > self.max_blocks { + return Err(MLError::ModelError(format!( + "curiosity trainer: num_blocks={num_blocks} exceeds max_blocks={}", + self.max_blocks + ))); + } - // Step 2a: Forward + backward (separate kernel — no inter-block sync needed) let fwd_bwd_cfg = LaunchConfig { - grid_dim: (((n_train as u32) + 255) / 256, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, + grid_dim: (num_blocks, 1, 1), + block_dim: (FWD_BWD_BLOCK_SIZE, 1, 1), + shared_mem_bytes: FWD_BWD_BLOCK_SIZE * std::mem::size_of::() as u32, }; unsafe { self.stream - .launch_builder(&self.fwd_bwd_func) + .launch_builder(&self.fwd_bwd_per_block_func) .arg(states) .arg(actions) .arg(&self.next_states_buf) @@ -395,19 +378,42 @@ impl GpuCuriosityTrainer { .arg(&weights.b1) .arg(&weights.w2) .arg(&weights.b2) + .arg(&mut self.partial_grads) + .arg(&n_i32) + .arg(&sd_i32) + .launch(fwd_bwd_cfg) + .map_err(|e| MLError::ModelError(format!("curiosity fwd_bwd_per_block launch: {e}")))?; + } + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("curiosity fwd_bwd_per_block CRASHED: {e}")))?; + + // ---- Launch 3/4: Deterministic gradient reduction ---- + let total_params_i32 = CUR_TOTAL_PARAMS as i32; + let num_blocks_i32 = num_blocks as i32; + let reduce_cfg = LaunchConfig { + grid_dim: (((CUR_TOTAL_PARAMS as u32) + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + // Suppress unused variable warning — total_params_i32 used only for clarity + let _ = total_params_i32; + unsafe { + self.stream + .launch_builder(&self.grad_reduce_func) + .arg(&self.partial_grads) .arg(&mut self.grad_w1) .arg(&mut self.grad_b1) .arg(&mut self.grad_w2) .arg(&mut self.grad_b2) - .arg(&n_i32) - .arg(&sd_i32) - .launch(fwd_bwd_cfg) - .map_err(|e| MLError::ModelError(format!("curiosity fwd_bwd launch: {e}")))?; + .arg(&num_blocks_i32) + .launch(reduce_cfg) + .map_err(|e| MLError::ModelError(format!("curiosity grad_reduce launch: {e}")))?; } - self.stream.synchronize() - .map_err(|e| MLError::ModelError(format!("curiosity fwd_bwd CRASHED: {e}")))?; - // Step 2b: Adam optimizer update (4 separate launches, one per param group) + // ---- Launch 4/4: Adam optimizer update (4 separate launches) ---- + self.step += 1; + let step = self.step; + launch_adam_step( &self.stream, &self.adam_func, &mut weights.w1, &self.grad_w1, CUR_W1_LEN, @@ -433,7 +439,7 @@ impl GpuCuriosityTrainer { n_train, step, )?; - debug!(step, n_train, "curiosity GPU training step complete (fused)"); + debug!(step, n_train, num_blocks, "curiosity GPU training step complete (deterministic)"); Ok(()) }