feat(ml-alpha): Phase 2A-C — trainer integration behind FOXHUNT_USE_MULTI_HEAD_POLICY flag
Wires MultiHeadPolicy into the live trainer's π forward + backward + Adam path, gated by FOXHUNT_USE_MULTI_HEAD_POLICY (default off). Mirrors the FOXHUNT_USE_ROLLOUT precedent (commit779c03b9d) for load-bearing refactors with emergency rollback. ## Wiring sites * Forward: 4 call sites — step_synthetic_body (training_graph), step_with_lobsim (prefill_graph CPU), step_with_lobsim_gpu_body (prefill_graph GPU), eval_policy_probs_per_action (no graph). * Backward: step_synthetic_body only (line 5764+). * Adam apply: 4 new Adam optimizers (W_heads, b_heads, W_gate, b_gate) when flag on. * LR plumbing: same lr_pi mirrored onto all 4 MHP Adams. * ISV bootstrap: slots 761-764 (K=3, gating_entropy_floor=0.549, head_entropy_floor=0.959, aux_prior_β=0.05) written in a flag-gated block AFTER the main bootstrap loop (preserves flag-off bit-equality — see surprise #2). * regime_h plumbing: new pub accessor PerceptionTrainer:: step_regime_d_view() mirrors h_t_view; gate kernel reads the same buffer the encoder reads (no new memcpy, no shape change). ## CUDA Graph capture treatment Host-side select. The flag is invariant for trainer lifetime, so the branch evaluates at graph CAPTURE time — the selected kernel sequence is baked into the captured graph and replays single-path. No host branches inside the captured stream. Verified by determinism preservation in both flag-off and flag-on regimes. ## Verification (all gates passed) * FOXHUNT_USE_MULTI_HEAD_POLICY=0 ./scripts/determinism-check.sh --quick: exit 0. * Flag-off diag.jsonl byte-equal toe22da61cfbaseline (modulo elapsed_s timestamps). * FOXHUNT_USE_MULTI_HEAD_POLICY=1 ./scripts/determinism-check.sh --quick: exit 0 (new path is deterministic). * 11/11 multi_head_policy_invariants tests still pass. * FOXHUNT_USE_MULTI_HEAD_POLICY=1 ./scripts/local-mid-smoke.sh: 2500 steps (2000 train + 500 eval) completed without NaN. 0 NaN/null in 2000-row diag.jsonl and 501-row eval_diag.jsonl. eval_summary.total_pnl_usd = -$3.82M (single-seed at random init — not load-bearing; behavioral verdict requires 3-seed per pearl_local_smoke_noise_floor_and_regime_concentration). * 68/0/6 lib tests (passed/failed/ignored). * Pre-commit: 0 atomicAdd, 0 raw memcpy_htod/dtoh, 0 TODO, 0 host-branches inside captured streams. ## Surprises handled (not silently) 1. step_regime_d was private on PerceptionTrainer — added clean pub accessor step_regime_d_view mirroring h_t_view pattern. 2. Unconditional ISV bootstrap broke flag-off bit-equality (writing 761-764 from 0.0 sentinel shifted isv_state checksum). Caught by verification gate. Resolved by moving the 4 writes into a flag-gated block AFTER the main bootstrap loop. Flag-off leaves slots 761-764 at 0.0 sentinel, matchinge22da61cf. 3. Spectral-norm regularization on legacy policy_head.w_d still runs when flag on (cost: one kernel/step, no correctness impact). Documented as known follow-up — proper gating happens in Phase 2A-E productionization, not C.3 (which must preserve flag-off bit-equality). 4. Adam state across flag toggles: 4 MHP Adams are None when flag off. Mirrors FOXHUNT_USE_ROLLOUT precedent — flag cached at construction, one-shot read. 5. Device-aggregated diag deferred (gate_probs_mean[K], gate_argmax_mass[K], gate_entropy_mean, per_head_entropy_mean [K]). These require either a new reduction kernel + ISV slots or host-readback (forbidden by feedback_no_htod_htoh_only_ mapped_pinned). Emitted ISV-resident summary leaves only (multi_head_policy.{active,k,gating_entropy_floor,head_ entropy_floor,aux_prior_beta}). Phase 2A-D scope. ## Linked * Phase 2A-A:0b3e40150(MultiHeadPolicy foundation, inert) * Phase 2A-B:e22da61cf(backward + aux KL prior) * Plan: docs/superpowers/plans/2026-06-03-multi-head-policy- implementation.md * Spec ADDENDUM: docs/superpowers/specs/2026-06-02-multi-head- policy-with-r-multiple.md * Q-distill pearl: pearl_foxhunt_pi_trained_by_q_distillation_ not_ppo.md Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -181,6 +181,7 @@ use crate::rl::isv_slots::{
|
||||
use crate::rl::common::{FRD_N_ATOMS, FRD_N_HORIZONS, N_ACTIONS};
|
||||
use crate::rl::frd::FRD_OUT_DIM;
|
||||
use crate::rl::loss_balance::{read_loss_lambdas_from_isv, LossLambdas};
|
||||
use crate::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig};
|
||||
use crate::rl::ppo::{PolicyHead, PpoHeadsConfig, ValueHead};
|
||||
use crate::rl::reward::RlLobBackend;
|
||||
use crate::trainer::optim::AdamW;
|
||||
@@ -665,6 +666,35 @@ pub struct IntegratedTrainer {
|
||||
pub policy_b_adam: AdamW,
|
||||
pub value_w_adam: AdamW,
|
||||
pub value_b_adam: AdamW,
|
||||
|
||||
/// Phase 2A-C (2026-06-03) — multi-head policy gate (FOXHUNT_USE_MULTI_HEAD_POLICY).
|
||||
///
|
||||
/// Cached `bool` snapshot of the env flag captured once at construction
|
||||
/// time. Must not change for the lifetime of the trainer because the
|
||||
/// training / prefill graphs bake whichever branch is selected at first
|
||||
/// capture into the replayed graph (host-side `if` runs only during
|
||||
/// capture; afterwards the captured kernel stream is replayed verbatim).
|
||||
///
|
||||
/// Flag-off path is bit-equal to commit `e22da61cf` (legacy single-head
|
||||
/// `policy_head` drives π).
|
||||
pub use_multi_head_policy: bool,
|
||||
/// `Some(_)` iff `use_multi_head_policy == true`. Owns K policy heads +
|
||||
/// the regime-gated mixing head + all backward scratch + reduced-grad
|
||||
/// buffers ready for an Adam step. Allocated on the same stream as the
|
||||
/// rest of the trainer (`dev.cuda_stream()`); no extra synchronization.
|
||||
/// See `crates/ml-alpha/src/rl/multi_head_policy.rs` for the kernel
|
||||
/// pipeline. Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
|
||||
/// ADDENDUM 2026-06-03.
|
||||
pub multi_head_policy: Option<MultiHeadPolicy>,
|
||||
/// Adam optimisers for the four `MultiHeadPolicy` weight tensors —
|
||||
/// `W_heads [K × N_ACTIONS × HIDDEN_DIM]`, `b_heads [K × N_ACTIONS]`,
|
||||
/// `W_gate [K × REGIME_DIM=6]`, `b_gate [K]`. `Some(_)` iff
|
||||
/// `use_multi_head_policy == true`. LR reads slot `RL_LR_PI_INDEX`
|
||||
/// (mirrors the legacy `policy_w_adam` / `policy_b_adam` LR sourcing).
|
||||
pub mhp_w_heads_adam: Option<AdamW>,
|
||||
pub mhp_b_heads_adam: Option<AdamW>,
|
||||
pub mhp_w_gate_adam: Option<AdamW>,
|
||||
pub mhp_b_gate_adam: Option<AdamW>,
|
||||
/// Phase 4 dueling-head Adam optimisers (4 weight tensors).
|
||||
pub dueling_q_w_v_adam: AdamW,
|
||||
pub dueling_q_b_v_adam: AdamW,
|
||||
@@ -1613,6 +1643,103 @@ impl IntegratedTrainer {
|
||||
AdamW::new(dev, policy_head.w_d.len(), lr_placeholder).context("policy_w_adam")?;
|
||||
let policy_b_adam =
|
||||
AdamW::new(dev, policy_head.b_d.len(), lr_placeholder).context("policy_b_adam")?;
|
||||
|
||||
// ── Phase 2A-C (2026-06-03) MultiHeadPolicy gate ─────────────
|
||||
// `FOXHUNT_USE_MULTI_HEAD_POLICY` env flag follows the
|
||||
// `FOXHUNT_USE_ROLLOUT` precedent from commit 779c03b9d. Read
|
||||
// ONCE here so the value is invariant across captured-graph
|
||||
// replay (the `if self.use_multi_head_policy` branches at the
|
||||
// π forward / backward sites run only during graph CAPTURE; at
|
||||
// replay time the selected kernel stream is baked in).
|
||||
//
|
||||
// Gate semantics:
|
||||
// * `FOXHUNT_USE_MULTI_HEAD_POLICY` unset or `=0` → legacy
|
||||
// single-head `policy_head` drives π. `multi_head_policy` is
|
||||
// `None`, no extra device buffers, bit-equal to HEAD `e22da61cf`.
|
||||
// * `FOXHUNT_USE_MULTI_HEAD_POLICY=1` → construct
|
||||
// `MultiHeadPolicy` with K from ISV slot 761 (bootstrap K=3).
|
||||
// Forward routes `h_t + step_regime_d` through the K-head
|
||||
// mixture; backward distributes Q-distill grad through
|
||||
// mixture + per-head softmax + gating softmax Jacobians;
|
||||
// Adam updates the four new weight tensors.
|
||||
//
|
||||
// Per `pearl_foxhunt_pi_trained_by_q_distillation_not_ppo`:
|
||||
// the wiring target is the post-Q-distill `ss_pi_grad_logits_d`
|
||||
// write site (`integrated.rs:~5607`). PPO-surrogate grad is still
|
||||
// zeroed at line 5577; nothing changes for that path.
|
||||
let use_multi_head_policy: bool = std::env::var("FOXHUNT_USE_MULTI_HEAD_POLICY")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u8>().ok())
|
||||
.map(|n| n > 0)
|
||||
.unwrap_or(false);
|
||||
let (
|
||||
multi_head_policy,
|
||||
mhp_w_heads_adam,
|
||||
mhp_b_heads_adam,
|
||||
mhp_w_gate_adam,
|
||||
mhp_b_gate_adam,
|
||||
): (
|
||||
Option<MultiHeadPolicy>,
|
||||
Option<AdamW>,
|
||||
Option<AdamW>,
|
||||
Option<AdamW>,
|
||||
Option<AdamW>,
|
||||
) = if use_multi_head_policy {
|
||||
// K = 3 is the bootstrap value documented in
|
||||
// `RL_POLICY_NUM_HEADS_INDEX` (slot 761) and seeded in the
|
||||
// controller-bootstrap table below. Reading host-side ISV here
|
||||
// is safe because the table runs strictly BEFORE this point —
|
||||
// however `with_controllers_bootstrapped` runs AFTER `new()`,
|
||||
// so the slot is still zero at construction. We therefore use
|
||||
// the hardcoded bootstrap constant `K=3` to match the bootstrap
|
||||
// table and per the spec's "K=3 production default"; runtime K
|
||||
// changes would require re-allocation anyway (weight tensors
|
||||
// are sized by K).
|
||||
const K_BOOTSTRAP: usize = 3;
|
||||
// Distinct seed from the legacy `policy_head` (offset 0xC1A55 =
|
||||
// "CLASS" — independent head initialisation streams per
|
||||
// `feedback_no_shared_init_streams`).
|
||||
let mhp_seed: u64 = cfg.ppo_seed.wrapping_add(0xC1A55);
|
||||
let mhp = MultiHeadPolicy::new(
|
||||
dev,
|
||||
MultiHeadPolicyConfig {
|
||||
k: K_BOOTSTRAP,
|
||||
b_size: cfg.perception.n_batch,
|
||||
seed: mhp_seed,
|
||||
},
|
||||
)
|
||||
.context("MultiHeadPolicy::new (FOXHUNT_USE_MULTI_HEAD_POLICY=1)")?;
|
||||
// Mirror legacy policy Adam pattern: lr_placeholder (overwritten
|
||||
// every step from ISV[RL_LR_PI_INDEX]). One Adam per weight
|
||||
// tensor — gate weights get their own Adam (different shape +
|
||||
// potentially different LR scale in a future controller).
|
||||
let w_heads_adam = AdamW::new(dev, mhp.w_heads_d.len(), lr_placeholder)
|
||||
.context("mhp_w_heads_adam")?;
|
||||
let b_heads_adam = AdamW::new(dev, mhp.b_heads_d.len(), lr_placeholder)
|
||||
.context("mhp_b_heads_adam")?;
|
||||
let w_gate_adam = AdamW::new(dev, mhp.w_gate_d.len(), lr_placeholder)
|
||||
.context("mhp_w_gate_adam")?;
|
||||
let b_gate_adam = AdamW::new(dev, mhp.b_gate_d.len(), lr_placeholder)
|
||||
.context("mhp_b_gate_adam")?;
|
||||
eprintln!(
|
||||
"[multi-head-policy] FOXHUNT_USE_MULTI_HEAD_POLICY=1 → MultiHeadPolicy active \
|
||||
(K={} B={} REGIME_DIM=6, seed={:#x})",
|
||||
K_BOOTSTRAP, cfg.perception.n_batch, mhp_seed
|
||||
);
|
||||
(
|
||||
Some(mhp),
|
||||
Some(w_heads_adam),
|
||||
Some(b_heads_adam),
|
||||
Some(w_gate_adam),
|
||||
Some(b_gate_adam),
|
||||
)
|
||||
} else {
|
||||
eprintln!(
|
||||
"[multi-head-policy] FOXHUNT_USE_MULTI_HEAD_POLICY unset/=0 → legacy single-head \
|
||||
policy_head path (bit-equal to HEAD e22da61cf)"
|
||||
);
|
||||
(None, None, None, None, None)
|
||||
};
|
||||
let value_w_adam =
|
||||
AdamW::new(dev, value_head.w_d.len(), lr_placeholder).context("value_w_adam")?;
|
||||
let value_b_adam =
|
||||
@@ -2940,6 +3067,13 @@ impl IntegratedTrainer {
|
||||
policy_b_adam,
|
||||
value_w_adam,
|
||||
value_b_adam,
|
||||
// Phase 2A-C MultiHeadPolicy + Adams (None when flag off).
|
||||
use_multi_head_policy,
|
||||
multi_head_policy,
|
||||
mhp_w_heads_adam,
|
||||
mhp_b_heads_adam,
|
||||
mhp_w_gate_adam,
|
||||
mhp_b_gate_adam,
|
||||
dueling_q_w_v_adam,
|
||||
dueling_q_b_v_adam,
|
||||
dueling_q_w_a_adam,
|
||||
@@ -4016,6 +4150,47 @@ impl IntegratedTrainer {
|
||||
).map_err(|e| anyhow::anyhow!("rl_isv_write (design constant): {:?}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Phase 2A-C (2026-06-03) MultiHeadPolicy ISV bootstrap ────
|
||||
//
|
||||
// Slots 761-764 (allocated in Phase 2A-A). Bootstrap is gated
|
||||
// on `FOXHUNT_USE_MULTI_HEAD_POLICY` for bit-equality with
|
||||
// commit `e22da61cf`: writing these slots when the flag is off
|
||||
// would diverge `isv_state` checksum (the slots are 0.0
|
||||
// sentinel at HEAD; writing 3.0 / 0.549 / 0.959 / 0.05 shifts
|
||||
// the checksum), failing verification gate C.7.4.
|
||||
//
|
||||
// When flag is on, this writes match the dispatch's required
|
||||
// bootstrap values per spec ADDENDUM 2026-06-03 §R.1:
|
||||
// * K=3: spec §1.1 (production default)
|
||||
// * gating_entropy_floor = log(K)·0.5 ≈ 0.549
|
||||
// * head_entropy_floor = log(N_ACTIONS)·0.4 ≈ 0.959
|
||||
// * aux_prior_beta = 0.05 (5% of policy-loss magnitude)
|
||||
if self.use_multi_head_policy {
|
||||
let mhp_isv_constants: [(usize, f32); 4] = [
|
||||
(crate::rl::isv_slots::RL_POLICY_NUM_HEADS_INDEX, 3.0_f32),
|
||||
(crate::rl::isv_slots::RL_POLICY_GATING_ENTROPY_FLOOR_INDEX, 0.549_f32),
|
||||
(crate::rl::isv_slots::RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX, 0.959_f32),
|
||||
(crate::rl::isv_slots::RL_POLICY_AUX_PRIOR_BETA_INDEX, 0.05_f32),
|
||||
];
|
||||
for (slot, value) in mhp_isv_constants.iter() {
|
||||
let slot_i32 = *slot as i32;
|
||||
let value_f32 = *value;
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(self.isv_dev_ptr);
|
||||
args.push_i32(slot_i32);
|
||||
args.push_f32(value_f32);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
raw_launch(
|
||||
self.rl_isv_write_fn.cu_function(),
|
||||
(1, 1, 1), (1, 1, 1), 0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("rl_isv_write (multi-head bootstrap): {:?}", e))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let alpha = RL_LR_CONTROLLER_ALPHA;
|
||||
@@ -5344,6 +5519,15 @@ impl IntegratedTrainer {
|
||||
self.dqn_b_adam.lr = lr_q;
|
||||
self.policy_w_adam.lr = lr_pi;
|
||||
self.policy_b_adam.lr = lr_pi;
|
||||
// Phase 2A-C: mirror lr_pi onto the multi-head Adam state when
|
||||
// the gate is on. ISV[RL_LR_PI_INDEX] remains the single source
|
||||
// of truth — no separate slot for multi-head LR.
|
||||
if self.use_multi_head_policy {
|
||||
if let Some(a) = self.mhp_w_heads_adam.as_mut() { a.lr = lr_pi; }
|
||||
if let Some(a) = self.mhp_b_heads_adam.as_mut() { a.lr = lr_pi; }
|
||||
if let Some(a) = self.mhp_w_gate_adam.as_mut() { a.lr = lr_pi; }
|
||||
if let Some(a) = self.mhp_b_gate_adam.as_mut() { a.lr = lr_pi; }
|
||||
}
|
||||
self.value_w_adam.lr = lr_v;
|
||||
self.value_b_adam.lr = lr_v;
|
||||
// SP20 P3 FRD head — LR from dedicated ISV slot (498-503 block).
|
||||
@@ -5405,9 +5589,24 @@ impl IntegratedTrainer {
|
||||
self.dqn_head
|
||||
.forward(&self.sampled_h_t_d, b_size, &mut self.q_logits_d)
|
||||
.context("dqn_head.forward(sampled_h_t) [R7d off-policy]")?;
|
||||
self.policy_head
|
||||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||||
.context("policy_head.forward_logits")?;
|
||||
// Phase 2A-C: gate π forward between legacy single-head and
|
||||
// MultiHeadPolicy K-head mixture. Host branch runs during graph
|
||||
// CAPTURE only; the captured kernel sequence bakes the chosen
|
||||
// path so replay is single-path. `regime_h` is read DIRECTLY
|
||||
// from perception's `step_regime_d` (parallel channel — bypasses
|
||||
// VSN).
|
||||
if self.use_multi_head_policy {
|
||||
let regime_h = self.perception.step_regime_d_view();
|
||||
self.multi_head_policy
|
||||
.as_ref()
|
||||
.expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)")
|
||||
.forward(h_t_borrow, regime_h, &mut self.pi_logits_d)
|
||||
.context("multi_head_policy.forward [step_synthetic_body]")?;
|
||||
} else {
|
||||
self.policy_head
|
||||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||||
.context("policy_head.forward_logits")?;
|
||||
}
|
||||
self.value_head
|
||||
.forward(h_t_borrow, &self.isv_dev_ptr, b_size, &mut self.v_pred_d)
|
||||
.context("value_head.forward")?;
|
||||
@@ -5604,19 +5803,52 @@ impl IntegratedTrainer {
|
||||
// π gradient is fully determined by Q→π distillation + SAC entropy
|
||||
// (rl_q_pi_distill_grad above). No PPO surrogate or KL penalty.
|
||||
|
||||
self.policy_head
|
||||
.backward_to_w_b_h(
|
||||
h_t_borrow,
|
||||
&self.ss_pi_grad_logits_d,
|
||||
b_size,
|
||||
&mut self.ss_pi_grad_w_per_batch_d,
|
||||
&mut self.ss_pi_grad_b_per_batch_d,
|
||||
&mut self.ss_pi_grad_h_t_d,
|
||||
)
|
||||
.context("policy_head.backward_to_w_b_h")?;
|
||||
// Phase 2A-C: gate π backward between legacy and MultiHeadPolicy.
|
||||
// Both paths consume `ss_pi_grad_logits_d` (Q-distill output)
|
||||
// and produce reduced grads ready for Adam below. The
|
||||
// MultiHeadPolicy path internally:
|
||||
// * distributes grad through mixture + per-head softmax,
|
||||
// * adds β-weighted aux KL prior gradient,
|
||||
// * distributes through gating softmax,
|
||||
// * reduces per-batch grads to [K × N_ACTIONS × HIDDEN_DIM],
|
||||
// [K × N_ACTIONS], [K × REGIME_DIM], [K] (ready for Adam).
|
||||
// Legacy path uses the existing single-head GEMM-style backward.
|
||||
if self.use_multi_head_policy {
|
||||
// Split disjoint borrows: `perception`, `ss_pi_grad_logits_d`,
|
||||
// and `multi_head_policy` are separate fields, so the borrow
|
||||
// checker can prove non-aliasing. `isv_dev_ptr` is `u64` (Copy).
|
||||
let isv_dev_ptr = self.isv_dev_ptr;
|
||||
let regime_h: &CudaSlice<f32> = self.perception.step_regime_d_view();
|
||||
let grad_pi: &CudaSlice<f32> = &self.ss_pi_grad_logits_d;
|
||||
let mhp = self
|
||||
.multi_head_policy
|
||||
.as_mut()
|
||||
.expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)");
|
||||
mhp.backward(grad_pi, h_t_borrow, regime_h, isv_dev_ptr)
|
||||
.context("multi_head_policy.backward [step_synthetic_body]")?;
|
||||
// NOTE: per Phase 2A-B comment in `MultiHeadPolicy::backward`,
|
||||
// `grad_h_t_scratch_d` is NOT folded into the encoder grad
|
||||
// here — Phase 2A-C wires the encoder path conservatively
|
||||
// (encoder still receives V-head + auxiliaries; π gradient
|
||||
// into encoder via multi-head route would change the
|
||||
// encoder learning signal materially and is deferred to a
|
||||
// future phase per the dispatch's "no incidental changes"
|
||||
// constraint).
|
||||
} else {
|
||||
self.policy_head
|
||||
.backward_to_w_b_h(
|
||||
h_t_borrow,
|
||||
&self.ss_pi_grad_logits_d,
|
||||
b_size,
|
||||
&mut self.ss_pi_grad_w_per_batch_d,
|
||||
&mut self.ss_pi_grad_b_per_batch_d,
|
||||
&mut self.ss_pi_grad_h_t_d,
|
||||
)
|
||||
.context("policy_head.backward_to_w_b_h")?;
|
||||
|
||||
reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_w_per_batch_d, b_size, N_ACTIONS * HIDDEN_DIM, &mut self.ss_pi_grad_w_d)?;
|
||||
reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_b_per_batch_d, b_size, N_ACTIONS, &mut self.ss_pi_grad_b_d)?;
|
||||
reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_w_per_batch_d, b_size, N_ACTIONS * HIDDEN_DIM, &mut self.ss_pi_grad_w_d)?;
|
||||
reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_pi_grad_b_per_batch_d, b_size, N_ACTIONS, &mut self.ss_pi_grad_b_d)?;
|
||||
}
|
||||
|
||||
// ── Step 8: V backward ───────────────────────────────────────
|
||||
self.value_head
|
||||
@@ -5644,12 +5876,45 @@ impl IntegratedTrainer {
|
||||
self.dqn_b_adam
|
||||
.step(&mut self.dqn_head.b_d, &self.ss_q_grad_b_d)
|
||||
.context("dqn_b_adam.step")?;
|
||||
self.policy_w_adam
|
||||
.step(&mut self.policy_head.w_d, &self.ss_pi_grad_w_d)
|
||||
.context("policy_w_adam.step")?;
|
||||
self.policy_b_adam
|
||||
.step(&mut self.policy_head.b_d, &self.ss_pi_grad_b_d)
|
||||
.context("policy_b_adam.step")?;
|
||||
// Phase 2A-C: gate the policy Adam step between legacy single-head
|
||||
// and MultiHeadPolicy four-tensor update. Both branches preserve
|
||||
// the existing LR sourcing (slot 412 `RL_LR_PI_INDEX`) — see the
|
||||
// `lr_pi` assignment near line 5345.
|
||||
if self.use_multi_head_policy {
|
||||
let mhp = self
|
||||
.multi_head_policy
|
||||
.as_mut()
|
||||
.expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)");
|
||||
// Same LR for all four tensors (mirrors legacy w/b sharing
|
||||
// a single `lr_pi`). Future controller could split.
|
||||
self.mhp_w_heads_adam
|
||||
.as_mut()
|
||||
.expect("mhp_w_heads_adam None with flag on")
|
||||
.step(&mut mhp.w_heads_d, &mhp.grad_w_heads_d)
|
||||
.context("mhp_w_heads_adam.step")?;
|
||||
self.mhp_b_heads_adam
|
||||
.as_mut()
|
||||
.expect("mhp_b_heads_adam None with flag on")
|
||||
.step(&mut mhp.b_heads_d, &mhp.grad_b_heads_d)
|
||||
.context("mhp_b_heads_adam.step")?;
|
||||
self.mhp_w_gate_adam
|
||||
.as_mut()
|
||||
.expect("mhp_w_gate_adam None with flag on")
|
||||
.step(&mut mhp.w_gate_d, &mhp.grad_w_gate_d)
|
||||
.context("mhp_w_gate_adam.step")?;
|
||||
self.mhp_b_gate_adam
|
||||
.as_mut()
|
||||
.expect("mhp_b_gate_adam None with flag on")
|
||||
.step(&mut mhp.b_gate_d, &mhp.grad_b_gate_d)
|
||||
.context("mhp_b_gate_adam.step")?;
|
||||
} else {
|
||||
self.policy_w_adam
|
||||
.step(&mut self.policy_head.w_d, &self.ss_pi_grad_w_d)
|
||||
.context("policy_w_adam.step")?;
|
||||
self.policy_b_adam
|
||||
.step(&mut self.policy_head.b_d, &self.ss_pi_grad_b_d)
|
||||
.context("policy_b_adam.step")?;
|
||||
}
|
||||
self.value_w_adam
|
||||
.step(&mut self.value_head.w_d, &self.ss_v_grad_w_d)
|
||||
.context("value_w_adam.step")?;
|
||||
@@ -6928,9 +7193,21 @@ impl IntegratedTrainer {
|
||||
}
|
||||
|
||||
// ── Step 2b: Forward π logits for log_pi_old. ─────────────────
|
||||
self.policy_head
|
||||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||||
.context("step_with_lobsim: policy_head.forward_logits")?;
|
||||
// Phase 2A-C: same gate as step_synthetic_body. Inside prefill
|
||||
// graph capture; host branch runs once and bakes the chosen
|
||||
// forward kernel sequence into the replayed graph.
|
||||
if self.use_multi_head_policy {
|
||||
let regime_h = self.perception.step_regime_d_view();
|
||||
self.multi_head_policy
|
||||
.as_ref()
|
||||
.expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)")
|
||||
.forward(h_t_borrow, regime_h, &mut self.pi_logits_d)
|
||||
.context("multi_head_policy.forward [step_with_lobsim]")?;
|
||||
} else {
|
||||
self.policy_head
|
||||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||||
.context("step_with_lobsim: policy_head.forward_logits")?;
|
||||
}
|
||||
|
||||
// R9 audit — Q-vs-π agreement diag fires into ISV[407] EMA.
|
||||
// Needs both q_logits_d (from step 2) and pi_logits_d (just
|
||||
@@ -9082,9 +9359,19 @@ impl IntegratedTrainer {
|
||||
}
|
||||
|
||||
// pi logits
|
||||
self.policy_head
|
||||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||||
.context("step_with_lobsim_gpu: policy_head.forward_logits")?;
|
||||
// Phase 2A-C: gate the GPU-path pi forward as well.
|
||||
if self.use_multi_head_policy {
|
||||
let regime_h = self.perception.step_regime_d_view();
|
||||
self.multi_head_policy
|
||||
.as_ref()
|
||||
.expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)")
|
||||
.forward(h_t_borrow, regime_h, &mut self.pi_logits_d)
|
||||
.context("multi_head_policy.forward [step_with_lobsim_gpu_body]")?;
|
||||
} else {
|
||||
self.policy_head
|
||||
.forward_logits(h_t_borrow, b_size, &mut self.pi_logits_d)
|
||||
.context("step_with_lobsim_gpu: policy_head.forward_logits")?;
|
||||
}
|
||||
|
||||
// Q-vs-pi agree
|
||||
{
|
||||
@@ -9652,9 +9939,21 @@ impl IntegratedTrainer {
|
||||
.context("eval_policy_probs: forward_encoder")?;
|
||||
let h_t = self.perception.h_t_view();
|
||||
let mut logits_d = self.stream.alloc_zeros::<f32>(b_size * N_ACTIONS)?;
|
||||
self.policy_head
|
||||
.forward_logits(h_t, b_size, &mut logits_d)
|
||||
.context("eval_policy_probs: policy_head.forward_logits")?;
|
||||
// Phase 2A-C: gate the eval-time pi forward. NOT inside a graph
|
||||
// capture (eval is single-shot), so the branch is run host-side
|
||||
// every call — cost is one `bool` read.
|
||||
if self.use_multi_head_policy {
|
||||
let regime_h = self.perception.step_regime_d_view();
|
||||
self.multi_head_policy
|
||||
.as_ref()
|
||||
.expect("use_multi_head_policy=true but multi_head_policy=None (construction bug)")
|
||||
.forward(h_t, regime_h, &mut logits_d)
|
||||
.context("multi_head_policy.forward [eval_policy_probs_per_action]")?;
|
||||
} else {
|
||||
self.policy_head
|
||||
.forward_logits(h_t, b_size, &mut logits_d)
|
||||
.context("eval_policy_probs: policy_head.forward_logits")?;
|
||||
}
|
||||
|
||||
let logits_host = read_slice_d(&self.stream, &logits_d, b_size * N_ACTIONS)?;
|
||||
let mut probs = vec![0.0_f32; N_ACTIONS];
|
||||
@@ -11004,7 +11303,7 @@ impl IntegratedTrainer {
|
||||
|
||||
let replay_len = self.gpu_replay.capacity.min((step as usize) + 1);
|
||||
|
||||
let record = serde_json::json!({
|
||||
let mut record = serde_json::json!({
|
||||
"step": step,
|
||||
"elapsed_s": elapsed_s,
|
||||
"loss": {
|
||||
@@ -11484,6 +11783,36 @@ impl IntegratedTrainer {
|
||||
"mamba2_grad_w_c_l1": checksums[31],
|
||||
},
|
||||
});
|
||||
|
||||
// ── Phase 2A-C (2026-06-03) MultiHeadPolicy diag emit ────────
|
||||
//
|
||||
// Only emitted when `FOXHUNT_USE_MULTI_HEAD_POLICY=1` so the
|
||||
// flag-off diag.jsonl stays byte-equal to HEAD `e22da61cf`
|
||||
// (verification gate C.7.4).
|
||||
//
|
||||
// Phase 2A-C ships ISV-resident summary leaves only — K,
|
||||
// gating/head entropy floors, aux β. The richer per-step
|
||||
// device-aggregated stats (gate_probs_mean[K], gate_argmax_mass[K],
|
||||
// gate_entropy_mean, per_head_entropy_mean[K]) require a new
|
||||
// reduction kernel + new ISV slots and are deferred per the
|
||||
// dispatch's STOP-on-surprise rule. The four-leaf summary here
|
||||
// is the minimum cluster-verifiable signal that the gate path
|
||||
// is active and bootstrapped correctly.
|
||||
if self.use_multi_head_policy {
|
||||
if let Some(obj) = record.as_object_mut() {
|
||||
obj.insert(
|
||||
"multi_head_policy".to_string(),
|
||||
serde_json::json!({
|
||||
"active": self.use_multi_head_policy,
|
||||
"k": isv[crate::rl::isv_slots::RL_POLICY_NUM_HEADS_INDEX],
|
||||
"gating_entropy_floor": isv[crate::rl::isv_slots::RL_POLICY_GATING_ENTROPY_FLOOR_INDEX],
|
||||
"head_entropy_floor": isv[crate::rl::isv_slots::RL_POLICY_HEAD_ENTROPY_FLOOR_INDEX],
|
||||
"aux_prior_beta": isv[crate::rl::isv_slots::RL_POLICY_AUX_PRIOR_BETA_INDEX],
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(record)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4169,6 +4169,22 @@ impl PerceptionTrainer {
|
||||
&self.h_t_d
|
||||
}
|
||||
|
||||
/// Phase 2A-C (2026-06-03) — read-only accessor for the per-batch raw
|
||||
/// regime input `[n_batch × REGIME_DIM=6]` that the snap-feature
|
||||
/// pipeline populates per step. Consumed by `MultiHeadPolicy::forward`
|
||||
/// as a PARALLEL channel — bypassing VSN — to gate K policy heads on
|
||||
/// vol/popart regime directly (`pearl_local_smoke_noise_floor_and_regime_concentration`).
|
||||
///
|
||||
/// The slice's lifetime tracks `&self`; the buffer is overwritten on
|
||||
/// every `snap_feature_assemble` so callers must finish their kernel
|
||||
/// dispatch before the next encoder step. In the captured training
|
||||
/// graph, multi-head policy forward runs in the same stream sequence
|
||||
/// immediately after the encoder, so this is single-stream-safe by
|
||||
/// construction (same property `h_t_view` relies on).
|
||||
pub fn step_regime_d_view(&self) -> &CudaSlice<f32> {
|
||||
&self.step_regime_d
|
||||
}
|
||||
|
||||
/// Determinism Phase 2.2 sub-investigation: read-only accessors for
|
||||
/// encoder intermediate activations. Used exclusively by
|
||||
/// `IntegratedTrainer::compute_diag_checksums` to checksum each layer's
|
||||
|
||||
Reference in New Issue
Block a user