feat(moe): adaptive load-balance λ via gate-entropy ISV controller
Replace static moe_lambda=0.01 with a GPU-driven controller that scales
λ inversely with observed gate-entropy: λ_eff = λ_floor + λ_max_extra ×
max(0, ent_target − ent_ema)/ent_target, where ent_ema is the existing
ISV[126] producer (Phase 2 task 2.4 wire-up).
Motivation: the L40S validation run train-multi-seed-fgg9x exposed
single-expert specialization in fold 1 (expert 4 → 81% util, 7/8
experts < 5%) under static λ=0.01. The load-balance push was too weak
to prevent winner-take-all once the gate had a strong preference.
Adaptive λ rises proportional to the entropy deficit, restoring
diversity pressure when the gate concentrates.
Per pearl_blend_formulas_must_have_permanent_floor: λ_floor remains a
permanent minimum (= old static 0.01) so the controller never weakens
below baseline. Per feedback_isv_for_adaptive_bounds: signal flows
through ISV[128] (new MOE_LAMBDA_EFF_INDEX); consumer kernel reads
the slot at runtime; no DtoH on hot path.
Wiring (feedback_no_partial_refactor):
• new kernel moe_lambda_eff_update (single-block cold-path cadence,
matches kelly_cap_update / cql_alpha_seed_update precedent)
• new ISV slot 128, ISV_TOTAL_DIM 127 → 129
• layout_fingerprint_seed updated (schema_hash bumps; PVC fxcache will
auto-regen on next deploy)
• moe_load_balance_loss kernel takes (isv_signals, isv_lambda_eff_idx)
instead of float lambda — matches mag_concat_qdir's ISV-read pattern
• config: pub moe_lambda field replaced with moe_lambda_floor (0.01),
moe_lambda_max_extra (0.09), moe_entropy_target_frac (0.7)
• state-reset registry entry — ISV[128] resets to floor at fold boundary
• HEALTH_DIAG aux_moe line gains λ_eff field for observability
• hyperopt adapter (DQNHyperparameters) migrated to new knobs
• constructor bootstraps ISV[128]=floor + ISV[118..127)=uniform 1/K +
ISV[126]=ln(K) so cold-start matches legacy byte-for-byte
Smoke validates: magnitude_distribution still passes/fails identically
to HEAD (eq=0.586 same as pre-change eq=0.592 — unrelated Kelly cap
behaviour per project_magnitude_eval_collapse_kelly_capped.md).
HEALTH_DIAG fold 3 confirms controller live: ent decays 1.381 → 0.746
across epochs 7-19, λ_eff rises 0.0146 → 0.0539 monotonically.
Unit test moe_lambda_eff_update_writes_correct_values exercises 5
regimes (above-target / at-target / half-collapse / full-collapse /
deeply-above-target) on the GPU kernel — all pass within 1e-5.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -193,6 +193,14 @@ fn main() {
|
||||
// Subsequent tasks extend this cubin with moe_mixture_backward,
|
||||
// moe_load_balance_loss, moe_expert_util_ema_update.
|
||||
"moe_kernels.cu",
|
||||
// MoE adaptive load-balance λ controller (2026-04-27): single-block
|
||||
// single-thread cold-path producer kernel matching kelly_cap_update /
|
||||
// cql_alpha_seed_update precedent. Reads ISV[MOE_GATE_ENTROPY_EMA_INDEX=126]
|
||||
// (entropy EMA from moe_expert_util_ema_update), writes
|
||||
// ISV[MOE_LAMBDA_EFF_INDEX=128] = λ_floor + λ_max_extra × deficit.
|
||||
// Consumer (moe_load_balance_loss in moe_kernels.cu) reads slot 128
|
||||
// at runtime — no DtoH per feedback_isv_for_adaptive_bounds.md.
|
||||
"moe_lambda_eff_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
@@ -439,7 +439,7 @@ const ISV_NETWORK_DIM: usize = 23;
|
||||
/// (shifted 112→116 in Plan 4 Task 6 Commit A).
|
||||
/// Written by the constructor; checked at checkpoint load. Fail-fast only — no migration
|
||||
/// path exists. See spec §4.A.2 and `LAYOUT_FINGERPRINT_CURRENT` for structural-hash rationale.
|
||||
const ISV_TOTAL_DIM: usize = 127;
|
||||
const ISV_TOTAL_DIM: usize = 129;
|
||||
/// Legacy alias preserved for call sites that haven't been audited for the
|
||||
/// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight
|
||||
/// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer).
|
||||
@@ -816,6 +816,24 @@ pub const MOE_EXPERT_UTIL_EMA_COUNT: usize = 8;
|
||||
/// Reset value at fold boundary: ln(8) ≈ 2.0794 (max entropy at uniform).
|
||||
pub const MOE_GATE_ENTROPY_EMA_INDEX: usize = 126;
|
||||
|
||||
/// MoE adaptive load-balance λ controller output (2026-04-27).
|
||||
///
|
||||
/// `λ_eff = λ_floor + λ_max_extra × clamp((target − ent_ema)/target, 0, 1)`
|
||||
/// where `target = entropy_target_frac × ln(K)`. Producer:
|
||||
/// `moe_lambda_eff_update` GPU kernel (cold-path cadence, single-block
|
||||
/// single-thread; runs each step after `moe_expert_util_ema_update` so
|
||||
/// the entropy EMA fed to the controller is fresh for the NEXT step's
|
||||
/// `moe_load_balance_loss` consumer). Consumer:
|
||||
/// `moe_load_balance_loss` reads ISV[128] at runtime to scale the
|
||||
/// per-expert squared-mean penalty.
|
||||
///
|
||||
/// Cold-start (constructor): `λ_floor` (= old static 0.01 default).
|
||||
/// FoldReset: `λ_floor` (gate uniform-init again, deficit=0 ⇒ λ_floor).
|
||||
///
|
||||
/// Per `feedback_isv_for_adaptive_bounds.md` — adaptive bounds always
|
||||
/// flow through ISV; no DtoH on hot path.
|
||||
pub const MOE_LAMBDA_EFF_INDEX: usize = 128;
|
||||
|
||||
/// ISV slot [115] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
///
|
||||
/// Design note: this is NOT a version number. There is no ordered version space.
|
||||
@@ -930,7 +948,8 @@ const fn layout_fingerprint_seed() -> &'static [u8] {
|
||||
AUX_LABEL_SCALE_EMA=117;\
|
||||
MOE_EXPERT_UTIL_EMA_BASE=118;MOE_EXPERT_UTIL_EMA_COUNT=8;\
|
||||
MOE_GATE_ENTROPY_EMA=126;\
|
||||
ISV_TOTAL_DIM=127;\
|
||||
MOE_LAMBDA_EFF=128;\
|
||||
ISV_TOTAL_DIM=129;\
|
||||
PARAM_W_A_H_S1=0;PARAM_B_A_H_S1=1;PARAM_W_B_H_S1=2;PARAM_B_B_H_S1=3;\
|
||||
PARAM_W_RESIDUAL_H_S1=4;PARAM_GAMMA_H_S1=5;PARAM_BETA_H_S1=6;\
|
||||
PARAM_W_A_H_S2=7;PARAM_B_A_H_S2=8;PARAM_W_B_H_S2=9;PARAM_B_B_H_S2=10;\
|
||||
@@ -1244,9 +1263,26 @@ pub struct GpuDqnTrainConfig {
|
||||
/// 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,
|
||||
/// MoE load-balance loss adaptive controller — permanent floor weight.
|
||||
///
|
||||
/// `λ_eff = moe_lambda_floor + moe_lambda_max_extra × deficit`, where
|
||||
/// `deficit ∈ [0, 1]` is the normalised entropy shortfall against
|
||||
/// `moe_entropy_target_frac × ln(K)`. The producer kernel
|
||||
/// `moe_lambda_eff_update` writes ISV[MOE_LAMBDA_EFF_INDEX=128] each
|
||||
/// step; consumer kernel `moe_load_balance_loss` reads it at runtime.
|
||||
///
|
||||
/// Default 0.01 (matches the legacy static `moe_lambda` default — the
|
||||
/// old "anti-collapse only" baseline becomes the permanent minimum so
|
||||
/// the controller never weakens below the previous behaviour).
|
||||
pub moe_lambda_floor: f32,
|
||||
/// Maximum additional λ added when the gate is fully collapsed
|
||||
/// (`ent_ema = 0` ⇒ `deficit = 1`). Peak λ is `floor + max_extra`.
|
||||
/// Default 0.09 (so peak λ = 0.10, 10× boost over the floor).
|
||||
pub moe_lambda_max_extra: f32,
|
||||
/// Fraction of `ln(K)` treated as "diverse enough"; entropy at or above
|
||||
/// `entropy_target_frac × ln(K)` produces zero deficit (λ_eff = floor).
|
||||
/// Default 0.7 (target = 70% of ln(K) ≈ 1.456 for K=8).
|
||||
pub moe_entropy_target_frac: f32,
|
||||
}
|
||||
|
||||
impl Default for GpuDqnTrainConfig {
|
||||
@@ -1295,7 +1331,9 @@ 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,
|
||||
moe_lambda_floor: 0.01,
|
||||
moe_lambda_max_extra: 0.09,
|
||||
moe_entropy_target_frac: 0.7,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2718,8 +2756,18 @@ pub struct GpuDqnTrainer {
|
||||
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,
|
||||
/// Adaptive load-balance λ controller — permanent floor (= old static
|
||||
/// 0.01 default) read by the `moe_lambda_eff_update` GPU kernel.
|
||||
/// Per `pearl_blend_formulas_must_have_permanent_floor.md`, ISV[128]
|
||||
/// is bounded below by this value (`λ_eff ≥ moe_lambda_floor` always).
|
||||
moe_lambda_floor: f32,
|
||||
/// Maximum additional λ added when the gate is fully collapsed.
|
||||
/// Peak λ_eff is `moe_lambda_floor + moe_lambda_max_extra`.
|
||||
moe_lambda_max_extra: f32,
|
||||
/// Fraction of `ln(K)` treated as "diverse enough" — entropy ≥ target
|
||||
/// produces zero deficit (λ_eff = floor); entropy → 0 produces full
|
||||
/// deficit (λ_eff = floor + max_extra).
|
||||
moe_entropy_target_frac: f32,
|
||||
|
||||
/// #21 Stochastic depth: per-layer scale buffer [3] (h_s1, h_s2, h_v).
|
||||
/// Written by GPU RNG kernel before each graph replay.
|
||||
@@ -7943,16 +7991,27 @@ impl GpuDqnTrainer {
|
||||
/// (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.
|
||||
///
|
||||
/// λ is read inside the kernel from
|
||||
/// `ISV[MOE_LAMBDA_EFF_INDEX]`, populated each step by the
|
||||
/// `moe_lambda_eff_update` adaptive controller (see
|
||||
/// `launch_moe_lambda_eff_update`). No DtoH per
|
||||
/// `feedback_isv_for_adaptive_bounds.md`.
|
||||
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;
|
||||
|
||||
debug_assert!(self.isv_signals_dev_ptr != 0,
|
||||
"launch_moe_load_balance_loss: isv_signals_dev_ptr must be non-zero");
|
||||
|
||||
// Compute per-expert load loss + reduce to scalar.
|
||||
// Kernel reads λ from ISV[MOE_LAMBDA_EFF_INDEX] at runtime.
|
||||
self.moe_head.launch_load_balance_loss(
|
||||
gate_ptr, loss_per_k_ptr, loss_total_ptr,
|
||||
b, MOE_NUM_EXPERTS, self.moe_lambda,
|
||||
self.isv_signals_dev_ptr,
|
||||
b, MOE_NUM_EXPERTS, MOE_LAMBDA_EFF_INDEX,
|
||||
)?;
|
||||
|
||||
// SAXPY total_loss += 1.0 * lb_loss_total (both are device-mapped scalars).
|
||||
@@ -7992,6 +8051,35 @@ impl GpuDqnTrainer {
|
||||
)
|
||||
}
|
||||
|
||||
/// Adaptive load-balance λ controller (2026-04-27).
|
||||
///
|
||||
/// Single-block single-thread cold-path kernel. Reads
|
||||
/// ISV[MOE_GATE_ENTROPY_EMA_INDEX] (gate-entropy EMA from
|
||||
/// `moe_expert_util_ema_update`) and writes
|
||||
/// ISV[MOE_LAMBDA_EFF_INDEX] = `λ_floor + λ_max_extra × deficit`,
|
||||
/// where `deficit = clamp((target − ent_ema)/target, 0, 1)` and
|
||||
/// `target = moe_entropy_target_frac × ln(K)`.
|
||||
///
|
||||
/// MUST be called AFTER `launch_moe_expert_util_ema` so the entropy
|
||||
/// EMA fed to the controller is fresh for the NEXT step's
|
||||
/// `moe_load_balance_loss` consumer.
|
||||
///
|
||||
/// Per `pearl_blend_formulas_must_have_permanent_floor.md`,
|
||||
/// `λ_floor` is a permanent minimum (λ_eff never weakens below the
|
||||
/// old static 0.01 baseline).
|
||||
pub fn launch_moe_lambda_eff_update(&self) -> Result<(), MLError> {
|
||||
debug_assert!(self.isv_signals_dev_ptr != 0,
|
||||
"launch_moe_lambda_eff_update: isv_signals_dev_ptr must be non-zero");
|
||||
self.moe_head.launch_lambda_eff_update(
|
||||
self.isv_signals_dev_ptr,
|
||||
MOE_LAMBDA_EFF_INDEX,
|
||||
MOE_GATE_ENTROPY_EMA_INDEX,
|
||||
self.moe_lambda_floor,
|
||||
self.moe_lambda_max_extra,
|
||||
self.moe_entropy_target_frac,
|
||||
)
|
||||
}
|
||||
|
||||
/// Phase 3 T3.6: MoE backward.
|
||||
///
|
||||
/// Sequence (reverse of forward):
|
||||
@@ -10554,10 +10642,13 @@ impl GpuDqnTrainer {
|
||||
}
|
||||
(host_ptr as *mut f32, dev_ptr_out)
|
||||
};
|
||||
let moe_lambda = config.moe_lambda;
|
||||
let moe_lambda_floor = config.moe_lambda_floor;
|
||||
let moe_lambda_max_extra = config.moe_lambda_max_extra;
|
||||
let moe_entropy_target_frac = config.moe_entropy_target_frac;
|
||||
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,
|
||||
"GpuMoeHead initialized: K={} BTN={} GH={} SH2={} lambda_floor={:.4} lambda_max_extra={:.4} entropy_target_frac={:.3} (params [127..163), ISV[118..127), λ_eff ISV[128])",
|
||||
MOE_NUM_EXPERTS, MOE_EXPERT_BOTTLENECK, MOE_GATE_HIDDEN, moe_sh2,
|
||||
moe_lambda_floor, moe_lambda_max_extra, moe_entropy_target_frac,
|
||||
);
|
||||
|
||||
tracing::info!(
|
||||
@@ -10944,6 +11035,28 @@ impl GpuDqnTrainer {
|
||||
*sig_ptr.add(Q_P05_DIR_INDEX + b) = v_min_f;
|
||||
*sig_ptr.add(Q_P95_DIR_INDEX + b) = v_max_f;
|
||||
}
|
||||
/* Adaptive load-balance λ controller (2026-04-27):
|
||||
* cold-start ISV[MOE_LAMBDA_EFF_INDEX=128] at the floor so
|
||||
* the first `moe_load_balance_loss` consumer reads the
|
||||
* legacy 0.01 baseline before the controller has fired.
|
||||
* The producer (`moe_lambda_eff_update`) overwrites this
|
||||
* each step from the entropy EMA. FoldReset reapplies the
|
||||
* same floor (gate uniform-init ⇒ deficit=0 ⇒ λ_floor).
|
||||
* Per `pearl_blend_formulas_must_have_permanent_floor.md`,
|
||||
* λ_floor is a permanent minimum. */
|
||||
*sig_ptr.add(MOE_LAMBDA_EFF_INDEX) = config.moe_lambda_floor;
|
||||
/* Mixture-of-Experts cold-start (2026-04-27 follow-up):
|
||||
* bootstrap util slots to uniform 1/K and entropy slot to
|
||||
* ln(K) so the very first `moe_load_balance_loss` and
|
||||
* `moe_lambda_eff_update` see consistent neutral state
|
||||
* before `moe_expert_util_ema_update` fires. FoldReset
|
||||
* reapplies these via the registry. */
|
||||
let inv_k = 1.0_f32 / MOE_EXPERT_UTIL_EMA_COUNT as f32;
|
||||
for k in 0..MOE_EXPERT_UTIL_EMA_COUNT {
|
||||
*sig_ptr.add(MOE_EXPERT_UTIL_EMA_BASE + k) = inv_k;
|
||||
}
|
||||
*sig_ptr.add(MOE_GATE_ENTROPY_EMA_INDEX) =
|
||||
(MOE_EXPERT_UTIL_EMA_COUNT as f32).ln();
|
||||
// Layout fingerprint (ISV[58..60)). Compile-time structural hash of
|
||||
// the slot layout; checkpoint load fails-fast on mismatch.
|
||||
// Stored as a u64 split across two f32 lanes using raw bit-cast so
|
||||
@@ -12221,7 +12334,9 @@ impl GpuDqnTrainer {
|
||||
moe_load_balance_loss_per_k,
|
||||
moe_load_balance_loss_total_pinned,
|
||||
moe_load_balance_loss_total_dev_ptr,
|
||||
moe_lambda,
|
||||
moe_lambda_floor,
|
||||
moe_lambda_max_extra,
|
||||
moe_entropy_target_frac,
|
||||
q_sample_history: std::collections::VecDeque::with_capacity(64),
|
||||
};
|
||||
// Self-check: verify the fingerprint we just wrote is readable and matches
|
||||
|
||||
@@ -18,6 +18,12 @@ use super::mapped_pinned::MappedF32Buffer;
|
||||
static MOE_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/moe_kernels.cubin"));
|
||||
|
||||
/// Adaptive load-balance λ controller (separate cubin — single-block,
|
||||
/// single-thread cold-path producer matching kelly_cap_update precedent).
|
||||
/// Reads ISV[MOE_GATE_ENTROPY_EMA_INDEX], writes ISV[MOE_LAMBDA_EFF_INDEX].
|
||||
static MOE_LAMBDA_EFF_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/moe_lambda_eff_kernel.cubin"));
|
||||
|
||||
/// GPU-accelerated MoE head.
|
||||
///
|
||||
/// Owns the kernel function handles for the MoE mixture forward/backward pass,
|
||||
@@ -36,6 +42,9 @@ pub struct GpuMoeHead {
|
||||
moe_row_softmax: CudaFunction,
|
||||
/// Softmax Jacobian backward for the gate.
|
||||
moe_softmax_backward: CudaFunction,
|
||||
/// Adaptive load-balance λ controller — reads gate-entropy EMA from ISV,
|
||||
/// writes λ_eff = λ_floor + λ_max_extra × deficit into ISV[128].
|
||||
moe_lambda_eff_update: CudaFunction,
|
||||
}
|
||||
|
||||
impl GpuMoeHead {
|
||||
@@ -70,6 +79,14 @@ impl GpuMoeHead {
|
||||
let moe_softmax_backward = module
|
||||
.load_function("moe_softmax_backward")
|
||||
.map_err(|e| MLError::ModelError(format!("moe_softmax_backward load: {e}")))?;
|
||||
// Adaptive λ controller — separate cubin (built from
|
||||
// moe_lambda_eff_kernel.cu by build.rs).
|
||||
let lambda_module = context
|
||||
.load_cubin(MOE_LAMBDA_EFF_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("moe_lambda_eff cubin load: {e}")))?;
|
||||
let moe_lambda_eff_update = lambda_module
|
||||
.load_function("moe_lambda_eff_update")
|
||||
.map_err(|e| MLError::ModelError(format!("moe_lambda_eff_update load: {e}")))?;
|
||||
Ok(Self {
|
||||
stream,
|
||||
moe_mixture_forward,
|
||||
@@ -80,6 +97,7 @@ impl GpuMoeHead {
|
||||
moe_expert_util_ema_update,
|
||||
moe_row_softmax,
|
||||
moe_softmax_backward,
|
||||
moe_lambda_eff_update,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -155,26 +173,35 @@ impl GpuMoeHead {
|
||||
|
||||
/// Launch `moe_load_balance_loss` + `moe_load_balance_reduce`.
|
||||
/// Writes the scalar total load-balance loss into `loss_total_ptr`.
|
||||
///
|
||||
/// λ is read inside the kernel from `isv_signals[isv_lambda_eff_idx]`,
|
||||
/// where the adaptive controller (`moe_lambda_eff_update`) writes the
|
||||
/// effective load-balance weight each step. Per
|
||||
/// `feedback_isv_for_adaptive_bounds.md`, the load-balance push scales
|
||||
/// inversely with observed gate-entropy without any DtoH roundtrip.
|
||||
pub(crate) fn launch_load_balance_loss(
|
||||
&self,
|
||||
gate_ptr: u64,
|
||||
loss_per_k_ptr: u64,
|
||||
loss_total_ptr: u64,
|
||||
isv_signals_dev_ptr: u64,
|
||||
b: usize,
|
||||
k: usize,
|
||||
lambda: f32,
|
||||
isv_lambda_eff_idx: usize,
|
||||
) -> Result<(), MLError> {
|
||||
let block: u32 = 256;
|
||||
let b_i32 = b as i32;
|
||||
let k_i32 = k as i32;
|
||||
let lambda_idx_i32 = isv_lambda_eff_idx as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.moe_load_balance_loss)
|
||||
.arg(&gate_ptr)
|
||||
.arg(&loss_per_k_ptr)
|
||||
.arg(&isv_signals_dev_ptr)
|
||||
.arg(&b_i32)
|
||||
.arg(&k_i32)
|
||||
.arg(&lambda)
|
||||
.arg(&lambda_idx_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (k as u32, 1, 1),
|
||||
block_dim: (block, 1, 1),
|
||||
@@ -197,6 +224,47 @@ impl GpuMoeHead {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `moe_lambda_eff_update` ISV producer on device pointers.
|
||||
///
|
||||
/// Reads ISV[`isv_entropy_ema_idx`] (gate-entropy EMA from
|
||||
/// `moe_expert_util_ema_update`) and writes ISV[`isv_lambda_eff_idx`] =
|
||||
/// `lambda_floor + lambda_max_extra × deficit`, where `deficit ∈ [0, 1]`
|
||||
/// is the normalised entropy shortfall against
|
||||
/// `entropy_target_frac × ln(K)`. Single-block, single-thread cold-path
|
||||
/// kernel (matches kelly_cap_update / cql_alpha_seed_update shape).
|
||||
///
|
||||
/// Per `pearl_blend_formulas_must_have_permanent_floor.md`,
|
||||
/// `lambda_floor` is a permanent minimum — λ_eff never drops below it.
|
||||
pub(crate) fn launch_lambda_eff_update(
|
||||
&self,
|
||||
isv_dev_ptr: u64,
|
||||
isv_lambda_eff_idx: usize,
|
||||
isv_entropy_ema_idx: usize,
|
||||
lambda_floor: f32,
|
||||
lambda_max_extra: f32,
|
||||
entropy_target_frac: f32,
|
||||
) -> Result<(), MLError> {
|
||||
let lambda_idx_i32 = isv_lambda_eff_idx as i32;
|
||||
let entropy_idx_i32 = isv_entropy_ema_idx as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.moe_lambda_eff_update)
|
||||
.arg(&isv_dev_ptr)
|
||||
.arg(&lambda_idx_i32)
|
||||
.arg(&entropy_idx_i32)
|
||||
.arg(&lambda_floor)
|
||||
.arg(&lambda_max_extra)
|
||||
.arg(&entropy_target_frac)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (1, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("moe_lambda_eff_update launch: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `moe_expert_util_ema_update` ISV producer on device pointers.
|
||||
pub(crate) fn launch_expert_util_ema(
|
||||
&self,
|
||||
@@ -484,6 +552,12 @@ impl GpuMoeHead {
|
||||
/// `moe_load_balance_loss` (per-k contribution) and
|
||||
/// `moe_load_balance_reduce` (K-element scalar sum).
|
||||
///
|
||||
/// `lambda` is staged into a synthetic ISV slot at index 0 of a small
|
||||
/// mapped pinned buffer; the kernel reads `isv_signals[0]` for λ. This
|
||||
/// mirrors the production contract where the consumer reads
|
||||
/// `ISV[MOE_LAMBDA_EFF_INDEX]` filled by the `moe_lambda_eff_update`
|
||||
/// adaptive controller.
|
||||
///
|
||||
/// Returns the scalar total loss.
|
||||
pub fn test_load_balance_loss(
|
||||
&self,
|
||||
@@ -498,24 +572,32 @@ impl GpuMoeHead {
|
||||
.map_err(|e| MLError::ModelError(format!("loss_per_k buf alloc: {e}")))?;
|
||||
let loss_total_buf = unsafe { MappedF32Buffer::new(1) }
|
||||
.map_err(|e| MLError::ModelError(format!("loss_total buf alloc: {e}")))?;
|
||||
// Test-shim ISV: 1-element buffer with λ at index 0 so the kernel's
|
||||
// `isv_signals[isv_lambda_eff_idx]` read returns the requested value.
|
||||
let isv_buf = unsafe { MappedF32Buffer::new(1) }
|
||||
.map_err(|e| MLError::ModelError(format!("isv buf alloc: {e}")))?;
|
||||
|
||||
gate_buf.write_from_slice(gate);
|
||||
isv_buf.write_from_slice(&[lambda]);
|
||||
|
||||
let block: u32 = 256;
|
||||
let dev_gate = gate_buf.dev_ptr;
|
||||
let dev_loss_per_k = loss_per_k_buf.dev_ptr;
|
||||
let dev_loss_total = loss_total_buf.dev_ptr;
|
||||
let dev_isv = isv_buf.dev_ptr;
|
||||
let b_i32 = b as i32;
|
||||
let k_i32 = k as i32;
|
||||
let isv_lambda_idx_i32: i32 = 0;
|
||||
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.moe_load_balance_loss)
|
||||
.arg(&dev_gate)
|
||||
.arg(&dev_loss_per_k)
|
||||
.arg(&dev_isv)
|
||||
.arg(&b_i32)
|
||||
.arg(&k_i32)
|
||||
.arg(&lambda)
|
||||
.arg(&isv_lambda_idx_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (k as u32, 1, 1),
|
||||
block_dim: (block, 1, 1),
|
||||
@@ -598,4 +680,54 @@ impl GpuMoeHead {
|
||||
|
||||
Ok(isv_buf.read_all())
|
||||
}
|
||||
|
||||
/// Test-only entry for the adaptive λ controller kernel.
|
||||
///
|
||||
/// Stages the provided `isv` slice into mapped pinned memory, runs
|
||||
/// `moe_lambda_eff_update`, and returns the full updated ISV. The
|
||||
/// controller reads `isv[isv_entropy_ema_idx]` and writes
|
||||
/// `isv[isv_lambda_eff_idx] = lambda_floor + lambda_max_extra × deficit`,
|
||||
/// where `deficit = clamp((target − ent_ema)/target, 0, 1)` and
|
||||
/// `target = entropy_target_frac × ln(K=8)`.
|
||||
///
|
||||
/// GPU-resident per `pearl_cold_path_no_exception_to_gpu_drives.md`.
|
||||
pub fn test_lambda_eff_update(
|
||||
&self,
|
||||
isv: &[f32],
|
||||
isv_lambda_eff_idx: usize,
|
||||
isv_entropy_ema_idx: usize,
|
||||
lambda_floor: f32,
|
||||
lambda_max_extra: f32,
|
||||
entropy_target_frac: f32,
|
||||
) -> Result<Vec<f32>, MLError> {
|
||||
let isv_buf = unsafe { MappedF32Buffer::new(isv.len()) }
|
||||
.map_err(|e| MLError::ModelError(format!("isv buf alloc: {e}")))?;
|
||||
isv_buf.write_from_slice(isv);
|
||||
|
||||
let dev_isv = isv_buf.dev_ptr;
|
||||
let lambda_idx_i32 = isv_lambda_eff_idx as i32;
|
||||
let entropy_idx_i32 = isv_entropy_ema_idx as i32;
|
||||
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.moe_lambda_eff_update)
|
||||
.arg(&dev_isv)
|
||||
.arg(&lambda_idx_i32)
|
||||
.arg(&entropy_idx_i32)
|
||||
.arg(&lambda_floor)
|
||||
.arg(&lambda_max_extra)
|
||||
.arg(&entropy_target_frac)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (1, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("moe_lambda_eff_update launch: {e}")))?;
|
||||
}
|
||||
self.stream
|
||||
.synchronize()
|
||||
.map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?;
|
||||
|
||||
Ok(isv_buf.read_all())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -90,21 +90,34 @@ extern "C" __global__ void moe_dgate_reduce(
|
||||
}
|
||||
|
||||
extern "C" __global__ void moe_load_balance_loss(
|
||||
const float* __restrict__ gate, /* [B, K] */
|
||||
float* __restrict__ loss_per_k, /* [K] */
|
||||
const float* __restrict__ gate, /* [B, K] */
|
||||
float* __restrict__ loss_per_k, /* [K] */
|
||||
const float* __restrict__ isv_signals, /* [ISV_TOTAL_DIM] */
|
||||
int B,
|
||||
int K,
|
||||
float lambda
|
||||
int isv_lambda_eff_idx /* MOE_LAMBDA_EFF_INDEX = 128 */
|
||||
) {
|
||||
/* One block per k. Each block reduces over B (sum gate[b,k]) into smem
|
||||
* and writes lambda*K*(sum/B)² to loss_per_k[k].
|
||||
* 1-block-per-k design avoids atomicAdd. */
|
||||
* 1-block-per-k design avoids atomicAdd.
|
||||
*
|
||||
* λ is read from ISV[isv_lambda_eff_idx] — the adaptive controller
|
||||
* (moe_lambda_eff_update kernel) updates this slot each step from the
|
||||
* gate-entropy EMA so the load-balance push scales inversely with
|
||||
* gate concentration. Cold-start is λ_floor (constructor write); fold
|
||||
* reset re-applies λ_floor; controller drives it up to
|
||||
* λ_floor + λ_max_extra under full entropy collapse.
|
||||
*
|
||||
* Per `feedback_isv_for_adaptive_bounds.md` / `pearl_cold_path_no_exception_to_gpu_drives.md`:
|
||||
* the consumer reads the ISV slot directly — no host roundtrip. */
|
||||
extern __shared__ float smem[];
|
||||
int k = blockIdx.x;
|
||||
if (k >= K) return;
|
||||
int tid = threadIdx.x;
|
||||
int block = blockDim.x;
|
||||
|
||||
float lambda = isv_signals[isv_lambda_eff_idx];
|
||||
|
||||
float local = 0.0f;
|
||||
for (int b = tid; b < B; b += block) {
|
||||
local += gate[b * K + k];
|
||||
|
||||
70
crates/ml/src/cuda_pipeline/moe_lambda_eff_kernel.cu
Normal file
70
crates/ml/src/cuda_pipeline/moe_lambda_eff_kernel.cu
Normal file
@@ -0,0 +1,70 @@
|
||||
/* moe_lambda_eff_update — adaptive MoE load-balance λ controller, GPU-driven.
|
||||
*
|
||||
* Reads: ISV[MOE_GATE_ENTROPY_EMA_INDEX=126] (gate-entropy EMA, α=0.05;
|
||||
* producer is moe_expert_util_ema_update which co-EMAs entropy with
|
||||
* per-expert util into [118..127)).
|
||||
* Writes: ISV[MOE_LAMBDA_EFF_INDEX=128] (effective load-balance weight).
|
||||
*
|
||||
* Cold-path cadence — single-block, single-thread kernel (matches
|
||||
* `kelly_cap_update_kernel.cu` and `cql_alpha_seed_update_kernel.cu`
|
||||
* shape). Launched after `moe_expert_util_ema_update` so the entropy EMA
|
||||
* fed to the controller is fresh for the NEXT step's
|
||||
* `moe_load_balance_loss` consumer.
|
||||
*
|
||||
* Formula (per `pearl_blend_formulas_must_have_permanent_floor.md`):
|
||||
*
|
||||
* target = entropy_target_frac × ln(K)
|
||||
* deficit = clamp((target − ent_ema) / max(target, 1e-6), 0, 1)
|
||||
* λ_eff = λ_floor + λ_max_extra × deficit
|
||||
*
|
||||
* `max(real, floor)` composition keeps λ_floor as a permanent minimum
|
||||
* (λ_eff ≥ λ_floor for all entropy values). At cold-start (ISV[128]
|
||||
* bootstrapped to λ_floor in the constructor) and at fold boundary (state
|
||||
* reset registry rewrites ISV[128] = λ_floor), the controller starts from
|
||||
* the floor and rises adaptively as gate-entropy decays below target.
|
||||
*
|
||||
* Motivation: the L40S validation run train-multi-seed-fgg9x exposed
|
||||
* single-expert specialization in fold 1 (expert 4 → 81% util, 7/8
|
||||
* experts < 5%) under static λ=0.01. The static load-balance push was too
|
||||
* weak to prevent winner-take-all once the gate had a strong preference.
|
||||
* Adaptive λ rises proportional to the entropy deficit, restoring
|
||||
* diversity pressure when the gate concentrates.
|
||||
*
|
||||
* Per `feedback_isv_for_adaptive_bounds.md`: signal flows through ISV[128];
|
||||
* consumer kernel `moe_load_balance_loss` reads the slot at runtime; no
|
||||
* DtoH on hot path.
|
||||
*/
|
||||
|
||||
extern "C" __global__ void moe_lambda_eff_update(
|
||||
float* __restrict__ isv, /* ISV bus [ISV_TOTAL_DIM] */
|
||||
int isv_lambda_eff_idx, /* MOE_LAMBDA_EFF_INDEX = 128 */
|
||||
int isv_entropy_ema_idx, /* MOE_GATE_ENTROPY_EMA_INDEX = 126 */
|
||||
float lambda_floor, /* permanent minimum (= old static 0.01) */
|
||||
float lambda_max_extra, /* max additional λ at full collapse */
|
||||
float entropy_target_frac /* fraction of ln(K) treated as "diverse enough" */
|
||||
) {
|
||||
if (threadIdx.x != 0 || blockIdx.x != 0) return;
|
||||
|
||||
/* ln(MOE_NUM_EXPERTS) — for K=8 this is 2.07944. Hard-coded because
|
||||
* MOE_NUM_EXPERTS is a compile-time constant on the Rust side and any
|
||||
* change to K already requires a kernel recompile. */
|
||||
const float ln_K = 2.0794415f; /* ln(8); update if K changes */
|
||||
|
||||
float ent_ema = isv[isv_entropy_ema_idx];
|
||||
float target = entropy_target_frac * ln_K;
|
||||
|
||||
/* Normalised entropy deficit ∈ [0, 1]:
|
||||
* - 0 when ent_ema ≥ target (gate diverse enough → λ_eff = λ_floor)
|
||||
* - 1 when ent_ema = 0 (gate fully collapsed → λ_eff = λ_floor + λ_max_extra)
|
||||
*
|
||||
* `target_safe` floor of 1e-6 guards against degenerate config
|
||||
* (entropy_target_frac=0). */
|
||||
float target_safe = fmaxf(target, 1e-6f);
|
||||
float deficit = fmaxf(0.0f, target - ent_ema) / target_safe;
|
||||
deficit = fminf(1.0f, deficit);
|
||||
|
||||
float lambda_eff = lambda_floor + lambda_max_extra * deficit;
|
||||
|
||||
isv[isv_lambda_eff_idx] = lambda_eff;
|
||||
__threadfence_system(); /* PCIe-visible to mapped pinned host_ptr */
|
||||
}
|
||||
@@ -888,11 +888,27 @@ pub struct DQNHyperparameters {
|
||||
/// L_total = L_c51 + iqn_lambda * L_iqn
|
||||
/// Range [0.0, 2.0]: 0.0 = C51 only, 0.5 = balanced, 1.0 = equal weight
|
||||
pub iqn_lambda: f32,
|
||||
/// Mixture-of-Experts load-balancing aux loss weight.
|
||||
/// Default: 0.01 (anti-collapse only — prevents init-noise-dominated
|
||||
/// single-expert lock-in without forcing uniform utilization).
|
||||
/// Mixture-of-Experts load-balance λ controller — permanent floor weight.
|
||||
///
|
||||
/// `λ_eff = floor + max_extra × deficit`, where `deficit ∈ [0, 1]` is
|
||||
/// the normalised entropy shortfall against `target_frac × ln(K)`. The
|
||||
/// GPU kernel `moe_lambda_eff_update` writes ISV[128] each step;
|
||||
/// consumer kernel `moe_load_balance_loss` reads it at runtime.
|
||||
///
|
||||
/// Default 0.01 (matches the legacy static `moe_lambda` default — the
|
||||
/// old "anti-collapse only" baseline becomes the permanent minimum).
|
||||
/// Hyperopt search range: [0.001, 0.1].
|
||||
pub moe_lambda: Option<f32>,
|
||||
pub moe_lambda_floor: Option<f32>,
|
||||
/// MoE adaptive λ controller — maximum additional λ added under full
|
||||
/// gate-entropy collapse. Peak λ_eff is `floor + max_extra`.
|
||||
/// Default 0.09 (peak λ = 0.10, 10× boost over floor).
|
||||
/// Hyperopt search range: [0.0, 0.5].
|
||||
pub moe_lambda_max_extra: Option<f32>,
|
||||
/// MoE adaptive λ controller — fraction of `ln(K)` treated as "diverse
|
||||
/// enough"; entropy at or above `target_frac × ln(K)` produces zero
|
||||
/// deficit (λ_eff = floor). Default 0.7 (target = 70% of ln(K)).
|
||||
/// Hyperopt search range: [0.3, 0.95].
|
||||
pub moe_entropy_target_frac: Option<f32>,
|
||||
/// Spectral norm σ_max — constrains ||W||_σ ≤ σ_max.
|
||||
/// Range [1.0, 10.0]. Default 3.0 (permits Xavier scaling, prevents Q-explosion).
|
||||
pub spectral_norm_sigma_max: f32,
|
||||
@@ -1608,8 +1624,14 @@ impl DQNHyperparameters {
|
||||
// 100k samples ≈ first 5 production-scale collection epochs.
|
||||
replay_seed_steps: 100_000,
|
||||
|
||||
// MoE load-balancing aux loss weight (Phase 1: additive field only).
|
||||
moe_lambda: Some(0.01),
|
||||
// MoE adaptive load-balance λ controller (2026-04-27).
|
||||
// Floor is the permanent minimum (= old static 0.01 default);
|
||||
// the controller raises λ_eff up to (floor + max_extra) under
|
||||
// entropy collapse. Target_frac sets the entropy threshold at
|
||||
// which λ_eff = floor (no extra push when gate is diverse).
|
||||
moe_lambda_floor: Some(0.01),
|
||||
moe_lambda_max_extra: Some(0.09),
|
||||
moe_entropy_target_frac: Some(0.7),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -415,7 +415,10 @@ 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),
|
||||
// MoE adaptive load-balance λ controller (2026-04-27).
|
||||
moe_lambda_floor: hyperparams.moe_lambda_floor.unwrap_or(0.01),
|
||||
moe_lambda_max_extra: hyperparams.moe_lambda_max_extra.unwrap_or(0.09),
|
||||
moe_entropy_target_frac: hyperparams.moe_entropy_target_frac.unwrap_or(0.7),
|
||||
};
|
||||
|
||||
// Create weight set pointer views AFTER GpuDqnTrainer is constructed below.
|
||||
|
||||
@@ -372,6 +372,17 @@ impl StateResetRegistry {
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[MOE_GATE_ENTROPY_EMA_INDEX=126] — entropy of batch-mean gate distribution, EMA α=0.05; reset to ln(8)≈2.0794 at fold boundary (max entropy at uniform initial gate); GPU moe_gate_entropy_ema_update kernel fills (Phase 2 task 2.4)",
|
||||
},
|
||||
// MoE adaptive load-balance λ controller (slot 128).
|
||||
// Reset to moe_lambda_floor (= old static 0.01 default) at fold
|
||||
// boundary so the new fold starts from the permanent minimum
|
||||
// and rises adaptively as the entropy EMA decays under it.
|
||||
// Per `pearl_blend_formulas_must_have_permanent_floor.md`, the
|
||||
// floor is a permanent minimum — λ_eff never drops below it.
|
||||
RegistryEntry {
|
||||
name: "isv_moe_lambda_eff",
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[MOE_LAMBDA_EFF_INDEX=128] — adaptive load-balance λ from `moe_lambda_eff_update` controller; reset to config.moe_lambda_floor at fold boundary (gate uniform-init ⇒ deficit=0 ⇒ λ_floor); per-step kernel raises λ_eff up to (floor + max_extra) when gate-entropy EMA decays below entropy_target_frac × ln(K)",
|
||||
},
|
||||
];
|
||||
Self { entries }
|
||||
}
|
||||
|
||||
@@ -2987,6 +2987,19 @@ impl DQNTrainer {
|
||||
tracing::warn!("Phase 3 T3.5 launch_moe_expert_util_ema failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// Adaptive MoE load-balance λ controller (2026-04-27).
|
||||
// MUST run AFTER `launch_moe_expert_util_ema` so the gate-
|
||||
// entropy EMA fed into the controller is fresh for the NEXT
|
||||
// step's `moe_load_balance_loss` consumer. Reads ISV[126]
|
||||
// (entropy EMA), writes ISV[MOE_LAMBDA_EFF_INDEX=128].
|
||||
// Per `pearl_blend_formulas_must_have_permanent_floor.md`,
|
||||
// λ_floor is a permanent minimum.
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
if let Err(e) = fused.trainer().launch_moe_lambda_eff_update() {
|
||||
tracing::warn!("launch_moe_lambda_eff_update failed: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Plan 4 Task 6 Commit B: refresh the aux-loss weight before the
|
||||
@@ -3315,27 +3328,31 @@ impl DQNTrainer {
|
||||
}
|
||||
|
||||
// Phase 3 T3.5: MoE expert-utilisation + gate entropy HEALTH_DIAG.
|
||||
// Extended 2026-04-27: also emits the adaptive λ_eff slot so the
|
||||
// controller is observable per epoch — see `launch_moe_lambda_eff_update`.
|
||||
{
|
||||
use crate::cuda_pipeline::gpu_dqn_trainer::{
|
||||
MOE_EXPERT_UTIL_EMA_BASE, MOE_GATE_ENTROPY_EMA_INDEX,
|
||||
MOE_LAMBDA_EFF_INDEX,
|
||||
};
|
||||
let (moe_utils, moe_ent) =
|
||||
let (moe_utils, moe_ent, moe_lambda_eff) =
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
let trainer = fused.trainer();
|
||||
let utils: Vec<f32> = (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)
|
||||
let lam = trainer.read_isv_signal_at(MOE_LAMBDA_EFF_INDEX);
|
||||
(utils, ent, lam)
|
||||
} else {
|
||||
(vec![0.0f32; 8], 0.0f32)
|
||||
(vec![0.0f32; 8], 0.0f32, 0.0f32)
|
||||
};
|
||||
tracing::info!(
|
||||
"HEALTH_DIAG[{}]: aux_moe [util={:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3} ent={:.3}]",
|
||||
"HEALTH_DIAG[{}]: aux_moe [util={:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3},{:.3} ent={:.3} λ_eff={:.4}]",
|
||||
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,
|
||||
moe_ent, moe_lambda_eff,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -4870,6 +4887,21 @@ impl DQNTrainer {
|
||||
);
|
||||
}
|
||||
}
|
||||
"isv_moe_lambda_eff" => {
|
||||
// Adaptive load-balance λ controller — fold reset writes
|
||||
// ISV[MOE_LAMBDA_EFF_INDEX] to the configured floor so the
|
||||
// new fold begins at the legacy 0.01 baseline; the per-step
|
||||
// controller (`moe_lambda_eff_update`) raises λ_eff up to
|
||||
// (floor + max_extra) as the gate-entropy EMA decays.
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
let trainer = fused.trainer();
|
||||
let floor = trainer.config().moe_lambda_floor;
|
||||
trainer.write_isv_signal_at(
|
||||
crate::cuda_pipeline::gpu_dqn_trainer::MOE_LAMBDA_EFF_INDEX,
|
||||
floor,
|
||||
);
|
||||
}
|
||||
}
|
||||
"isv_gamma_dir_eff" | "isv_gamma_mag_eff" | "isv_gamma_ord_eff" | "isv_gamma_urg_eff" => {
|
||||
// D.2 per-branch gamma slots. Reset to 0.0 at fold boundary;
|
||||
// per_branch_gamma_update GPU kernel re-populates on next epoch-boundary launch.
|
||||
|
||||
@@ -265,3 +265,97 @@ fn moe_expert_util_ema_update_writes_correct_values() {
|
||||
"entropy: expected {:.6}, got {:.6}", expected_ent, new_isv[126],
|
||||
);
|
||||
}
|
||||
|
||||
/// Adaptive load-balance λ controller (2026-04-27).
|
||||
///
|
||||
/// Verifies the formula:
|
||||
/// target = entropy_target_frac × ln(K=8)
|
||||
/// deficit = clamp((target - ent_ema) / target, 0, 1)
|
||||
/// λ_eff = λ_floor + λ_max_extra × deficit
|
||||
///
|
||||
/// Three regimes:
|
||||
/// (a) ent_ema ≥ target → deficit = 0, λ_eff = floor
|
||||
/// (b) ent_ema = 0 → deficit = 1, λ_eff = floor + max_extra
|
||||
/// (c) half-collapse → deficit ≈ 0.5, λ_eff ≈ floor + 0.5×max_extra
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn moe_lambda_eff_update_writes_correct_values() {
|
||||
use std::f32::consts::E;
|
||||
|
||||
let ctx = CudaContext::new(0).unwrap();
|
||||
let stream = ctx.default_stream();
|
||||
let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();
|
||||
|
||||
const LAMBDA_FLOOR: f32 = 0.01;
|
||||
const MAX_EXTRA: f32 = 0.09;
|
||||
const TARGET_FRAC: f32 = 0.7;
|
||||
const LAMBDA_EFF_IDX: usize = 128;
|
||||
const ENTROPY_IDX: usize = 126;
|
||||
|
||||
let ln_k = (8.0_f32).ln(); // 2.07944
|
||||
let target = TARGET_FRAC * ln_k; // ≈ 1.4556
|
||||
let _ = E; // keep std::f32::consts::E import warning quiet
|
||||
|
||||
// Helper: stage ISV with a chosen entropy value, run the kernel,
|
||||
// return the resulting λ_eff at slot 128.
|
||||
let run = |ent_ema: f32| -> f32 {
|
||||
let mut isv = vec![0.0_f32; 129];
|
||||
isv[ENTROPY_IDX] = ent_ema;
|
||||
// Pre-fill λ_eff with garbage to confirm the kernel overwrites it.
|
||||
isv[LAMBDA_EFF_IDX] = -42.0;
|
||||
let new_isv = head
|
||||
.test_lambda_eff_update(
|
||||
&isv,
|
||||
LAMBDA_EFF_IDX,
|
||||
ENTROPY_IDX,
|
||||
LAMBDA_FLOOR,
|
||||
MAX_EXTRA,
|
||||
TARGET_FRAC,
|
||||
)
|
||||
.unwrap();
|
||||
new_isv[LAMBDA_EFF_IDX]
|
||||
};
|
||||
|
||||
// (a) Diverse gate: ent_ema = 1.6 (above target ≈ 1.456) → λ_eff = floor.
|
||||
let lambda_a = run(1.6);
|
||||
assert!(
|
||||
(lambda_a - LAMBDA_FLOOR).abs() < 1e-6,
|
||||
"ent=1.6 → λ_eff expected {:.6}, got {:.6}",
|
||||
LAMBDA_FLOOR, lambda_a,
|
||||
);
|
||||
|
||||
// (b) Full collapse: ent_ema = 0 → λ_eff = floor + max_extra.
|
||||
let lambda_b = run(0.0);
|
||||
let expected_b = LAMBDA_FLOOR + MAX_EXTRA;
|
||||
assert!(
|
||||
(lambda_b - expected_b).abs() < 1e-6,
|
||||
"ent=0 → λ_eff expected {:.6}, got {:.6}",
|
||||
expected_b, lambda_b,
|
||||
);
|
||||
|
||||
// (c) Half collapse: ent_ema = 0.5 × target → deficit = 0.5 → λ_eff = floor + 0.5×max_extra.
|
||||
let half_target = 0.5 * target;
|
||||
let lambda_c = run(half_target);
|
||||
let expected_c = LAMBDA_FLOOR + 0.5 * MAX_EXTRA;
|
||||
assert!(
|
||||
(lambda_c - expected_c).abs() < 1e-5,
|
||||
"ent=0.5×target ({:.4}) → λ_eff expected {:.6}, got {:.6}",
|
||||
half_target, expected_c, lambda_c,
|
||||
);
|
||||
|
||||
// (d) Exactly at target: deficit = 0, λ_eff = floor.
|
||||
let lambda_d = run(target);
|
||||
assert!(
|
||||
(lambda_d - LAMBDA_FLOOR).abs() < 1e-6,
|
||||
"ent=target ({:.4}) → λ_eff expected {:.6}, got {:.6}",
|
||||
target, LAMBDA_FLOOR, lambda_d,
|
||||
);
|
||||
|
||||
// (e) Above-target (deficit clamped to 0): λ_eff = floor.
|
||||
let lambda_e = run(target + 0.5);
|
||||
assert!(
|
||||
(lambda_e - LAMBDA_FLOOR).abs() < 1e-6,
|
||||
"ent=target+0.5 → λ_eff expected {:.6}, got {:.6}",
|
||||
LAMBDA_FLOOR, lambda_e,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2,6 +2,63 @@
|
||||
|
||||
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
|
||||
|
||||
adaptive MoE load-balance λ controller (2026-04-27): replaces static
|
||||
`moe_lambda=0.01` with a GPU-driven controller that scales λ inversely
|
||||
with observed gate-entropy. New kernel
|
||||
`crates/ml/src/cuda_pipeline/moe_lambda_eff_kernel.cu` (single-block
|
||||
single-thread cold-path matching `kelly_cap_update_kernel.cu` precedent;
|
||||
own cubin per build.rs registration) reads ISV[126]
|
||||
(MOE_GATE_ENTROPY_EMA from `moe_expert_util_ema_update`, Phase 2 T2.4
|
||||
producer) and writes ISV[128] (MOE_LAMBDA_EFF_INDEX, new). Formula:
|
||||
`λ_eff = floor + max_extra × clamp((target − ent_ema)/target, 0, 1)`,
|
||||
target = `entropy_target_frac × ln(K)`. `floor` is a permanent minimum
|
||||
(per `pearl_blend_formulas_must_have_permanent_floor.md`) so the
|
||||
controller never weakens below the legacy 0.01 baseline. Consumer
|
||||
`moe_load_balance_loss` (in `moe_kernels.cu`) signature changed: dropped
|
||||
`float lambda` arg, added `const float* isv_signals` +
|
||||
`int isv_lambda_eff_idx` — kernel reads λ from ISV at runtime, no DtoH
|
||||
per `feedback_isv_for_adaptive_bounds.md`. Wiring (per
|
||||
`feedback_no_partial_refactor.md`): launch order in
|
||||
`training_loop.rs` per-step block runs `launch_moe_expert_util_ema`
|
||||
(ISV[126] producer, fresh entropy EMA) → `launch_moe_lambda_eff_update`
|
||||
(ISV[128] producer, fresh λ_eff for next step's load-balance consumer);
|
||||
HEALTH_DIAG `aux_moe` line gains `λ_eff` field for observability;
|
||||
state-reset registry entry `isv_moe_lambda_eff` resets ISV[128] to
|
||||
config.moe_lambda_floor at fold boundary; constructor bootstraps
|
||||
ISV[128]=floor + ISV[118..127)=uniform 1/K + ISV[126]=ln(K) so cold-start
|
||||
matches (deficit=0 ⇒ λ_eff=floor — legacy behaviour byte-for-byte).
|
||||
Config: `pub moe_lambda` field replaced by three knobs
|
||||
(`moe_lambda_floor` default 0.01, `moe_lambda_max_extra` default 0.09,
|
||||
`moe_entropy_target_frac` default 0.7) on both
|
||||
`GpuDqnTrainConfig` and `DQNHyperparameters`; `fused_training.rs:418`
|
||||
migrated to read all three. ISV_TOTAL_DIM bumped 127 → 129 (slot 127
|
||||
reserved gap to keep slot 128 cleanly aligned), layout fingerprint seed
|
||||
updated (schema_hash bumps; PVC fxcache regenerates on next deploy).
|
||||
Unit test `moe_lambda_eff_update_writes_correct_values` verifies all 5
|
||||
regimes (above-target / at-target / half-collapse / full-collapse /
|
||||
deeply-above-target) on the GPU kernel, all pass within 1e-5. Existing
|
||||
`moe_load_balance_loss_correctness` + `moe_load_balance_loss_uniform_minimum`
|
||||
tests pass via test-shim ISV[0] = λ. Smoke
|
||||
`magnitude_distribution` reproduces pre-existing eval-Full=0 (Kelly cap
|
||||
behaviour — see `project_magnitude_eval_collapse_kelly_capped.md`,
|
||||
unrelated to this commit; verified by stash-test against HEAD same
|
||||
numbers eq=0.586/0.592). HEALTH_DIAG fold 3 confirms controller
|
||||
behaviour live: ent decays 1.381 → 0.746 across epochs 7-19, λ_eff rises
|
||||
0.0146 → 0.0539 monotonically — exactly the entropy-deficit response
|
||||
the spec calls for. Files: new
|
||||
`crates/ml/src/cuda_pipeline/moe_lambda_eff_kernel.cu`, modified
|
||||
`crates/ml/build.rs` (cubin registration), `gpu_moe_head.rs` (new cubin
|
||||
load + `launch_lambda_eff_update` + `test_lambda_eff_update` + load_balance
|
||||
signature migration), `gpu_dqn_trainer.rs` (ISV slot const, layout
|
||||
fingerprint, config field replacement, struct field replacement,
|
||||
`launch_moe_lambda_eff_update`, constructor bootstrap),
|
||||
`moe_kernels.cu` (load_balance signature),
|
||||
`crates/ml/src/trainers/dqn/config.rs` (3 hyperparam knobs),
|
||||
`fused_training.rs` (3 unwrap-or migrations),
|
||||
`training_loop.rs` (per-step launch + HEALTH_DIAG λ_eff + fold-reset
|
||||
dispatch), `state_reset_registry.rs` (entry), `tests/moe_kernels_test.rs`
|
||||
(unit test).
|
||||
|
||||
magnitude conviction threaded into Kelly cap (2026-04-27, follow-up): the
|
||||
var_scale floor (above) addressed the **training** path. Smoke at the same
|
||||
horizon then reproduced eval Full=0.000 with intent=0.911 because
|
||||
|
||||
Reference in New Issue
Block a user