diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a695c60e0..2c5efc93f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -54,6 +54,7 @@ use cudarc::cublaslt::result as cublaslt_result; use cudarc::cublaslt::sys as cublaslt_sys; use crate::MLError; +use ml_dqn::{MOE_EXPERT_BOTTLENECK, MOE_GATE_HIDDEN, MOE_NUM_EXPERTS}; use super::gpu_attention::GpuAttention; use super::gpu_weights::{DuelingWeightSet, BranchingWeightSet}; use super::batched_forward::{CublasForward, CublasGemmSet, f32_weight_ptrs_from_base}; @@ -977,7 +978,16 @@ const fn layout_fingerprint_seed() -> &'static [u8] { PARAM_VSN_W1_G5=115;PARAM_VSN_B1_G5=116;PARAM_VSN_W2_G5=117;PARAM_VSN_B2_G5=118;\ PARAM_AUX_NB_W1=119;PARAM_AUX_NB_B1=120;PARAM_AUX_NB_W2=121;PARAM_AUX_NB_B2=122;\ PARAM_AUX_RG_W1=123;PARAM_AUX_RG_B1=124;PARAM_AUX_RG_W2=125;PARAM_AUX_RG_B2=126;\ - PARAM_TOTAL_TENSORS=127" + PARAM_MOE_GATE_W1=127;PARAM_MOE_GATE_B1=128;PARAM_MOE_GATE_W2=129;PARAM_MOE_GATE_B2=130;\ + PARAM_MOE_EXPERT_0_W1=131;PARAM_MOE_EXPERT_0_B1=132;PARAM_MOE_EXPERT_0_W2=133;PARAM_MOE_EXPERT_0_B2=134;\ + PARAM_MOE_EXPERT_1_W1=135;PARAM_MOE_EXPERT_1_B1=136;PARAM_MOE_EXPERT_1_W2=137;PARAM_MOE_EXPERT_1_B2=138;\ + PARAM_MOE_EXPERT_2_W1=139;PARAM_MOE_EXPERT_2_B1=140;PARAM_MOE_EXPERT_2_W2=141;PARAM_MOE_EXPERT_2_B2=142;\ + PARAM_MOE_EXPERT_3_W1=143;PARAM_MOE_EXPERT_3_B1=144;PARAM_MOE_EXPERT_3_W2=145;PARAM_MOE_EXPERT_3_B2=146;\ + PARAM_MOE_EXPERT_4_W1=147;PARAM_MOE_EXPERT_4_B1=148;PARAM_MOE_EXPERT_4_W2=149;PARAM_MOE_EXPERT_4_B2=150;\ + PARAM_MOE_EXPERT_5_W1=151;PARAM_MOE_EXPERT_5_B1=152;PARAM_MOE_EXPERT_5_W2=153;PARAM_MOE_EXPERT_5_B2=154;\ + PARAM_MOE_EXPERT_6_W1=155;PARAM_MOE_EXPERT_6_B1=156;PARAM_MOE_EXPERT_6_W2=157;PARAM_MOE_EXPERT_6_B2=158;\ + PARAM_MOE_EXPERT_7_W1=159;PARAM_MOE_EXPERT_7_B1=160;PARAM_MOE_EXPERT_7_W2=161;PARAM_MOE_EXPERT_7_B2=162;\ + PARAM_TOTAL_TENSORS=163" } /// Compile-time layout fingerprint. Burned into the binary at build time. @@ -1387,7 +1397,11 @@ impl Default for CausalInterventionConfig { /// tensors will receive zero-gradient SAXPYs each step until Commit B /// activates the backward — Adam keeps them at Xavier init meanwhile, /// which is the intended dormant state. -pub(crate) const NUM_WEIGHT_TENSORS: usize = 127; +/// + 36 MoE params (Phase 1 Task 1.6): 4 gate tensors + 8 experts × 4 tensors +/// → indices [127..163). Gate: gate_w1[SD,64], gate_b1[64], gate_w2[64,8], +/// gate_b2[8]. Per expert: w1[SH2,64], b1[64], w2[64,SH2], b2[SH2]. +/// Allocated but not yet wired into forward/backward — Phase 3 wire-up. +pub(crate) const NUM_WEIGHT_TENSORS: usize = 163; /// Compute the size (element count) of each weight tensor. /// @@ -1581,8 +1595,56 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT sizes[125] = crate::cuda_pipeline::gpu_aux_heads::AUX_REGIME_K * crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; // [125] aux_rg_w2 [5, 32] sizes[126] = crate::cuda_pipeline::gpu_aux_heads::AUX_REGIME_K; // [126] aux_rg_b2 [5] - debug_assert!(NUM_WEIGHT_TENSORS == 127, - "compute_param_sizes: NUM_WEIGHT_TENSORS expected 127 after aux-heads, got {}", + debug_assert!( + { + // VSN fill should have ended at 119 and aux-heads fill indices [119..127). + // Check the last aux-head slot was written before proceeding to MoE. + sizes[126] > 0 || crate::cuda_pipeline::gpu_aux_heads::AUX_REGIME_K == 0 + }, + "compute_param_sizes: aux-head tail sanity failed before MoE extension" + ); + + // ── Phase 1 Task 1.6: MoE params (36 tensors, indices [127..163)) ───────── + // Gate subnetwork (4 tensors at [127..131)): + // gate_w1 [STATE_DIM, MOE_GATE_HIDDEN] — zero-init so g(s) = uniform 1/K initially + // gate_b1 [MOE_GATE_HIDDEN] + // gate_w2 [MOE_GATE_HIDDEN, MOE_NUM_EXPERTS] + // gate_b2 [MOE_NUM_EXPERTS] + // 8 experts × 4 tensors (4 per expert at [131..163)): + // expert_k_w1 [cfg.shared_h2, MOE_EXPERT_BOTTLENECK] — Xavier init + // expert_k_b1 [MOE_EXPERT_BOTTLENECK] + // expert_k_w2 [MOE_EXPERT_BOTTLENECK, cfg.shared_h2] + // expert_k_b2 [cfg.shared_h2] + // Not yet wired into forward/backward — Phase 3 wire-up. + // Adam SAXPY iterates 0..NUM_WEIGHT_TENSORS so these tensors receive + // zero-gradient SAXPYs until Phase 3 activates the backward chain, + // keeping them at init state meanwhile (correct dormant behaviour). + let sd = ml_core::state_layout::STATE_DIM; + // Gate tensors [127..131) + sizes[127] = sd * MOE_GATE_HIDDEN; // gate_w1 [SD, MOE_GATE_HIDDEN] + sizes[128] = MOE_GATE_HIDDEN; // gate_b1 [MOE_GATE_HIDDEN] + sizes[129] = MOE_GATE_HIDDEN * MOE_NUM_EXPERTS; // gate_w2 [MOE_GATE_HIDDEN, K] + sizes[130] = MOE_NUM_EXPERTS; // gate_b2 [MOE_NUM_EXPERTS] + // Expert tensors [131..163) — 4 per expert, 8 experts + { + let sh2 = cfg.shared_h2; + let btn = MOE_EXPERT_BOTTLENECK; + let mut eidx = 131usize; + let mut k = 0usize; + while k < MOE_NUM_EXPERTS { + sizes[eidx] = sh2 * btn; // expert_k_w1 [SH2, BTN] + sizes[eidx + 1] = btn; // expert_k_b1 [BTN] + sizes[eidx + 2] = btn * sh2; // expert_k_w2 [BTN, SH2] + sizes[eidx + 3] = sh2; // expert_k_b2 [SH2] + eidx += 4; + k += 1; + } + debug_assert!(eidx == 163, + "compute_param_sizes: MoE expert fill ended at {} but expected 163", + eidx); + } + debug_assert!(NUM_WEIGHT_TENSORS == 163, + "compute_param_sizes: NUM_WEIGHT_TENSORS expected 163 after MoE extension, got {}", NUM_WEIGHT_TENSORS); sizes @@ -17740,6 +17802,32 @@ impl GpuDqnTrainer { fan_dims[126] = (0, 0); // aux_rg_b2 (zero) } + // ── Phase 1 Task 1.6: MoE fan dims (36 tensors at [127..163)) ───────── + // Gate: zero-init (fan_out=0,fan_in=0) so g(s) = uniform 1/K at cold start. + // Experts: Xavier init on w1/w2 for initial differentiation; zero on biases. + { + let sd = ml_core::state_layout::STATE_DIM; + let btn = MOE_EXPERT_BOTTLENECK; + let sh2 = cfg.shared_h2; + // Gate tensors [127..131) — zero-init: softmax of zero-logits = uniform 1/K + fan_dims[127] = (0, 0); // gate_w1 [SD, MOE_GATE_HIDDEN] (zero-init per spec) + fan_dims[128] = (0, 0); // gate_b1 [MOE_GATE_HIDDEN] + fan_dims[129] = (0, 0); // gate_w2 [MOE_GATE_HIDDEN, K] (zero-init per spec) + fan_dims[130] = (0, 0); // gate_b2 [MOE_NUM_EXPERTS] + // Expert tensors [131..163) — Xavier on w1/w2, zero on b1/b2 + let _ = sd; // used indirectly via gate tensors above + let mut eidx = 131usize; + let mut k = 0usize; + while k < MOE_NUM_EXPERTS { + fan_dims[eidx] = (btn, sh2); // expert_k_w1 [SH2→BTN] (Xavier fan_out=BTN,fan_in=SH2) + fan_dims[eidx + 1] = (0, 0); // expert_k_b1 [BTN] (zero) + fan_dims[eidx + 2] = (sh2, btn); // expert_k_w2 [BTN→SH2] (Xavier fan_out=SH2,fan_in=BTN) + fan_dims[eidx + 3] = (0, 0); // expert_k_b2 [SH2] (zero) + eidx += 4; + k += 1; + } + } + // Build flat host buffer: Xavier init for weights, zeros for biases + padding. let cutlass_tile_pad = 32 * cfg.adv_h.max(cfg.value_h); let buf_len = total + cutlass_tile_pad; diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 23a11e27e..7f36720fd 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -767,3 +767,14 @@ DQN struct gains optional `moe_gate: Option` and `moe_experts: Option>` fields, initialized to `None` — actual wire-up into the forward graph in Phase 3. --> +