Files
foxhunt/docs
jgrusewski a52d996135 feat(moe): wire MoE forward + backward + load-balance + monitoring
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) <noreply@anthropic.com>
2026-04-27 19:47:19 +02:00
..