diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index c1d38f134..04a40d4af 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -60,6 +60,7 @@ const KERNELS: &[&str] = &[ "rl_kl_approx_b", // Schulman-style KL = mean(log π_old − log π_new) → kl_pi_ema (ISV[419]) feeding rl_ppo_clip "rl_l2_diff_norm", // ‖W_online − W_target‖₂ → q_divergence_ema (ISV[418]) feeding rl_target_tau "rl_step_counter_update", // per-batch trade-duration counter + done-gated emit → mean_trade_duration_ema (ISV[417]) feeding rl_gamma + "rl_l2_norm", // ‖x‖₂ single-buffer reduction → q/pi/v grad-norm EMAs (ISV[424..427]) feeding rl_lr_controller // (entropy_observed_ema, ISV[420], feeds rl_entropy_coef directly via ema_update_per_step's internal mean reduce on entropy_d — no separate kernel needed.) ]; diff --git a/crates/ml-alpha/cuda/rl_l2_norm.cu b/crates/ml-alpha/cuda/rl_l2_norm.cu new file mode 100644 index 000000000..8540eccfc --- /dev/null +++ b/crates/ml-alpha/cuda/rl_l2_norm.cu @@ -0,0 +1,40 @@ +// rl_l2_norm.cu — `‖x‖₂ = sqrt(Σ x²)` over n floats → scalar. +// +// Single-buffer L2 norm reduction. Used by the LR controller's input +// path to compute per-head gradient norms (Q, π, V) for the +// signal-modulated learning-rate target. +// +// Same shape as rl_l2_diff_norm (single block, grid-stride loop, tree +// reduce, no atomics) but operating on one buffer rather than two. +// The trainer allocates one launch per head (Q grad_w 24k floats, π +// grad_w 1k floats, V grad_w 128 floats) and feeds the scalar output +// to ema_update_per_step at the corresponding ISV slot +// (RL_Q_GRAD_NORM_EMA_INDEX = 424, π = 425, V = 426). + +extern "C" __global__ void rl_l2_norm( + const float* __restrict__ x, + float* __restrict__ out, + int n +) { + extern __shared__ float s[]; + const int tid = threadIdx.x; + + float acc = 0.0f; + for (int i = tid; i < n; i += blockDim.x) { + const float v = x[i]; + acc += v * v; + } + s[tid] = acc; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + s[tid] += s[tid + stride]; + } + __syncthreads(); + } + + if (tid == 0) { + out[0] = sqrtf(s[0]); + } +} diff --git a/crates/ml-alpha/cuda/rl_lr_controller.cu b/crates/ml-alpha/cuda/rl_lr_controller.cu index 466ffd159..b3b326c4c 100644 --- a/crates/ml-alpha/cuda/rl_lr_controller.cu +++ b/crates/ml-alpha/cuda/rl_lr_controller.cu @@ -1,28 +1,45 @@ -// rl_lr_controller.cu — emits per-head learning rates to ISV[412..417]. -// -// Phase E.2-DEFER scope (2026-05-22): controller emits constant 1e-3 to -// all 5 slots on first observation (sentinel-zero → bootstrap value); -// subsequent emits Wiener-α blend with the same bootstrap target. A -// future enhancement (Phase E.3+) lifts the target from constant -// `LR_BOOTSTRAP` to a signal-modulated function of per-head gradient-norm -// EMA divergence — the controller already accepts the diagnostic inputs -// (`*_signal` args) so the upgrade is a pure kernel body change with no -// caller-side refactor. +// rl_lr_controller.cu — emits per-head learning rates to ISV[412..417] +// with signal-modulated targets (per-head gradient-norm EMA divergence +// from a target magnitude). // // Per `pearl_controller_anchors_isv_driven` + `feedback_isv_for_adaptive_bounds`: -// every learning rate that was previously a Rust constant -// (PHASE_E2_DEFAULT_LR) now lives in ISV. The IntegratedTrainer DtoH -// refreshes ISV each step before composing Adam args. +// every learning rate is ISV-resident and driven by an observable +// signal (grad-norm EMA) rather than a hardcoded constant. The +// trainer: +// 1. Computes per-head L2 grad-norm (rl_l2_norm) after each backward. +// 2. Feeds it through ema_update_per_step into the relevant grad-norm +// EMA slot (RL_Q_GRAD_NORM_EMA_INDEX, RL_PI_GRAD_NORM_EMA_INDEX, +// RL_V_GRAD_NORM_EMA_INDEX). +// 3. Launches this kernel passing those slot indices. +// This kernel reads each EMA and derives a per-head LR target: +// +// target_lr = lr_prev × (TARGET_GRAD_NORM / max(observed_grad_norm, ε)) +// +// Multiplicative pattern (same shape as rl_target_tau / rl_ppo_clip): +// observed > target → ratio < 1 → shrink lr (calm the gradients); +// observed < target → ratio > 1 → grow lr (push more aggressive +// updates). Clamped to [LR_MIN, LR_MAX]. // // Per `pearl_first_observation_bootstrap`: sentinel zero at the slot's // address means "uninitialised — emit the bootstrap value on the next -// controller fire." The IntegratedTrainer zero-allocates the ISV buffer -// at construction. +// controller fire." Cold-start gate: when the input EMA is sentinel +// zero (no backward has fired yet), the controller holds prev at +// bootstrap — same anti-pattern fix applied to τ/ε/n_roll/scale. +// Replace-directly on first warm observation: when prev is exactly +// the hardcoded bootstrap value, write target directly instead of +// Wiener-blending (60% bootstrap contamination would otherwise +// distort the first emit). // // Per `pearl_wiener_alpha_floor_for_nonstationary`: the α blend is // floored at 0.4 since the target drives a co-adapting closed loop // (parameter LR feeds back into gradient magnitude via Adam). // +// BCE and AUX heads are emitted at their hardcoded bootstrap. Passing +// a `signal_slot < 0` for these heads skips the signal-driven path +// and falls back to `target_lr = LR_BOOTSTRAP` — the perception +// trainer owns those heads, not the RL trainer, so signal wiring is +// out of scope here. +// // Block layout: // grid = (1, 1, 1) // block = (1, 1, 1) @@ -34,48 +51,77 @@ #define RL_LR_V_INDEX 415 #define RL_LR_AUX_INDEX 416 -#define LR_BOOTSTRAP 1e-3f -#define LR_MIN 1e-5f -#define LR_MAX 1e-2f +#define LR_BOOTSTRAP 1e-3f +#define LR_MIN 1e-5f +#define LR_MAX 1e-2f +// Target grad-norm magnitude. Chosen so a typical post-backward +// grad-norm at LR=LR_BOOTSTRAP lands near this value; observed > +// target → controller shrinks LR; observed < target → grows LR. +// Magnitude derived empirically from the 50k smoke (Q grad_w +// L2 norms typically O(1) at LR=1e-3). +#define TARGET_GRAD_NORM 1.0f +#define EMA_EPS 1e-6f -__device__ __forceinline__ void update_lr( - float* isv, int idx, float alpha, float lr_target +__device__ __forceinline__ void update_lr_with_signal( + float* isv, int lr_idx, float alpha, int signal_slot ) { - const float lr_prev = isv[idx]; + const float lr_prev = isv[lr_idx]; + + // Bootstrap on sentinel 0.0 per pearl_first_observation_bootstrap. if (lr_prev == 0.0f) { - // First-observation bootstrap per pearl_first_observation_bootstrap. - isv[idx] = LR_BOOTSTRAP; + isv[lr_idx] = LR_BOOTSTRAP; return; } + + // No signal wiring for this head (BCE / AUX caller-pass -1) — + // hold at bootstrap, Wiener-blend toward LR_BOOTSTRAP (no-op since + // prev is already there in steady state, but the blend gives the + // controller authority to override an externally-edited slot). + if (signal_slot < 0) { + const float a = fmaxf(alpha, 0.4f); + isv[lr_idx] = (1.0f - a) * lr_prev + a * LR_BOOTSTRAP; + return; + } + + const float observed = isv[signal_slot]; + + // Cold-start gate: no grad-norm signal observed yet → hold at + // bootstrap. Same fix that closed the cold-start migration in + // τ/ε/n_roll/scale. + if (observed == 0.0f) return; + + // Multiplicative target derivation. + float target_lr = lr_prev * (TARGET_GRAD_NORM / fmaxf(observed, EMA_EPS)); + target_lr = fminf(LR_MAX, fmaxf(LR_MIN, target_lr)); + + // First-observation replace-directly: avoid Wiener-blend's 60% + // bootstrap contamination on the first warm step. + if (lr_prev == LR_BOOTSTRAP) { + isv[lr_idx] = target_lr; + return; + } + // Wiener-α floor 0.4 per pearl_wiener_alpha_floor_for_nonstationary. const float a = fmaxf(alpha, 0.4f); - const float lr_target_clamped = fminf(LR_MAX, fmaxf(LR_MIN, lr_target)); - isv[idx] = (1.0f - a) * lr_prev + a * lr_target_clamped; + float lr_new = (1.0f - a) * lr_prev + a * target_lr; + lr_new = fminf(LR_MAX, fmaxf(LR_MIN, lr_new)); + isv[lr_idx] = lr_new; } extern "C" __global__ void rl_lr_controller( float* __restrict__ isv, float alpha, - float bce_signal, // diagnostic input (e.g. per-head grad-norm - float q_signal, // EMA) — reserved for Phase E.3+'s - float pi_signal, // signal-modulated target. Currently - float v_signal, // unused; the controller emits the - float aux_signal // constant bootstrap. See header. + int bce_signal_slot, // <0 → no signal (BCE owned by perception) + int q_signal_slot, // RL_Q_GRAD_NORM_EMA_INDEX + int pi_signal_slot, // RL_PI_GRAD_NORM_EMA_INDEX + int v_signal_slot, // RL_V_GRAD_NORM_EMA_INDEX + int aux_signal_slot // <0 → no signal (AUX owned by perception) ) { if (threadIdx.x != 0 || blockIdx.x != 0) return; - // Silence unused-arg warnings until Phase E.3 lifts these into the - // target derivation. Reads of the same address are free. - (void)bce_signal; - (void)q_signal; - (void)pi_signal; - (void)v_signal; - (void)aux_signal; - - update_lr(isv, RL_LR_BCE_INDEX, alpha, LR_BOOTSTRAP); - update_lr(isv, RL_LR_Q_INDEX, alpha, LR_BOOTSTRAP); - update_lr(isv, RL_LR_PI_INDEX, alpha, LR_BOOTSTRAP); - update_lr(isv, RL_LR_V_INDEX, alpha, LR_BOOTSTRAP); - update_lr(isv, RL_LR_AUX_INDEX, alpha, LR_BOOTSTRAP); - // Phase E.3+: replace LR_BOOTSTRAP with `derive_target(*_signal)`. + update_lr_with_signal(isv, RL_LR_BCE_INDEX, alpha, bce_signal_slot); + update_lr_with_signal(isv, RL_LR_Q_INDEX, alpha, q_signal_slot); + update_lr_with_signal(isv, RL_LR_PI_INDEX, alpha, pi_signal_slot); + update_lr_with_signal(isv, RL_LR_V_INDEX, alpha, v_signal_slot); + update_lr_with_signal(isv, RL_LR_AUX_INDEX, alpha, aux_signal_slot); } diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index 3b6beaea7..f1a5bc3ff 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -133,6 +133,19 @@ pub const RL_TD_KURTOSIS_EMA_INDEX: usize = 422; /// `rl_reward_scale_controller`. pub const RL_MEAN_ABS_PNL_EMA_INDEX: usize = 423; +/// EMA of L2 norm of Q-head weight gradient `‖grad_w_q‖₂`. Input to +/// `rl_lr_controller`'s Q-LR target derivation (multiplicative: +/// observed > target → shrink LR; observed < target → grow LR). +pub const RL_Q_GRAD_NORM_EMA_INDEX: usize = 424; + +/// EMA of L2 norm of policy-head weight gradient `‖grad_w_pi‖₂`. Input +/// to `rl_lr_controller`'s π-LR target derivation. +pub const RL_PI_GRAD_NORM_EMA_INDEX: usize = 425; + +/// EMA of L2 norm of value-head weight gradient `‖grad_w_v‖₂`. Input +/// to `rl_lr_controller`'s V-LR target derivation. +pub const RL_V_GRAD_NORM_EMA_INDEX: usize = 426; + /// Last RL-allocated slot index (exclusive). The integrated trainer /// extends `ISV_TOTAL_DIM` to at least this value at trainer init time. -pub const RL_SLOTS_END: usize = 424; +pub const RL_SLOTS_END: usize = 427; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 1043884da..c59211fda 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -185,6 +185,8 @@ const RL_L2_DIFF_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_l2_diff_norm.cubin")); const RL_STEP_COUNTER_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_step_counter_update.cubin")); +const RL_L2_NORM_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_l2_norm.cubin")); /// Per-head LR controller Wiener-α target. The controller kernel floors /// this at 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` so the @@ -384,6 +386,8 @@ pub struct IntegratedTrainer { rl_l2_diff_norm_fn: CudaFunction, _rl_step_counter_update_module: Arc, rl_step_counter_update_fn: CudaFunction, + _rl_l2_norm_module: Arc, + rl_l2_norm_fn: CudaFunction, /// Per-batch step counter for `trade_duration_ema` — increments /// every step, resets on done. Zero-init at construction. pub steps_since_done_d: CudaSlice, @@ -759,6 +763,12 @@ impl IntegratedTrainer { let rl_step_counter_update_fn = rl_step_counter_update_module .load_function("rl_step_counter_update") .context("load rl_step_counter_update")?; + let rl_l2_norm_module = ctx + .load_cubin(RL_L2_NORM_CUBIN.to_vec()) + .context("load rl_l2_norm cubin")?; + let rl_l2_norm_fn = rl_l2_norm_module + .load_function("rl_l2_norm") + .context("load rl_l2_norm")?; // Per-batch PRNG state for the Thompson sampler. Seeded // deterministically from cfg.dqn_seed via ChaCha8 host RNG so @@ -978,6 +988,8 @@ impl IntegratedTrainer { rl_l2_diff_norm_fn, _rl_step_counter_update_module: rl_step_counter_update_module, rl_step_counter_update_fn, + _rl_l2_norm_module: rl_l2_norm_module, + rl_l2_norm_fn, steps_since_done_d, trade_duration_emit_d, prev_realized_pnl_d, @@ -1955,6 +1967,45 @@ impl IntegratedTrainer { ) .context("ema_update_per_step(kl_pi)")?; + // Per-head gradient-norm EMAs (Q, π, V) — feed the + // signal-modulated LR controller (rl_lr_controller's target + // formula: `lr_prev × TARGET_GRAD_NORM / observed_grad_norm`). + // grad_w_*_d buffers are still in scope here (allocated at + // top of step_synthetic, consumed by the Adam steps in Step 9 + // but not freed until function return); the l2_norm launches + // here read them after the Adam step has already used them. + let k_dqn_hidden = (N_ACTIONS * Q_N_ATOMS * HIDDEN_DIM) as usize; + self.launch_l2_norm(&q_grad_w_d.clone(), k_dqn_hidden) + .context("launch_l2_norm(q_grad_w_d)")?; + self.launch_ema_update_per_step( + crate::rl::isv_slots::RL_Q_GRAD_NORM_EMA_INDEX, + RL_LR_CONTROLLER_ALPHA, + &self.ema_input_scratch_d.clone(), + 1, + ) + .context("ema_update_per_step(q_grad_norm)")?; + + let n_actions_hidden = (N_ACTIONS * HIDDEN_DIM) as usize; + self.launch_l2_norm(&pi_grad_w_d.clone(), n_actions_hidden) + .context("launch_l2_norm(pi_grad_w_d)")?; + self.launch_ema_update_per_step( + crate::rl::isv_slots::RL_PI_GRAD_NORM_EMA_INDEX, + RL_LR_CONTROLLER_ALPHA, + &self.ema_input_scratch_d.clone(), + 1, + ) + .context("ema_update_per_step(pi_grad_norm)")?; + + self.launch_l2_norm(&v_grad_w_d.clone(), HIDDEN_DIM) + .context("launch_l2_norm(v_grad_w_d)")?; + self.launch_ema_update_per_step( + crate::rl::isv_slots::RL_V_GRAD_NORM_EMA_INDEX, + RL_LR_CONTROLLER_ALPHA, + &self.ema_input_scratch_d.clone(), + 1, + ) + .context("ema_update_per_step(v_grad_norm)")?; + // ── Step 12: compose stats ─────────────────────────────────── // BCE / aux losses are NOT read this phase — perception is driven // separately by callers via its existing step_batched path. The @@ -2798,6 +2849,39 @@ impl IntegratedTrainer { Ok(()) } + /// Reduce `‖x‖₂ = sqrt(Σ x²)` over `input_d[n]` into + /// `self.ema_input_scratch_d[1]`. Single block, 256 threads, + /// grid-stride loop covers arbitrary n via per-thread partial + /// accumulation + shared-mem tree-reduce. Used by the LR + /// controller's grad-norm input path (per-head grad_w_d L2 norm + /// fed to ema_update_per_step → ISV[424..426]). + fn launch_l2_norm( + &mut self, + input_d: &CudaSlice, + n: usize, + ) -> Result<()> { + debug_assert!(input_d.len() >= n); + if n == 0 { + return Ok(()); + } + let block_dim: u32 = 256; + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: (block_dim as usize * std::mem::size_of::()) as u32, + }; + let n_i = n as i32; + let mut launch = self.stream.launch_builder(&self.rl_l2_norm_fn); + launch + .arg(input_d) + .arg(&mut self.ema_input_scratch_d) + .arg(&n_i); + unsafe { + launch.launch(cfg).context("rl_l2_norm launch")?; + } + Ok(()) + } + /// Reduce kurtosis `E[(x-μ)⁴]/σ⁴` over `input_d[b_size]` into /// `self.ema_input_scratch_d[1]`. Returns 0 at b_size < 2 /// (kurtosis undefined; the per_α controller's cold-start gate @@ -2830,10 +2914,12 @@ impl IntegratedTrainer { } /// Launch `rl_lr_controller` to emit per-head learning rates into - /// `ISV[412..417]`. Phase E.2-DEFER item 3. Diagnostic-input args - /// are currently zero (the controller emits the bootstrap target - /// regardless); Phase E.3+ wires gradient-norm EMAs into the - /// `*_signal` args to drive a signal-modulated target. + /// `ISV[412..417]`. Reads per-head grad-norm EMAs from + /// `ISV[RL_Q_GRAD_NORM_EMA_INDEX]` etc. and derives multiplicative + /// targets `lr_prev × (TARGET_GRAD_NORM / observed)`. BCE and AUX + /// heads pass sentinel `-1` for their signal slots (those heads + /// are owned by the perception trainer, not the RL trainer; their + /// LRs hold at LR_BOOTSTRAP). fn launch_rl_lr_controller(&self) -> Result<()> { let cfg = LaunchConfig { grid_dim: (1, 1, 1), @@ -2841,16 +2927,23 @@ impl IntegratedTrainer { shared_mem_bytes: 0, }; let alpha = RL_LR_CONTROLLER_ALPHA; - let zero = 0.0_f32; + let bce_signal_slot: i32 = -1; // owned by perception + let q_signal_slot: i32 = + crate::rl::isv_slots::RL_Q_GRAD_NORM_EMA_INDEX as i32; + let pi_signal_slot: i32 = + crate::rl::isv_slots::RL_PI_GRAD_NORM_EMA_INDEX as i32; + let v_signal_slot: i32 = + crate::rl::isv_slots::RL_V_GRAD_NORM_EMA_INDEX as i32; + let aux_signal_slot: i32 = -1; // owned by perception let mut launch = self.stream.launch_builder(&self.rl_lr_controller_fn); launch .arg(&self.isv_d) .arg(&alpha) - .arg(&zero) // bce_signal — reserved for Phase E.3+ - .arg(&zero) // q_signal - .arg(&zero) // pi_signal - .arg(&zero) // v_signal - .arg(&zero); // aux_signal + .arg(&bce_signal_slot) + .arg(&q_signal_slot) + .arg(&pi_signal_slot) + .arg(&v_signal_slot) + .arg(&aux_signal_slot); unsafe { launch.launch(cfg).context("rl_lr_controller launch")?; }