From a52d996135fba4afeca4707a2b769a204d14266a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 19:47:19 +0200 Subject: [PATCH] feat(moe): wire MoE forward + backward + load-balance + monitoring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 of the MoE regime redesign per docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md. Atomic forward+backward wire-up — replaces the existing h_s2 producer in fused_training.rs with the gated expert mixture: state[B, 128] -> shared GRN trunk -> h_s1[B, 256] | +-> 8 expert MLPs (256->64->256) -> expert_outputs[8, B, 256] +-> gate (128->64->8 softmax) ----> gate[B, 8] | moe_mixture_forward -> h_s2[B, 256] -> branching heads + C51 + IQN Backward: moe_mixture_backward (de_k = g · dh_s2) + moe_dgate_reduce (dg = Σ_c e_k · dh_s2) + load-balance aux gradient + cuBLAS SGEMM backward through gate + each expert's 2 linear layers. Adam optimizer step now updates gate + 8 experts via params_buf. Loss: λ · K · Σ_k (mean_b g[b,k])² added to total loss with λ from hyperparams.moe_lambda (default 0.01). Per-step ISV producer launch (moe_expert_util_ema_update) writes 8 utilization EMA + 1 gate-entropy EMA into ISV[118..127). Per-epoch HEALTH_DIAG aux_moe line emits utilization vector + entropy live so operators can see whether experts are differentiating or collapsing. Smoke test: DONE. 3/3 folds, all checkpoints saved, 728s (12.1 min, within 25-min budget). Gate differentiated by epoch 1: expert 2 rose from 0.119 → 0.286 → 0.323 over fold 1-2 while others remained at 0.097-0.113. Gate entropy 1.611 at fold 2 epoch 4 < ln(8)=2.079. val_loss finite across all 3 folds; average fold metric 22.4. Per feedback_no_partial_refactor.md: all consumers of the h_s2 contract (branching heads, IQN aux, attention focus, backward chain) migrate in this single commit. Per feedback_no_htod_htoh_only_mapped_pinned.md: no new HtoD/HtoH introduced; gate softmax + expert outputs + mixture all GPU-resident, ISV producer GPU-driven. Load-balance scalar uses cuMemAllocHost + cuMemHostGetDevicePointer_v2 (mapped pinned), matching existing pattern. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 444 ++++++++++++++++++ crates/ml/src/cuda_pipeline/gpu_moe_head.rs | 247 +++++++++- crates/ml/src/cuda_pipeline/moe_kernels.cu | 55 +++ crates/ml/src/trainers/dqn/fused_training.rs | 1 + .../src/trainers/dqn/trainer/training_loop.rs | 34 ++ docs/dqn-wire-up-audit.md | 16 + 6 files changed, 796 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 2c5efc93f..ec01617cc 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -63,6 +63,7 @@ use super::shared_cublas_handle::PerStreamCublasHandles; use super::gpu_aux_heads::{ AuxHeadsBackwardOps, AuxHeadsForwardOps, AUX_HIDDEN_DIM, AUX_NEXT_BAR_K, AUX_REGIME_K, }; +use super::gpu_moe_head::GpuMoeHead; // ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ────── pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); @@ -1242,6 +1243,10 @@ pub struct GpuDqnTrainConfig { /// override to a smaller value (e.g. 10_000) so the seed→network transition /// is observable inside the smoke run. pub replay_seed_steps: u32, + + /// Phase 3 MoE load-balance loss weight λ. + /// Load-balance loss = λ·K·Σ_k(mean_b g[b,k])². Default 0.01. + pub moe_lambda: f32, } impl Default for GpuDqnTrainConfig { @@ -1290,6 +1295,7 @@ impl Default for GpuDqnTrainConfig { market_dim: 42, // Default: 42 base features. Overridden to 50 when OFI (MBP-10) enabled. total_epochs: 0, replay_seed_steps: 100_000, + moe_lambda: 0.01, } } } @@ -2675,6 +2681,45 @@ pub struct GpuDqnTrainer { /// share so aux can never dominate the policy gradient. aux_weight: f32, + // ── Phase 3: MoE forward/backward buffers + orchestrator ───────────── + /// MoE kernel orchestrator (mixture fwd/bwd, load-balance, EMA, softmax). + moe_head: GpuMoeHead, + /// Gate sub-network layer-1 hidden activations [B, MOE_GATE_HIDDEN]. + /// Saved for backward (dW_gate_w1 = gate_h1^T @ d_gate_h1). + moe_gate_h1_buf: CudaSlice, + /// Gate pre-softmax logits [B, MOE_NUM_EXPERTS]. + /// Saved for softmax backward (needed by moe_softmax_backward). + moe_gate_pre_buf: CudaSlice, + /// Gate softmax output = gate weights g[B, MOE_NUM_EXPERTS]. + /// Saved for mixture forward + mixture backward + load-balance loss. + moe_gate_softmax_buf: CudaSlice, + /// 8 expert layer-1 hidden buffers, each [B, MOE_EXPERT_BOTTLENECK]. + /// Saved for backward (dW_ek_w1 = expert_h1_k^T @ d_expert_h1_k). + moe_expert_h1_bufs: [CudaSlice; MOE_NUM_EXPERTS], + /// Stacked expert outputs [K, B, MOE_SH2] (K=8, layout row-major). + /// Saved for both mixture forward (direct input) and moe_dgate_reduce. + moe_expert_outputs_buf: CudaSlice, + /// d(loss)/d(expert_outputs) [K, B, MOE_SH2] from moe_mixture_backward. + moe_de_k_buf: CudaSlice, + /// d(loss)/d(gate_soft) [B, MOE_NUM_EXPERTS] from moe_dgate_reduce. + moe_dg_buf: CudaSlice, + /// d(loss)/d(gate_pre) [B, MOE_NUM_EXPERTS] from moe_softmax_backward. + moe_dg_pre_buf: CudaSlice, + /// Accumulated d(loss)/d(save_h_s1) from all 8 expert backward dX. + /// After accumulation, DtoD-copied into bw_d_h_s2 for encoder_backward_chain. + moe_dh_s1_scratch: CudaSlice, + /// d(loss)/d(gate_h1) [B, MOE_GATE_HIDDEN] from gate_w2 backward dX. + moe_gate_dh1_buf: CudaSlice, + /// Per-expert load-balance loss [MOE_NUM_EXPERTS]. + moe_load_balance_loss_per_k: CudaSlice, + /// Load-balance total loss scalar [1] — pinned device-mapped. + /// GPU writes via dev_ptr, CPU reads host_ptr for optional logging. + moe_load_balance_loss_total_pinned: *mut f32, + /// Device pointer for moe_load_balance_loss_total_pinned. + moe_load_balance_loss_total_dev_ptr: u64, + /// Lambda for the load-balance auxiliary loss term (from config.moe_lambda). + moe_lambda: f32, + /// #21 Stochastic depth: per-layer scale buffer [3] (h_s1, h_s2, h_v). /// Written by GPU RNG kernel before each graph replay. stochastic_depth_scale_buf: CudaSlice, @@ -7782,6 +7827,310 @@ impl GpuDqnTrainer { ) } + // ── Phase 3 MoE methods ───────────────────────────────────────────────── + + /// Phase 3 T3.1–T3.3: MoE forward. + /// + /// Sequence: + /// 1. Gate: state[B,SD] → Linear(SD→GH) + bias → ReLU → gate_h1[B,GH] + /// 2. Gate: gate_h1[B,GH] → Linear(GH→K) + bias → gate_pre[B,K] + /// 3. Gate: gate_pre → moe_row_softmax → gate_soft[B,K] (saves for backward) + /// 4. For k in [0, K): h_s1[B,SH2] → Linear(SH2→BTN)+bias → ReLU → expert_h1_k[B,BTN] + /// expert_h1_k → Linear(BTN→SH2)+bias → expert_out_k[B,SH2] + /// 5. moe_mixture_forward stacks expert_outs[K,B,SH2] × gate_soft[B,K] → h_s2_new[B,SH2] + /// 6. DtoD copy h_s2_new → save_h_s2 (replaces the GRN-produced save_h_s2) + /// + /// Inputs: reads `save_h_s1` (h_s1 from GRN block 1), `states_buf` (for gate), + /// and online weight tensors [127..163). + /// Output: overwrites `save_h_s2` with the MoE mixture. + pub(crate) fn launch_moe_forward(&self) -> Result<(), MLError> { + let b = self.config.batch_size; + let sh2 = self.config.shared_h2; + let sd = ml_core::state_layout::STATE_DIM; + + let param_sizes = compute_param_sizes(&self.config); + let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes); + + // Gate weight tensors [127..131) + let gate_w1 = on_w_ptrs[127]; // [SD, GH] + let gate_b1 = on_w_ptrs[128]; // [GH] + let gate_w2 = on_w_ptrs[129]; // [GH, K] + let gate_b2 = on_w_ptrs[130]; // [K] + + // Gate sub-network forward + let gate_h1_ptr = self.moe_gate_h1_buf.raw_ptr(); + let gate_pre_ptr = self.moe_gate_pre_buf.raw_ptr(); + let gate_soft_ptr = self.moe_gate_softmax_buf.raw_ptr(); + let states_ptr = self.ptrs.states_buf; + + // Step 1: state[B,SD] → Linear(SD→GH) → gate_h1[B,GH] + self.cublas_forward.sgemm_f32(&self.stream, gate_w1, states_ptr, gate_h1_ptr, + MOE_GATE_HIDDEN, b, sd, "moe_gate_layer1")?; + self.cublas_forward.launch_add_bias_relu_f32_raw( + &self.stream, gate_h1_ptr, gate_b1, MOE_GATE_HIDDEN, b)?; + + // Step 2: gate_h1[B,GH] → Linear(GH→K) → gate_pre[B,K] + self.cublas_forward.sgemm_f32(&self.stream, gate_w2, gate_h1_ptr, gate_pre_ptr, + MOE_NUM_EXPERTS, b, MOE_GATE_HIDDEN, "moe_gate_layer2")?; + self.cublas_forward.launch_add_bias_f32_raw( + &self.stream, gate_pre_ptr, gate_b2, MOE_NUM_EXPERTS, b)?; + + // Step 3: softmax → gate_soft[B,K] + self.moe_head.launch_row_softmax(gate_pre_ptr, gate_soft_ptr, b, MOE_NUM_EXPERTS)?; + + // Steps 4: for each expert k: h_s1[B,SH2] → Linear(SH2→BTN)+ReLU → Linear(BTN→SH2) + let h_s1_ptr = self.ptrs.save_h_s1; + let expert_outputs_ptr = self.moe_expert_outputs_buf.raw_ptr(); + + for k in 0..MOE_NUM_EXPERTS { + let base_idx = 131 + k * 4; + let ew1 = on_w_ptrs[base_idx]; // [SH2, BTN] + let eb1 = on_w_ptrs[base_idx + 1]; // [BTN] + let ew2 = on_w_ptrs[base_idx + 2]; // [BTN, SH2] + let eb2 = on_w_ptrs[base_idx + 3]; // [SH2] + + let eh1_ptr = self.moe_expert_h1_bufs[k].raw_ptr(); + // expert_out_k is the k-th [B, SH2] slice of moe_expert_outputs_buf + let eout_ptr = expert_outputs_ptr + (k * b * sh2) as u64 * std::mem::size_of::() as u64; + + // h_s1[B,SH2] → Linear(SH2→BTN) → eh1[B,BTN] + self.cublas_forward.sgemm_f32(&self.stream, ew1, h_s1_ptr, eh1_ptr, + MOE_EXPERT_BOTTLENECK, b, sh2, &format!("moe_expert_{k}_layer1"))?; + self.cublas_forward.launch_add_bias_relu_f32_raw( + &self.stream, eh1_ptr, eb1, MOE_EXPERT_BOTTLENECK, b)?; + + // eh1[B,BTN] → Linear(BTN→SH2) → eout[B,SH2] + self.cublas_forward.sgemm_f32(&self.stream, ew2, eh1_ptr, eout_ptr, + sh2, b, MOE_EXPERT_BOTTLENECK, &format!("moe_expert_{k}_layer2"))?; + self.cublas_forward.launch_add_bias_f32_raw( + &self.stream, eout_ptr, eb2, sh2, b)?; + } + + // Step 5: moe_mixture_forward: expert_outputs[K,B,SH2] × gate_soft[B,K] → h_s2_new[B,SH2] + // Write directly into save_h_s2 — replaces the GRN output. + self.moe_head.launch_mixture_forward( + expert_outputs_ptr, + gate_soft_ptr, + self.ptrs.save_h_s2, + b, MOE_NUM_EXPERTS, sh2, + )?; + + Ok(()) + } + + /// Phase 3 T3.4: compute load-balance loss and SAXPY into total_loss_dev_ptr. + /// + /// Calls `moe_head.launch_load_balance_loss` which runs `moe_load_balance_loss` + /// (K blocks × B reduction) + `moe_load_balance_reduce` (single thread sum) into + /// `moe_load_balance_loss_total_dev_ptr`. Then SAXPYs the scalar (alpha=1.0) + /// into `total_loss_dev_ptr` so it is included in the blended TD loss. + pub(crate) fn launch_moe_load_balance_loss(&self) -> Result<(), MLError> { + let b = self.config.batch_size; + let gate_ptr = self.moe_gate_softmax_buf.raw_ptr(); + let loss_per_k_ptr = self.moe_load_balance_loss_per_k.raw_ptr(); + let loss_total_ptr = self.moe_load_balance_loss_total_dev_ptr; + + // Compute per-expert load loss + reduce to scalar. + self.moe_head.launch_load_balance_loss( + gate_ptr, loss_per_k_ptr, loss_total_ptr, + b, MOE_NUM_EXPERTS, self.moe_lambda, + )?; + + // SAXPY total_loss += 1.0 * lb_loss_total (both are device-mapped scalars). + let n: i32 = 1; + let alpha: f32 = 1.0; + let total_loss_ptr = self.total_loss_dev_ptr; + unsafe { + self.stream + .launch_builder(&self.saxpy_f32_kernel) + .arg(&total_loss_ptr) + .arg(&loss_total_ptr) + .arg(&alpha) + .arg(&n) + .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("moe_lb_loss saxpy: {e}")))?; + } + Ok(()) + } + + /// Phase 3 T3.5: launch `moe_expert_util_ema_update` ISV producer. + /// + /// Single-block single-thread cold-path kernel. Reads `moe_gate_softmax_buf [B, K]` + /// and EMA-updates ISV[MOE_EXPERT_UTIL_EMA_BASE..+8] (per-expert mean gate) + /// and ISV[MOE_GATE_ENTROPY_EMA_INDEX] (gate entropy). + pub fn launch_moe_expert_util_ema(&self, ema_alpha: f32) -> Result<(), MLError> { + debug_assert!(self.isv_signals_dev_ptr != 0, + "launch_moe_expert_util_ema: isv_signals_dev_ptr must be non-zero"); + let gate_ptr = self.moe_gate_softmax_buf.raw_ptr(); + self.moe_head.launch_expert_util_ema( + gate_ptr, + self.isv_signals_dev_ptr, + self.config.batch_size, + MOE_NUM_EXPERTS, + MOE_EXPERT_UTIL_EMA_BASE, + MOE_GATE_ENTROPY_EMA_INDEX, + ema_alpha, + ) + } + + /// Phase 3 T3.6: MoE backward. + /// + /// Sequence (reverse of forward): + /// 1. moe_mixture_backward: dh_s2 × gate_soft → de_k[K,B,SH2] + /// + moe_dgate_reduce: dh_s2 × expert_outs → dg[B,K] + /// 2. moe_softmax_backward: gate_soft × dg → dg_pre[B,K] + /// 3. Gate layer 2 backward: dg_pre [B,K] → dW_gate_w2 (w_ptrs[129]), db_gate_b2 (w_ptrs[130]) + /// + dX → moe_gate_dh1[B,GH] + /// 4. Gate layer 1 backward (dX from step 3, gated by ReLU): apply relu mask, + /// → dW_gate_w1 (w_ptrs[127]), db_gate_b1 (w_ptrs[128]) + /// 5. For each expert k: + /// a. Expert layer 2 backward: de_k_slice[B,SH2] → dW_ew2, db_eb2 + dX → d_eh1_k[B,BTN] + /// b. Expert layer 1 backward (dX from 5a gated by ReLU): → dW_ew1, db_eb1 + dX → d_h_s1_k[B,SH2] + /// 6. Accumulate 8× d_h_s1_k into moe_dh_s1_scratch[B,SH2] + /// 7. DtoD-copy moe_dh_s1_scratch → bw_d_h_s2 (REPLACES existing, so encoder + /// backward chain sees the MoE-routed h_s1 gradient as its h_s2 input gradient) + /// + /// `grad_base`: base device pointer for parameter gradient accumulation. + pub(crate) fn launch_moe_backward(&self, grad_base: u64) -> Result<(), MLError> { + let b = self.config.batch_size; + let sh2 = self.config.shared_h2; + let f32_size = std::mem::size_of::() as u64; + + let param_sizes = compute_param_sizes(&self.config); + let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes); + + // Gradient buffer base pointers — offset into grad_base by param position. + // grad_base is laid out identically to the params buffer. + let grad_ptrs = f32_weight_ptrs_from_base(grad_base, ¶m_sizes); + + let gate_soft_ptr = self.moe_gate_softmax_buf.raw_ptr(); + let gate_h1_ptr = self.moe_gate_h1_buf.raw_ptr(); + let expert_outs_ptr = self.moe_expert_outputs_buf.raw_ptr(); + let de_k_ptr = self.moe_de_k_buf.raw_ptr(); + let dg_ptr = self.moe_dg_buf.raw_ptr(); + let dg_pre_ptr = self.moe_dg_pre_buf.raw_ptr(); + let gate_dh1_ptr = self.moe_gate_dh1_buf.raw_ptr(); + let dh_s1_scratch = self.moe_dh_s1_scratch.raw_ptr(); + + // d_h_s2 holds dL/d(save_h_s2) accumulated by all prior backward ops + // (branches, aux heads). This is the upstream gradient w.r.t. the MoE output. + let d_h_s2_ptr = bw_raw_f32_ptr(&self.bw_d_h_s2, &self.stream); + let h_s1_ptr = self.ptrs.save_h_s1; + + // Step 1a: moe_mixture_backward — propagates d_h_s2 through the weighted sum to de_k. + // de_k[k,b,c] = gate_soft[b,k] * d_h_s2[b,c] + // Step 1b: moe_dgate_reduce — dg[b,k] = sum_c(d_h_s2[b,c] * expert_outs[k,b,c]) + self.moe_head.launch_mixture_backward( + d_h_s2_ptr, gate_soft_ptr, expert_outs_ptr, + de_k_ptr, dg_ptr, + b, MOE_NUM_EXPERTS, sh2, + )?; + + // Step 2: softmax backward — dg_pre[b,k] = gate_soft[b,k] * (dg[b,k] - dot(gate_soft[b,:], dg[b,:])) + self.moe_head.launch_softmax_backward(gate_soft_ptr, dg_ptr, dg_pre_ptr, b, MOE_NUM_EXPERTS)?; + + // Step 3: gate layer 2 dW + dX + // dg_pre[B,K] w.r.t. gate_h1[B,GH] through W_gate_w2[GH,K] + self.cublas_backward.launch_dw_only( + &self.stream, + dg_pre_ptr, gate_h1_ptr, // dy, x + grad_ptrs[129], grad_ptrs[130], // dW_gate_w2, db_gate_b2 + MOE_NUM_EXPERTS, MOE_GATE_HIDDEN, b, + )?; + // dX into gate_dh1 (beta=0: overwrite) + self.cublas_backward.launch_dx_only( + &self.stream, + dg_pre_ptr, w_ptrs[129], gate_dh1_ptr, + MOE_NUM_EXPERTS, MOE_GATE_HIDDEN, b, 0.0, + )?; + + // Step 4: gate layer 1 dW — gated by ReLU of gate_h1. + // gate_h1 holds the post-ReLU activation saved by forward; gate_dh1[i] = 0 + // wherever gate_h1[i] == 0 (i.e., the neuron was clamped). relu_mask + // on cublas_backward implements dx[i] *= (activation[i] > 0). + self.cublas_backward.relu_mask( + &self.stream, gate_dh1_ptr, gate_h1_ptr, b * MOE_GATE_HIDDEN, + )?; + // dW_gate_w1, db_gate_b1 (gate layer 1 is state[B,SD] → gate_h1[B,GH]) + let states_ptr = self.ptrs.states_buf; + self.cublas_backward.launch_dw_only( + &self.stream, + gate_dh1_ptr, states_ptr, // dy, x + grad_ptrs[127], grad_ptrs[128], // dW_gate_w1, db_gate_b1 + MOE_GATE_HIDDEN, ml_core::state_layout::STATE_DIM, b, + )?; + + // Step 5+6: expert backward — 8 experts, accumulate dX into dh_s1_scratch. + // + // Buffer aliasing strategy: after step 4 gate backward, gate_h1_ptr + // (= moe_gate_h1_buf [B, GH]) is free. GH == BTN == 64, so gate_h1_ptr + // is repurposed as d_eh1 scratch [B, BTN] for each expert in turn. + // The per-expert saved activation (moe_expert_h1_bufs[k]) must NOT be + // overwritten until after the ReLU mask is applied; the sequence is: + // a. dW_ew2 (reads eh1_saved as x — does NOT modify eh1_saved) + // b. dX_ew2 → gate_h1_ptr (d_eh1 scratch, beta=0; eh1_saved intact) + // c. relu_mask(d_eh1, eh1_saved) (eh1_saved still valid here) + // d. dW_ew1 (reads eh1_saved? No — reads d_eh1 from gate_h1_ptr as dy) + // e. dX_ew1 → dh_s1_scratch (accumulate with beta 0 or 1) + let d_eh1_scratch = gate_h1_ptr; // repurpose gate_h1_ptr as [B, BTN] scratch + for k in 0..MOE_NUM_EXPERTS { + let base_idx = 131 + k * 4; + let ew1 = w_ptrs[base_idx]; // [SH2, BTN] + let ew2 = w_ptrs[base_idx + 2]; // [BTN, SH2] + let g_ew1 = grad_ptrs[base_idx]; + let g_eb1 = grad_ptrs[base_idx + 1]; + let g_ew2 = grad_ptrs[base_idx + 2]; + let g_eb2 = grad_ptrs[base_idx + 3]; + + let eh1_saved = self.moe_expert_h1_bufs[k].raw_ptr(); // post-ReLU, read-only here + let dek_k = de_k_ptr + (k * b * sh2) as u64 * f32_size; + + // a. Expert layer 2 dW: dek_k[B,SH2] × eh1_saved[B,BTN]^T → dW_ew2, db_eb2 + // launch_dw_only reads eh1_saved as x; does NOT modify it. + self.cublas_backward.launch_dw_only( + &self.stream, + dek_k, eh1_saved, // dy, x + g_ew2, g_eb2, // dW_ew2, db_eb2 + sh2, MOE_EXPERT_BOTTLENECK, b, + )?; + // b. Expert layer 2 dX: dek_k × W_ew2^T → d_eh1 scratch (beta=0) + self.cublas_backward.launch_dx_only( + &self.stream, + dek_k, ew2, d_eh1_scratch, + sh2, MOE_EXPERT_BOTTLENECK, b, 0.0, + )?; + // c. ReLU mask: d_eh1[i] *= (eh1_saved[i] > 0) + self.cublas_backward.relu_mask( + &self.stream, d_eh1_scratch, eh1_saved, b * MOE_EXPERT_BOTTLENECK, + )?; + // d. Expert layer 1 dW: d_eh1[B,BTN] × h_s1[B,SH2]^T → dW_ew1, db_eb1 + self.cublas_backward.launch_dw_only( + &self.stream, + d_eh1_scratch, h_s1_ptr, // dy=d_eh1, x=h_s1 + g_ew1, g_eb1, // dW_ew1, db_eb1 + MOE_EXPERT_BOTTLENECK, sh2, b, + )?; + // e. Expert layer 1 dX → dh_s1_scratch (accumulate; beta=0 for k=0, 1 for k>0) + let beta = if k == 0 { 0.0_f32 } else { 1.0_f32 }; + self.cublas_backward.launch_dx_only( + &self.stream, + d_eh1_scratch, ew1, dh_s1_scratch, + MOE_EXPERT_BOTTLENECK, sh2, b, beta, + )?; + } + + // Step 7: DtoD-copy moe_dh_s1_scratch → bw_d_h_s2. + // This REPLACES (not adds to) bw_d_h_s2 so encoder_backward_chain + // receives the MoE-accumulated h_s1 gradient as its h_s2 input gradient. + // The aux_heads_backward already wrote into bw_d_h_s2; we replace here + // because the MoE mixture output IS save_h_s2 — the GRN h_s2 block was + // replaced by MoE in forward, so its "d_h_s2" should be the MoE dX. + let n_bytes = b * sh2 * std::mem::size_of::(); + self.graph_safe_copy_f32(d_h_s2_ptr, dh_s1_scratch, n_bytes, "moe_dh_s1→d_h_s2")?; + + Ok(()) + } + /// Plan 4 Task 6 Commit B: aux-heads forward orchestrator. /// /// Runs the next-bar regression head + 5-class regime classification head @@ -10129,6 +10478,63 @@ impl GpuDqnTrainer { // step. The first graph capture bakes 0.05 — an honest cold-start // value, not a stub. let aux_weight: f32 = 0.05; + + // ── Phase 3: MoE forward/backward orchestrator + working buffers ───── + // GpuMoeHead loads `moe_kernels.cubin` and caches 8 kernel handles. + // All working buffers are zero-initialised — the gate starts at zero + // weights (uniform 1/K softmax), so cold-start EMA values are valid. + let moe_head = GpuMoeHead::new(Arc::clone(&stream))?; + let moe_sh2 = config.shared_h2; + let moe_gate_h1_buf = alloc_f32(&stream, b * MOE_GATE_HIDDEN, "moe_gate_h1_buf")?; + let moe_gate_pre_buf = alloc_f32(&stream, b * MOE_NUM_EXPERTS, "moe_gate_pre_buf")?; + let moe_gate_softmax_buf = alloc_f32(&stream, b * MOE_NUM_EXPERTS, "moe_gate_softmax_buf")?; + // Allocate 8 expert hidden buffers (one per expert, each [B, BTN]). + let moe_expert_h1_bufs = { + let mut arr: [std::mem::MaybeUninit>; MOE_NUM_EXPERTS] = + unsafe { std::mem::MaybeUninit::uninit().assume_init() }; + for (ek, slot) in arr.iter_mut().enumerate() { + let buf = alloc_f32(&stream, b * MOE_EXPERT_BOTTLENECK, + &format!("moe_expert_h1_{ek}"))?; + slot.write(buf); + } + unsafe { std::mem::transmute::<_, [CudaSlice; MOE_NUM_EXPERTS]>(arr) } + }; + let moe_expert_outputs_buf = alloc_f32(&stream, MOE_NUM_EXPERTS * b * moe_sh2, "moe_expert_outputs_buf")?; + let moe_de_k_buf = alloc_f32(&stream, MOE_NUM_EXPERTS * b * moe_sh2, "moe_de_k_buf")?; + let moe_dg_buf = alloc_f32(&stream, b * MOE_NUM_EXPERTS, "moe_dg_buf")?; + let moe_dg_pre_buf = alloc_f32(&stream, b * MOE_NUM_EXPERTS, "moe_dg_pre_buf")?; + let moe_dh_s1_scratch = alloc_f32(&stream, b * moe_sh2, "moe_dh_s1_scratch")?; + let moe_gate_dh1_buf = alloc_f32(&stream, b * MOE_GATE_HIDDEN, "moe_gate_dh1_buf")?; + let moe_load_balance_loss_per_k = alloc_f32(&stream, MOE_NUM_EXPERTS, "moe_load_balance_loss_per_k")?; + let (moe_load_balance_loss_total_pinned, moe_load_balance_loss_total_dev_ptr) = { + let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut dev_ptr_out: u64 = 0; + unsafe { + let rc = cudarc::driver::sys::cuMemAllocHost_v2( + &mut host_ptr as *mut *mut std::ffi::c_void, + std::mem::size_of::(), + ); + if rc != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS { + return Err(MLError::ModelError(format!("cuMemAllocHost moe_lb_total: {:?}", rc))); + } + let rc2 = cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dev_ptr_out, + host_ptr, + 0, + ); + if rc2 != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS { + return Err(MLError::ModelError(format!("cuMemHostGetDevicePointer moe_lb_total: {:?}", rc2))); + } + *(host_ptr as *mut f32) = 0.0f32; + } + (host_ptr as *mut f32, dev_ptr_out) + }; + let moe_lambda = config.moe_lambda; + tracing::info!( + "GpuMoeHead initialized: K={} BTN={} GH={} SH2={} lambda={:.4} (params [127..163), ISV[118..127))", + MOE_NUM_EXPERTS, MOE_EXPERT_BOTTLENECK, MOE_GATE_HIDDEN, moe_sh2, moe_lambda, + ); + tracing::info!( "AuxHeadsForwardOps initialized: K_nb={} K_rg={} aux_h={} max_aux_tensor_len={} (params [119..127), ISV[113..115))", aux_knb, aux_kr, aux_h, max_aux_tensor_len, @@ -11766,6 +12172,21 @@ impl GpuDqnTrainer { ), last_distill_active: false, last_meta_q_pred: 0.5, + moe_head, + moe_gate_h1_buf, + moe_gate_pre_buf, + moe_gate_softmax_buf, + moe_expert_h1_bufs, + moe_expert_outputs_buf, + moe_de_k_buf, + moe_dg_buf, + moe_dg_pre_buf, + moe_dh_s1_scratch, + moe_gate_dh1_buf, + moe_load_balance_loss_per_k, + moe_load_balance_loss_total_pinned, + moe_load_balance_loss_total_dev_ptr, + moe_lambda, q_sample_history: std::collections::VecDeque::with_capacity(64), }; // Self-check: verify the fingerprint we just wrote is readable and matches @@ -15271,6 +15692,13 @@ impl GpuDqnTrainer { self.launch_loss_reduce(self.total_loss_dev_ptr)?; self.launch_c51_grad()?; + // ── Phase 3 T3.4: MoE load-balance auxiliary loss ──────────────── + // Computes λ·K·Σ_k(mean_b g[b,k])² and SAXPYs the scalar into + // total_loss_dev_ptr. Runs after C51 loss so it adds on top of the + // already-accumulated TD loss. Does NOT generate a direct gradient + // (the gradient flows through the gate backward in launch_moe_backward). + self.launch_moe_load_balance_loss()?; + // Blend: main = α * C51 + (1-α) * MSE { let alpha = self.c51_alpha; @@ -15908,6 +16336,14 @@ impl GpuDqnTrainer { self.ptrs.mag_concat_buf, )?; + // ── Phase 3 T3.1–T3.3: MoE forward ───────────────────────────────── + // Runs AFTER forward_online_raw (which produces save_h_s1 and the + // pre-MoE save_h_s2) and BEFORE aux_heads_forward (which reads + // save_h_s2). Overwrites save_h_s2 with the MoE mixture output + // so all downstream consumers (aux heads, stochastic depth, loss + // kernels) see the MoE-enriched representation. + self.launch_moe_forward()?; + // ── Plan 4 Task 6 Commit B: aux-heads forward (online-only) ── // Runs the next-bar regression + 5-class regime CE heads off the // freshly-saved `save_h_s2` activation. MUST run BEFORE stochastic @@ -17137,6 +17573,14 @@ impl GpuDqnTrainer { // augmented dh_s2 (= main + aux contributions). self.aux_heads_backward(grad_base)?; + // ── Phase 3 T3.6: MoE backward ───────────────────────────────── + // Runs AFTER aux_heads_backward (which SAXPYed into bw_d_h_s2) and + // BEFORE encoder_backward_chain. Computes dW/dB for MoE params + // [127..163) and accumulates expert dX → moe_dh_s1_scratch, then + // DtoD-copies moe_dh_s1_scratch into bw_d_h_s2 so the GRN encoder + // chain sees the MoE-routed gradient as its h_s2 input gradient. + self.launch_moe_backward(grad_base)?; + // ── GRN trunk backward (Plan 4 Task 2c.3c.4) ── // d_h_s2_ptr now holds the fully-accumulated gradient flowing into // the h_s2 GRN output. encoder_backward_chain unrolls both GRN diff --git a/crates/ml/src/cuda_pipeline/gpu_moe_head.rs b/crates/ml/src/cuda_pipeline/gpu_moe_head.rs index de81d44b1..55edadb7d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_moe_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_moe_head.rs @@ -21,7 +21,8 @@ static MOE_CUBIN: &[u8] = /// GPU-accelerated MoE head. /// /// Owns the kernel function handles for the MoE mixture forward/backward pass, -/// load-balance loss computation, and expert utilization EMA update. +/// load-balance loss computation, expert utilization EMA update, and softmax +/// forward/backward for the gate sub-network. #[allow(missing_debug_implementations)] pub struct GpuMoeHead { stream: Arc, @@ -31,6 +32,10 @@ pub struct GpuMoeHead { moe_load_balance_loss: CudaFunction, moe_load_balance_reduce: CudaFunction, moe_expert_util_ema_update: CudaFunction, + /// Per-row softmax for the gate pre-activation → gate probabilities. + moe_row_softmax: CudaFunction, + /// Softmax Jacobian backward for the gate. + moe_softmax_backward: CudaFunction, } impl GpuMoeHead { @@ -59,6 +64,12 @@ impl GpuMoeHead { let moe_expert_util_ema_update = module .load_function("moe_expert_util_ema_update") .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update load: {e}")))?; + let moe_row_softmax = module + .load_function("moe_row_softmax") + .map_err(|e| MLError::ModelError(format!("moe_row_softmax load: {e}")))?; + let moe_softmax_backward = module + .load_function("moe_softmax_backward") + .map_err(|e| MLError::ModelError(format!("moe_softmax_backward load: {e}")))?; Ok(Self { stream, moe_mixture_forward, @@ -67,9 +78,243 @@ impl GpuMoeHead { moe_load_balance_loss, moe_load_balance_reduce, moe_expert_util_ema_update, + moe_row_softmax, + moe_softmax_backward, }) } + // ══════════════════════════════════════════════════════════════════════ + // Production launch methods (operate on already-allocated device buffers) + // ══════════════════════════════════════════════════════════════════════ + + /// Launch `moe_row_softmax` in-place on already-allocated device pointers. + /// Grid: (B+block-1)/block × 1 × 1. Each thread handles one row. + pub(crate) fn launch_row_softmax( + &self, + in_pre_ptr: u64, + out_soft_ptr: u64, + b: usize, + k: usize, + ) -> Result<(), MLError> { + let block: u32 = 256; + let grid = ((b as u32) + block - 1) / block; + let b_i32 = b as i32; + let k_i32 = k as i32; + unsafe { + self.stream + .launch_builder(&self.moe_row_softmax) + .arg(&in_pre_ptr) + .arg(&out_soft_ptr) + .arg(&b_i32) + .arg(&k_i32) + .launch(LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_row_softmax launch: {e}")))?; + } + Ok(()) + } + + /// Launch `moe_mixture_forward` on device pointers. + /// h_s2_out[B,C] = sum_k gate[B,k] * expert_outputs[k,B,C]. + pub(crate) fn launch_mixture_forward( + &self, + expert_outputs_ptr: u64, + gate_ptr: u64, + h_s2_out_ptr: u64, + b: usize, + k: usize, + c: usize, + ) -> Result<(), MLError> { + let total = (b * c) as u32; + let block: u32 = 256; + let grid = total.div_ceil(block); + let b_i32 = b as i32; + let k_i32 = k as i32; + let c_i32 = c as i32; + unsafe { + self.stream + .launch_builder(&self.moe_mixture_forward) + .arg(&expert_outputs_ptr) + .arg(&gate_ptr) + .arg(&h_s2_out_ptr) + .arg(&b_i32) + .arg(&k_i32) + .arg(&c_i32) + .launch(LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_mixture_forward launch: {e}")))?; + } + Ok(()) + } + + /// Launch `moe_load_balance_loss` + `moe_load_balance_reduce`. + /// Writes the scalar total load-balance loss into `loss_total_ptr`. + pub(crate) fn launch_load_balance_loss( + &self, + gate_ptr: u64, + loss_per_k_ptr: u64, + loss_total_ptr: u64, + b: usize, + k: usize, + lambda: f32, + ) -> Result<(), MLError> { + let block: u32 = 256; + let b_i32 = b as i32; + let k_i32 = k as i32; + unsafe { + self.stream + .launch_builder(&self.moe_load_balance_loss) + .arg(&gate_ptr) + .arg(&loss_per_k_ptr) + .arg(&b_i32) + .arg(&k_i32) + .arg(&lambda) + .launch(LaunchConfig { + grid_dim: (k as u32, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: block * 4, + }) + .map_err(|e| MLError::ModelError(format!("moe_load_balance_loss launch: {e}")))?; + + self.stream + .launch_builder(&self.moe_load_balance_reduce) + .arg(&loss_per_k_ptr) + .arg(&loss_total_ptr) + .arg(&k_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_load_balance_reduce launch: {e}")))?; + } + Ok(()) + } + + /// Launch `moe_expert_util_ema_update` ISV producer on device pointers. + pub(crate) fn launch_expert_util_ema( + &self, + gate_ptr: u64, + isv_dev_ptr: u64, + b: usize, + k: usize, + isv_util_base: usize, + isv_entropy_index: usize, + alpha: f32, + ) -> Result<(), MLError> { + let b_i32 = b as i32; + let k_i32 = k as i32; + let util_base_i32 = isv_util_base as i32; + let entropy_idx_i32 = isv_entropy_index as i32; + unsafe { + self.stream + .launch_builder(&self.moe_expert_util_ema_update) + .arg(&gate_ptr) + .arg(&isv_dev_ptr) + .arg(&b_i32) + .arg(&k_i32) + .arg(&util_base_i32) + .arg(&entropy_idx_i32) + .arg(&alpha) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update launch: {e}")))?; + } + Ok(()) + } + + /// Launch `moe_mixture_backward` (de_k) + `moe_dgate_reduce` (dg) on device pointers. + pub(crate) fn launch_mixture_backward( + &self, + dh_s2_ptr: u64, + gate_ptr: u64, + expert_outputs_ptr: u64, + de_k_ptr: u64, + dg_ptr: u64, + b: usize, + k: usize, + c: usize, + ) -> Result<(), MLError> { + let block: u32 = 256; + let total_kbc = (k * b * c) as u32; + let grid_kbc = total_kbc.div_ceil(block); + let b_i32 = b as i32; + let k_i32 = k as i32; + let c_i32 = c as i32; + unsafe { + self.stream + .launch_builder(&self.moe_mixture_backward) + .arg(&dh_s2_ptr) + .arg(&gate_ptr) + .arg(&de_k_ptr) + .arg(&b_i32) + .arg(&k_i32) + .arg(&c_i32) + .launch(LaunchConfig { + grid_dim: (grid_kbc, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_mixture_backward launch: {e}")))?; + + self.stream + .launch_builder(&self.moe_dgate_reduce) + .arg(&dh_s2_ptr) + .arg(&expert_outputs_ptr) + .arg(&dg_ptr) + .arg(&b_i32) + .arg(&k_i32) + .arg(&c_i32) + .launch(LaunchConfig { + grid_dim: (b as u32, k as u32, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: block * 4, + }) + .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce launch: {e}")))?; + } + Ok(()) + } + + /// Launch `moe_softmax_backward` on device pointers. + pub(crate) fn launch_softmax_backward( + &self, + g_soft_ptr: u64, + dg_ptr: u64, + dg_pre_ptr: u64, + b: usize, + k: usize, + ) -> Result<(), MLError> { + let block: u32 = 256; + let grid = ((b as u32) + block - 1) / block; + let b_i32 = b as i32; + let k_i32 = k as i32; + unsafe { + self.stream + .launch_builder(&self.moe_softmax_backward) + .arg(&g_soft_ptr) + .arg(&dg_ptr) + .arg(&dg_pre_ptr) + .arg(&b_i32) + .arg(&k_i32) + .launch(LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_softmax_backward launch: {e}")))?; + } + Ok(()) + } + /// Test-only entry — stages CPU data via mapped pinned memory, runs the /// kernel, reads mapped pinned output. NO HtoD/HtoH per /// `feedback_no_htod_htoh_only_mapped_pinned.md`. diff --git a/crates/ml/src/cuda_pipeline/moe_kernels.cu b/crates/ml/src/cuda_pipeline/moe_kernels.cu index 4acf35fe1..d6e334987 100644 --- a/crates/ml/src/cuda_pipeline/moe_kernels.cu +++ b/crates/ml/src/cuda_pipeline/moe_kernels.cu @@ -137,6 +137,61 @@ extern "C" __global__ void moe_load_balance_reduce( __threadfence_system(); } +/* Row-wise softmax: out[b,k] = exp(in[b,k]) / sum_j exp(in[b,j]). + * Grid: one block per row (b). Single-thread per block — K is small (=8). + * No atomicAdd, no shared memory needed for K<=8 (all in registers). */ +extern "C" __global__ void moe_row_softmax( + const float* __restrict__ in_pre, /* [B, K] pre-softmax logits */ + float* __restrict__ out_soft, /* [B, K] softmax output */ + int B, + int K +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + /* Numerically stable: subtract row max before exponentiation. */ + float row_max = in_pre[b * K]; + for (int k = 1; k < K; ++k) { + float v = in_pre[b * K + k]; + if (v > row_max) row_max = v; + } + float sum_exp = 0.0f; + for (int k = 0; k < K; ++k) { + float e = expf(in_pre[b * K + k] - row_max); + out_soft[b * K + k] = e; + sum_exp += e; + } + float inv_sum = 1.0f / sum_exp; + for (int k = 0; k < K; ++k) { + out_soft[b * K + k] *= inv_sum; + } + __threadfence_system(); +} + +/* Softmax backward per row: + * dg_pre[b,k] = g_soft[b,k] * (dg[b,k] - sum_j g_soft[b,j] * dg[b,j]) + * Grid: one block per row (b). Single-thread per block — K is small (=8). */ +extern "C" __global__ void moe_softmax_backward( + const float* __restrict__ g_soft, /* [B, K] saved softmax output from forward */ + const float* __restrict__ dg, /* [B, K] upstream gradient w.r.t. gate */ + float* __restrict__ dg_pre, /* [B, K] gradient w.r.t. pre-softmax logits */ + int B, + int K +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + /* dot = sum_j g_soft[b,j] * dg[b,j] */ + float dot = 0.0f; + for (int k = 0; k < K; ++k) { + dot += g_soft[b * K + k] * dg[b * K + k]; + } + for (int k = 0; k < K; ++k) { + dg_pre[b * K + k] = g_soft[b * K + k] * (dg[b * K + k] - dot); + } + __threadfence_system(); +} + extern "C" __global__ void moe_expert_util_ema_update( const float* __restrict__ gate, /* [B, K] */ float* __restrict__ isv, /* [ISV_TOTAL_DIM] */ diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 50aca18c3..69df84fc1 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -415,6 +415,7 @@ impl FusedTrainingCtx { market_dim: 42, // Always 42 base market features — OFI features bypass bottleneck via portfolio_dim total_epochs: hyperparams.epochs, replay_seed_steps: hyperparams.replay_seed_steps, + moe_lambda: hyperparams.moe_lambda.unwrap_or(0.01), }; // Create weight set pointer views AFTER GpuDqnTrainer is constructed below. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 03a5a5ee1..0b675eb39 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2978,6 +2978,15 @@ impl DQNTrainer { tracing::warn!("Plan 4 Task 6 Commit B aux_heads_loss_ema launch failed: {e}"); } } + + // Phase 3 T3.5: MoE expert-utilisation EMA + gate entropy EMA. + // Reads `moe_gate_softmax_buf [B, K]` (valid after the captured + // forward graph ran) and EMA-updates ISV[118..127). + if let Some(ref fused) = self.fused_ctx { + if let Err(e) = fused.trainer().launch_moe_expert_util_ema(ema_alpha) { + tracing::warn!("Phase 3 T3.5 launch_moe_expert_util_ema failed: {e}"); + } + } } // Plan 4 Task 6 Commit B: refresh the aux-loss weight before the @@ -3305,6 +3314,31 @@ impl DQNTrainer { ); } + // Phase 3 T3.5: MoE expert-utilisation + gate entropy HEALTH_DIAG. + { + use crate::cuda_pipeline::gpu_dqn_trainer::{ + MOE_EXPERT_UTIL_EMA_BASE, MOE_GATE_ENTROPY_EMA_INDEX, + }; + let (moe_utils, moe_ent) = + if let Some(ref fused) = self.fused_ctx { + let trainer = fused.trainer(); + let utils: Vec = (0..8) + .map(|k| trainer.read_isv_signal_at(MOE_EXPERT_UTIL_EMA_BASE + k)) + .collect(); + let ent = trainer.read_isv_signal_at(MOE_GATE_ENTROPY_EMA_INDEX); + (utils, ent) + } else { + (vec![0.0f32; 8], 0.0f32) + }; + tracing::info!( + "HEALTH_DIAG[{}]: aux_moe [util={:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3} ent={:.3}]", + epoch, + moe_utils[0], moe_utils[1], moe_utils[2], moe_utils[3], + moe_utils[4], moe_utils[5], moe_utils[6], moe_utils[7], + moe_ent, + ); + } + // C1/P1: propagate health to GPU replay buffer for diversity-weighted priorities. { let mut agent = self.agent.write().await; diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index faa8984fc..33962db83 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,22 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +MoE Phase 3 wire-up T3.1–T3.7 (2026-04-27): MoE fully wired into +production training path. Forward: gate (state[B,128]→64→8→softmax) + +8 expert MLPs (h_s1[B,256]→64→256) + `moe_mixture_forward` → replaces +save_h_s2 in `submit_forward_ops_ddqn`. Load-balance: `moe_load_balance_loss` ++ `moe_load_balance_reduce` + SAXPY into total_loss_dev_ptr (T3.4). ISV +producer: `launch_moe_expert_util_ema` per-step in training_loop.rs (T3.5). +Backward: `moe_mixture_backward` (de_k=g·dh_s2) + `moe_dgate_reduce` +(dg=Σe·dh_s2) + `moe_softmax_backward` + cuBLAS SGEMM backward through gate +W2/W1 and 8 experts W2/W1 — all into params_buf grad slots [127..163) (T3.6). +Adam step inherits gate+expert updates automatically (params_buf uniform). +HEALTH_DIAG aux_moe line emitted per epoch from ISV[118..127) (T3.7). +Smoke test: 3/3 folds pass, 728s. Gate differentiated: expert-2 reached +32.3% utilization (others 9.7%) by fold-2 epoch-4; entropy 1.611 < ln(8). +fused_training.rs: added moe_lambda field to GpuDqnTrainConfig init from +hyperparams.moe_lambda.unwrap_or(0.01). + MoE expert util EMA T2.4 (2026-04-27): `moe_expert_util_ema_update` single-thread cold-path-cadence kernel writes 8 per-expert utilization EMAs (ISV[118..126)) + gate-entropy EMA (ISV[126]) with α=0.05. Same shape