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}")) })?;