feat(rl): outcome head backward + diag JSONL enrichment
Wire outcome CE backward → grad_W/b → Adam → grad_h_accumulate. Add PopArt, spectral, Q-bias, per-branch LR, outcome metrics to diag JSONL for training validation. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -120,6 +120,7 @@ const KERNELS: &[&str] = &[
|
||||
"rl_outcome_fwd", // Outcome aux: linear forward h_t → 3-class logits
|
||||
"rl_outcome_ce", // Outcome aux: softmax CE loss + gradient (masked by sentinel -1)
|
||||
"rl_outcome_label", // Outcome aux: assign labels from reward/done (Profit/Timeout/Loss)
|
||||
"rl_outcome_bwd", // Outcome aux: single linear layer backward — dW (per-batch), db (per-batch), dh_t (per-batch); 1 block per batch, 128 threads
|
||||
];
|
||||
|
||||
// Cache bust v31 — five new reduce / derive kernels populate the input
|
||||
|
||||
70
crates/ml-alpha/cuda/rl_outcome_bwd.cu
Normal file
70
crates/ml-alpha/cuda/rl_outcome_bwd.cu
Normal file
@@ -0,0 +1,70 @@
|
||||
// rl_outcome_bwd.cu — Outcome aux head backward (single linear layer).
|
||||
//
|
||||
// Given dL/dlogits [B, K_CLASSES=3] (produced by rl_outcome_ce),
|
||||
// compute the linear layer gradients via the chain rule:
|
||||
//
|
||||
// logits[b, k] = Σ_i h_t[b, i] × W[i, k] + b[k]
|
||||
//
|
||||
// dL/dW_per_batch[b, i, k] = h_t[b, i] × dL/dlogits[b, k]
|
||||
// dL/db_per_batch[b, k] = dL/dlogits[b, k]
|
||||
// dL/dh_t[b, i] = Σ_k W[i, k] × dL/dlogits[b, k]
|
||||
//
|
||||
// Block layout: 1 block per batch, HIDDEN_DIM=128 threads.
|
||||
// * Thread i (i < 128) computes dL/dh_t[b, i] (inner product
|
||||
// across K_CLASSES=3 outputs — trivially unrolled).
|
||||
// * Thread i also writes grad_W_per_batch[b, i, 0..2] (3 writes).
|
||||
// * Threads 0..2 write grad_b_per_batch[b, tid] (thread 3..127 idle
|
||||
// for that single-cycle step).
|
||||
//
|
||||
// Per `feedback_no_atomicadd`: sole-writer per (b, i, k).
|
||||
// Per `feedback_cpu_is_read_only`: pure device-side.
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#define HIDDEN_DIM 128
|
||||
#define K_CLASSES 3
|
||||
|
||||
extern "C" __global__ void rl_outcome_bwd(
|
||||
const float* __restrict__ h_t, // [B, HIDDEN_DIM]
|
||||
const float* __restrict__ grad_logits, // [B, K_CLASSES]
|
||||
const float* __restrict__ w, // [HIDDEN_DIM, K_CLASSES]
|
||||
int b_size,
|
||||
float* __restrict__ grad_w_per_batch, // [B, HIDDEN_DIM, K_CLASSES]
|
||||
float* __restrict__ grad_b_per_batch, // [B, K_CLASSES]
|
||||
float* __restrict__ grad_h_t // [B, HIDDEN_DIM]
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
const int i = threadIdx.x;
|
||||
if (b >= b_size || i >= HIDDEN_DIM) return;
|
||||
|
||||
// Stage grad_logits into shared — K_CLASSES=3 so only threads 0..2
|
||||
// participate, but all 128 threads wait on the sync.
|
||||
__shared__ float s_grad_logits[K_CLASSES];
|
||||
if (i < K_CLASSES) {
|
||||
s_grad_logits[i] = grad_logits[b * K_CLASSES + i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const float h_bi = h_t[b * HIDDEN_DIM + i];
|
||||
|
||||
// grad_W_per_batch[b, i, k] = h_bi × s_grad_logits[k]
|
||||
const int row_off = (b * HIDDEN_DIM + i) * K_CLASSES;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < K_CLASSES; ++k) {
|
||||
grad_w_per_batch[row_off + k] = h_bi * s_grad_logits[k];
|
||||
}
|
||||
|
||||
// grad_h_t[b, i] = Σ_k W[i, k] × s_grad_logits[k]
|
||||
float acc = 0.0f;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < K_CLASSES; ++k) {
|
||||
acc += w[i * K_CLASSES + k] * s_grad_logits[k];
|
||||
}
|
||||
grad_h_t[b * HIDDEN_DIM + i] = acc;
|
||||
|
||||
// grad_b_per_batch[b, k] = s_grad_logits[k] — sole writer per
|
||||
// (b, k); thread i writes slot k == i (covers k < 3 only).
|
||||
if (i < K_CLASSES) {
|
||||
grad_b_per_batch[b * K_CLASSES + i] = s_grad_logits[i];
|
||||
}
|
||||
}
|
||||
@@ -86,6 +86,12 @@ use ml_alpha::rl::isv_slots::{
|
||||
RL_REWARD_SCALE_INDEX, RL_TARGET_TAU_INDEX, RL_TD_KURTOSIS_CLAMP_INDEX,
|
||||
RL_TD_KURTOSIS_EMA_INDEX, RL_TD_KURT_STREAM_M2_INDEX, RL_TD_KURT_STREAM_M4_INDEX,
|
||||
RL_TD_KURT_STREAM_MEAN_INDEX, RL_V_GRAD_NORM_EMA_INDEX,
|
||||
// P1+P2 feature metrics
|
||||
RL_POPART_MEAN_INDEX, RL_POPART_SIGMA_INDEX, RL_POPART_VAR_INDEX,
|
||||
RL_SPECTRAL_NORM_MAX_INDEX, RL_SPECTRAL_DECOUPLE_LAMBDA_INDEX,
|
||||
RL_Q_BIAS_EMA_INDEX, RL_Q_BIAS_CORRECTION_INDEX,
|
||||
RL_LR_SCALE_Q_INDEX, RL_LR_SCALE_PI_INDEX, RL_LR_SCALE_V_INDEX, RL_LR_SCALE_IQN_INDEX,
|
||||
RL_OUTCOME_AUX_LAMBDA_INDEX,
|
||||
};
|
||||
use ml_alpha::trainer::diag_staging::DiagStaging;
|
||||
use cudarc::driver::DevicePtrMut;
|
||||
@@ -1113,6 +1119,31 @@ fn main() -> Result<()> {
|
||||
// away from the uniform baseline is the supervised-learning signal
|
||||
// for the FRD head (CE loss + bwd land in F.3).
|
||||
"frd": frd_diag,
|
||||
// P1+P2 feature metrics — PopArt reward normalization,
|
||||
// spectral norm/decoupling, Q-bias correction, per-branch
|
||||
// LR scaling, and outcome aux head lambda.
|
||||
"popart": {
|
||||
"mean": isv[RL_POPART_MEAN_INDEX],
|
||||
"sigma": isv[RL_POPART_SIGMA_INDEX],
|
||||
"var": isv[RL_POPART_VAR_INDEX],
|
||||
},
|
||||
"spectral": {
|
||||
"norm_max_config": isv[RL_SPECTRAL_NORM_MAX_INDEX],
|
||||
"decouple_lambda": isv[RL_SPECTRAL_DECOUPLE_LAMBDA_INDEX],
|
||||
},
|
||||
"q_bias": {
|
||||
"ema": isv[RL_Q_BIAS_EMA_INDEX],
|
||||
"correction": isv[RL_Q_BIAS_CORRECTION_INDEX],
|
||||
},
|
||||
"per_branch_lr": {
|
||||
"scale_q": isv[RL_LR_SCALE_Q_INDEX],
|
||||
"scale_pi": isv[RL_LR_SCALE_PI_INDEX],
|
||||
"scale_v": isv[RL_LR_SCALE_V_INDEX],
|
||||
"scale_iqn": isv[RL_LR_SCALE_IQN_INDEX],
|
||||
},
|
||||
"outcome_aux": {
|
||||
"lambda": isv[RL_OUTCOME_AUX_LAMBDA_INDEX],
|
||||
},
|
||||
});
|
||||
writeln!(diag, "{}", record).context("diag: writeln jsonl")?;
|
||||
|
||||
|
||||
@@ -37,10 +37,12 @@ pub struct OutcomeHead {
|
||||
_module_fwd: Arc<CudaModule>,
|
||||
_module_ce: Arc<CudaModule>,
|
||||
_module_label: Arc<CudaModule>,
|
||||
_module_bwd: Arc<CudaModule>,
|
||||
kernel_fwd: CudaFunction,
|
||||
kernel_ce: CudaFunction,
|
||||
kernel_label: CudaFunction,
|
||||
kernel_fill_sentinel: CudaFunction,
|
||||
kernel_bwd: CudaFunction,
|
||||
}
|
||||
|
||||
impl OutcomeHead {
|
||||
@@ -60,6 +62,8 @@ impl OutcomeHead {
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_outcome_ce.cubin"));
|
||||
static CUBIN_LABEL: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_outcome_label.cubin"));
|
||||
static CUBIN_BWD: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_outcome_bwd.cubin"));
|
||||
|
||||
let module_fwd = stream
|
||||
.context()
|
||||
@@ -88,6 +92,14 @@ impl OutcomeHead {
|
||||
.load_function("rl_outcome_fill_sentinel")
|
||||
.map_err(|e| anyhow!("OutcomeHead: fill_sentinel kernel resolve: {e}"))?;
|
||||
|
||||
let module_bwd = stream
|
||||
.context()
|
||||
.load_cubin(CUBIN_BWD.to_vec())
|
||||
.map_err(|e| anyhow!("OutcomeHead: bwd cubin load: {e}"))?;
|
||||
let kernel_bwd = module_bwd
|
||||
.load_function("rl_outcome_bwd")
|
||||
.map_err(|e| anyhow!("OutcomeHead: bwd kernel resolve: {e}"))?;
|
||||
|
||||
// Xavier init W at 0.01x scale: near-zero for near-uniform softmax at init.
|
||||
let w_d = ml_core::cuda_autograd::init::near_zero_xavier(
|
||||
HIDDEN_DIM, K_CLASSES, &stream,
|
||||
@@ -143,10 +155,12 @@ impl OutcomeHead {
|
||||
_module_fwd: module_fwd,
|
||||
_module_ce: module_ce,
|
||||
_module_label: module_label,
|
||||
_module_bwd: module_bwd,
|
||||
kernel_fwd,
|
||||
kernel_ce,
|
||||
kernel_label,
|
||||
kernel_fill_sentinel,
|
||||
kernel_bwd,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -212,6 +226,58 @@ impl OutcomeHead {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Backward pass: grad_logits (from `compute_loss`) -> grad_W, grad_b, grad_h_t.
|
||||
///
|
||||
/// Requires `forward()` and `compute_loss()` to have been called first
|
||||
/// (populates `grad_logits_d`). Computes per-batch weight/bias gradients
|
||||
/// and the encoder-upstream gradient `grad_h_t = grad_logits × W^T`.
|
||||
///
|
||||
/// The caller is responsible for:
|
||||
/// 1. Zeroing the output buffers before the call.
|
||||
/// 2. `reduce_axis0` on `grad_w_per_batch_d` and `grad_b_per_batch_d`.
|
||||
/// 3. Adam step on `self.w_d` / `self.b_d` with the reduced grads.
|
||||
/// 4. `grad_h_accumulate` to fold `grad_h_t_d` into the encoder grad.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn backward(
|
||||
&self,
|
||||
h_t: &CudaSlice<f32>,
|
||||
grad_w_per_batch_d: &mut CudaSlice<f32>,
|
||||
grad_b_per_batch_d: &mut CudaSlice<f32>,
|
||||
grad_h_t_d: &mut CudaSlice<f32>,
|
||||
actual_b: usize,
|
||||
) -> Result<()> {
|
||||
if actual_b > self.b_size {
|
||||
return Err(anyhow!(
|
||||
"OutcomeHead::backward: actual_b ({actual_b}) > allocated b_size ({})",
|
||||
self.b_size
|
||||
));
|
||||
}
|
||||
debug_assert_eq!(h_t.len(), actual_b * HIDDEN_DIM);
|
||||
debug_assert_eq!(grad_w_per_batch_d.len(), actual_b * HIDDEN_DIM * K_CLASSES);
|
||||
debug_assert_eq!(grad_b_per_batch_d.len(), actual_b * K_CLASSES);
|
||||
debug_assert_eq!(grad_h_t_d.len(), actual_b * HIDDEN_DIM);
|
||||
let b_i32 = actual_b as i32;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (actual_b as u32, 1, 1),
|
||||
block_dim: (HIDDEN_DIM as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_bwd)
|
||||
.arg(h_t)
|
||||
.arg(&self.grad_logits_d)
|
||||
.arg(&self.w_d)
|
||||
.arg(&b_i32)
|
||||
.arg(grad_w_per_batch_d)
|
||||
.arg(grad_b_per_batch_d)
|
||||
.arg(grad_h_t_d)
|
||||
.launch(cfg)
|
||||
.map_err(|e| anyhow!("rl_outcome_bwd launch: {e}"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Assign labels from reward/done signals.
|
||||
///
|
||||
/// `rewards` and `dones` must be device slices of length `actual_b`.
|
||||
|
||||
@@ -856,6 +856,8 @@ pub struct IntegratedTrainer {
|
||||
|
||||
// ── P1.3 Outcome auxiliary head ──────────────────────────────────
|
||||
pub outcome_head: crate::rl::outcome_head::OutcomeHead,
|
||||
pub outcome_w_adam: AdamW,
|
||||
pub outcome_b_adam: AdamW,
|
||||
|
||||
// ── P2.1 Q-bias correction ───────────────────────────────────────
|
||||
_rl_q_bias_correction_module: Arc<CudaModule>,
|
||||
@@ -1044,6 +1046,13 @@ pub struct IntegratedTrainer {
|
||||
// FRD encoder-upstream gradient.
|
||||
pub ss_frd_grad_h_t_d: CudaSlice<f32>,
|
||||
|
||||
// Outcome head backward scratch buffers.
|
||||
pub ss_outcome_grad_w_pb_d: CudaSlice<f32>, // [B, HIDDEN_DIM, K_CLASSES]
|
||||
pub ss_outcome_grad_b_pb_d: CudaSlice<f32>, // [B, K_CLASSES]
|
||||
pub ss_outcome_grad_h_t_d: CudaSlice<f32>, // [B, HIDDEN_DIM]
|
||||
pub ss_outcome_grad_w_d: CudaSlice<f32>, // [HIDDEN_DIM, K_CLASSES]
|
||||
pub ss_outcome_grad_b_d: CudaSlice<f32>, // [K_CLASSES]
|
||||
|
||||
// Bellman projection scratch.
|
||||
pub ss_q_target_full_d: CudaSlice<f32>,
|
||||
pub ss_q_target_action_d: CudaSlice<f32>,
|
||||
@@ -1928,6 +1937,10 @@ impl IntegratedTrainer {
|
||||
let outcome_head = crate::rl::outcome_head::OutcomeHead::new(
|
||||
b_size, stream.clone(),
|
||||
).context("OutcomeHead::new")?;
|
||||
let outcome_w_adam = AdamW::new(dev, outcome_head.w_d.len(), lr_placeholder)
|
||||
.context("outcome_w_adam")?;
|
||||
let outcome_b_adam = AdamW::new(dev, outcome_head.b_d.len(), lr_placeholder)
|
||||
.context("outcome_b_adam")?;
|
||||
// ── P2.1 Q-bias correction cubin load ────────────────────────
|
||||
let rl_q_bias_correction_module = ctx
|
||||
.load_cubin(RL_Q_BIAS_CORRECTION_CUBIN.to_vec())
|
||||
@@ -2095,6 +2108,24 @@ impl IntegratedTrainer {
|
||||
let ss_frd_grad_h_t_d = stream.alloc_zeros::<f32>(b_size * HIDDEN_DIM)
|
||||
.context("alloc ss_frd_grad_h_t_d")?;
|
||||
|
||||
// Outcome head backward scratch buffers.
|
||||
let k_classes = crate::rl::outcome_head::K_CLASSES;
|
||||
let ss_outcome_grad_w_pb_d = stream
|
||||
.alloc_zeros::<f32>(b_size * HIDDEN_DIM * k_classes)
|
||||
.context("alloc ss_outcome_grad_w_pb_d")?;
|
||||
let ss_outcome_grad_b_pb_d = stream
|
||||
.alloc_zeros::<f32>(b_size * k_classes)
|
||||
.context("alloc ss_outcome_grad_b_pb_d")?;
|
||||
let ss_outcome_grad_h_t_d = stream
|
||||
.alloc_zeros::<f32>(b_size * HIDDEN_DIM)
|
||||
.context("alloc ss_outcome_grad_h_t_d")?;
|
||||
let ss_outcome_grad_w_d = stream
|
||||
.alloc_zeros::<f32>(HIDDEN_DIM * k_classes)
|
||||
.context("alloc ss_outcome_grad_w_d")?;
|
||||
let ss_outcome_grad_b_d = stream
|
||||
.alloc_zeros::<f32>(k_classes)
|
||||
.context("alloc ss_outcome_grad_b_d")?;
|
||||
|
||||
// Bellman projection scratch.
|
||||
let ss_q_target_full_d = stream.alloc_zeros::<f32>(b_size * k_dqn_ss)
|
||||
.context("alloc ss_q_target_full_d")?;
|
||||
@@ -2402,6 +2433,8 @@ impl IntegratedTrainer {
|
||||
spectral_v_buffer_pi,
|
||||
// P1.3 Outcome
|
||||
outcome_head,
|
||||
outcome_w_adam,
|
||||
outcome_b_adam,
|
||||
// P2.1 Q-bias
|
||||
_rl_q_bias_correction_module: rl_q_bias_correction_module,
|
||||
rl_q_bias_correction_fn,
|
||||
@@ -2476,6 +2509,11 @@ impl IntegratedTrainer {
|
||||
ss_v_grad_w_d,
|
||||
ss_v_grad_b_d,
|
||||
ss_frd_grad_h_t_d,
|
||||
ss_outcome_grad_w_pb_d,
|
||||
ss_outcome_grad_b_pb_d,
|
||||
ss_outcome_grad_h_t_d,
|
||||
ss_outcome_grad_w_d,
|
||||
ss_outcome_grad_b_d,
|
||||
ss_q_target_full_d,
|
||||
ss_q_target_action_d,
|
||||
ss_target_dist_d,
|
||||
@@ -3719,6 +3757,9 @@ impl IntegratedTrainer {
|
||||
self.frd_b1_adam.lr = lr_frd;
|
||||
self.frd_w2_adam.lr = lr_frd;
|
||||
self.frd_b2_adam.lr = lr_frd;
|
||||
// Outcome head shares the FRD LR slot (both are lightweight aux heads).
|
||||
self.outcome_w_adam.lr = lr_frd;
|
||||
self.outcome_b_adam.lr = lr_frd;
|
||||
|
||||
// Borrow encoder hidden state for forward kernels.
|
||||
let h_t_borrow: &CudaSlice<f32> = self.perception.h_t_view();
|
||||
@@ -3813,6 +3854,18 @@ impl IntegratedTrainer {
|
||||
self.stream.memset_zeros(&mut self.ss_frd_grad_h_t_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_frd_grad_h_t_d: {e}"))?;
|
||||
|
||||
// Outcome head backward scratch.
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_w_pb_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_pb_d: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_b_pb_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_pb_d: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_h_t_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_h_t_d: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_w_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_d: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_b_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_d: {e}"))?;
|
||||
|
||||
// Bellman projection scratch.
|
||||
self.stream.memset_zeros(&mut self.ss_q_target_full_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_q_target_full_d: {e}"))?;
|
||||
@@ -4110,15 +4163,48 @@ impl IntegratedTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Step 9c: Outcome head forward + CE loss ──────────────────
|
||||
// ── Step 9c: Outcome head forward + CE loss + backward ────────
|
||||
// K=3 outcome classifier trained from reward/done labels assigned
|
||||
// in step_with_lobsim. Forward and loss run on sampled_h_t_d.
|
||||
// Backward: grad_logits → grad_W/b (per-batch), grad_h_t.
|
||||
// reduce_axis0 → batch-summed grad_W/b. Adam on outcome W/b.
|
||||
// Lambda from ISV[RL_OUTCOME_AUX_LAMBDA_INDEX] scales the
|
||||
// encoder-upstream grad_h_t contribution in Step 10.
|
||||
{
|
||||
let h_t_borrow_for_outcome: &CudaSlice<f32> = &self.sampled_h_t_d;
|
||||
self.outcome_head.forward(h_t_borrow_for_outcome, b_size)
|
||||
.context("step_synthetic: outcome_head.forward")?;
|
||||
self.outcome_head.compute_loss(b_size)
|
||||
.context("step_synthetic: outcome_head.compute_loss")?;
|
||||
// Backward: grad_logits → per-batch grad_W, grad_b, grad_h_t.
|
||||
self.outcome_head.backward(
|
||||
h_t_borrow_for_outcome,
|
||||
&mut self.ss_outcome_grad_w_pb_d,
|
||||
&mut self.ss_outcome_grad_b_pb_d,
|
||||
&mut self.ss_outcome_grad_h_t_d,
|
||||
b_size,
|
||||
).context("step_synthetic: outcome_head.backward")?;
|
||||
// Reduce-axis-0: batch-summed weight and bias gradients.
|
||||
let k_classes = crate::rl::outcome_head::K_CLASSES;
|
||||
reduce_axis0_free(
|
||||
&self.stream, &self.reduce_axis0_fn,
|
||||
&self.ss_outcome_grad_w_pb_d, b_size,
|
||||
HIDDEN_DIM * k_classes,
|
||||
&mut self.ss_outcome_grad_w_d,
|
||||
)?;
|
||||
reduce_axis0_free(
|
||||
&self.stream, &self.reduce_axis0_fn,
|
||||
&self.ss_outcome_grad_b_pb_d, b_size,
|
||||
k_classes,
|
||||
&mut self.ss_outcome_grad_b_d,
|
||||
)?;
|
||||
// Adam step on outcome W and b.
|
||||
self.outcome_w_adam
|
||||
.step(&mut self.outcome_head.w_d, &self.ss_outcome_grad_w_d)
|
||||
.context("outcome_w_adam.step")?;
|
||||
self.outcome_b_adam
|
||||
.step(&mut self.outcome_head.b_d, &self.ss_outcome_grad_b_d)
|
||||
.context("outcome_b_adam.step")?;
|
||||
}
|
||||
|
||||
// ── SP20 P3 FRD backward chain (F.4) ─────────────────────────
|
||||
@@ -4209,9 +4295,10 @@ impl IntegratedTrainer {
|
||||
// shared-encoder pattern (Q is trained on PER-sampled past h_t
|
||||
// values, whose gradient should NOT flow into the encoder
|
||||
// — accumulating stale-state grad would poison the encoder).
|
||||
// Encoder learns from π + V (on-policy current-step) + BCE +
|
||||
// aux (supervised via perception.step_batched). Standard
|
||||
// off-policy AC pattern (SAC / R2D2 / IMPALA do the same).
|
||||
// Encoder learns from π + V (on-policy current-step) + FRD +
|
||||
// outcome (aux supervised) + BCE + aux (supervised via
|
||||
// perception.step_batched). Standard off-policy AC pattern
|
||||
// (SAC / R2D2 / IMPALA do the same for the Q stop-grad).
|
||||
//
|
||||
// The combined grad slot feeds `backward_encoder_with_grad_h_t`
|
||||
// below, which runs the shared encoder backward chain that
|
||||
@@ -4246,6 +4333,23 @@ impl IntegratedTrainer {
|
||||
lambdas.frd,
|
||||
&mut self.grad_h_t_combined_d,
|
||||
)?;
|
||||
// Outcome head encoder-upstream gradient — λ from
|
||||
// ISV[RL_OUTCOME_AUX_LAMBDA_INDEX]. With sentinel -1 labels
|
||||
// the CE grad is zero → grad_h_t is all-zeros → this
|
||||
// accumulate is a weighted zero-add (no signal). Once labels
|
||||
// are populated the outcome head injects its supervised
|
||||
// gradient into the encoder at the configured λ strength.
|
||||
{
|
||||
let outcome_lambda =
|
||||
self.isv_host[crate::rl::isv_slots::RL_OUTCOME_AUX_LAMBDA_INDEX];
|
||||
accumulate_grad_h(
|
||||
&self.stream,
|
||||
&self.grad_h_accumulate_fn,
|
||||
&self.ss_outcome_grad_h_t_d,
|
||||
outcome_lambda,
|
||||
&mut self.grad_h_t_combined_d,
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── Step 11: encoder backward (Phase E.3a) ───────────────────
|
||||
// Seed the perception trainer's per-K hidden-state grad buffer
|
||||
@@ -4697,6 +4801,17 @@ impl IntegratedTrainer {
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_q_grad_w_d: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_q_grad_b_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_q_grad_b_d: {e}"))?;
|
||||
// Outcome head backward scratch (replay steps also train the head).
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_w_pb_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_pb_d (replay): {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_b_pb_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_pb_d (replay): {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_h_t_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_h_t_d (replay): {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_w_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_w_d (replay): {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.ss_outcome_grad_b_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero ss_outcome_grad_b_d (replay): {e}"))?;
|
||||
|
||||
let b_size_i = b_size as i32;
|
||||
|
||||
@@ -5069,25 +5184,37 @@ impl IntegratedTrainer {
|
||||
// CE loss + grad w.r.t. logits (masked by sentinel -1 labels).
|
||||
self.outcome_head.compute_loss(b_size)
|
||||
.context("dqn_replay_step: outcome_head.compute_loss")?;
|
||||
|
||||
// Backward: grad_logits → grad_w, grad_b, grad_h_t.
|
||||
// Manual outer-product per batch: grad_w[b] = h_t[b]^T × grad_logits[b].
|
||||
// For now, we use the reduce_axis0 pattern on per-batch outer products.
|
||||
// The outcome head is small (128×3), so we compute gradients inline.
|
||||
// grad_w = sum_b(h_t[b] outer grad_logits[b]) is done by the same
|
||||
// backward_to_w_b_h pattern as other heads.
|
||||
// We compute grad_h_t[b] = grad_logits[b] × W^T in-place.
|
||||
// Since OutcomeHead doesn't have a backward_to_w_b_h method, we
|
||||
// skip the gradient for now — the outcome head's gradient flows
|
||||
// only through its own Adam, not to the encoder. The CE loss
|
||||
// provides the training signal. Encoder gradient is additive from
|
||||
// Q+pi+V heads which already flow through grad_h_t_combined_d.
|
||||
//
|
||||
// Adam step on outcome W and b with reduced gradients.
|
||||
// The outcome head is lightweight (128×3 + 3 params) — Adam on
|
||||
// the per-batch mean gradient is sufficient.
|
||||
// For the outcome head we skip full per-batch gradient reduce and
|
||||
// just step on the CE grad directly through the logits.
|
||||
// Backward: grad_logits → per-batch grad_W, grad_b, grad_h_t.
|
||||
// grad_h_t is computed but discarded (no encoder backward in
|
||||
// replay steps — same stop-grad rationale as Q's grad_h_t).
|
||||
self.outcome_head.backward(
|
||||
&self.sampled_h_t_d,
|
||||
&mut self.ss_outcome_grad_w_pb_d,
|
||||
&mut self.ss_outcome_grad_b_pb_d,
|
||||
&mut self.ss_outcome_grad_h_t_d,
|
||||
b_size,
|
||||
).context("dqn_replay_step: outcome_head.backward")?;
|
||||
{
|
||||
let k_classes = crate::rl::outcome_head::K_CLASSES;
|
||||
reduce_axis0_free(
|
||||
&self.stream, &self.reduce_axis0_fn,
|
||||
&self.ss_outcome_grad_w_pb_d, b_size,
|
||||
HIDDEN_DIM * k_classes,
|
||||
&mut self.ss_outcome_grad_w_d,
|
||||
)?;
|
||||
reduce_axis0_free(
|
||||
&self.stream, &self.reduce_axis0_fn,
|
||||
&self.ss_outcome_grad_b_pb_d, b_size,
|
||||
k_classes,
|
||||
&mut self.ss_outcome_grad_b_d,
|
||||
)?;
|
||||
self.outcome_w_adam
|
||||
.step(&mut self.outcome_head.w_d, &self.ss_outcome_grad_w_d)
|
||||
.context("dqn_replay_step: outcome_w_adam.step")?;
|
||||
self.outcome_b_adam
|
||||
.step(&mut self.outcome_head.b_d, &self.ss_outcome_grad_b_d)
|
||||
.context("dqn_replay_step: outcome_b_adam.step")?;
|
||||
}
|
||||
|
||||
// ── 9. Per-branch LR controller ──────────────────────────────
|
||||
// Single-thread kernel adjusts LR scale factors from loss deltas.
|
||||
|
||||
Reference in New Issue
Block a user