From 6cb2257163e44c5f5f35d36616da9d7d6aedd4b7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 23 Apr 2026 09:16:11 +0200 Subject: [PATCH] fix(dqn): adaptive per-branch gradient-norm balancer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Caps any branch's weight-gradient L2 norm at num_branches × median (median across the 4 branches' norms). Scales the offending branch's gradient down to the cap; healthy branches pass through unchanged. Fixes the observed pathology in L40S train-mdh86: grad_ratio_mag_dir was 15k–26k× for 6 consecutive epochs, then collapsed to ~100× in a single step at epoch 7 and destabilised learning (Sharpe flipped +34 → -67, never recovered). Symmetric per-branch capping at `num_branches × median` prevents the swing at both ends without requiring a global ratio bound. No tuned knobs: `num_branches = 4` is architectural (factored action space: direction × magnitude × order × urgency), `median_branch_norm` is a per-step statistical reference that tracks the current gradient regime, and the product is fully adaptive. Per feedback_adaptive_not_tuned.md, the only static value is the architectural axis count; medians and derived caps are signal-driven. Implementation — two CUDA kernel launches in `branch_grad_balance_kernel.cu`: branch_grad_norm_reduce: grid=(4,1,1), block=(256,1,1). One block per branch; sum-of-squares via shared-mem tree reduce writes `branch_norms_dev[4]`. No atomicAdd (one-block-per-branch, single writer per slot). branch_grad_rescale: grid=(max_blocks, 4, 1), block=(256,1,1). Each block caches the 4 branch norms into shared memory, computes the median via a 5-comparator sorting network + two-element average (branch-deterministic, no reduction primitive), derives the 4 per-branch scales `scale[d] = min(1, 4×median/norm[d])`, then threads multiply their slice element by the owning branch's scale. No atomicAdd (each thread writes one distinct element). Insertion point: inside the `adam_grad_child` graph between the aux phase and `compute_grad_norm_for_adam`, so Adam's global clip and the Adam update both observe the rebalanced gradient. Also wired into the ungraphed fallback paths so no code path can skip the cap. The kernels have fixed launch configs, no host syncs, no dynamic allocations — safe to capture. Per-branch slice metadata (starts/lens for each of the 4 contiguous 4-tensor branch slices in `grad_buf`) is precomputed from `compute_param_sizes` at trainer construction and uploaded once to device i32 buffers, matching the existing `grad_decomp_kernel` layout convention. Smoke tests (local RTX 3050 Ti, 4 GB): magnitude_distribution: PASS (MAG_DIST Q=0.637 H=0.114 F=0.249, EVAL_DIST Q=0.153 H=0.255 F=0.592) multi_fold_convergence: PASS (3/3 folds produce best-checkpoint; fold Sharpes +57.8 / +55.6 / +119.3) grad_ratio_mag_dir trajectory (mag_dist smoke, first fold, first 5 epochs) — pre-fix values from /tmp/l40s_diag/health.log (L40S train-mdh86): pre-fix: 14793, 11934, 12858, 16406, 15141 (×1000 regime) post-fix: 55, 78, 35, 11, 10 (×10-100 regime) Three+ orders of magnitude reduction. The residual ratio can still exceed `num_branches = 4` when the direction branch's norm sits below the median — the cap bounds each branch's absolute norm (≤ 4×median), not the pairwise ratio, by design (direction-outlier smallness is a separate pathology that would be masked by a ratio bound). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 1 + .../branch_grad_balance_kernel.cu | 177 ++++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 254 ++++++++++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 16 ++ 4 files changed, 448 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 8c0637d3e..a4d91764b 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -82,6 +82,7 @@ fn main() { "mamba2_temporal_kernel.cu", "graph_utility_kernels.cu", "grad_decomp_kernel.cu", + "branch_grad_balance_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu b/crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu new file mode 100644 index 000000000..13638792a --- /dev/null +++ b/crates/ml/src/cuda_pipeline/branch_grad_balance_kernel.cu @@ -0,0 +1,177 @@ +/** + * branch_grad_balance_kernel — adaptive per-branch gradient-norm balancer. + * + * Caps each of the 4 factored-action branches' gradient L2 norm at + * `num_branches × median_branch_norm` + * (median across the 4 branches' norms at the current step). Branches + * within the cap pass through unchanged. + * + * Why this shape (documented architecturally, NOT tuned): + * * `num_branches` (= 4) is architectural — the count of factored action + * axes (direction, magnitude, order, urgency). It is NOT a tuning + * knob; changing the action factorisation changes this number. + * * `median_branch_norm` is a per-step statistical reference that scales + * naturally with the training regime (Q-magnitude, loss blend, etc.). + * There is no fixed target scale, no schedule, and no hyperparameter. + * * Their product enforces the defensible bound: "no single branch can + * dominate the gradient by more than `num_branches` × the median." + * Geometrically, a branch at the cap contributes as much L2 mass as + * all other branches combined evaluated at the median. + * * Scale-only-when-needed preserves branches that are proportional; + * the mechanism never forces artificial equality. + * + * Motivation (L40S train-mdh86, epochs 0-6): + * HEALTH_DIAG reported `grad_ratio_mag_dir` ≈ 1.5e4 – 2.6e4×. Magnitude + * gradients were three-to-four orders of magnitude larger than + * direction for the entire warm-up, which destabilised the dueling + + * factored-advantage decomposition. On a single step at epoch 7 the + * ratio collapsed to ~100× and Sharpe went from +34 to -67 without + * recovering. Per-branch capping at `num_branches × median` is a + * symmetric, signal-driven bound that prevents either end of the + * swing. + * + * Pipeline (one kernel, two launches): + * branch_grad_norm_reduce: grid=(4,1,1), block=(256,1,1). One block + * per branch; reduces the branch's + * sum-of-squares via shared-mem tree and + * writes sqrt to branch_norms_dev[branch]. + * No atomicAdd (one-block-per-branch). + * + * branch_grad_rescale: grid=(max_blocks, 4, 1), block=(256,1,1). + * blockIdx.y = branch id; blockIdx.x = 256- + * element sub-chunk. First 4 threads cache + * the 4 branch norms into shared memory, + * compute the median + per-branch scales, + * then each thread multiplies one element + * of its branch's slice by its branch scale. + * No atomicAdd (each thread writes a + * distinct element). + * + * CUDA Graph safety: + * All launches are plain kernel launches with fixed grid/block config. + * No dynamic allocations, no host syncs, no atomicAdd. Safe to capture + * inside `adam_grad_child` between the post_aux phase and + * `compute_grad_norm_for_adam`. + */ + +extern "C" __global__ void branch_grad_norm_reduce( + const float* __restrict__ grad_buf, + const int* __restrict__ branch_starts, /* [4] element offsets into grad_buf */ + const int* __restrict__ branch_lens, /* [4] element counts (padded, zero fill) */ + float* __restrict__ branch_norms_dev /* [4] output L2 norms */ +) { + const int branch = blockIdx.x; /* 0..3 */ + if (branch >= 4) return; + const int tid = threadIdx.x; + const int start = branch_starts[branch]; + const int len = branch_lens[branch]; + + __shared__ float sh[256]; + float acc = 0.0f; + for (int i = tid; i < len; i += blockDim.x) { + float g = grad_buf[start + i]; + acc += g * g; + } + sh[tid] = acc; + __syncthreads(); + + /* Block tree reduction 256→1 */ + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + sh[tid] += sh[tid + s]; + } + __syncthreads(); + } + if (tid == 0) { + branch_norms_dev[branch] = sqrtf(sh[0]); + } +} + +/* Device helper: median of 4 non-negative floats (sort-based, fully + * branch-deterministic, no reduction primitives). Returns the average + * of the middle two after sorting — the standard even-length median. + */ +__device__ __forceinline__ float branch_grad_balance_median4( + float a, float b, float c, float d +) { + /* 5-comparator sorting network for 4 elements (Bose-Nelson). */ + #define SWAP_MIN_MAX(x, y) do { \ + float mn = fminf((x), (y)); \ + float mx = fmaxf((x), (y)); \ + (x) = mn; \ + (y) = mx; \ + } while (0) + SWAP_MIN_MAX(a, c); + SWAP_MIN_MAX(b, d); + SWAP_MIN_MAX(a, b); + SWAP_MIN_MAX(c, d); + SWAP_MIN_MAX(b, c); + #undef SWAP_MIN_MAX + /* a ≤ b ≤ c ≤ d; median = (b + c) / 2 */ + return 0.5f * (b + c); +} + +extern "C" __global__ void branch_grad_rescale( + float* __restrict__ grad_buf, + const int* __restrict__ branch_starts, /* [4] element offsets */ + const int* __restrict__ branch_lens, /* [4] element counts */ + const float* __restrict__ branch_norms_dev, /* [4] L2 norms from reduce pass */ + int num_branches, /* architectural constant (= 4); cap multiplier */ + float* __restrict__ branch_scales_dev /* [4] optional diagnostic output — scale applied per branch */ +) { + const int branch = blockIdx.y; + if (branch >= 4) return; + + __shared__ float sh_norms[4]; + __shared__ float sh_scales[4]; + + const int tid = threadIdx.x; + if (tid < 4) { + sh_norms[tid] = branch_norms_dev[tid]; + sh_scales[tid] = 1.0f; + } + __syncthreads(); + + /* Thread 0 computes median + per-branch scales once per block. + * The cap is `num_branches × median` — a branch at the cap has + * L2 mass equal to the sum of `num_branches` median-branches, + * i.e. the sum of the other branches evaluated at the median. */ + if (tid == 0) { + float n0 = sh_norms[0]; + float n1 = sh_norms[1]; + float n2 = sh_norms[2]; + float n3 = sh_norms[3]; + float med = branch_grad_balance_median4(n0, n1, n2, n3); + float cap = (float)num_branches * med; + /* Zero/near-zero median means all branches are quiescent — no + * rescaling is defensible (cap would be 0, divide-by-zero). + * Guard with the same epsilon the existing grad_ratio_mag_dir + * accessor uses (1e-9). */ + if (med > 1e-9f) { + sh_scales[0] = (n0 > cap) ? (cap / n0) : 1.0f; + sh_scales[1] = (n1 > cap) ? (cap / n1) : 1.0f; + sh_scales[2] = (n2 > cap) ? (cap / n2) : 1.0f; + sh_scales[3] = (n3 > cap) ? (cap / n3) : 1.0f; + } + /* Publish scale to the optional diagnostic output (only block + * (0, branch, 0) writes to avoid redundant stores across + * element-chunk blocks). The grad_buf writes below are + * always idempotent across blockIdx.x so any block can emit. */ + if (blockIdx.x == 0 && branch_scales_dev != nullptr) { + branch_scales_dev[branch] = sh_scales[branch]; + } + } + __syncthreads(); + + const float scale = sh_scales[branch]; + /* Skip the multiply entirely when scale == 1.0f — preserves the + * original gradient bit-identically on healthy branches. */ + if (scale == 1.0f) return; + + const int start = branch_starts[branch]; + const int len = branch_lens[branch]; + const int idx = blockIdx.x * blockDim.x + tid; + if (idx < len) { + grad_buf[start + idx] *= scale; + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 752f30e5e..55ae264bb 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -76,6 +76,7 @@ static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_gra static MAMBA2_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin")); pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/graph_utility_kernels.cubin")); static GRAD_DECOMP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad_decomp_kernel.cubin")); +static BRANCH_GRAD_BALANCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/branch_grad_balance_kernel.cubin")); /// Mamba2 temporal scan configuration. const MAMBA2_HISTORY_K: usize = 8; // Rolling history length @@ -1115,6 +1116,41 @@ pub struct GpuDqnTrainer { grad_decomp_snapshot_len: usize, /// Reduction kernel (one-block 256-thread launch, no atomics). grad_decomp_kernel: CudaFunction, + /// Adaptive per-branch gradient-norm balancer — caps any branch whose + /// L2 norm exceeds `num_branches × median_branch_norm` (see + /// `branch_grad_balance_kernel.cu`). `num_branches = 4` is the + /// architectural factored-action axis count (direction, magnitude, + /// order, urgency); the median is a per-step statistical reference, + /// so the cap scales with the current training regime without tuned + /// knobs. Launched inside the `adam_grad_child` graph after all aux + /// grad writes complete and BEFORE `compute_grad_norm_for_adam`, so + /// Adam and the global clip observe the rebalanced gradient. + branch_grad_balance_reduce: CudaFunction, + branch_grad_balance_rescale: CudaFunction, + /// `[4]` element offsets into `grad_buf` for branches 0-3 + /// (direction, magnitude, order, urgency). Each branch covers the + /// 4 contiguous tensors `(8+4b)..(12+4b)`: fc W, fc b, out W, out b. + /// Stored device-side so both balancer kernels read them directly + /// (graph-safe; values are written once at construction). + branch_slice_starts_dev: CudaSlice, + /// `[4]` element lengths (padded to align4) matching + /// `branch_slice_starts_dev`. Padding bytes in `grad_buf` are + /// always zero (memset at step start; cuBLAS/aux kernels write only + /// valid indices) so reading the padded range yields the same L2 as + /// the unpadded range. + branch_slice_lens_dev: CudaSlice, + /// `[4]` scratch slot receiving per-branch L2 norms from the reduce + /// kernel. Consumed by the rescale kernel in the next launch. + branch_grad_norms_dev: CudaSlice, + /// `[4]` diagnostic slot — the scale each branch received on the + /// last rescale launch (`1.0` when no scaling was needed). Written + /// by block (0, branch, 0) of `branch_grad_rescale`; readable at + /// epoch boundary for HEALTH_DIAG auditing. + branch_grad_scales_dev: CudaSlice, + /// Maximum of the 4 branch slice lengths (element count). Used to + /// size the x-dimension of the rescale launch so every branch gets + /// enough blocks to cover all its elements. + branch_slice_max_len: usize, /// Host-side cached per-component norms populated at epoch boundary /// from the pinned result slot. Index order matches `grad_mag_*_ratio` /// accessors: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`. @@ -6489,6 +6525,75 @@ impl GpuDqnTrainer { grad_decomp_mag_start, grad_decomp_mag_len, ); + // ── Adaptive per-branch gradient-norm balancer (Task: L40S grad imbalance fix) ── + // Caps any of the 4 factored-action branches' weight-gradient L2 norm at + // `num_branches × median_branch_norm` + // where `num_branches = 4` is architectural (direction, magnitude, + // order, urgency — the action factorisation) and `median_branch_norm` + // is a per-step statistical reference. Scale-only-when-needed, no + // tuned constants. + let (branch_slice_starts_dev, branch_slice_lens_dev, branch_slice_max_len) = { + let param_sizes = compute_param_sizes(&config); + let elem_size = std::mem::size_of::(); + let mut starts = [0_i32; 4]; + let mut lens = [0_i32; 4]; + let mut max_len = 0_usize; + // Branch b owns tensors (8+4b)..(12+4b), which are contiguous in + // the padded flat layout. Reading the padded range is safe + // because grad_buf is memset-zeroed at step start and writes + // never touch padding bytes — so padding always contributes 0 + // to the L2 sum and no-ops under scalar multiplication. + for b in 0..4_usize { + let first = 8 + b * 4; + let start_byte = padded_byte_offset(¶m_sizes, first) as usize; + let end_byte = padded_byte_offset(¶m_sizes, first + 4) as usize; + let start = start_byte / elem_size; + let end = end_byte / elem_size; + let len = end - start; + starts[b] = start as i32; + lens[b] = len as i32; + if len > max_len { max_len = len; } + } + let mut starts_dev = stream.alloc_zeros::(4) + .map_err(|e| MLError::ModelError(format!("alloc branch_slice_starts_dev: {e}")))?; + stream.memcpy_htod(&starts, &mut starts_dev) + .map_err(|e| MLError::ModelError(format!("htod branch_slice_starts_dev: {e}")))?; + let mut lens_dev = stream.alloc_zeros::(4) + .map_err(|e| MLError::ModelError(format!("alloc branch_slice_lens_dev: {e}")))?; + stream.memcpy_htod(&lens, &mut lens_dev) + .map_err(|e| MLError::ModelError(format!("htod branch_slice_lens_dev: {e}")))?; + (starts_dev, lens_dev, max_len) + }; + let branch_grad_norms_dev = stream.alloc_zeros::(4) + .map_err(|e| MLError::ModelError(format!("alloc branch_grad_norms_dev: {e}")))?; + let branch_grad_scales_dev = { + // Initialised to 1.0 so HEALTH_DIAG reads a sensible "no-op" + // value before the first rescale launch writes real scales. + let ones = [1.0_f32; 4]; + let mut buf = stream.alloc_zeros::(4) + .map_err(|e| MLError::ModelError(format!("alloc branch_grad_scales_dev: {e}")))?; + stream.memcpy_htod(&ones, &mut buf) + .map_err(|e| MLError::ModelError(format!("htod branch_grad_scales_dev: {e}")))?; + buf + }; + // Load reduce + rescale kernels from a dedicated CUmodule so their + // CUfunction handles are isolated from all other graphs (Hopper + // CUfunction-isolation rule — shared handles corrupt captured + // kernel state). + let (branch_grad_balance_reduce, branch_grad_balance_rescale) = { + let module = stream.context().load_cubin(BRANCH_GRAD_BALANCE_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("branch_grad_balance cubin load: {e}")))?; + let reduce = module.load_function("branch_grad_norm_reduce") + .map_err(|e| MLError::ModelError(format!("branch_grad_norm_reduce load: {e}")))?; + let rescale = module.load_function("branch_grad_rescale") + .map_err(|e| MLError::ModelError(format!("branch_grad_rescale load: {e}")))?; + (reduce, rescale) + }; + info!( + "GpuDqnTrainer: branch_grad_balance kernels loaded — max_slice_len={} num_branches=4 (architectural)", + branch_slice_max_len, + ); + // eval_v_range — pinned device-mapped (zero-copy write from CPU). let eval_v_range_pinned: *mut f32 = unsafe { let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; @@ -8460,6 +8565,13 @@ impl GpuDqnTrainer { grad_decomp_mag_len, grad_decomp_snapshot_len, grad_decomp_kernel, + branch_grad_balance_reduce, + branch_grad_balance_rescale, + branch_slice_starts_dev, + branch_slice_lens_dev, + branch_grad_norms_dev, + branch_grad_scales_dev, + branch_slice_max_len, grad_component_norms_mag: [0.0_f32; 9], grad_component_norms_dir: [0.0_f32; 9], grad_component_norms_trunk: [0.0_f32; 9], @@ -9546,6 +9658,136 @@ impl GpuDqnTrainer { /// Replay graph_adam + readback scalars. Call AFTER injecting auxiliary /// gradients into grad_buf (IQN, attention, ensemble). + /// Adaptive per-branch gradient-norm balancer. + /// + /// For each factored-action branch `d ∈ {0..4}` (direction, magnitude, + /// order, urgency) computes the branch's weight-gradient L2 norm, + /// and scales the branch's gradient down to the cap + /// + /// `cap = num_branches × median(branch_norms)` + /// + /// whenever its norm exceeds the cap. Branches within the cap pass + /// through unchanged. + /// + /// Why this shape is fully adaptive with no tuned knobs: + /// * `num_branches = 4` is the architectural count of factored + /// action axes. Not a hyperparameter. + /// * `median(branch_norms)` is a per-step statistical reference + /// that tracks the current gradient regime (Q-magnitude, loss + /// blend, ISV-driven scales). + /// * Their product gives the defensible bound: "no branch may + /// dominate by more than `num_branches` × the median" — i.e. + /// no branch can carry more L2 mass than all other branches + /// would, evaluated at the median. + /// + /// Fixes the pathology observed on L40S (train-mdh86, epochs 0-6): + /// `grad_ratio_mag_dir ≈ 1.5e4-2.6e4×` for the warm-up, then a + /// single-step collapse to ~100× at epoch 7 that destabilised + /// learning (Sharpe +34 → -67). Symmetric capping prevents the + /// swing at both ends. + /// + /// Insertion point: inside the `adam_grad_child` graph, AFTER all + /// aux grad writes complete and BEFORE `compute_grad_norm_for_adam` + /// — so Adam's global clip and the Adam update both observe the + /// rebalanced gradient. Two kernel launches; no atomicAdd; no host + /// syncs; no dynamic allocations — safe to capture. + pub(crate) fn launch_branch_grad_balance(&self) -> Result<(), MLError> { + let grad_ptr = self.ptrs.grad_buf; + let starts_ptr = self.branch_slice_starts_dev.raw_ptr(); + let lens_ptr = self.branch_slice_lens_dev.raw_ptr(); + let norms_ptr = self.branch_grad_norms_dev.raw_ptr(); + let scales_ptr = self.branch_grad_scales_dev.raw_ptr(); + + // ── Pass 1: per-branch L2 reduction (grid=(4,), block=(256,)) ─ + unsafe { + self.stream + .launch_builder(&self.branch_grad_balance_reduce) + .arg(&grad_ptr) + .arg(&starts_ptr) + .arg(&lens_ptr) + .arg(&norms_ptr) + .launch(LaunchConfig { + grid_dim: (4, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError( + format!("branch_grad_norm_reduce launch: {e}") + ))?; + } + + // ── Pass 2: cap + rescale (grid=(blocks_x, 4,), block=(256,)) ─ + // blocks_x covers the largest branch; shorter branches early-exit + // via the `idx < len` guard in the kernel. Architectural value + // `num_branches = 4` is passed as the cap multiplier. + let blocks_x = ((self.branch_slice_max_len as u32 + 255) / 256).max(1); + let num_branches: i32 = 4; + unsafe { + self.stream + .launch_builder(&self.branch_grad_balance_rescale) + .arg(&grad_ptr) + .arg(&starts_ptr) + .arg(&lens_ptr) + .arg(&norms_ptr) + .arg(&num_branches) + .arg(&scales_ptr) + .launch(LaunchConfig { + grid_dim: (blocks_x, 4, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError( + format!("branch_grad_rescale launch: {e}") + ))?; + } + + Ok(()) + } + + /// Read back the per-branch scales applied on the most recent + /// `launch_branch_grad_balance` call — `[dir, mag, ord, urg]`. + /// Values in `[0, 1]`; `1.0` means the branch passed through + /// unchanged; `< 1.0` means the branch exceeded + /// `num_branches × median(branch_norms)` and was capped to that + /// ratio. + /// + /// Epoch-boundary cold path (stream-syncs once, same pattern as + /// `per_branch_grad_norms`). Useful for HEALTH_DIAG auditing — a + /// scale of `~1.0/K` on the magnitude branch confirms the cap + /// fired at `num_branches × median / K × norm_mag = cap`. + pub fn branch_grad_scales(&self) -> Result<[f32; 4], MLError> { + self.stream.synchronize().map_err(|e| { + MLError::ModelError(format!("branch_grad_scales sync: {e}")) + })?; + let mut host = [1.0_f32; 4]; + self.stream + .memcpy_dtoh(&self.branch_grad_scales_dev.slice(..4), &mut host) + .map_err(|e| MLError::ModelError(format!("branch_grad_scales dtoh: {e}")))?; + self.stream.synchronize().map_err(|e| { + MLError::ModelError(format!("branch_grad_scales post-dtoh sync: {e}")) + })?; + Ok(host) + } + + /// Read back the per-branch gradient L2 norms computed by the most + /// recent `launch_branch_grad_balance` reduce pass — `[dir, mag, + /// ord, urg]`. These are the norms the rescale pass observed + /// BEFORE capping; compare against `branch_grad_scales` to + /// reconstruct the post-cap ratio (`norm × scale ≤ cap`). + pub fn branch_grad_norms_device(&self) -> Result<[f32; 4], MLError> { + self.stream.synchronize().map_err(|e| { + MLError::ModelError(format!("branch_grad_norms_device sync: {e}")) + })?; + let mut host = [0.0_f32; 4]; + self.stream + .memcpy_dtoh(&self.branch_grad_norms_dev.slice(..4), &mut host) + .map_err(|e| MLError::ModelError(format!("branch_grad_norms_device dtoh: {e}")))?; + self.stream.synchronize().map_err(|e| { + MLError::ModelError(format!("branch_grad_norms_device post-dtoh sync: {e}")) + })?; + Ok(host) + } + /// Compute gradient norm for adam_child graph capture. /// Uses grad_norm_kernel (NOT standalone) — the standalone is already captured /// in forward_child (d_logits clipping). Sharing the same CUfunction across @@ -9677,6 +9919,14 @@ impl GpuDqnTrainer { self.update_stochastic_depth_mask()?; self.submit_forward_ops_main()?; self.submit_forward_ops_ddqn()?; + // Adaptive per-branch gradient-norm balancer — caps each branch's + // weight-gradient L2 norm at `num_branches × median_branch_norm`. + // Runs AFTER all aux grad writes complete and BEFORE the global + // grad_norm reduction so Adam's clip and the update both see the + // rebalanced gradient. (No aux path in this ungraphed fallback — + // but the method is the sole guard against the L40S pathology so + // it must be unconditional on every step, graphed or not.) + self.launch_branch_grad_balance()?; self.compute_grad_norm_for_adam()?; self.submit_adam_ops(online_dueling, online_branching)?; @@ -10261,6 +10511,10 @@ impl GpuDqnTrainer { // ── Execute training step ungraphed (graph capture in FusedTrainingCtx) ── self.submit_forward_ops_main()?; self.submit_forward_ops_ddqn()?; + // Adaptive per-branch gradient-norm balancer — see the detailed + // rationale on `launch_branch_grad_balance`. Kept unconditional + // so no code path can skip the cap. + self.launch_branch_grad_balance()?; self.compute_grad_norm_for_adam()?; self.submit_adam_ops(online_dueling, online_branching)?; diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2cf9254d1..3a2ca15aa 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1194,6 +1194,17 @@ impl FusedTrainingCtx { self.step_ofi_embed_adam()?; self.trainer.apply_pruning_mask() .map_err(|e| anyhow::anyhow!("pruning mask: {e}"))?; + // Adaptive per-branch gradient-norm balancer — caps each + // branch's weight-gradient L2 norm at + // `num_branches × median_branch_norm`. Must run AFTER all + // aux grad writes complete and BEFORE `compute_grad_norm_for_adam` + // so the global clip + Adam update both observe the rebalanced + // gradient. See `launch_branch_grad_balance` for the full + // architectural rationale (num_branches = 4 is the factored + // action-axis count, median is a per-step reference — no + // tuned knobs). + self.trainer.launch_branch_grad_balance() + .map_err(|e| anyhow::anyhow!("branch_grad_balance: {e}"))?; self.trainer.compute_grad_norm_for_adam() .map_err(|e| anyhow::anyhow!("grad_norm: {e}"))?; self.trainer.submit_adam_ops(&self.online_dueling, &self.online_branching) @@ -1963,6 +1974,11 @@ impl FusedTrainingCtx { s.step_ofi_embed_adam()?; s.trainer.apply_pruning_mask() .map_err(|e| anyhow::anyhow!("{e}"))?; + // Adaptive per-branch gradient-norm balancer — captured inside + // adam_grad_child so every graph replay rebalances before the + // global grad_norm reduction. See `launch_branch_grad_balance`. + s.trainer.launch_branch_grad_balance() + .map_err(|e| anyhow::anyhow!("branch_grad_balance: {e}"))?; s.trainer.compute_grad_norm_for_adam() .map_err(|e| anyhow::anyhow!("{e}")) })?;