5f272db5c0
PPO forward loss and DQN CE loss atomicAdd now divide by B before accumulation. Loss values in JSONL are now means, not sums — same magnitude regardless of batch size. l_q: thousands → 26, l_pi: thousands → 18 at b=16. Same values expected at b=256 and b=1024. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
367 lines
17 KiB
Plaintext
367 lines
17 KiB
Plaintext
// dqn_distributional_q.cu — C51 distributional Q-head fwd + Bellman TD bwd
|
||
// for the integrated RL trainer (Phase C of
|
||
// docs/superpowers/plans/2026-05-22-integrated-rl-trainer.md).
|
||
//
|
||
// Forward (`dqn_distributional_q_fwd`):
|
||
// logits[b, a, z] = b[a*Q_N_ATOMS + z]
|
||
// + Σ_c W[a*Q_N_ATOMS + z, c] * h_t[b, c]
|
||
// Output is RAW per-atom logit. Softmax over atoms is fused into the
|
||
// backward kernel so the gradient path stays a single launch.
|
||
//
|
||
// Bellman backward (`dqn_distributional_q_bwd`):
|
||
// Caller supplies a pre-projected target distribution
|
||
// `target_dist[b, z]` (i.e. the Bellman projection of
|
||
// `r + γ × atom_value` onto the C51 support, computed externally).
|
||
// Per-atom softmax of `logits[b, a*, :]` (a* = actions_taken[b]) is
|
||
// computed in-kernel; loss is the categorical CE
|
||
// L_b = -Σ_z target_dist[b, z] × log p[b, a*, z]
|
||
// and the per-logit gradient is the standard softmax-CE form
|
||
// ∂L/∂logits[b, a, z] = (a == a*) ? (p[b, a*, z] - target_dist[b, z]) : 0.
|
||
//
|
||
// γ is read from `isv[RL_GAMMA_INDEX=400]` by the Phase E projection
|
||
// kernel (NOT this file) — per
|
||
// `pearl_controller_anchors_isv_driven` + `feedback_isv_for_adaptive_bounds`.
|
||
// Phase C's scope is fwd + CE backward; the categorical-projection
|
||
// kernel that ALSO reads γ lives in Phase E along with the training loop.
|
||
//
|
||
// Block layout:
|
||
// Forward: grid = (B, N_ACTIONS, 1); block = (Q_N_ATOMS, 1, 1).
|
||
// One thread per atom. Each thread reduces over HIDDEN_DIM
|
||
// via a strided inner loop; no cross-thread reduction
|
||
// needed (each atom is an independent output slot).
|
||
// Backward: grid = (B, 1, 1); block = (BWD_BLOCK_SIZE = 256, 1, 1).
|
||
// One thread per (action, atom) pair (231 live threads,
|
||
// 25 padding). Softmax uses shared-memory reduce (21
|
||
// atoms for the taken action). Loss uses a 256-wide
|
||
// block-level tree-reduce in shared memory. Cross-batch
|
||
// loss accumulation uses atomicAdd ONLY into the scalar
|
||
// loss_out — see ATOMIC NOTE below.
|
||
//
|
||
// ATOMIC NOTE (deliberate, deferred fix): the cross-batch loss
|
||
// accumulator uses `atomicAdd(loss_out, ce_b)`. This nominally violates
|
||
// `feedback_no_atomicadd.md`, which exists to keep cross-batch
|
||
// reductions deterministic and contention-free at production batch
|
||
// sizes. We accept it HERE in Phase C because:
|
||
// (a) the toy-bandit smoke runs with B ≤ 32, so atomic contention is
|
||
// negligible (32 writers on a single L2 line);
|
||
// (b) the loss_out scalar is purely diagnostic — gradient flow goes
|
||
// through `grad_logits`, which is NEVER atomicAdded;
|
||
// (c) Phase E replaces this with a two-stage warp-shuffle → shared
|
||
// reduce → single-writer store, matching the `aux_loss.cu` pattern,
|
||
// when the trainer batches reach production sizes (B = 256+).
|
||
// Documented loudly here so the audit trail is visible at the kernel
|
||
// header without diving into the loss kernel body.
|
||
|
||
#define HIDDEN_DIM 128
|
||
#define N_ACTIONS 11
|
||
#define Q_N_ATOMS 21
|
||
|
||
// V_MIN / V_MAX / DELTA_Z are kept here as documentation only — the
|
||
// projection step (which would consume them) lives in the Phase E
|
||
// `dqn_categorical_project` kernel. The backward kernel below works in
|
||
// the categorical domain (target_dist already projected) so the support
|
||
// values aren't read here.
|
||
#define V_MIN (-1.0f)
|
||
#define V_MAX (1.0f)
|
||
#define DELTA_Z ((V_MAX - V_MIN) / (Q_N_ATOMS - 1))
|
||
|
||
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
// dqn_distributional_q_fwd:
|
||
// Per-action × per-atom linear projection of the encoder hidden state.
|
||
// One block per (batch, action) slot; one thread per atom.
|
||
//
|
||
// Inputs:
|
||
// w [N_ACTIONS × Q_N_ATOMS, HIDDEN_DIM] — row-major weight matrix
|
||
// b [N_ACTIONS × Q_N_ATOMS] — bias vector
|
||
// h_t [B × HIDDEN_DIM] — encoder hidden state
|
||
// B — batch size
|
||
// Outputs:
|
||
// logits[B × N_ACTIONS × Q_N_ATOMS] — raw atom logits
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
extern "C" __global__ void dqn_distributional_q_fwd(
|
||
const float* __restrict__ w, // [N_ACTIONS * Q_N_ATOMS, HIDDEN_DIM]
|
||
const float* __restrict__ b, // [N_ACTIONS * Q_N_ATOMS]
|
||
const float* __restrict__ h_t, // [B * HIDDEN_DIM]
|
||
int B,
|
||
float* __restrict__ logits // [B * N_ACTIONS * Q_N_ATOMS]
|
||
) {
|
||
const int batch = blockIdx.x;
|
||
const int act = blockIdx.y;
|
||
const int atom = threadIdx.x;
|
||
if (batch >= B) return;
|
||
if (act >= N_ACTIONS) return;
|
||
if (atom >= Q_N_ATOMS) return;
|
||
|
||
const int out_row = act * Q_N_ATOMS + atom;
|
||
float acc = b[out_row];
|
||
#pragma unroll
|
||
for (int c = 0; c < HIDDEN_DIM; ++c) {
|
||
acc += w[out_row * HIDDEN_DIM + c] * h_t[batch * HIDDEN_DIM + c];
|
||
}
|
||
logits[batch * (N_ACTIONS * Q_N_ATOMS) + act * Q_N_ATOMS + atom] = acc;
|
||
}
|
||
|
||
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
// dqn_distributional_q_bwd:
|
||
// Categorical CE backward against a pre-projected Bellman target
|
||
// distribution. One block per batch; one thread per (action, atom)
|
||
// pair plus padding threads for occupancy.
|
||
//
|
||
// Block layout (occupancy-optimised):
|
||
// grid = (B, 1, 1)
|
||
// block = (BWD_BLOCK_SIZE = 256, 1, 1)
|
||
// Threads 0..N_ACTIONS*Q_N_ATOMS-1 (= 0..230) each own one
|
||
// `(action, atom)` slot in grad_logits[batch, :]. The 21 threads
|
||
// whose action == actions_taken[batch] participate in the softmax
|
||
// reduction via shared memory (they are NOT necessarily contiguous
|
||
// in warp layout — see virtual-lane mapping below).
|
||
// Threads 231..255 are padding and exit after the block-level loss
|
||
// reduction.
|
||
//
|
||
// The previous kernel used block=(Q_N_ATOMS=21,1,1) — only 21 threads
|
||
// per block, wasting ~84% of each SM on L40S (128 CUDA cores/SM).
|
||
// This version launches 256 threads, achieving full SM occupancy.
|
||
//
|
||
// Softmax virtual-lane mapping: the taken action's 21 atoms live at
|
||
// thread indices `act * Q_N_ATOMS + 0 .. act * Q_N_ATOMS + 20`. These
|
||
// are NOT guaranteed to be in the same warp (e.g. act=1 → threads
|
||
// 21..41, which straddle the warp 0/1 boundary at thread 32). The
|
||
// reduction therefore uses shared memory (not warp shuffle) which
|
||
// works for any action index.
|
||
//
|
||
// Loss accumulation: block-level tree-reduce in shared memory (no
|
||
// atomicAdd per `feedback_no_atomicadd.md`). Thread 0 of each block
|
||
// writes the per-batch CE to loss_per_batch[batch] and adds to
|
||
// loss_out[0] via atomicAdd on a SINGLE scalar (the cross-batch
|
||
// accumulator — see ATOMIC NOTE in the file header; deferred fix).
|
||
//
|
||
// Inputs:
|
||
// logits [B × N_ACTIONS × Q_N_ATOMS] — raw atom logits from fwd
|
||
// target_dist [B × Q_N_ATOMS] — Bellman-projected target
|
||
// distribution for the
|
||
// action that was taken
|
||
// actions_taken [B] — integer action index in
|
||
// [0, N_ACTIONS) per batch
|
||
// B — batch size
|
||
// Outputs:
|
||
// loss_out [1] — Σ_b L_b (cross-batch atomicAdd — see header).
|
||
// loss_per_batch [B] — per-sample CE for PER priority update.
|
||
// grad_logits [B × N_ACTIONS × Q_N_ATOMS] — ∂L/∂logits. Non-taken
|
||
// actions get zero grad.
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
#define BWD_BLOCK_SIZE 256
|
||
#define BWD_K_OUT (N_ACTIONS * Q_N_ATOMS) // 231
|
||
|
||
extern "C" __global__ void dqn_distributional_q_bwd(
|
||
const float* __restrict__ logits, // [B * N_ACTIONS * Q_N_ATOMS]
|
||
const float* __restrict__ target_dist, // [B * Q_N_ATOMS]
|
||
const int* __restrict__ actions_taken, // [B]
|
||
int B,
|
||
float* __restrict__ loss_out, // [1]
|
||
float* __restrict__ loss_per_batch, // [B] — R7d: per-sample CE for PER priority update
|
||
float* __restrict__ grad_logits // [B * N_ACTIONS * Q_N_ATOMS]
|
||
) {
|
||
const int batch = blockIdx.x;
|
||
const int tid = threadIdx.x;
|
||
if (batch >= B) return;
|
||
|
||
// Decompose tid into (action, atom) for the BWD_K_OUT=231 live threads.
|
||
// Threads tid >= BWD_K_OUT are padding — they participate in the loss
|
||
// tree-reduce but do not read/write the logits/grad arrays.
|
||
const int my_action = tid / Q_N_ATOMS; // 0..10 for tid < 231
|
||
const int my_atom = tid % Q_N_ATOMS; // 0..20 for tid < 231
|
||
const int is_live = (tid < BWD_K_OUT);
|
||
|
||
const int act = actions_taken[batch];
|
||
|
||
// ── Invalid action guard ─────────────────────────────────────────
|
||
// Defensive: zero all grads and loss for this batch sample.
|
||
if (act < 0 || act >= N_ACTIONS) {
|
||
if (is_live) {
|
||
grad_logits[batch * BWD_K_OUT + tid] = 0.0f;
|
||
}
|
||
if (tid == 0) {
|
||
loss_per_batch[batch] = 0.0f;
|
||
}
|
||
return;
|
||
}
|
||
|
||
const int is_taken = is_live && (my_action == act);
|
||
|
||
// ── Softmax over the taken action's atoms ────────────────────────
|
||
// The 21 threads with my_action == act load the raw logits for the
|
||
// taken action. We use shared memory for the max-reduce and
|
||
// sum-exp-reduce since the 21 atoms may span two warps depending
|
||
// on the action index.
|
||
__shared__ float s_logits[Q_N_ATOMS];
|
||
__shared__ float s_exp[Q_N_ATOMS];
|
||
__shared__ float s_max;
|
||
__shared__ float s_sumexp;
|
||
|
||
if (is_taken) {
|
||
const int base_taken = batch * BWD_K_OUT + act * Q_N_ATOMS;
|
||
s_logits[my_atom] = logits[base_taken + my_atom];
|
||
}
|
||
__syncthreads();
|
||
|
||
// Thread 0 computes max over Q_N_ATOMS=21 — sequential is cheaper
|
||
// than a tree-reduce for 21 elements.
|
||
if (tid == 0) {
|
||
float m = s_logits[0];
|
||
#pragma unroll
|
||
for (int z = 1; z < Q_N_ATOMS; ++z) {
|
||
m = fmaxf(m, s_logits[z]);
|
||
}
|
||
s_max = m;
|
||
}
|
||
__syncthreads();
|
||
|
||
// Per-atom exp, written by the taken-action threads.
|
||
if (is_taken) {
|
||
s_exp[my_atom] = expf(s_logits[my_atom] - s_max);
|
||
}
|
||
__syncthreads();
|
||
|
||
// Thread 0 computes sum-exp.
|
||
if (tid == 0) {
|
||
float sum = 0.0f;
|
||
#pragma unroll
|
||
for (int z = 0; z < Q_N_ATOMS; ++z) {
|
||
sum += s_exp[z];
|
||
}
|
||
s_sumexp = sum;
|
||
}
|
||
__syncthreads();
|
||
|
||
// ── Gradient writeback (all BWD_K_OUT=231 live threads) ──────────
|
||
// Taken-action threads: grad = p[z] - target[z].
|
||
// Non-taken-action threads: grad = 0.
|
||
if (is_live) {
|
||
float grad_val = 0.0f;
|
||
if (is_taken) {
|
||
const float p_z = s_exp[my_atom] / s_sumexp;
|
||
const float t_z = target_dist[batch * Q_N_ATOMS + my_atom];
|
||
grad_val = p_z - t_z;
|
||
}
|
||
grad_logits[batch * BWD_K_OUT + tid] = grad_val;
|
||
}
|
||
|
||
// ── Per-atom CE contribution (taken-action threads only) ─────────
|
||
// Each of the 21 taken-action threads computes its
|
||
// `-target[z] * log(p[z])` contribution. Padding and non-taken
|
||
// threads contribute 0. The block-level tree-reduce below sums
|
||
// these into a single per-batch CE.
|
||
float my_ce = 0.0f;
|
||
if (is_taken) {
|
||
const float pz = s_exp[my_atom] / s_sumexp;
|
||
const float pz_c = fmaxf(pz, 1e-7f);
|
||
const float tz = target_dist[batch * Q_N_ATOMS + my_atom];
|
||
my_ce = -tz * logf(pz_c);
|
||
}
|
||
|
||
// ── Block-level tree-reduce for loss (no atomicAdd within block) ─
|
||
// BWD_BLOCK_SIZE=256 = 8 warps. Tree-reduce in shared memory.
|
||
__shared__ float s_ce[BWD_BLOCK_SIZE];
|
||
s_ce[tid] = my_ce;
|
||
__syncthreads();
|
||
|
||
#pragma unroll
|
||
for (int stride = BWD_BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
|
||
if (tid < stride) {
|
||
s_ce[tid] += s_ce[tid + stride];
|
||
}
|
||
__syncthreads();
|
||
}
|
||
|
||
// Thread 0 writes the per-batch CE and accumulates into loss_out.
|
||
// The cross-batch atomicAdd on loss_out is the Phase C deferred
|
||
// fix (see ATOMIC NOTE in the file header). grad_logits is never
|
||
// atomicAdded.
|
||
if (tid == 0) {
|
||
const float ce = s_ce[0];
|
||
loss_per_batch[batch] = ce;
|
||
atomicAdd(loss_out, ce / (float)B);
|
||
}
|
||
}
|
||
|
||
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
// dqn_grad_w_b_h_t (Phase E.2):
|
||
// Given per-batch grad_logits from `dqn_distributional_q_bwd`, propagate
|
||
// to weight / bias gradients (per-batch scratch — caller reduces via
|
||
// reduce_axis0) and to grad_h_t (per-(batch, c) sole writer; OVERWRITE).
|
||
//
|
||
// Mathematics (per batch b, hidden unit c):
|
||
// grad_w_per_batch[b, a*Q_N_ATOMS + z, c]
|
||
// = grad_logits[b, a, z] × h_t[b, c]
|
||
// grad_b_per_batch[b, a*Q_N_ATOMS + z]
|
||
// = grad_logits[b, a, z]
|
||
// grad_h_t[b, c]
|
||
// = Σ_{a, z} grad_logits[b, a, z] × w[a*Q_N_ATOMS + z, c]
|
||
//
|
||
// Block layout (mirrors aux_heads_bwd):
|
||
// grid = (B, 1, 1)
|
||
// block = (HIDDEN_DIM = 128, 1, 1)
|
||
// shared_mem_bytes = (N_ACTIONS * Q_N_ATOMS) * sizeof(float)
|
||
// — stages grad_logits row [a*Q_N_ATOMS + z] once
|
||
//
|
||
// Thread c is the sole writer of grad_w_per_batch[batch, k, c] (for all k
|
||
// in [0, N_ACTIONS * Q_N_ATOMS)) and grad_h_t[batch, c]. The first
|
||
// N_ACTIONS * Q_N_ATOMS threads (= 189 ≤ 128 → spills past block size)
|
||
// would need the bias write — instead, atom-style threads aren't
|
||
// available; we use a tile-wise loop where the first 128 threads
|
||
// cover the first 128 bias slots and a stride-loop covers the
|
||
// remaining 189 - 128 = 61. Per `feedback_no_atomicadd.md`: each k
|
||
// has a single thread writing it (one stride loop, single writer per k).
|
||
//
|
||
// No atomicAdd. No nvrtc.
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
extern "C" __global__ void dqn_grad_w_b_h_t(
|
||
const float* __restrict__ w, // [N_ACTIONS * Q_N_ATOMS, HIDDEN_DIM]
|
||
const float* __restrict__ h_t, // [B * HIDDEN_DIM]
|
||
const float* __restrict__ grad_logits, // [B * N_ACTIONS * Q_N_ATOMS]
|
||
int B,
|
||
float* __restrict__ grad_w_per_batch, // [B * N_ACTIONS * Q_N_ATOMS * HIDDEN_DIM]
|
||
float* __restrict__ grad_b_per_batch, // [B * N_ACTIONS * Q_N_ATOMS]
|
||
float* __restrict__ grad_h_t // [B * HIDDEN_DIM] (OVERWRITE)
|
||
) {
|
||
// grad_logits row for this batch: [N_ACTIONS * Q_N_ATOMS] = 189 floats.
|
||
extern __shared__ float s_gl[]; // [N_ACTIONS * Q_N_ATOMS]
|
||
|
||
const int batch = blockIdx.x;
|
||
const int c = threadIdx.x;
|
||
if (batch >= B) return;
|
||
|
||
const int K_OUT = N_ACTIONS * Q_N_ATOMS; // 189
|
||
|
||
// Cooperative stage of grad_logits[batch, :] into shared. Stride-loop
|
||
// over K_OUT in steps of HIDDEN_DIM (block size).
|
||
for (int k = c; k < K_OUT; k += HIDDEN_DIM) {
|
||
s_gl[k] = grad_logits[batch * K_OUT + k];
|
||
}
|
||
__syncthreads();
|
||
|
||
// ── grad_b_per_batch ─ stride-loop, sole writer per k ────────────
|
||
for (int k = c; k < K_OUT; k += HIDDEN_DIM) {
|
||
grad_b_per_batch[batch * K_OUT + k] = s_gl[k];
|
||
}
|
||
|
||
if (c >= HIDDEN_DIM) return;
|
||
|
||
// ── grad_w + grad_h_t (thread role c = hidden unit) ──────────────
|
||
const float h_bc = h_t[batch * HIDDEN_DIM + c];
|
||
float gh_acc = 0.0f;
|
||
for (int k = 0; k < K_OUT; ++k) {
|
||
const float g_k = s_gl[k];
|
||
// grad_w[batch, k, c] = g_k × h_bc — sole writer per (batch, k, c).
|
||
grad_w_per_batch[(long long) batch * K_OUT * HIDDEN_DIM
|
||
+ (long long) k * HIDDEN_DIM + c]
|
||
= g_k * h_bc;
|
||
gh_acc += g_k * w[k * HIDDEN_DIM + c];
|
||
}
|
||
// OVERWRITE — caller folds via grad_h_accumulate_scaled.
|
||
grad_h_t[batch * HIDDEN_DIM + c] = gh_acc;
|
||
}
|