feat(isv): target network — ISV weights online-only
EMA sync restricted to indices 0-67 (non-ISV params). ISV weights (68-77) exist only in online params_buf. Target forward uses online gamma_buf (drift-conditioned) and online branch_gate_buf — correct for Bellman projection. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -83,6 +83,9 @@ const MAMBA2_STATE_DIM: usize = 16; // SSM state dimension
|
||||
/// Introspective State Vector (ISV) configuration.
|
||||
const ISV_K: usize = 4; // Temporal ISV history length
|
||||
const ISV_DIM: usize = 8; // ISV signal count
|
||||
/// First ISV tensor index in the flat param buffer.
|
||||
/// ISV weights (68-77) are online-only — NOT synced to the target network.
|
||||
const FIRST_ISV_TENSOR: usize = 68;
|
||||
const NEW_COMPONENT_WARMUP_STEPS: u32 = 500;
|
||||
|
||||
/// Homeostatic regularizer: number of observable signals.
|
||||
@@ -501,6 +504,12 @@ pub(crate) fn compute_total_params(cfg: &GpuDqnTrainConfig) -> usize {
|
||||
compute_param_sizes(cfg).iter().map(|&s| align4(s)).sum()
|
||||
}
|
||||
|
||||
/// Non-ISV parameter count (indices 0..FIRST_ISV_TENSOR).
|
||||
/// EMA target sync is restricted to this range — ISV weights are online-only.
|
||||
pub(crate) fn compute_non_isv_params(cfg: &GpuDqnTrainConfig) -> usize {
|
||||
compute_param_sizes(cfg)[..FIRST_ISV_TENSOR].iter().map(|&s| align4(s)).sum()
|
||||
}
|
||||
|
||||
/// Cumulative byte offset to tensor `idx` in the padded flat buffer.
|
||||
pub(crate) fn padded_byte_offset(param_sizes: &[usize], idx: usize) -> u64 {
|
||||
param_sizes[..idx].iter()
|
||||
@@ -894,6 +903,8 @@ pub struct GpuDqnTrainer {
|
||||
// ── Training state ──────────────────────────────────────────────
|
||||
pub(crate) adam_step: i32,
|
||||
total_params: usize,
|
||||
/// Non-ISV parameter count: EMA target sync restricted to this range.
|
||||
non_isv_params: usize,
|
||||
pub(crate) params_initialized: bool,
|
||||
target_params_initialized: bool,
|
||||
attention_initialized: bool,
|
||||
@@ -3800,6 +3811,7 @@ impl GpuDqnTrainer {
|
||||
) -> Result<Self, MLError> {
|
||||
let b = config.batch_size;
|
||||
let total_params = compute_total_params(&config);
|
||||
let non_isv_params = compute_non_isv_params(&config);
|
||||
|
||||
// Event tracking: kept ENABLED during normal ops (buffer allocation, DtoD).
|
||||
// DISABLED only during CUDA Graph capture (in capture_training_graphs).
|
||||
@@ -5520,6 +5532,7 @@ impl GpuDqnTrainer {
|
||||
q_div_ema: 0.0,
|
||||
adam_step: 0,
|
||||
total_params,
|
||||
non_isv_params,
|
||||
params_initialized: false,
|
||||
target_params_initialized: false,
|
||||
attention_initialized: false,
|
||||
@@ -9109,7 +9122,9 @@ impl GpuDqnTrainer {
|
||||
/// matches the GOFF_* defines exactly, with `align4()` padding per tensor.
|
||||
///
|
||||
/// Both params_buf and target_params_buf receive identical weights — the EMA
|
||||
/// kernel will diverge them during training.
|
||||
/// kernel will diverge them during training. ISV weights (indices 68-77) in
|
||||
/// target_params_buf are inert: EMA skips them, and the target forward never
|
||||
/// reads them.
|
||||
pub(crate) fn xavier_init_params_buf(&mut self) -> Result<(), MLError> {
|
||||
let cfg = &self.config;
|
||||
let sizes = compute_param_sizes(cfg);
|
||||
@@ -9401,11 +9416,13 @@ impl GpuDqnTrainer {
|
||||
|
||||
/// GPU-native Polyak EMA: `target[i] = (1-tau)*target[i] + tau*online[i]`
|
||||
///
|
||||
/// Fused single-kernel update over flat parameter buffers. On first call,
|
||||
/// flattens target weights into `target_params_buf` (same GOFF_* layout as
|
||||
/// `params_buf`). Then launches ONE EMA kernel over the entire flat buffer
|
||||
/// instead of 20 per-tensor launches. After the kernel, scatters the updated
|
||||
/// flat target weights back to individual tensors via `unflatten_target_weights()`.
|
||||
/// Fused single-kernel update over flat parameter buffers, restricted to
|
||||
/// non-ISV params (indices 0-67). ISV weights (68-77) are online-only and
|
||||
/// never synced to the target network. On first call, flattens target weights
|
||||
/// into `target_params_buf` (same GOFF_* layout as `params_buf`). Then
|
||||
/// launches ONE EMA kernel over the non-ISV portion instead of 20 per-tensor
|
||||
/// launches. After the kernel, scatters the updated flat target weights back
|
||||
/// to individual tensors via `unflatten_target_weights()`.
|
||||
///
|
||||
/// Runs OUTSIDE the captured CUDA Graph -- device pointers are stable so
|
||||
/// the graph stays valid.
|
||||
@@ -9432,8 +9449,11 @@ impl GpuDqnTrainer {
|
||||
tau
|
||||
};
|
||||
|
||||
let n = self.total_params as i32;
|
||||
let blocks = ((self.total_params + 255) / 256) as u32;
|
||||
// EMA only over non-ISV params (indices 0-67). ISV weights (68-77) are
|
||||
// online-only — the target network uses neutral defaults (base_gamma,
|
||||
// uniform gates) instead of ISV-modulated values.
|
||||
let n = self.non_isv_params as i32;
|
||||
let blocks = ((self.non_isv_params + 255) / 256) as u32;
|
||||
let launch_cfg = LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
|
||||
Reference in New Issue
Block a user