feat(rl): signal-driven LR controller from per-head grad-norm EMAs
The rl_lr_controller emitted a hardcoded `LR_BOOTSTRAP = 1e-3` for
every step regardless of training dynamics. The kernel accepted 5
`*_signal` scalar args but ignored them via `(void)signal;` — a
stub. This commit makes the LR genuinely signal-driven per
`pearl_controller_anchors_isv_driven` + `feedback_isv_for_adaptive_bounds`.
## Architecture
Per-head grad-norm EMA → LR target derivation:
observed_grad_norm = EMA(‖grad_w_head‖₂)
target_lr = lr_prev × (TARGET_GRAD_NORM / max(observed, ε))
Wiener-α blend (floor 0.4) + clamp to [LR_MIN, LR_MAX].
Multiplicative pattern — same shape as rl_target_tau / rl_ppo_clip
controllers. High observed gradient (model thrashing) shrinks LR
(calm updates); low observed gradient (model coasting) grows LR
(push more aggressive learning).
## Components
1. **rl_l2_norm.cu** (new) — single-buffer L2 norm `‖x‖₂` via
grid-stride loop + shared-mem tree reduce. Used for per-head
grad_w_*_d reductions.
2. **rl_lr_controller.cu** (rewrite) — kernel signature changes
from 5 scalar `*_signal` args to 5 `int *_signal_slot` args
(ISV slot indices). The kernel reads each signal from
`isv[slot]`, derives target multiplicatively, and applies the
cold-start gate + replace-directly pattern (same R9-audit fixes
that closed the dead-zones in the other multiplicative
controllers). BCE and AUX heads pass sentinel `-1` for their
signal slot (those heads are owned by the perception trainer);
the kernel falls back to LR_BOOTSTRAP for those.
3. **ISV slot extension** (`isv_slots.rs`):
* `RL_Q_GRAD_NORM_EMA_INDEX = 424`
* `RL_PI_GRAD_NORM_EMA_INDEX = 425`
* `RL_V_GRAD_NORM_EMA_INDEX = 426`
* `RL_SLOTS_END = 427` (was 424).
4. **Trainer wiring** (`integrated.rs`):
* New `rl_l2_norm` module + fn fields + load in `new()`.
* New `launch_l2_norm` helper (256-thread single-block reduce).
* After-encoder-backward block in `step_synthetic` gains 3
grad-norm + EMA launches (Q grad_w 24,192 floats, π grad_w
1,152 floats, V grad_w 128 floats) alongside the existing
entropy / td_kurtosis / kl_pi EMAs.
* `launch_rl_lr_controller` updated to pass i32 slot indices
instead of f32 scalars.
## What's NOT in this commit
* BCE and AUX LR signals — those heads' gradients live in the
perception trainer, not the RL trainer. A future commit can
wire `perception.bce_grad_w_d` → ISV slot if the BCE/AUX LRs
need to adapt for cross-trainer alignment.
* Production tuning of `TARGET_GRAD_NORM = 1.0`. Empirical from
the 50k smoke (Q grad_w L2 norm landed near 1 at LR=1e-3); the
smoke at this commit will confirm whether the LR controller
drives the grad-norm to this anchor.
## Verified gates (local sm_86)
G1 isv_bootstrap ✅ (per_α, γ, etc. — unchanged)
G3 controllers_emit ✅ (test pre-seeds inputs)
G4 target_soft_update ✅
G6 r7d_per_wiring ✅
smoke ✅ all losses finite
## Expected effect
Prior 50k run showed l_q oscillating in 2.7-4.4 range without
visible convergence at LR=1e-3 constant. With LR now adaptive,
the trainer should:
* Shrink Q LR when Q grad-norm spikes (large per-sample CE
after a big trade close).
* Grow Q LR when grad-norm stays small (steady-state coasting).
* Same logic for π and V.
Next cluster smoke at 50k steps will produce a diag.jsonl where
ISV[413..415] (lr_q, lr_pi, lr_v) AND ISV[424..426] (grad-norm
EMAs) both evolve over time — observable convergence dynamics
that previously didn't exist.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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.)
|
||||
];
|
||||
|
||||
|
||||
40
crates/ml-alpha/cuda/rl_l2_norm.cu
Normal file
40
crates/ml-alpha/cuda/rl_l2_norm.cu
Normal file
@@ -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]);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<CudaModule>,
|
||||
rl_step_counter_update_fn: CudaFunction,
|
||||
_rl_l2_norm_module: Arc<CudaModule>,
|
||||
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<i32>,
|
||||
@@ -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<f32>,
|
||||
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::<f32>()) 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")?;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user