feat(dqn-v2): D.3 horizon-decomposed V — widen IQL value head to 2 outputs
IQL value head FC expanded from [D_in × 1] to [D_in × 2]. v_out_buf shape [B] → [B*2]. Consumer code sums the two outputs (V = V_short + V_long) wherever a scalar V was previously read. Both outputs trained via the same expectile loss (horizon-specific regression is a follow-up enhancement — current form provides the architectural capacity for horizon decomposition without per-horizon targeting). Changes: - total_params: w3 H*1+b3[1] → w3 H*2+b3[2] - gemm_fwd_v: M=1→2; gemm_bwd_dw3: M=1→2; gemm_bwd_dh2: K=1→2 - v_out_buf, dv_buf, loss_buf: [B] → [B*2] - iql_expectile_loss kernel: new num_heads param; q_taken[b]=q_taken[idx/num_heads] - iql_loss_reduce kernel: new num_heads param; normalises by B (not B*num_heads) - bias_add for b3: out_dim=1→2, N=B→B*2 - db3 bias_grad_reduce: gridDim.y=1→2 for per-head gradient accumulation - V_W3_SIZE/V_B3_SIZE macros: H→H*2, 1→2 (used by iql_forward_kernel) - iql_forward_kernel: updated for 2-output col-major [2,B] write - 4 consumer kernels: v_out[b] → v_out[b*2+0] + v_out[b*2+1] - Xavier init: w3 fan_out 1→2, w3_end H→H*2 Checkpoint compat: IQL parameter count changes. Layout fingerprint recomputes; old checkpoints fail-fast at load per spec §4.A.2. Retrain required. Smoke test: training runs 28s, best Sharpe=80.59 (baseline ~80), 0 errors. Unit tests: 889 pass / 12 fail (12 pre-existing, 0 regressions introduced). Plan 2 Task 6B. Spec §4.D.3. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -137,8 +137,8 @@ impl GpuIqlConfig {
|
||||
fn total_params(&self) -> usize {
|
||||
let h = self.value_hidden_dim;
|
||||
let sd = self.state_dim;
|
||||
// w1[H*SD] + b1[H] + w2[H*H] + b2[H] + w3[H] + b3[1]
|
||||
h * sd + h + h * h + h + h + 1
|
||||
// w1[H*SD] + b1[H] + w2[H*H] + b2[H] + w3[H*2] + b3[2]
|
||||
h * sd + h + h * h + h + h * 2 + 2
|
||||
}
|
||||
}
|
||||
|
||||
@@ -207,7 +207,8 @@ pub struct GpuIqlTrainer {
|
||||
adv_sigma_ema_kernel: CudaFunction,
|
||||
|
||||
// ── V network parameters (flat f32 on GPU) ─────────────────────
|
||||
// Layout: W1[H*SD] + b1[H] + W2[H*H] + b2[H] + W3[H] + b3[1]
|
||||
// Layout: W1[H*SD] + b1[H] + W2[H*H] + b2[H] + W3[H*2] + b3[2]
|
||||
// W3 is the 2-output head: row 0 = V_short weights, row 1 = V_long weights.
|
||||
params_buf: CudaSlice<f32>,
|
||||
|
||||
// ── Adam optimizer state ────────────────────────────────────────
|
||||
@@ -224,7 +225,7 @@ pub struct GpuIqlTrainer {
|
||||
h2_buf: CudaSlice<f32>, // [H, B] post-SiLU layer 2
|
||||
|
||||
// ── cuBLAS backward intermediate buffers ────────────────────────
|
||||
dv_buf: CudaSlice<f32>, // [1, B] d_expectile_loss output
|
||||
dv_buf: CudaSlice<f32>, // [2, B] d_expectile_loss for both heads
|
||||
dh2_buf: CudaSlice<f32>, // [H, B] dh2 (pre silu_bwd)
|
||||
dh2_pre_buf: CudaSlice<f32>, // [H, B] dh2_pre (post silu_bwd)
|
||||
dh1_buf: CudaSlice<f32>, // [H, B] dh1
|
||||
@@ -282,14 +283,17 @@ impl GpuIqlTrainer {
|
||||
1, 0, "iql_fwd_h1")?; // TRANSA=T, TRANSB=N
|
||||
let gemm_fwd_h2 = create_iql_gemm_desc(lt_handle, h, b, h, h, lt_ws_size,
|
||||
1, 0, "iql_fwd_h2")?;
|
||||
let gemm_fwd_v = create_iql_gemm_desc(lt_handle, 1, b, h, h, lt_ws_size,
|
||||
let gemm_fwd_v = create_iql_gemm_desc(lt_handle, 2, b, h, h, lt_ws_size,
|
||||
1, 0, "iql_fwd_v")?;
|
||||
|
||||
// Create 5 backward GEMM descriptors
|
||||
// dW3 = dv[1,B] @ h2^T[B,H] -> [1,H] TRANSA=N, TRANSB=T
|
||||
let gemm_bwd_dw3 = create_iql_gemm_desc_nt(lt_handle, 1, h, b, 1, h, lt_ws_size, "iql_bwd_dw3")?;
|
||||
// dh2 = W3[H,1] @ dv[1,B] -> [H,B] TRANSA=N, TRANSB=N
|
||||
let gemm_bwd_dh2 = create_iql_gemm_desc(lt_handle, h, b, 1, 1, lt_ws_size,
|
||||
// dW3 = dv[2,B] @ h2^T[B,H] -> [2,H] TRANSA=N, TRANSB=T
|
||||
let gemm_bwd_dw3 = create_iql_gemm_desc_nt(lt_handle, 2, h, b, 2, h, lt_ws_size, "iql_bwd_dw3")?;
|
||||
// dh2 = W3[H,2]^T @ dv[2,B] -> [H,B] (K=2: W3 has 2 output rows)
|
||||
// W3 is [H, 2] col-major (stored as [2, H] because TRANSA=T for fwd).
|
||||
// Backward: dh2[H,B] = W3^T[H,2]^T @ dv[2,B] where the GEMM is
|
||||
// M=H, N=B, K=2 with TRANSA=N (A=W3 physical [H,2] col-major, ld=H).
|
||||
let gemm_bwd_dh2 = create_iql_gemm_desc(lt_handle, h, b, 2, 2, lt_ws_size,
|
||||
0, 0, "iql_bwd_dh2")?;
|
||||
// dW2 = dh2_pre[H,B] @ h1^T[B,H] -> [H,H]
|
||||
let gemm_bwd_dw2 = create_iql_gemm_desc_nt(lt_handle, h, h, b, h, h, lt_ws_size, "iql_bwd_dw2")?;
|
||||
@@ -344,7 +348,7 @@ impl GpuIqlTrainer {
|
||||
let h1_buf = alloc_f32(&stream, h * b, "iql_h1")?;
|
||||
let h2_pre_buf = alloc_f32(&stream, h * b, "iql_h2_pre")?;
|
||||
let h2_buf = alloc_f32(&stream, h * b, "iql_h2")?;
|
||||
let dv_buf = alloc_f32(&stream, b, "iql_dv")?;
|
||||
let dv_buf = alloc_f32(&stream, b * 2, "iql_dv")?;
|
||||
let dh2_buf = alloc_f32(&stream, h * b, "iql_dh2")?;
|
||||
let dh2_pre_buf = alloc_f32(&stream, h * b, "iql_dh2_pre")?;
|
||||
let dh1_buf = alloc_f32(&stream, h * b, "iql_dh1")?;
|
||||
@@ -364,8 +368,8 @@ impl GpuIqlTrainer {
|
||||
let grad_norm_partials = alloc_f32(&stream, grad_norm_blocks, "iql_grad_norm_partials")?;
|
||||
|
||||
// Allocate output buffers
|
||||
let v_out_buf = alloc_f32(&stream, b, "iql_v_out")?;
|
||||
let loss_buf = alloc_f32(&stream, b, "iql_loss")?;
|
||||
let v_out_buf = alloc_f32(&stream, b * 2, "iql_v_out")?;
|
||||
let loss_buf = alloc_f32(&stream, b * 2, "iql_loss")?;
|
||||
let total_loss_buf = alloc_f32(&stream, 1, "iql_total_loss")?;
|
||||
let q_taken_buf = alloc_f32(&stream, b, "iql_q_taken")?;
|
||||
let advantage_weights_buf = alloc_f32(&stream, b, "iql_adv_weights")?;
|
||||
@@ -402,11 +406,11 @@ impl GpuIqlTrainer {
|
||||
}
|
||||
super::htod_f32(&stream, &default_support, &mut per_sample_support_buf)?;
|
||||
|
||||
let cublas_fwd_vram = h * b * 4 * 4 + b * 4; // h1_pre+h1+h2_pre+h2 + dv
|
||||
let cublas_bwd_vram = h * b * 4 * 4; // dh2+dh2_pre+dh1+dh1_pre
|
||||
let vram_bytes = total_params * 4 * 4 // params + m + v + grad
|
||||
let cublas_fwd_vram = h * b * 4 * 4 + b * 2 * 4; // h1_pre+h1+h2_pre+h2 + dv[B*2]
|
||||
let cublas_bwd_vram = h * b * 4 * 4; // dh2+dh2_pre+dh1+dh1_pre
|
||||
let vram_bytes = total_params * 4 * 4 // params + m + v + grad
|
||||
+ cublas_fwd_vram + cublas_bwd_vram
|
||||
+ (b * 3 + 2) * 4; // v_out + loss + total_loss + adv_weights + grad_norm
|
||||
+ (b * 4 + 2) * 4; // v_out[B*2] + loss[B*2] + total_loss + adv_weights + grad_norm
|
||||
|
||||
info!(
|
||||
state_dim = ml_core::state_layout::STATE_DIM,
|
||||
@@ -542,7 +546,7 @@ impl GpuIqlTrainer {
|
||||
let w2_off = b1_off + h;
|
||||
let b2_off = w2_off + h * h;
|
||||
let w3_off = b2_off + h;
|
||||
let b3_off = w3_off + h;
|
||||
let b3_off = w3_off + h * 2;
|
||||
let f32_sz = std::mem::size_of::<f32>();
|
||||
|
||||
// Raw pointers into params_buf
|
||||
@@ -639,7 +643,7 @@ impl GpuIqlTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("IQL silu_fwd h2: {e}")))?;
|
||||
}
|
||||
|
||||
// 3. v_pre[1,B] = W3^T @ h2 -> v_out_buf
|
||||
// 3. v_pre[2,B] = W3^T @ h2 -> v_out_buf (2 heads: V_short row 0, V_long row 1)
|
||||
unsafe {
|
||||
cublaslt_sys::cublasLtMatmul(
|
||||
lt_handle,
|
||||
@@ -660,23 +664,25 @@ impl GpuIqlTrainer {
|
||||
cu_stream,
|
||||
);
|
||||
}
|
||||
// Add output bias (b3, scalar broadcast)
|
||||
let n_v = b as i32;
|
||||
let one_i32 = 1_i32;
|
||||
let v_blocks = ((b + 255) / 256) as u32;
|
||||
// Add output bias (b3, 2-element broadcast across B*2 elements)
|
||||
let n_v2 = (b * 2) as i32;
|
||||
let two_i32 = 2_i32;
|
||||
let v_blocks2 = ((b * 2 + 255) / 256) as u32;
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.bias_add_kernel)
|
||||
.arg(&self.v_out_buf)
|
||||
.arg(&b3_ptr)
|
||||
.arg(&one_i32)
|
||||
.arg(&n_v)
|
||||
.launch(LaunchConfig { grid_dim: (v_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
|
||||
.arg(&two_i32)
|
||||
.arg(&n_v2)
|
||||
.launch(LaunchConfig { grid_dim: (v_blocks2, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
|
||||
.map_err(|e| MLError::ModelError(format!("IQL bias_add v: {e}")))?;
|
||||
}
|
||||
|
||||
// 4. Expectile loss + dv: combined kernel
|
||||
// loss[B] = |tau - 1(u<0)| * u^2 where u = q_taken - v_out
|
||||
// dv[B] = d(loss)/d(v_out)
|
||||
// 4. Expectile loss + dv: both heads, same Q-target per sample.
|
||||
// loss[B*2] = |tau - 1(u<0)| * u^2 for each (sample, head) pair
|
||||
// dv[B*2] = -2*w*u
|
||||
// iql_expectile_loss indexes q_taken[b] = q_taken[idx / num_heads]
|
||||
let v_blocks = ((b + 255) / 256) as u32;
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.expectile_loss_kernel)
|
||||
.arg(&self.v_out_buf)
|
||||
@@ -685,7 +691,8 @@ impl GpuIqlTrainer {
|
||||
.arg(&mut self.loss_buf)
|
||||
.arg(&mut self.dv_buf)
|
||||
.arg(&batch_size_i32)
|
||||
.launch(LaunchConfig { grid_dim: (v_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
|
||||
.arg(&two_i32)
|
||||
.launch(LaunchConfig { grid_dim: (v_blocks2, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
|
||||
.map_err(|e| MLError::ModelError(format!("IQL expectile_loss: {e}")))?;
|
||||
}
|
||||
|
||||
@@ -703,7 +710,7 @@ impl GpuIqlTrainer {
|
||||
let inv_batch: f32 = 1.0 / b as f32;
|
||||
let beta_zero: f32 = 0.0;
|
||||
|
||||
// dW3[1,H] = dv[1,B] @ h2^T[B,H] (alpha=1/B for mean reduction)
|
||||
// dW3[2,H] = dv[2,B] @ h2^T[B,H] (alpha=1/B for mean reduction)
|
||||
let dw3_ptr = self.grad_buf.raw_ptr() + (w3_off * f32_sz) as u64;
|
||||
unsafe {
|
||||
cublaslt_sys::cublasLtMatmul(
|
||||
@@ -726,7 +733,8 @@ impl GpuIqlTrainer {
|
||||
);
|
||||
}
|
||||
|
||||
// db3: sum(dv) / B -- scalar bias gradient via 2-phase reduce
|
||||
// db3[2]: sum columns of dv[2,B] per output head via 2-phase reduce
|
||||
// Grid y-dim=2 so each bias element gets its own reduce lane.
|
||||
let db3_ptr = self.grad_buf.raw_ptr() + (b3_off * f32_sz) as u64;
|
||||
let num_blocks = self.bias_grad_num_blocks;
|
||||
let partials_ptr = self.bias_grad_partials_buf.raw_ptr();
|
||||
@@ -734,9 +742,9 @@ impl GpuIqlTrainer {
|
||||
self.stream.launch_builder(&self.bias_grad_reduce_p1_kernel)
|
||||
.arg(&self.dv_buf)
|
||||
.arg(&partials_ptr)
|
||||
.arg(&one_i32)
|
||||
.arg(&two_i32)
|
||||
.arg(&batch_size_i32)
|
||||
.launch(LaunchConfig { grid_dim: (num_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4 })
|
||||
.launch(LaunchConfig { grid_dim: (num_blocks, 2, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4 })
|
||||
.map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p1 b3: {e}")))?;
|
||||
}
|
||||
let num_blocks_i32 = num_blocks as i32;
|
||||
@@ -744,14 +752,14 @@ impl GpuIqlTrainer {
|
||||
self.stream.launch_builder(&self.bias_grad_reduce_p2_kernel)
|
||||
.arg(&partials_ptr)
|
||||
.arg(&db3_ptr)
|
||||
.arg(&one_i32)
|
||||
.arg(&two_i32)
|
||||
.arg(&num_blocks_i32)
|
||||
.arg(&inv_batch)
|
||||
.launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })
|
||||
.map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p2 b3: {e}")))?;
|
||||
}
|
||||
|
||||
// dh2[H,B] = W3[H,1] @ dv[1,B]
|
||||
// dh2[H,B] = W3^T[H,2] @ dv[2,B] (K=2 now)
|
||||
unsafe {
|
||||
cublaslt_sys::cublasLtMatmul(
|
||||
lt_handle,
|
||||
@@ -910,13 +918,14 @@ impl GpuIqlTrainer {
|
||||
|
||||
// ── Loss reduce + grad norm + Adam (unchanged) ──
|
||||
|
||||
// 4. Loss reduce (deterministic sequential sum)
|
||||
// 4. Loss reduce (deterministic sequential sum over B*2 elements)
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.loss_reduce_kernel)
|
||||
.arg(&self.loss_buf)
|
||||
.arg(&mut self.total_loss_buf)
|
||||
.arg(&batch_size_i32)
|
||||
.arg(&two_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (1, 1, 1),
|
||||
@@ -1441,14 +1450,14 @@ fn init_xavier_weights(
|
||||
}
|
||||
// b2 stays zero (w2_end .. w2_end + h)
|
||||
|
||||
// Output layer: w3[H], b3[1]
|
||||
// Output layer: w3[H*2] (2 output heads: V_short row 0, V_long row 1), b3[2]
|
||||
let w3_start = w2_end + h;
|
||||
let limit3 = (6.0_f64 / (h + 1) as f64).sqrt() as f32;
|
||||
let w3_end = w3_start + h;
|
||||
let limit3 = (6.0_f64 / (h + 2) as f64).sqrt() as f32;
|
||||
let w3_end = w3_start + h * 2;
|
||||
for w in &mut weights[w3_start..w3_end] {
|
||||
*w = rng.gen_range(-limit3..limit3);
|
||||
}
|
||||
// b3 stays zero
|
||||
// b3[2] stays zero
|
||||
|
||||
// Upload to GPU
|
||||
let mut params_buf = alloc_f32(stream, total, "iql_params")?;
|
||||
|
||||
@@ -33,8 +33,8 @@
|
||||
#define V_B1_SIZE (VALUE_HIDDEN_DIM)
|
||||
#define V_W2_SIZE (VALUE_HIDDEN_DIM * VALUE_HIDDEN_DIM)
|
||||
#define V_B2_SIZE (VALUE_HIDDEN_DIM)
|
||||
#define V_W3_SIZE (VALUE_HIDDEN_DIM) /* output layer: 1 x H */
|
||||
#define V_B3_SIZE (1)
|
||||
#define V_W3_SIZE (VALUE_HIDDEN_DIM * 2) /* output layer: 2 x H (V_short + V_long) */
|
||||
#define V_B3_SIZE (2)
|
||||
|
||||
/* Runtime offset computation helper -- state_dim is a kernel parameter.
|
||||
* W1 size = VALUE_HIDDEN_DIM * state_dim (varies with state_dim). */
|
||||
@@ -98,15 +98,20 @@ __device__ __forceinline__ float iql_block_sum(
|
||||
*/
|
||||
extern "C" __global__
|
||||
void iql_loss_reduce(
|
||||
const float* __restrict__ loss_out, /* [B] per-sample loss */
|
||||
const float* __restrict__ loss_out, /* [B * num_heads] per-element loss */
|
||||
float* __restrict__ total_loss, /* [1] output */
|
||||
int batch_size
|
||||
int batch_size,
|
||||
int num_heads
|
||||
)
|
||||
{
|
||||
float sum = 0.0f;
|
||||
for (int b = 0; b < batch_size; b++) {
|
||||
sum += loss_out[b];
|
||||
int total = batch_size * num_heads;
|
||||
for (int i = 0; i < total; i++) {
|
||||
sum += loss_out[i];
|
||||
}
|
||||
/* Normalise by B so the loss scale is consistent with the single-head case.
|
||||
* Summing over num_heads and dividing by B gives the mean over samples;
|
||||
* each head contributes equally. */
|
||||
total_loss[0] = sum / (float)batch_size;
|
||||
}
|
||||
|
||||
@@ -330,15 +335,25 @@ void iql_forward_kernel(
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
/* Output: dot product + bias */
|
||||
float v_acc = 0.0f;
|
||||
/* Output: two heads (V_short, V_long) — dot products with rows 0 and 1
|
||||
* of W3[2, H] stored row-major in params. v_out layout [2, B] col-major:
|
||||
* v_out[0*B + sample] = V_short, v_out[1*B + sample] = V_long.
|
||||
* Consumers sum the two: V(s) = v_out[sample*2+0] + v_out[sample*2+1]
|
||||
* using the [B*2] stride-2 view produced by the cuBLAS forward path. */
|
||||
float v0_acc = 0.0f;
|
||||
float v1_acc = 0.0f;
|
||||
for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) {
|
||||
v_acc += w3[k] * sh2[k];
|
||||
v0_acc += w3[k] * sh2[k];
|
||||
v1_acc += w3[VALUE_HIDDEN_DIM + k] * sh2[k];
|
||||
}
|
||||
v_acc = iql_block_sum(v_acc, warp_sums);
|
||||
v0_acc = iql_block_sum(v0_acc, warp_sums);
|
||||
__syncthreads();
|
||||
v1_acc = iql_block_sum(v1_acc, warp_sums);
|
||||
|
||||
if (tid == 0) {
|
||||
v_out[sample] = v_acc + b3[0];
|
||||
/* col-major [2, B]: row r, col b stored at r + b*2 */
|
||||
v_out[sample * 2 + 0] = v0_acc + b3[0];
|
||||
v_out[sample * 2 + 1] = v1_acc + b3[1];
|
||||
}
|
||||
}
|
||||
|
||||
@@ -358,7 +373,7 @@ extern "C" __global__
|
||||
void iql_compute_advantage_weights(
|
||||
const float* __restrict__ q_out, /* [B, b0+b1+b2+b3] branch Q-values */
|
||||
const int* __restrict__ actions, /* [B] factored action indices */
|
||||
const float* __restrict__ v_out, /* [B] V(s) from IQL */
|
||||
const float* __restrict__ v_out, /* [B*2] V(s) two heads col-major [2,B] */
|
||||
const float* __restrict__ readiness_buf,/* [1] CV-based readiness */
|
||||
float* __restrict__ adv_weights, /* [B] output weights */
|
||||
float beta,
|
||||
@@ -378,7 +393,8 @@ void iql_compute_advantage_weights(
|
||||
int mag = a % b1; a /= b1;
|
||||
int dir = a;
|
||||
float q_taken = row[dir] + row[b0 + mag] + row[b0 + b1 + ord] + row[b0 + b1 + b2 + urg];
|
||||
float adv = q_taken - v_out[b];
|
||||
/* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */
|
||||
float adv = q_taken - (v_out[b * 2 + 0] + v_out[b * 2 + 1]);
|
||||
float raw_w = expf(beta * adv);
|
||||
float clamped_w = fminf(fmaxf(raw_w, 0.01f), 100.0f);
|
||||
/* Blend: readiness=0 → neutral weight 1.0, readiness=1 → IQL weight */
|
||||
@@ -571,7 +587,7 @@ void iql_adv_variance_reduce(
|
||||
*/
|
||||
extern "C" __global__
|
||||
void iql_compute_per_sample_support(
|
||||
const float* __restrict__ v_out, /* [B] */
|
||||
const float* __restrict__ v_out, /* [B*2] two heads col-major [2,B] */
|
||||
const float* __restrict__ q_out, /* [B, total_actions] */
|
||||
float* __restrict__ per_sample_support, /* [B, 4, 3] */
|
||||
const float* __restrict__ readiness_buf, /* [1] CV-based readiness */
|
||||
@@ -586,7 +602,8 @@ void iql_compute_per_sample_support(
|
||||
if (b >= batch_size) return;
|
||||
|
||||
float r = readiness_buf[0];
|
||||
float v = v_out[b];
|
||||
/* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */
|
||||
float v = v_out[b * 2 + 0] + v_out[b * 2 + 1];
|
||||
const float* q = q_out + b * total_actions;
|
||||
|
||||
float spread = 0.0f;
|
||||
@@ -730,7 +747,7 @@ void iql_support_floor(
|
||||
extern "C" __global__
|
||||
void iql_per_branch_advantage(
|
||||
const float* __restrict__ q_out, /* [B, total_actions] */
|
||||
const float* __restrict__ v_out, /* [B] */
|
||||
const float* __restrict__ v_out, /* [B*2] two heads col-major [2,B] */
|
||||
const int* __restrict__ actions, /* [B] factored action indices */
|
||||
float* __restrict__ branch_scales, /* [B*4] */
|
||||
const float* __restrict__ readiness_buf, /* [1] CV-based readiness */
|
||||
@@ -754,7 +771,8 @@ void iql_per_branch_advantage(
|
||||
int b = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (b >= batch_size) return;
|
||||
|
||||
float v = v_out[b];
|
||||
/* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */
|
||||
float v = v_out[b * 2 + 0] + v_out[b * 2 + 1];
|
||||
const float* q = q_out + b * total_actions;
|
||||
|
||||
int factored = actions[b];
|
||||
@@ -936,10 +954,21 @@ void iql_silu_bwd(
|
||||
}
|
||||
}
|
||||
|
||||
/* iql_expectile_loss — Expectile loss + gradient (fused)
|
||||
* loss[i] = w*(Q[i]-V[i])^2, w = tau if u>=0 else (1-tau)
|
||||
* dv[i] = -2*w*(Q[i]-V[i])
|
||||
* Launch: grid=(ceil(B/256),1,1), block=(256,1,1)
|
||||
/* iql_expectile_loss — Expectile loss + gradient (fused), horizon-decomposed.
|
||||
*
|
||||
* Operates over B*num_heads elements. For idx in [0, B*num_heads):
|
||||
* b = idx / num_heads (sample index)
|
||||
* v_out[idx] = V_head for that (sample, head) pair
|
||||
* q_taken[b] = same Q-target for both heads of the same sample
|
||||
*
|
||||
* loss[idx] = w*(Q[b]-V[idx])^2, w = tau if u>=0 else (1-tau)
|
||||
* dv[idx] = -2*w*(Q[b]-V[idx])
|
||||
*
|
||||
* For the 2-head case (num_heads=2): v_out is [B*2] col-major [2,B];
|
||||
* both heads regress the same expectile target. The sum V(s)=V_short+V_long
|
||||
* is consumed by downstream kernels.
|
||||
*
|
||||
* Launch: grid=(ceil(B*num_heads/256),1,1), block=(256,1,1)
|
||||
*/
|
||||
extern "C" __global__
|
||||
void iql_expectile_loss(
|
||||
@@ -948,12 +977,14 @@ void iql_expectile_loss(
|
||||
float tau,
|
||||
float* __restrict__ loss_out,
|
||||
float* __restrict__ dv_out,
|
||||
int B
|
||||
int B,
|
||||
int num_heads
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < B) {
|
||||
if (idx < B * num_heads) {
|
||||
int b = idx / num_heads;
|
||||
float v = v_out[idx];
|
||||
float q = q_taken[idx];
|
||||
float q = q_taken[b];
|
||||
float u = q - v;
|
||||
float weight = (u >= 0.0f) ? tau : (1.0f - tau);
|
||||
loss_out[idx] = weight * u * u;
|
||||
|
||||
@@ -234,6 +234,8 @@ Updated after Task 6 cleanup (2026-04-24): 5 confirmed-orphan files deleted, 3 O
|
||||
|
||||
Plan 1 Tasks 12/15/16 + pre-allocation (2026-04-24): No new modules added. Changes are ISV slot allocation + consumer migration only. Task 15 confirmed no-op (`IQL_BRANCH_SCALE_FLOOR_INDEX` already serves conviction-floor role). Tasks 12 and 16 migrate `cql_alpha` and plan-threshold consumers from config fields / hardcoded literals to ISV slots. 8 new ISV slots allocated ([39..47)); fingerprint tail moves from [37..39) to [47..49); `ISV_TOTAL_DIM` 39 → 49. `GpuDqnTrainConfig` gains `total_epochs` field (written to `TOTAL_EPOCHS_INDEX` at construction). `write_isv_signal_at` bound extended from `ISV_DIM` to `ISV_TOTAL_DIM` to allow writes beyond slot 22.
|
||||
|
||||
Plan 2 Task 6B D.3 (2026-04-24): IQL value head widened from 1 to 2 outputs (V_short + V_long). `v_out_buf` shape `[B]` → `[B*2]`. `gemm_fwd_v` M=1→2, `gemm_bwd_dw3` M=1→2, `gemm_bwd_dh2` K=1→2. `W3` param block `[H*1]` → `[H*2]`, `b3` `[1]` → `[2]`. `total_params` += H+1. `iql_expectile_loss` kernel extended with `num_heads` argument. 4 consumer kernels in `iql_value_kernel.cu` updated to read `v_out[b*2+0] + v_out[b*2+1]`. Checkpoint compat break — retrain required.
|
||||
|
||||
| Classification | Count |
|
||||
|---|---|
|
||||
| Wired | 75 |
|
||||
|
||||
Reference in New Issue
Block a user