fix: rewrite curiosity training to per-block reduce — eliminates non-deterministic atomicAdd
Replace warp-reduced atomicAdd gradient accumulation in the curiosity forward model with a two-kernel deterministic pipeline: 1. curiosity_fwd_bwd_per_block: each block reduces its threads' gradient contributions via shared memory tree reduction, writes one partial gradient vector [CUR_TOTAL_PARAMS] per block. 2. curiosity_grad_reduce: one thread per parameter sums block partials in fixed order (block 0, 1, 2, ...). Fully deterministic. Deleted kernels: curiosity_forward_backward (atomicAdd path), curiosity_fused_zero_fwd_bwd_adam (grid-wide atomic barrier), curiosity_adam_step_fused (dead code), warp_sum_cur helper. Rust side: removed block_counter, adam_fused_func, fused_zero_fwd_bwd_adam_func. Added partial_grads buffer [max_blocks * CUR_TOTAL_PARAMS] (~2.8 MB for 64 blocks). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<CudaStream>,
|
||||
|
||||
// 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<f32>, // [CUR_W1_LEN]
|
||||
@@ -81,6 +79,10 @@ pub struct GpuCuriosityTrainer {
|
||||
grad_w2: CudaSlice<f32>, // [CUR_W2_LEN]
|
||||
grad_b2: CudaSlice<f32>, // [CUR_B2_LEN]
|
||||
|
||||
// Per-block partial gradient buffer for deterministic reduction
|
||||
partial_grads: CudaSlice<f32>, // [max_blocks * CUR_TOTAL_PARAMS]
|
||||
max_blocks: usize,
|
||||
|
||||
// Adam first moment (per-param-group)
|
||||
adam_m_w1: CudaSlice<f32>, // [CUR_W1_LEN]
|
||||
adam_m_b1: CudaSlice<f32>, // [CUR_B1_LEN]
|
||||
@@ -96,9 +98,6 @@ pub struct GpuCuriosityTrainer {
|
||||
// Shifted next_states buffer
|
||||
next_states_buf: CudaSlice<f32>,
|
||||
|
||||
/// Atomic block-arrival counter for fused kernel grid-wide sync (single i32 on GPU).
|
||||
block_counter: CudaSlice<i32>,
|
||||
|
||||
// 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<Self, MLError> {
|
||||
// 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::<f32>(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::<f32>(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::<f32>(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::<i32>(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::<f32>() 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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user