diff --git a/crates/ml/examples/evaluate_baseline.rs b/crates/ml/examples/evaluate_baseline.rs index 589f1028b..e3c50d028 100644 --- a/crates/ml/examples/evaluate_baseline.rs +++ b/crates/ml/examples/evaluate_baseline.rs @@ -1197,15 +1197,19 @@ fn evaluate_dqn_fold_gpu( dueling_config.advantage_hidden_dim, ); - let weights = extract_dueling_weights(dqn.get_q_network_vars(), &stream) + let (_d_backing, weights) = extract_dueling_weights(dqn.get_q_network_vars(), &stream) .with_context(|| format!("extract_dueling_weights failed for fold {}", fold))?; // Extract branching weights if branching DQN is active use ml::cuda_pipeline::gpu_weights::extract_branching_weights; - let branching_weights = dqn.branching_q_network.as_ref() + let branching_result = dqn.branching_q_network.as_ref() .map(|br| extract_branching_weights(br.vars(), &stream)) .transpose() .with_context(|| format!("extract_branching_weights failed for fold {}", fold))?; + let (_b_backing, branching_weights) = match branching_result { + Some((b, w)) => (Some(b), Some(w)), + None => (None, None), + }; let dqn_cfg = ml::cuda_pipeline::gpu_backtest_evaluator::DqnBacktestConfig::from_network_dims(network_dims); diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index d784a57f8..7861a4cf5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -1616,10 +1616,10 @@ impl GpuBacktestEvaluator { // Helper: convert a single f32 weight tensor to f32 into the flat buffer at current offset. macro_rules! copy_weight { - ($slice:expr, $idx:expr) => {{ + ($ptr:expr, $idx:expr) => {{ let n_elems = param_sizes[$idx]; if n_elems > 0 { - let src = raw_bf16_ptr($slice, &self.stream); + let src: u64 = $ptr; let dst = dst_base + f32_offset; let n_i32 = n_elems as i32; let blocks = ((n_elems + 255) / 256) as u32; @@ -1644,55 +1644,52 @@ impl GpuBacktestEvaluator { } // Shared trunk: w_s1, b_s1, w_s2, b_s2 - copy_weight!(&online.w_s1, 0); - copy_weight!(&online.b_s1, 1); - copy_weight!(&online.w_s2, 2); - copy_weight!(&online.b_s2, 3); + copy_weight!(online.w_s1, 0); + copy_weight!(online.b_s1, 1); + copy_weight!(online.w_s2, 2); + copy_weight!(online.b_s2, 3); // Value head: w_v1, b_v1, w_v2, b_v2 - copy_weight!(&online.w_v1, 4); - copy_weight!(&online.b_v1, 5); - copy_weight!(&online.w_v2, 6); - copy_weight!(&online.b_v2, 7); + copy_weight!(online.w_v1, 4); + copy_weight!(online.b_v1, 5); + copy_weight!(online.w_v2, 6); + copy_weight!(online.b_v2, 7); // Branch 0 (exposure — from DuelingWeightSet): w_b0fc, b_b0fc, w_b0out, b_b0out - copy_weight!(&online.w_a1, 8); - copy_weight!(&online.b_a1, 9); - copy_weight!(&online.w_a2, 10); - copy_weight!(&online.b_a2, 11); + copy_weight!(online.w_a1, 8); + copy_weight!(online.b_a1, 9); + copy_weight!(online.w_a2, 10); + copy_weight!(online.b_a2, 11); // Branch 1 (order), Branch 2 (urgency), Branch 3 (magnitude) if let Some(bw) = branching { // Branching mode: separate weights for each branch - copy_weight!(&bw.w_bo1, 12); - copy_weight!(&bw.b_bo1, 13); - copy_weight!(&bw.w_bo2, 14); - copy_weight!(&bw.b_bo2, 15); - copy_weight!(&bw.w_bu1, 16); - copy_weight!(&bw.b_bu1, 17); - copy_weight!(&bw.w_bu2, 18); - copy_weight!(&bw.b_bu2, 19); - copy_weight!(&bw.w_bg1, 20); - copy_weight!(&bw.b_bg1, 21); - copy_weight!(&bw.w_bg2, 22); - copy_weight!(&bw.b_bg2, 23); + copy_weight!(bw.w_bo1, 12); + copy_weight!(bw.b_bo1, 13); + copy_weight!(bw.w_bo2, 14); + copy_weight!(bw.b_bo2, 15); + copy_weight!(bw.w_bu1, 16); + copy_weight!(bw.b_bu1, 17); + copy_weight!(bw.w_bu2, 18); + copy_weight!(bw.b_bu2, 19); + copy_weight!(bw.w_bg1, 20); + copy_weight!(bw.b_bg1, 21); + copy_weight!(bw.w_bg2, 22); + copy_weight!(bw.b_bg2, 23); } else { // Non-branching: replicate branch 0 weights for branches 1, 2, and 3. - // The FC layers (w_b1fc/b_b1fc) have the same shape as branch 0 FC, - // and the out layers differ only in action count (but for non-branching - // DQN all branch sizes are the same = branch_0_size). - copy_weight!(&online.w_a1, 12); - copy_weight!(&online.b_a1, 13); - copy_weight!(&online.w_a2, 14); - copy_weight!(&online.b_a2, 15); - copy_weight!(&online.w_a1, 16); - copy_weight!(&online.b_a1, 17); - copy_weight!(&online.w_a2, 18); - copy_weight!(&online.b_a2, 19); - copy_weight!(&online.w_a1, 20); - copy_weight!(&online.b_a1, 21); - copy_weight!(&online.w_a2, 22); - copy_weight!(&online.b_a2, 23); + copy_weight!(online.w_a1, 12); + copy_weight!(online.b_a1, 13); + copy_weight!(online.w_a2, 14); + copy_weight!(online.b_a2, 15); + copy_weight!(online.w_a1, 16); + copy_weight!(online.b_a1, 17); + copy_weight!(online.w_a2, 18); + copy_weight!(online.b_a2, 19); + copy_weight!(online.w_a1, 20); + copy_weight!(online.b_a1, 21); + copy_weight!(online.w_a2, 22); + copy_weight!(online.b_a2, 23); } Ok(()) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 3842eb9b0..85c9c0a69 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -5851,19 +5851,19 @@ impl GpuDqnTrainer { let dst_base = self.params_buf.raw_ptr(); // Ordered to match GOFF_* layout (indices 0-23; 24-25 = bottleneck handled separately) - let slices: [&CudaSlice; 24] = [ - &online_d.w_s1, &online_d.b_s1, - &online_d.w_s2, &online_d.b_s2, - &online_d.w_v1, &online_d.b_v1, - &online_d.w_v2, &online_d.b_v2, - &online_d.w_a1, &online_d.b_a1, // branch 0 = direction - &online_d.w_a2, &online_d.b_a2, - &online_b.w_bo1, &online_b.b_bo1, // branch 1 = order - &online_b.w_bo2, &online_b.b_bo2, - &online_b.w_bu1, &online_b.b_bu1, // branch 2 = urgency - &online_b.w_bu2, &online_b.b_bu2, - &online_b.w_bg1, &online_b.b_bg1, // branch 3 = magnitude - &online_b.w_bg2, &online_b.b_bg2, + let ptrs: [(u64, usize); 24] = [ + (online_d.w_s1, online_d.w_s1_len), (online_d.b_s1, online_d.b_s1_len), + (online_d.w_s2, online_d.w_s2_len), (online_d.b_s2, online_d.b_s2_len), + (online_d.w_v1, online_d.w_v1_len), (online_d.b_v1, online_d.b_v1_len), + (online_d.w_v2, online_d.w_v2_len), (online_d.b_v2, online_d.b_v2_len), + (online_d.w_a1, online_d.w_a1_len), (online_d.b_a1, online_d.b_a1_len), + (online_d.w_a2, online_d.w_a2_len), (online_d.b_a2, online_d.b_a2_len), + (online_b.w_bo1, online_b.w_bo1_len), (online_b.b_bo1, online_b.b_bo1_len), + (online_b.w_bo2, online_b.w_bo2_len), (online_b.b_bo2, online_b.b_bo2_len), + (online_b.w_bu1, online_b.w_bu1_len), (online_b.b_bu1, online_b.b_bu1_len), + (online_b.w_bu2, online_b.w_bu2_len), (online_b.b_bu2, online_b.b_bu2_len), + (online_b.w_bg1, online_b.w_bg1_len), (online_b.b_bg1, online_b.b_bg1_len), + (online_b.w_bg2, online_b.w_bg2_len), (online_b.b_bg2, online_b.b_bg2_len), ]; let weight_names = [ @@ -5876,27 +5876,25 @@ impl GpuDqnTrainer { ]; let mut byte_offset: u64 = 0; - for (i, slice) in slices.iter().enumerate() { - let actual = slice.len(); + for (i, &(src, actual)) in ptrs.iter().enumerate() { if actual != sizes[i] { // When bottleneck is active, w_s1/b_s1 have different dimensions in VarStore // vs flat buffer (VarStore uses full state_dim, flat buffer uses reduced). // Skip these — the flat buffer is Xavier-initialized at the correct dimension. if self.config.bottleneck_dim > 0 && (i == 0 || i == 1) { tracing::debug!( - "flatten[{}] ({}) bottleneck skip: VarStore={}, flat={}", + "flatten[{}] ({}) bottleneck skip: src={}, flat={}", i, weight_names[i], actual, sizes[i] ); byte_offset += (sizes[i] * std::mem::size_of::()) as u64; continue; } return Err(MLError::ModelError(format!( - "flatten[{}] ({}) size mismatch: CudaSlice has {} elems, expected {}", + "flatten[{}] ({}) size mismatch: src has {} elems, expected {}", i, weight_names[i], actual, sizes[i], ))); } let num_bytes = sizes[i] * std::mem::size_of::(); - let src = slice.raw_ptr(); dtod_copy(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten")?; byte_offset += num_bytes as u64; } @@ -5904,43 +5902,46 @@ impl GpuDqnTrainer { Ok(()) } - /// Copy flat `params_buf` back to 24 individual f32 weight tensors. + /// No-op when weight sets are already pointer views into params_buf. /// - /// Called after each Adam step to sync the canonical weight storage. - /// Pure device-to-device copies — zero host roundtrip. - /// Indices 24-25 (bottleneck) are not backed by individual CudaSlice fields. - /// - /// During CUDA Graph capture, these d2d copies are recorded and replayed. - /// The device pointers are stable (CudaSlice allocations don't move). + /// When DuelingWeightSet/BranchingWeightSet point directly into params_buf + /// (created via `from_flat_buffer()`), unflatten is a no-op — the pointers + /// already reference the correct regions. When the weight sets point at + /// separate allocations (experience collector, ensemble), this copies the + /// flat buffer content back to the individual regions. pub(crate) fn unflatten_online_weights( &self, online_d: &DuelingWeightSet, online_b: &BranchingWeightSet, ) -> Result<(), MLError> { let sizes = compute_param_sizes(&self.config); - // Read from params_buf (single f32 buffer — no shadow). let src_base = self.params_buf.raw_ptr(); + // Check if the weight set points directly into params_buf (zero-copy path). + // If w_s1 points at the start of params_buf, all pointers are views — skip D2D. + if online_d.w_s1 == src_base { + return Ok(()); + } + // Must match GOFF_* order exactly (indices 0-23) - let slices: [&CudaSlice; 24] = [ - &online_d.w_s1, &online_d.b_s1, - &online_d.w_s2, &online_d.b_s2, - &online_d.w_v1, &online_d.b_v1, - &online_d.w_v2, &online_d.b_v2, - &online_d.w_a1, &online_d.b_a1, - &online_d.w_a2, &online_d.b_a2, - &online_b.w_bo1, &online_b.b_bo1, - &online_b.w_bo2, &online_b.b_bo2, - &online_b.w_bu1, &online_b.b_bu1, - &online_b.w_bu2, &online_b.b_bu2, - &online_b.w_bg1, &online_b.b_bg1, - &online_b.w_bg2, &online_b.b_bg2, + let ptrs: [u64; 24] = [ + online_d.w_s1, online_d.b_s1, + online_d.w_s2, online_d.b_s2, + online_d.w_v1, online_d.b_v1, + online_d.w_v2, online_d.b_v2, + online_d.w_a1, online_d.b_a1, + online_d.w_a2, online_d.b_a2, + online_b.w_bo1, online_b.b_bo1, + online_b.w_bo2, online_b.b_bo2, + online_b.w_bu1, online_b.b_bu1, + online_b.w_bu2, online_b.b_bu2, + online_b.w_bg1, online_b.b_bg1, + online_b.w_bg2, online_b.b_bg2, ]; let mut byte_offset: u64 = 0; - for (i, slice) in slices.iter().enumerate() { + for (i, &dst) in ptrs.iter().enumerate() { let num_bytes = sizes[i] * std::mem::size_of::(); - let dst = slice.raw_ptr(); dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten")?; byte_offset += num_bytes as u64; } @@ -5963,22 +5964,26 @@ impl GpuDqnTrainer { target_b: &BranchingWeightSet, ) -> Result<(), MLError> { let sizes = compute_param_sizes(&self.config); - // Flatten directly into target_params_buf (single buffer, no shadow). let dst_base = self.target_params_buf.raw_ptr(); - let slices: [&CudaSlice; 24] = [ - &target_d.w_s1, &target_d.b_s1, - &target_d.w_s2, &target_d.b_s2, - &target_d.w_v1, &target_d.b_v1, - &target_d.w_v2, &target_d.b_v2, - &target_d.w_a1, &target_d.b_a1, - &target_d.w_a2, &target_d.b_a2, - &target_b.w_bo1, &target_b.b_bo1, - &target_b.w_bo2, &target_b.b_bo2, - &target_b.w_bu1, &target_b.b_bu1, - &target_b.w_bu2, &target_b.b_bu2, - &target_b.w_bg1, &target_b.b_bg1, - &target_b.w_bg2, &target_b.b_bg2, + // If target already points into target_params_buf, no-op. + if target_d.w_s1 == dst_base { + return Ok(()); + } + + let ptrs: [u64; 24] = [ + target_d.w_s1, target_d.b_s1, + target_d.w_s2, target_d.b_s2, + target_d.w_v1, target_d.b_v1, + target_d.w_v2, target_d.b_v2, + target_d.w_a1, target_d.b_a1, + target_d.w_a2, target_d.b_a2, + target_b.w_bo1, target_b.b_bo1, + target_b.w_bo2, target_b.b_bo2, + target_b.w_bu1, target_b.b_bu1, + target_b.w_bu2, target_b.b_bu2, + target_b.w_bg1, target_b.b_bg1, + target_b.w_bg2, target_b.b_bg2, ]; // Sync stream to catch any deferred errors from graph replay @@ -5987,9 +5992,8 @@ impl GpuDqnTrainer { return Err(MLError::ModelError(format!("flatten_target pre-sync: {sync_r:?}"))); } let mut byte_offset: u64 = 0; - for (i, slice) in slices.iter().enumerate() { + for (i, &src) in ptrs.iter().enumerate() { let num_bytes = sizes[i] * std::mem::size_of::(); - let src = slice.raw_ptr(); let dst = dst_base + byte_offset; dtod_copy(dst, src, num_bytes, &self.stream, i, "flatten_target")?; byte_offset += num_bytes as u64; @@ -5998,40 +6002,40 @@ impl GpuDqnTrainer { Ok(()) } - /// Copy flat `target_params_buf` back to 24 individual target weight tensors. + /// Scatter flat target_params_buf back to individual target weight pointers. /// - /// Called after the fused EMA kernel to scatter updated flat target weights - /// back into the individual `DuelingWeightSet` + `BranchingWeightSet` tensors. - /// Indices 24-25 (bottleneck) are not backed by individual CudaSlice fields. - /// Pure device-to-device copies -- zero host roundtrip. + /// No-op when target weight set points directly into target_params_buf. fn unflatten_target_weights( &self, target_d: &DuelingWeightSet, target_b: &BranchingWeightSet, ) -> Result<(), MLError> { let sizes = compute_param_sizes(&self.config); - // Read from target_params_buf (single buffer, no shadow). let src_base = self.target_params_buf.raw_ptr(); - let slices: [&CudaSlice; 24] = [ - &target_d.w_s1, &target_d.b_s1, - &target_d.w_s2, &target_d.b_s2, - &target_d.w_v1, &target_d.b_v1, - &target_d.w_v2, &target_d.b_v2, - &target_d.w_a1, &target_d.b_a1, - &target_d.w_a2, &target_d.b_a2, - &target_b.w_bo1, &target_b.b_bo1, - &target_b.w_bo2, &target_b.b_bo2, - &target_b.w_bu1, &target_b.b_bu1, - &target_b.w_bu2, &target_b.b_bu2, - &target_b.w_bg1, &target_b.b_bg1, - &target_b.w_bg2, &target_b.b_bg2, + // No-op if target already points into target_params_buf. + if target_d.w_s1 == src_base { + return Ok(()); + } + + let ptrs: [u64; 24] = [ + target_d.w_s1, target_d.b_s1, + target_d.w_s2, target_d.b_s2, + target_d.w_v1, target_d.b_v1, + target_d.w_v2, target_d.b_v2, + target_d.w_a1, target_d.b_a1, + target_d.w_a2, target_d.b_a2, + target_b.w_bo1, target_b.b_bo1, + target_b.w_bo2, target_b.b_bo2, + target_b.w_bu1, target_b.b_bu1, + target_b.w_bu2, target_b.b_bu2, + target_b.w_bg1, target_b.b_bg1, + target_b.w_bg2, target_b.b_bg2, ]; let mut byte_offset: u64 = 0; - for (i, slice) in slices.iter().enumerate() { + for (i, &dst) in ptrs.iter().enumerate() { let num_bytes = sizes[i] * std::mem::size_of::(); - let dst = slice.raw_ptr(); dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?; byte_offset += num_bytes as u64; } @@ -6057,8 +6061,8 @@ impl GpuDqnTrainer { &mut self, _online_d: &DuelingWeightSet, _online_b: &BranchingWeightSet, - target_d: &mut DuelingWeightSet, - target_b: &mut BranchingWeightSet, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, tau: f32, ) -> Result<(), MLError> { // Disable cudarc event tracking for the duration of this method. diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 0e945dc5f..22e4c6335 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -540,7 +540,7 @@ pub struct GpuExperienceCollector { exp_v_logits: CudaSlice, // [N, num_atoms] f32 (output layer — no f32 truncation) exp_b_logits: CudaSlice, // [N, total_branch_atoms] f32 - // ── Legacy weight sets (for sync_online_weights/sync_target_weights compat) ── + // ── Weight sets (pointer views + backing storage) ── online_weights: DuelingWeightSet, target_weights: DuelingWeightSet, curiosity_weights: CuriosityWeightSet, @@ -548,6 +548,12 @@ pub struct GpuExperienceCollector { target_rmsnorm: RmsNormWeightSet, online_branching: BranchingWeightSet, target_branching: BranchingWeightSet, + /// Backing GPU allocations that keep the weight pointers valid. + #[allow(dead_code)] + _weight_backings: ( + super::gpu_weights::DuelingWeightBacking, super::gpu_weights::DuelingWeightBacking, + super::gpu_weights::BranchingWeightBacking, super::gpu_weights::BranchingWeightBacking, + ), /// GPU buffer for trade_stats_reduce output: [TRADE_STATS_FLOATS] = 6 floats. trade_stats_buf: CudaSlice, @@ -806,8 +812,8 @@ impl GpuExperienceCollector { ); // ── Step 4: Extract legacy weight sets ────────────────────────── - let online_weights = extract_dueling_weights_branching(online_vars, &stream)?; - let target_weights = extract_dueling_weights_branching(target_vars, &stream)?; + let (_online_d_backing, online_weights) = extract_dueling_weights_branching(online_vars, &stream)?; + let (_target_d_backing, target_weights) = extract_dueling_weights_branching(target_vars, &stream)?; let curiosity_weights = match curiosity_vars { Some(vars) => extract_curiosity_weights(vars, &stream)?, None => { @@ -832,10 +838,8 @@ impl GpuExperienceCollector { }; info!("GPU experience collector: extracting branching weights (order + urgency heads)"); - let (online_branching, target_branching) = ( - extract_branching_weights(online_vars, &stream)?, - extract_branching_weights(target_vars, &stream)?, - ); + let (_online_b_backing, online_branching) = extract_branching_weights(online_vars, &stream)?; + let (_target_b_backing, target_branching) = extract_branching_weights(target_vars, &stream)?; // ── Step 5: L2 cache persistence hints (Hopper GPUs) ──────────── { @@ -1150,6 +1154,7 @@ impl GpuExperienceCollector { target_rmsnorm, online_branching, target_branching, + _weight_backings: (_online_d_backing, _target_d_backing, _online_b_backing, _target_b_backing), portfolio_states, rng_states, episode_starts_buf, @@ -2264,28 +2269,28 @@ impl GpuExperienceCollector { /// Extracts per-tensor weights from VarStore into the legacy weight sets, /// then flattens them into the contiguous `online_params_flat` buffer for cuBLAS. pub fn sync_online_weights(&mut self, vars: &GpuVarStore) -> Result<(), MLError> { - sync_dueling_weights_branching(vars, &mut self.online_weights, &self.stream)?; + sync_dueling_weights_branching(vars, &self.online_weights, &self.stream)?; self.flatten_online_weights() } /// Sync target Q-network weights from a `GpuVarStore` (after soft update). /// - /// Updates the legacy target weight sets. The target network is not used + /// Updates the target weight views. The target network is not used /// by the cuBLAS timestep loop (only online network is used for action /// selection during experience collection). pub fn sync_target_weights(&mut self, vars: &GpuVarStore) -> Result<(), MLError> { - sync_dueling_weights_branching(vars, &mut self.target_weights, &self.stream) + sync_dueling_weights_branching(vars, &self.target_weights, &self.stream) } /// Sync online branching DQN extra head weights (branches 1+2). pub fn sync_online_branching(&mut self, vars: &GpuVarStore) -> Result<(), MLError> { - sync_branching_weights(vars, &mut self.online_branching, &self.stream)?; + sync_branching_weights(vars, &self.online_branching, &self.stream)?; self.flatten_online_weights() } /// Sync target branching DQN extra head weights (branches 1+2). pub fn sync_target_branching(&mut self, vars: &GpuVarStore) -> Result<(), MLError> { - sync_branching_weights(vars, &mut self.target_branching, &self.stream) + sync_branching_weights(vars, &self.target_branching, &self.stream) } /// Sync curiosity forward model weights from a `GpuVarStore`. @@ -2333,52 +2338,51 @@ impl GpuExperienceCollector { let mut byte_offset: u64 = 0; let dst_base = self.online_params_flat.raw_ptr(); - // Helper: copy a single weight tensor into the flat buffer at current offset. + // Helper: copy a single weight tensor from its device pointer into the flat buffer. macro_rules! copy_weight { - ($slice:expr, $idx:expr) => {{ + ($ptr:expr, $idx:expr) => {{ let n_elems = self.param_sizes[$idx]; let n_bytes = n_elems * std::mem::size_of::(); - let src = $slice.raw_ptr(); - dtod_copy(dst_base + byte_offset, src, n_bytes, &self.stream, $idx, "flatten")?; + dtod_copy(dst_base + byte_offset, $ptr, n_bytes, &self.stream, $idx, "flatten")?; byte_offset += n_bytes as u64; }}; } // Shared trunk: w_s1, b_s1, w_s2, b_s2 - copy_weight!(&self.online_weights.w_s1, 0); - copy_weight!(&self.online_weights.b_s1, 1); - copy_weight!(&self.online_weights.w_s2, 2); - copy_weight!(&self.online_weights.b_s2, 3); + copy_weight!(self.online_weights.w_s1, 0); + copy_weight!(self.online_weights.b_s1, 1); + copy_weight!(self.online_weights.w_s2, 2); + copy_weight!(self.online_weights.b_s2, 3); // Value head: w_v1, b_v1, w_v2, b_v2 - copy_weight!(&self.online_weights.w_v1, 4); - copy_weight!(&self.online_weights.b_v1, 5); - copy_weight!(&self.online_weights.w_v2, 6); - copy_weight!(&self.online_weights.b_v2, 7); + copy_weight!(self.online_weights.w_v1, 4); + copy_weight!(self.online_weights.b_v1, 5); + copy_weight!(self.online_weights.w_v2, 6); + copy_weight!(self.online_weights.b_v2, 7); // Branch 0 (direction — from DuelingWeightSet): w_b0fc, b_b0fc, w_b0out, b_b0out - copy_weight!(&self.online_weights.w_a1, 8); - copy_weight!(&self.online_weights.b_a1, 9); - copy_weight!(&self.online_weights.w_a2, 10); - copy_weight!(&self.online_weights.b_a2, 11); + copy_weight!(self.online_weights.w_a1, 8); + copy_weight!(self.online_weights.b_a1, 9); + copy_weight!(self.online_weights.w_a2, 10); + copy_weight!(self.online_weights.b_a2, 11); // Branch 1 (order — from BranchingWeightSet): w_b1fc, b_b1fc, w_b1out, b_b1out - copy_weight!(&self.online_branching.w_bo1, 12); - copy_weight!(&self.online_branching.b_bo1, 13); - copy_weight!(&self.online_branching.w_bo2, 14); - copy_weight!(&self.online_branching.b_bo2, 15); + copy_weight!(self.online_branching.w_bo1, 12); + copy_weight!(self.online_branching.b_bo1, 13); + copy_weight!(self.online_branching.w_bo2, 14); + copy_weight!(self.online_branching.b_bo2, 15); // Branch 2 (urgency — from BranchingWeightSet): w_b2fc, b_b2fc, w_b2out, b_b2out - copy_weight!(&self.online_branching.w_bu1, 16); - copy_weight!(&self.online_branching.b_bu1, 17); - copy_weight!(&self.online_branching.w_bu2, 18); - copy_weight!(&self.online_branching.b_bu2, 19); + copy_weight!(self.online_branching.w_bu1, 16); + copy_weight!(self.online_branching.b_bu1, 17); + copy_weight!(self.online_branching.w_bu2, 18); + copy_weight!(self.online_branching.b_bu2, 19); // Branch 3 (magnitude — from BranchingWeightSet): w_b3fc, b_b3fc, w_b3out, b_b3out - copy_weight!(&self.online_branching.w_bg1, 20); - copy_weight!(&self.online_branching.b_bg1, 21); - copy_weight!(&self.online_branching.w_bg2, 22); - copy_weight!(&self.online_branching.b_bg2, 23); + copy_weight!(self.online_branching.w_bg1, 20); + copy_weight!(self.online_branching.b_bg1, 21); + copy_weight!(self.online_branching.w_bg2, 22); + copy_weight!(self.online_branching.b_bg2, 23); Ok(()) } diff --git a/crates/ml/src/cuda_pipeline/gpu_weights.rs b/crates/ml/src/cuda_pipeline/gpu_weights.rs index dff63bbb7..724ca4352 100644 --- a/crates/ml/src/cuda_pipeline/gpu_weights.rs +++ b/crates/ml/src/cuda_pipeline/gpu_weights.rs @@ -1,18 +1,16 @@ //! Weight extraction and GPU upload for DQN experience collection kernel. //! -//! Extracts neural network weights from `GpuVarStore` objects and copies -//! them as flat `CudaSlice` GPU buffers for the CUDA experience -//! collection kernels (`experience_kernels.cu` + cuBLAS). -//! -//! Two weight sets are managed: -//! - **Dueling Q-Network** (online + target): shared layers, value head, advantage head (12 tensors each) -//! - **Curiosity Forward Model**: two fully-connected layers (4 tensors) +//! Two weight-set representations: +//! - **`DuelingWeightSet` / `BranchingWeightSet`**: raw `u64` device-pointer + `usize` +//! element-count pairs that can be **zero-copy views** into the flat `params_buf` +//! (training hot path) or **owned allocations** (experience collector, ensemble heads, +//! tests). No `CudaSlice` ownership — the pointed-to memory lives in `params_buf` +//! or in a separately-allocated `CudaSlice` held by the caller. +//! - **Curiosity / PPO / RMSNorm weight sets**: still use owned `CudaSlice` +//! because they are not part of the flat training buffer. //! //! Weights are stored row-major `[out_features, in_features]`. The CUDA //! kernel reads them in the same layout. -//! -//! All weights are `CudaSlice` on GPU. Extraction and sync use -//! device-to-device copy (zero CPU roundtrip). // CUDA FFI module — memcpy_dtod_async requires unsafe by design. #![allow(unsafe_code)] @@ -23,6 +21,7 @@ use ml_core::cuda_autograd::GpuVarStore; use tracing::info; use crate::MLError; +use super::gpu_dqn_trainer::{NUM_WEIGHT_TENSORS, compute_param_sizes, align4, GpuDqnTrainConfig}; // --------------------------------------------------------------------------- // Weight name constants @@ -102,7 +101,14 @@ const PPO_CRITIC_WEIGHT_NAMES: [&str; 12] = [ // Structs // --------------------------------------------------------------------------- -/// GPU buffers holding all 12 weight tensors of a Dueling Q-Network. +/// Raw device-pointer + element-count view into GPU memory holding 12 weight +/// tensors of a Dueling Q-Network. +/// +/// Fields are raw `u64` CUDA device pointers and `usize` element counts. +/// They may point into a flat `params_buf` (zero-copy training path) or into +/// independently-allocated `CudaSlice` buffers (experience collector, +/// ensemble heads, tests). The caller must ensure the underlying allocation +/// outlives this struct. /// /// When used with `BranchingDuelingQNetwork` (distributional C51 always enabled): /// - Shared layers: `[SHARED_H1, STATE_DIM]`, `[SHARED_H1]`, `[SHARED_H2, SHARED_H1]`, `[SHARED_H2]` @@ -110,30 +116,96 @@ const PPO_CRITIC_WEIGHT_NAMES: [&str; 12] = [ /// - Exposure head: `[ADV_H, SHARED_H2]`, `[ADV_H]`, `[BRANCH_0_SIZE*NUM_ATOMS, ADV_H]`, `[BRANCH_0_SIZE*NUM_ATOMS]` /// /// Defaults: SHARED_H1=SHARED_H2=256, VALUE_H=ADV_H=128, NUM_ATOMS=51, BRANCH_0_SIZE=5. -#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug +#[derive(Debug)] pub struct DuelingWeightSet { // Shared layers - pub w_s1: CudaSlice, // shared_0.weight [SHARED_H1, STATE_DIM] - pub b_s1: CudaSlice, // shared_0.bias [SHARED_H1] - pub w_s2: CudaSlice, // shared_1.weight [SHARED_H2, SHARED_H1] - pub b_s2: CudaSlice, // shared_1.bias [SHARED_H2] + pub w_s1: u64, pub w_s1_len: usize, // shared_0.weight [SHARED_H1, STATE_DIM] + pub b_s1: u64, pub b_s1_len: usize, // shared_0.bias [SHARED_H1] + pub w_s2: u64, pub w_s2_len: usize, // shared_1.weight [SHARED_H2, SHARED_H1] + pub b_s2: u64, pub b_s2_len: usize, // shared_1.bias [SHARED_H2] // Value head (distributional: NUM_ATOMS outputs, not 1) - pub w_v1: CudaSlice, // value_fc.weight [VALUE_H, SHARED_H2] - pub b_v1: CudaSlice, // value_fc.bias [VALUE_H] - pub w_v2: CudaSlice, // value_out.weight [NUM_ATOMS, VALUE_H] - pub b_v2: CudaSlice, // value_out.bias [NUM_ATOMS] + pub w_v1: u64, pub w_v1_len: usize, // value_fc.weight [VALUE_H, SHARED_H2] + pub b_v1: u64, pub b_v1_len: usize, // value_fc.bias [VALUE_H] + pub w_v2: u64, pub w_v2_len: usize, // value_out.weight [NUM_ATOMS, VALUE_H] + pub b_v2: u64, pub b_v2_len: usize, // value_out.bias [NUM_ATOMS] // Exposure head / branch 0 (distributional: BRANCH_0_SIZE*NUM_ATOMS outputs) - pub w_a1: CudaSlice, // branch_0_fc.weight [ADV_H, SHARED_H2] - pub b_a1: CudaSlice, // branch_0_fc.bias [ADV_H] - pub w_a2: CudaSlice, // branch_0_out.weight [BRANCH_0_SIZE*NUM_ATOMS, ADV_H] - pub b_a2: CudaSlice, // branch_0_out.bias [BRANCH_0_SIZE*NUM_ATOMS] + pub w_a1: u64, pub w_a1_len: usize, // branch_0_fc.weight [ADV_H, SHARED_H2] + pub b_a1: u64, pub b_a1_len: usize, // branch_0_fc.bias [ADV_H] + pub w_a2: u64, pub w_a2_len: usize, // branch_0_out.weight [BRANCH_0_SIZE*NUM_ATOMS, ADV_H] + pub b_a2: u64, pub b_a2_len: usize, // branch_0_out.bias [BRANCH_0_SIZE*NUM_ATOMS] } -/// GPU buffers holding the 3 extra advantage head weight sets for Branching DQN. +impl DuelingWeightSet { + /// Create pointer views into a flat `params_buf` at the GOFF_* layout offsets. + /// + /// `base_ptr` is the raw device pointer of the flat buffer. `sizes` is the + /// 26-element array from `compute_param_sizes()`. Only the first 12 entries + /// (indices 0..12) are used for the dueling weight set. Each tensor is + /// padded to `align4()` for 16-byte cuBLAS alignment. + pub fn from_flat_buffer(base_ptr: u64, sizes: &[usize; NUM_WEIGHT_TENSORS]) -> Self { + let mut offset = 0u64; + let mut ptr_at = |i: usize| -> (u64, usize) { + let p = base_ptr + offset; + let s = sizes[i]; + offset += (align4(s) * std::mem::size_of::()) as u64; + (p, s) + }; + let (w_s1, w_s1_len) = ptr_at(0); + let (b_s1, b_s1_len) = ptr_at(1); + let (w_s2, w_s2_len) = ptr_at(2); + let (b_s2, b_s2_len) = ptr_at(3); + let (w_v1, w_v1_len) = ptr_at(4); + let (b_v1, b_v1_len) = ptr_at(5); + let (w_v2, w_v2_len) = ptr_at(6); + let (b_v2, b_v2_len) = ptr_at(7); + let (w_a1, w_a1_len) = ptr_at(8); + let (b_a1, b_a1_len) = ptr_at(9); + let (w_a2, w_a2_len) = ptr_at(10); + let (b_a2, b_a2_len) = ptr_at(11); + Self { + w_s1, w_s1_len, b_s1, b_s1_len, + w_s2, w_s2_len, b_s2, b_s2_len, + w_v1, w_v1_len, b_v1, b_v1_len, + w_v2, w_v2_len, b_v2, b_v2_len, + w_a1, w_a1_len, b_a1, b_a1_len, + w_a2, w_a2_len, b_a2, b_a2_len, + } + } + + /// Create a DuelingWeightSet pointing at owned `CudaSlice` allocations. + /// + /// Extracts raw device pointers from each slice. The caller MUST keep the + /// slices alive for as long as this weight set is used. + pub fn from_slices( + w_s1: &CudaSlice, b_s1: &CudaSlice, + w_s2: &CudaSlice, b_s2: &CudaSlice, + w_v1: &CudaSlice, b_v1: &CudaSlice, + w_v2: &CudaSlice, b_v2: &CudaSlice, + w_a1: &CudaSlice, b_a1: &CudaSlice, + w_a2: &CudaSlice, b_a2: &CudaSlice, + ) -> Self { + Self { + w_s1: w_s1.raw_ptr(), w_s1_len: w_s1.len(), + b_s1: b_s1.raw_ptr(), b_s1_len: b_s1.len(), + w_s2: w_s2.raw_ptr(), w_s2_len: w_s2.len(), + b_s2: b_s2.raw_ptr(), b_s2_len: b_s2.len(), + w_v1: w_v1.raw_ptr(), w_v1_len: w_v1.len(), + b_v1: b_v1.raw_ptr(), b_v1_len: b_v1.len(), + w_v2: w_v2.raw_ptr(), w_v2_len: w_v2.len(), + b_v2: b_v2.raw_ptr(), b_v2_len: b_v2.len(), + w_a1: w_a1.raw_ptr(), w_a1_len: w_a1.len(), + b_a1: b_a1.raw_ptr(), b_a1_len: b_a1.len(), + w_a2: w_a2.raw_ptr(), w_a2_len: w_a2.len(), + b_a2: b_a2.raw_ptr(), b_a2_len: b_a2.len(), + } + } +} + +/// Raw device-pointer + element-count view for the 3 extra advantage head +/// weight sets of Branching DQN. /// -/// In branching mode, the standard advantage head (w_a1/b_a1/w_a2/b_a2 in `DuelingWeightSet`) -/// becomes the direction head (branch 0). This struct holds the order head (branch 1), -/// urgency head (branch 2), and magnitude head (branch 3). +/// Same pointer semantics as `DuelingWeightSet` — may be zero-copy views into +/// `params_buf` or point at independently-allocated `CudaSlice` buffers. /// /// Layout matches `BranchingDuelingQNetwork` in `crates/ml-dqn/src/branching.rs`: /// - Order FC: `[ADV_H, SHARED_H2]`, `[ADV_H]` @@ -144,23 +216,137 @@ pub struct DuelingWeightSet { /// - Magnitude Out: `[BRANCH_3_SIZE*NUM_ATOMS, ADV_H]`, `[BRANCH_3_SIZE*NUM_ATOMS]` /// /// Defaults: ADV_H=128, SHARED_H2=256, BRANCH_1/2/3_SIZE=3, NUM_ATOMS=51. -#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug +#[derive(Debug)] pub struct BranchingWeightSet { // Order head (branch 1) - pub w_bo1: CudaSlice, // branch_1_fc.weight [ADV_H, SHARED_H2] - pub b_bo1: CudaSlice, // branch_1_fc.bias [ADV_H] - pub w_bo2: CudaSlice, // branch_1_out.weight [BRANCH_1_SIZE*NUM_ATOMS, ADV_H] - pub b_bo2: CudaSlice, // branch_1_out.bias [BRANCH_1_SIZE*NUM_ATOMS] + pub w_bo1: u64, pub w_bo1_len: usize, // branch_1_fc.weight [ADV_H, SHARED_H2] + pub b_bo1: u64, pub b_bo1_len: usize, // branch_1_fc.bias [ADV_H] + pub w_bo2: u64, pub w_bo2_len: usize, // branch_1_out.weight [BRANCH_1_SIZE*NUM_ATOMS, ADV_H] + pub b_bo2: u64, pub b_bo2_len: usize, // branch_1_out.bias [BRANCH_1_SIZE*NUM_ATOMS] // Urgency head (branch 2) - pub w_bu1: CudaSlice, // branch_2_fc.weight [ADV_H, SHARED_H2] - pub b_bu1: CudaSlice, // branch_2_fc.bias [ADV_H] - pub w_bu2: CudaSlice, // branch_2_out.weight [BRANCH_2_SIZE*NUM_ATOMS, ADV_H] - pub b_bu2: CudaSlice, // branch_2_out.bias [BRANCH_2_SIZE*NUM_ATOMS] + pub w_bu1: u64, pub w_bu1_len: usize, // branch_2_fc.weight [ADV_H, SHARED_H2] + pub b_bu1: u64, pub b_bu1_len: usize, // branch_2_fc.bias [ADV_H] + pub w_bu2: u64, pub w_bu2_len: usize, // branch_2_out.weight [BRANCH_2_SIZE*NUM_ATOMS, ADV_H] + pub b_bu2: u64, pub b_bu2_len: usize, // branch_2_out.bias [BRANCH_2_SIZE*NUM_ATOMS] // Magnitude head (branch 3) — 4-branch refactor - pub w_bg1: CudaSlice, // branch_3_fc.weight [ADV_H, SHARED_H2] - pub b_bg1: CudaSlice, // branch_3_fc.bias [ADV_H] - pub w_bg2: CudaSlice, // branch_3_out.weight [BRANCH_3_SIZE*NUM_ATOMS, ADV_H] - pub b_bg2: CudaSlice, // branch_3_out.bias [BRANCH_3_SIZE*NUM_ATOMS] + pub w_bg1: u64, pub w_bg1_len: usize, // branch_3_fc.weight [ADV_H, SHARED_H2] + pub b_bg1: u64, pub b_bg1_len: usize, // branch_3_fc.bias [ADV_H] + pub w_bg2: u64, pub w_bg2_len: usize, // branch_3_out.weight [BRANCH_3_SIZE*NUM_ATOMS, ADV_H] + pub b_bg2: u64, pub b_bg2_len: usize, // branch_3_out.bias [BRANCH_3_SIZE*NUM_ATOMS] +} + +impl BranchingWeightSet { + /// Create pointer views into a flat `params_buf` at the GOFF_* layout offsets. + /// + /// `base_ptr` is the raw device pointer of the flat buffer. `sizes` is the + /// 26-element array from `compute_param_sizes()`. Branching weights start at + /// index 12 (indices 12..24). Each tensor is padded to `align4()`. + pub fn from_flat_buffer(base_ptr: u64, sizes: &[usize; NUM_WEIGHT_TENSORS]) -> Self { + // Skip indices 0..12 (dueling weights) + let mut offset = 0u64; + for i in 0..12 { + offset += (align4(sizes[i]) * std::mem::size_of::()) as u64; + } + let mut ptr_at = |i: usize| -> (u64, usize) { + let p = base_ptr + offset; + let s = sizes[i]; + offset += (align4(s) * std::mem::size_of::()) as u64; + (p, s) + }; + let (w_bo1, w_bo1_len) = ptr_at(12); + let (b_bo1, b_bo1_len) = ptr_at(13); + let (w_bo2, w_bo2_len) = ptr_at(14); + let (b_bo2, b_bo2_len) = ptr_at(15); + let (w_bu1, w_bu1_len) = ptr_at(16); + let (b_bu1, b_bu1_len) = ptr_at(17); + let (w_bu2, w_bu2_len) = ptr_at(18); + let (b_bu2, b_bu2_len) = ptr_at(19); + let (w_bg1, w_bg1_len) = ptr_at(20); + let (b_bg1, b_bg1_len) = ptr_at(21); + let (w_bg2, w_bg2_len) = ptr_at(22); + let (b_bg2, b_bg2_len) = ptr_at(23); + Self { + w_bo1, w_bo1_len, b_bo1, b_bo1_len, + w_bo2, w_bo2_len, b_bo2, b_bo2_len, + w_bu1, w_bu1_len, b_bu1, b_bu1_len, + w_bu2, w_bu2_len, b_bu2, b_bu2_len, + w_bg1, w_bg1_len, b_bg1, b_bg1_len, + w_bg2, w_bg2_len, b_bg2, b_bg2_len, + } + } + + /// Create a BranchingWeightSet pointing at owned `CudaSlice` allocations. + /// + /// Extracts raw device pointers from each slice. The caller MUST keep the + /// slices alive for as long as this weight set is used. + pub fn from_slices( + w_bo1: &CudaSlice, b_bo1: &CudaSlice, + w_bo2: &CudaSlice, b_bo2: &CudaSlice, + w_bu1: &CudaSlice, b_bu1: &CudaSlice, + w_bu2: &CudaSlice, b_bu2: &CudaSlice, + w_bg1: &CudaSlice, b_bg1: &CudaSlice, + w_bg2: &CudaSlice, b_bg2: &CudaSlice, + ) -> Self { + Self { + w_bo1: w_bo1.raw_ptr(), w_bo1_len: w_bo1.len(), + b_bo1: b_bo1.raw_ptr(), b_bo1_len: b_bo1.len(), + w_bo2: w_bo2.raw_ptr(), w_bo2_len: w_bo2.len(), + b_bo2: b_bo2.raw_ptr(), b_bo2_len: b_bo2.len(), + w_bu1: w_bu1.raw_ptr(), w_bu1_len: w_bu1.len(), + b_bu1: b_bu1.raw_ptr(), b_bu1_len: b_bu1.len(), + w_bu2: w_bu2.raw_ptr(), w_bu2_len: w_bu2.len(), + b_bu2: b_bu2.raw_ptr(), b_bu2_len: b_bu2.len(), + w_bg1: w_bg1.raw_ptr(), w_bg1_len: w_bg1.len(), + b_bg1: b_bg1.raw_ptr(), b_bg1_len: b_bg1.len(), + w_bg2: w_bg2.raw_ptr(), w_bg2_len: w_bg2.len(), + b_bg2: b_bg2.raw_ptr(), b_bg2_len: b_bg2.len(), + } + } +} + +/// Owned GPU allocations backing a `DuelingWeightSet`. +/// +/// When weight sets are created from VarStore extraction (experience collector, +/// ensemble heads), these allocations keep the pointed-to GPU memory alive. +/// When weight sets are created via `from_flat_buffer()` (training hot path), +/// no backing storage is needed — `params_buf` owns the memory. +#[allow(missing_debug_implementations)] +pub struct DuelingWeightBacking { + pub slices: [CudaSlice; 12], +} + +impl DuelingWeightBacking { + /// Build a `DuelingWeightSet` of pointer views into these owned slices. + pub fn weight_set(&self) -> DuelingWeightSet { + DuelingWeightSet::from_slices( + &self.slices[0], &self.slices[1], + &self.slices[2], &self.slices[3], + &self.slices[4], &self.slices[5], + &self.slices[6], &self.slices[7], + &self.slices[8], &self.slices[9], + &self.slices[10], &self.slices[11], + ) + } +} + +/// Owned GPU allocations backing a `BranchingWeightSet`. +#[allow(missing_debug_implementations)] +pub struct BranchingWeightBacking { + pub slices: [CudaSlice; 12], +} + +impl BranchingWeightBacking { + /// Build a `BranchingWeightSet` of pointer views into these owned slices. + pub fn weight_set(&self) -> BranchingWeightSet { + BranchingWeightSet::from_slices( + &self.slices[0], &self.slices[1], + &self.slices[2], &self.slices[3], + &self.slices[4], &self.slices[5], + &self.slices[6], &self.slices[7], + &self.slices[8], &self.slices[9], + &self.slices[10], &self.slices[11], + ) + } } /// GPU buffers holding all 4 weight tensors of the Curiosity Forward Model. @@ -408,32 +594,20 @@ impl KernelWeightPack { ) -> Self { Self { online: GpuDuelingPtrs { - w_s1: raw_device_ptr(&online.w_s1, stream), - b_s1: raw_device_ptr(&online.b_s1, stream), - w_s2: raw_device_ptr(&online.w_s2, stream), - b_s2: raw_device_ptr(&online.b_s2, stream), - w_v1: raw_device_ptr(&online.w_v1, stream), - b_v1: raw_device_ptr(&online.b_v1, stream), - w_v2: raw_device_ptr(&online.w_v2, stream), - b_v2: raw_device_ptr(&online.b_v2, stream), - w_a1: raw_device_ptr(&online.w_a1, stream), - b_a1: raw_device_ptr(&online.b_a1, stream), - w_a2: raw_device_ptr(&online.w_a2, stream), - b_a2: raw_device_ptr(&online.b_a2, stream), + w_s1: online.w_s1, b_s1: online.b_s1, + w_s2: online.w_s2, b_s2: online.b_s2, + w_v1: online.w_v1, b_v1: online.b_v1, + w_v2: online.w_v2, b_v2: online.b_v2, + w_a1: online.w_a1, b_a1: online.b_a1, + w_a2: online.w_a2, b_a2: online.b_a2, }, target: GpuDuelingPtrs { - w_s1: raw_device_ptr(&target.w_s1, stream), - b_s1: raw_device_ptr(&target.b_s1, stream), - w_s2: raw_device_ptr(&target.w_s2, stream), - b_s2: raw_device_ptr(&target.b_s2, stream), - w_v1: raw_device_ptr(&target.w_v1, stream), - b_v1: raw_device_ptr(&target.b_v1, stream), - w_v2: raw_device_ptr(&target.w_v2, stream), - b_v2: raw_device_ptr(&target.b_v2, stream), - w_a1: raw_device_ptr(&target.w_a1, stream), - b_a1: raw_device_ptr(&target.b_a1, stream), - w_a2: raw_device_ptr(&target.w_a2, stream), - b_a2: raw_device_ptr(&target.b_a2, stream), + w_s1: target.w_s1, b_s1: target.b_s1, + w_s2: target.w_s2, b_s2: target.b_s2, + w_v1: target.w_v1, b_v1: target.b_v1, + w_v2: target.w_v2, b_v2: target.b_v2, + w_a1: target.w_a1, b_a1: target.b_a1, + w_a2: target.w_a2, b_a2: target.b_a2, }, curiosity: GpuCuriosityPtrs { w1: raw_device_ptr(&curiosity.w1, stream), @@ -448,32 +622,20 @@ impl KernelWeightPack { a: raw_device_ptr(&rmsnorm.gamma_a, stream), }, online_br: GpuBranchPtrs { - w_o1: raw_device_ptr(&online_br.w_bo1, stream), - b_o1: raw_device_ptr(&online_br.b_bo1, stream), - w_o2: raw_device_ptr(&online_br.w_bo2, stream), - b_o2: raw_device_ptr(&online_br.b_bo2, stream), - w_u1: raw_device_ptr(&online_br.w_bu1, stream), - b_u1: raw_device_ptr(&online_br.b_bu1, stream), - w_u2: raw_device_ptr(&online_br.w_bu2, stream), - b_u2: raw_device_ptr(&online_br.b_bu2, stream), - w_g1: raw_device_ptr(&online_br.w_bg1, stream), - b_g1: raw_device_ptr(&online_br.b_bg1, stream), - w_g2: raw_device_ptr(&online_br.w_bg2, stream), - b_g2: raw_device_ptr(&online_br.b_bg2, stream), + w_o1: online_br.w_bo1, b_o1: online_br.b_bo1, + w_o2: online_br.w_bo2, b_o2: online_br.b_bo2, + w_u1: online_br.w_bu1, b_u1: online_br.b_bu1, + w_u2: online_br.w_bu2, b_u2: online_br.b_bu2, + w_g1: online_br.w_bg1, b_g1: online_br.b_bg1, + w_g2: online_br.w_bg2, b_g2: online_br.b_bg2, }, target_br: GpuBranchPtrs { - w_o1: raw_device_ptr(&target_br.w_bo1, stream), - b_o1: raw_device_ptr(&target_br.b_bo1, stream), - w_o2: raw_device_ptr(&target_br.w_bo2, stream), - b_o2: raw_device_ptr(&target_br.b_bo2, stream), - w_u1: raw_device_ptr(&target_br.w_bu1, stream), - b_u1: raw_device_ptr(&target_br.b_bu1, stream), - w_u2: raw_device_ptr(&target_br.w_bu2, stream), - b_u2: raw_device_ptr(&target_br.b_bu2, stream), - w_g1: raw_device_ptr(&target_br.w_bg1, stream), - b_g1: raw_device_ptr(&target_br.b_bg1, stream), - w_g2: raw_device_ptr(&target_br.w_bg2, stream), - b_g2: raw_device_ptr(&target_br.b_bg2, stream), + w_o1: target_br.w_bo1, b_o1: target_br.b_bo1, + w_o2: target_br.w_bo2, b_o2: target_br.b_bo2, + w_u1: target_br.w_bu1, b_u1: target_br.b_bu1, + w_u2: target_br.w_bu2, b_u2: target_br.b_bu2, + w_g1: target_br.w_bg1, b_g1: target_br.b_bg1, + w_g2: target_br.w_bg2, b_g2: target_br.b_bg2, }, } } @@ -545,6 +707,33 @@ fn sync_one( Ok(()) } +/// Re-upload a single weight from `GpuVarStore` into a raw device pointer destination. +/// Uses device-to-device copy (zero CPU roundtrip). +fn sync_one_to_ptr( + vars: &GpuVarStore, + name: &str, + dst_ptr: u64, + expected_len: usize, + stream: &Arc, +) -> Result<(), MLError> { + let param = vars + .get(name) + .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))?; + + let n_elems = param.data.len(); + if n_elems != expected_len { + return Err(MLError::ModelError(format!( + "sync_to_ptr {name}: VarStore has {n_elems} elems, expected {expected_len}" + ))); + } + + let src_ptr = param.data.raw_ptr(); + let num_bytes = n_elems * std::mem::size_of::(); + dtod_copy_checked(dst_ptr, src_ptr, num_bytes, stream, name, "sync_to_ptr")?; + + Ok(()) +} + // --------------------------------------------------------------------------- // Public API // --------------------------------------------------------------------------- @@ -560,7 +749,7 @@ fn sync_one( pub fn extract_dueling_weights( vars: &GpuVarStore, stream: &Arc, -) -> Result { +) -> Result<(DuelingWeightBacking, DuelingWeightSet), MLError> { let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) @@ -592,20 +781,11 @@ pub fn extract_dueling_weights( "Dueling Q-Network weights extracted and uploaded to GPU" ); - Ok(DuelingWeightSet { - w_s1, - b_s1, - w_s2, - b_s2, - w_v1, - b_v1, - w_v2, - b_v2, - w_a1, - b_a1, - w_a2, - b_a2, - }) + let backing = DuelingWeightBacking { + slices: [w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2, w_a1, b_a1, w_a2, b_a2], + }; + let ws = backing.weight_set(); + Ok((backing, ws)) } /// Extract all 12 Branching DQN extra head weight tensors from a `GpuVarStore` @@ -618,7 +798,7 @@ pub fn extract_dueling_weights( pub fn extract_branching_weights( vars: &GpuVarStore, stream: &Arc, -) -> Result { +) -> Result<(BranchingWeightBacking, BranchingWeightSet), MLError> { let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) @@ -650,20 +830,11 @@ pub fn extract_branching_weights( "Branching DQN weights extracted and uploaded to GPU" ); - Ok(BranchingWeightSet { - w_bo1, - b_bo1, - w_bo2, - b_bo2, - w_bu1, - b_bu1, - w_bu2, - b_bu2, - w_bg1, - b_bg1, - w_bg2, - b_bg2, - }) + let backing = BranchingWeightBacking { + slices: [w_bo1, b_bo1, w_bo2, b_bo2, w_bu1, b_bu1, w_bu2, b_bu2, w_bg1, b_bg1, w_bg2, b_bg2], + }; + let ws = backing.weight_set(); + Ok((backing, ws)) } /// Extract Dueling Q-Network weights using branching GpuVarStore key names. @@ -678,7 +849,7 @@ pub fn extract_branching_weights( pub fn extract_dueling_weights_branching( vars: &GpuVarStore, stream: &Arc, -) -> Result { +) -> Result<(DuelingWeightBacking, DuelingWeightSet), MLError> { let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) @@ -701,11 +872,11 @@ pub fn extract_dueling_weights_branching( info!("Branching DQN dueling weights (branch_0=exposure) extracted and uploaded to GPU"); - Ok(DuelingWeightSet { - w_s1, b_s1, w_s2, b_s2, - w_v1, b_v1, w_v2, b_v2, - w_a1, b_a1, w_a2, b_a2, - }) + let backing = DuelingWeightBacking { + slices: [w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2, w_a1, b_a1, w_a2, b_a2], + }; + let ws = backing.weight_set(); + Ok((backing, ws)) } /// Re-upload Dueling Q-Network weights using branching GpuVarStore key names. @@ -713,23 +884,23 @@ pub fn extract_dueling_weights_branching( /// Counterpart to [`extract_dueling_weights_branching`] for weight sync. pub fn sync_dueling_weights_branching( vars: &GpuVarStore, - weights: &mut DuelingWeightSet, + weights: &DuelingWeightSet, stream: &Arc, ) -> Result<(), MLError> { - sync_one(vars, "shared_0.weight", &mut weights.w_s1, stream)?; - sync_one(vars, "shared_0.bias", &mut weights.b_s1, stream)?; - sync_one(vars, "shared_1.weight", &mut weights.w_s2, stream)?; - sync_one(vars, "shared_1.bias", &mut weights.b_s2, stream)?; - sync_one(vars, "value_fc.weight", &mut weights.w_v1, stream)?; - sync_one(vars, "value_fc.bias", &mut weights.b_v1, stream)?; - sync_one(vars, "value_out.weight", &mut weights.w_v2, stream)?; - sync_one(vars, "value_out.bias", &mut weights.b_v2, stream)?; + sync_one_to_ptr(vars, "shared_0.weight", weights.w_s1, weights.w_s1_len, stream)?; + sync_one_to_ptr(vars, "shared_0.bias", weights.b_s1, weights.b_s1_len, stream)?; + sync_one_to_ptr(vars, "shared_1.weight", weights.w_s2, weights.w_s2_len, stream)?; + sync_one_to_ptr(vars, "shared_1.bias", weights.b_s2, weights.b_s2_len, stream)?; + sync_one_to_ptr(vars, "value_fc.weight", weights.w_v1, weights.w_v1_len, stream)?; + sync_one_to_ptr(vars, "value_fc.bias", weights.b_v1, weights.b_v1_len, stream)?; + sync_one_to_ptr(vars, "value_out.weight", weights.w_v2, weights.w_v2_len, stream)?; + sync_one_to_ptr(vars, "value_out.bias", weights.b_v2, weights.b_v2_len, stream)?; // Advantage/exposure: branch_0_* keys in branching GpuVarStore - sync_one(vars, "branch_0_fc.weight", &mut weights.w_a1, stream)?; - sync_one(vars, "branch_0_fc.bias", &mut weights.b_a1, stream)?; - sync_one(vars, "branch_0_out.weight", &mut weights.w_a2, stream)?; - sync_one(vars, "branch_0_out.bias", &mut weights.b_a2, stream)?; + sync_one_to_ptr(vars, "branch_0_fc.weight", weights.w_a1, weights.w_a1_len, stream)?; + sync_one_to_ptr(vars, "branch_0_fc.bias", weights.b_a1, weights.b_a1_len, stream)?; + sync_one_to_ptr(vars, "branch_0_out.weight", weights.w_a2, weights.w_a2_len, stream)?; + sync_one_to_ptr(vars, "branch_0_out.bias", weights.b_a2, weights.b_a2_len, stream)?; Ok(()) } @@ -776,22 +947,22 @@ pub fn extract_curiosity_weights( /// to keep the GPU weight copies in sync with the model's GpuVarStore. pub fn sync_dueling_weights( vars: &GpuVarStore, - weights: &mut DuelingWeightSet, + weights: &DuelingWeightSet, stream: &Arc, ) -> Result<(), MLError> { - sync_one(vars, "shared_0.weight", &mut weights.w_s1, stream)?; - sync_one(vars, "shared_0.bias", &mut weights.b_s1, stream)?; - sync_one(vars, "shared_1.weight", &mut weights.w_s2, stream)?; - sync_one(vars, "shared_1.bias", &mut weights.b_s2, stream)?; - sync_one(vars, "value_fc.weight", &mut weights.w_v1, stream)?; - sync_one(vars, "value_fc.bias", &mut weights.b_v1, stream)?; - sync_one(vars, "value_out.weight", &mut weights.w_v2, stream)?; - sync_one(vars, "value_out.bias", &mut weights.b_v2, stream)?; - sync_one(vars, "advantage_fc.weight", &mut weights.w_a1, stream)?; - sync_one(vars, "advantage_fc.bias", &mut weights.b_a1, stream)?; - sync_one(vars, "advantage_out.weight", &mut weights.w_a2, stream)?; - sync_one(vars, "advantage_out.bias", &mut weights.b_a2, stream)?; + sync_one_to_ptr(vars, "shared_0.weight", weights.w_s1, weights.w_s1_len, stream)?; + sync_one_to_ptr(vars, "shared_0.bias", weights.b_s1, weights.b_s1_len, stream)?; + sync_one_to_ptr(vars, "shared_1.weight", weights.w_s2, weights.w_s2_len, stream)?; + sync_one_to_ptr(vars, "shared_1.bias", weights.b_s2, weights.b_s2_len, stream)?; + sync_one_to_ptr(vars, "value_fc.weight", weights.w_v1, weights.w_v1_len, stream)?; + sync_one_to_ptr(vars, "value_fc.bias", weights.b_v1, weights.b_v1_len, stream)?; + sync_one_to_ptr(vars, "value_out.weight", weights.w_v2, weights.w_v2_len, stream)?; + sync_one_to_ptr(vars, "value_out.bias", weights.b_v2, weights.b_v2_len, stream)?; + sync_one_to_ptr(vars, "advantage_fc.weight", weights.w_a1, weights.w_a1_len, stream)?; + sync_one_to_ptr(vars, "advantage_fc.bias", weights.b_a1, weights.b_a1_len, stream)?; + sync_one_to_ptr(vars, "advantage_out.weight", weights.w_a2, weights.w_a2_len, stream)?; + sync_one_to_ptr(vars, "advantage_out.bias", weights.b_a2, weights.b_a2_len, stream)?; Ok(()) } @@ -803,22 +974,22 @@ pub fn sync_dueling_weights( /// to keep the GPU weight copies in sync with the model's GpuVarStore. pub fn sync_branching_weights( vars: &GpuVarStore, - weights: &mut BranchingWeightSet, + weights: &BranchingWeightSet, stream: &Arc, ) -> Result<(), MLError> { - sync_one(vars, "branch_1_fc.weight", &mut weights.w_bo1, stream)?; - sync_one(vars, "branch_1_fc.bias", &mut weights.b_bo1, stream)?; - sync_one(vars, "branch_1_out.weight", &mut weights.w_bo2, stream)?; - sync_one(vars, "branch_1_out.bias", &mut weights.b_bo2, stream)?; - sync_one(vars, "branch_2_fc.weight", &mut weights.w_bu1, stream)?; - sync_one(vars, "branch_2_fc.bias", &mut weights.b_bu1, stream)?; - sync_one(vars, "branch_2_out.weight", &mut weights.w_bu2, stream)?; - sync_one(vars, "branch_2_out.bias", &mut weights.b_bu2, stream)?; - sync_one(vars, "branch_3_fc.weight", &mut weights.w_bg1, stream)?; - sync_one(vars, "branch_3_fc.bias", &mut weights.b_bg1, stream)?; - sync_one(vars, "branch_3_out.weight", &mut weights.w_bg2, stream)?; - sync_one(vars, "branch_3_out.bias", &mut weights.b_bg2, stream)?; + sync_one_to_ptr(vars, "branch_1_fc.weight", weights.w_bo1, weights.w_bo1_len, stream)?; + sync_one_to_ptr(vars, "branch_1_fc.bias", weights.b_bo1, weights.b_bo1_len, stream)?; + sync_one_to_ptr(vars, "branch_1_out.weight", weights.w_bo2, weights.w_bo2_len, stream)?; + sync_one_to_ptr(vars, "branch_1_out.bias", weights.b_bo2, weights.b_bo2_len, stream)?; + sync_one_to_ptr(vars, "branch_2_fc.weight", weights.w_bu1, weights.w_bu1_len, stream)?; + sync_one_to_ptr(vars, "branch_2_fc.bias", weights.b_bu1, weights.b_bu1_len, stream)?; + sync_one_to_ptr(vars, "branch_2_out.weight", weights.w_bu2, weights.w_bu2_len, stream)?; + sync_one_to_ptr(vars, "branch_2_out.bias", weights.b_bu2, weights.b_bu2_len, stream)?; + sync_one_to_ptr(vars, "branch_3_fc.weight", weights.w_bg1, weights.w_bg1_len, stream)?; + sync_one_to_ptr(vars, "branch_3_fc.bias", weights.b_bg1, weights.b_bg1_len, stream)?; + sync_one_to_ptr(vars, "branch_3_out.weight", weights.w_bg2, weights.w_bg2_len, stream)?; + sync_one_to_ptr(vars, "branch_3_out.bias", weights.b_bg2, weights.b_bg2_len, stream)?; Ok(()) } @@ -827,13 +998,14 @@ pub fn sync_branching_weights( // Reverse sync: CudaSlice → GpuVarStore (write back GPU-updated weights) // --------------------------------------------------------------------------- -/// Write a single CudaSlice back to the corresponding GpuVarStore parameter. +/// Write a raw device pointer back to the corresponding GpuVarStore parameter. /// /// Uses device-to-device copy (zero CPU roundtrip). -fn reverse_sync_one( +fn reverse_sync_one_from_ptr( vars: &GpuVarStore, name: &str, - src: &CudaSlice, + src_ptr: u64, + src_len: usize, stream: &Arc, ) -> Result<(), MLError> { let param = vars @@ -841,25 +1013,23 @@ fn reverse_sync_one( .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))?; let n_elems = param.data.len(); - if src.len() != n_elems { + if src_len != n_elems { return Err(MLError::ModelError(format!( - "reverse_sync {name}: CudaSlice len {} != param elems {n_elems}", - src.len() + "reverse_sync {name}: src len {src_len} != param elems {n_elems}", ))); } - let (dst_ptr, _dst_sync) = param.data.device_ptr(stream); - let (src_ptr, _src_sync) = src.device_ptr(stream); + let dst_ptr = param.data.raw_ptr(); let num_bytes = n_elems * std::mem::size_of::(); dtod_copy_checked(dst_ptr, src_ptr, num_bytes, stream, name, "reverse_sync")?; Ok(()) } -/// Write back all 12 branching-dueling weight tensors from CudaSlice to GpuVarStore. +/// Write back all 12 branching-dueling weight tensors to GpuVarStore. /// /// Reverse of `sync_dueling_weights_branching`: copies GPU-updated weights -/// from the fused training kernel's CudaSlice buffers back to the GpuVarStore +/// from the fused training pointer views back to the GpuVarStore /// so that Polyak target updates can read current online weights. pub fn reverse_sync_dueling_weights_branching( vars: &GpuVarStore, @@ -867,23 +1037,23 @@ pub fn reverse_sync_dueling_weights_branching( stream: &Arc, ) -> Result<(), MLError> { - reverse_sync_one(vars, "shared_0.weight", &weights.w_s1, stream)?; - reverse_sync_one(vars, "shared_0.bias", &weights.b_s1, stream)?; - reverse_sync_one(vars, "shared_1.weight", &weights.w_s2, stream)?; - reverse_sync_one(vars, "shared_1.bias", &weights.b_s2, stream)?; - reverse_sync_one(vars, "value_fc.weight", &weights.w_v1, stream)?; - reverse_sync_one(vars, "value_fc.bias", &weights.b_v1, stream)?; - reverse_sync_one(vars, "value_out.weight", &weights.w_v2, stream)?; - reverse_sync_one(vars, "value_out.bias", &weights.b_v2, stream)?; - reverse_sync_one(vars, "branch_0_fc.weight", &weights.w_a1, stream)?; - reverse_sync_one(vars, "branch_0_fc.bias", &weights.b_a1, stream)?; - reverse_sync_one(vars, "branch_0_out.weight", &weights.w_a2, stream)?; - reverse_sync_one(vars, "branch_0_out.bias", &weights.b_a2, stream)?; + reverse_sync_one_from_ptr(vars, "shared_0.weight", weights.w_s1, weights.w_s1_len, stream)?; + reverse_sync_one_from_ptr(vars, "shared_0.bias", weights.b_s1, weights.b_s1_len, stream)?; + reverse_sync_one_from_ptr(vars, "shared_1.weight", weights.w_s2, weights.w_s2_len, stream)?; + reverse_sync_one_from_ptr(vars, "shared_1.bias", weights.b_s2, weights.b_s2_len, stream)?; + reverse_sync_one_from_ptr(vars, "value_fc.weight", weights.w_v1, weights.w_v1_len, stream)?; + reverse_sync_one_from_ptr(vars, "value_fc.bias", weights.b_v1, weights.b_v1_len, stream)?; + reverse_sync_one_from_ptr(vars, "value_out.weight", weights.w_v2, weights.w_v2_len, stream)?; + reverse_sync_one_from_ptr(vars, "value_out.bias", weights.b_v2, weights.b_v2_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_0_fc.weight", weights.w_a1, weights.w_a1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_0_fc.bias", weights.b_a1, weights.b_a1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_0_out.weight", weights.w_a2, weights.w_a2_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_0_out.bias", weights.b_a2, weights.b_a2_len, stream)?; Ok(()) } -/// Write back all 12 branching weight tensors (branches 1+2+3) from CudaSlice to GpuVarStore. +/// Write back all 12 branching weight tensors (branches 1+2+3) to GpuVarStore. /// /// Reverse of `sync_branching_weights`. pub fn reverse_sync_branching_weights( @@ -892,18 +1062,18 @@ pub fn reverse_sync_branching_weights( stream: &Arc, ) -> Result<(), MLError> { - reverse_sync_one(vars, "branch_1_fc.weight", &weights.w_bo1, stream)?; - reverse_sync_one(vars, "branch_1_fc.bias", &weights.b_bo1, stream)?; - reverse_sync_one(vars, "branch_1_out.weight", &weights.w_bo2, stream)?; - reverse_sync_one(vars, "branch_1_out.bias", &weights.b_bo2, stream)?; - reverse_sync_one(vars, "branch_2_fc.weight", &weights.w_bu1, stream)?; - reverse_sync_one(vars, "branch_2_fc.bias", &weights.b_bu1, stream)?; - reverse_sync_one(vars, "branch_2_out.weight", &weights.w_bu2, stream)?; - reverse_sync_one(vars, "branch_2_out.bias", &weights.b_bu2, stream)?; - reverse_sync_one(vars, "branch_3_fc.weight", &weights.w_bg1, stream)?; - reverse_sync_one(vars, "branch_3_fc.bias", &weights.b_bg1, stream)?; - reverse_sync_one(vars, "branch_3_out.weight", &weights.w_bg2, stream)?; - reverse_sync_one(vars, "branch_3_out.bias", &weights.b_bg2, stream)?; + reverse_sync_one_from_ptr(vars, "branch_1_fc.weight", weights.w_bo1, weights.w_bo1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_1_fc.bias", weights.b_bo1, weights.b_bo1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_1_out.weight", weights.w_bo2, weights.w_bo2_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_1_out.bias", weights.b_bo2, weights.b_bo2_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_2_fc.weight", weights.w_bu1, weights.w_bu1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_2_fc.bias", weights.b_bu1, weights.b_bu1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_2_out.weight", weights.w_bu2, weights.w_bu2_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_2_out.bias", weights.b_bu2, weights.b_bu2_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_3_fc.weight", weights.w_bg1, weights.w_bg1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_3_fc.bias", weights.b_bg1, weights.b_bg1_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_3_out.weight", weights.w_bg2, weights.w_bg2_len, stream)?; + reverse_sync_one_from_ptr(vars, "branch_3_out.bias", weights.b_bg2, weights.b_bg2_len, stream)?; Ok(()) } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 5799462ee..13d686f2a 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -169,6 +169,17 @@ pub(crate) struct FusedTrainingCtx { online_branching: BranchingWeightSet, target_dueling: DuelingWeightSet, target_branching: BranchingWeightSet, + /// Backing GPU allocations for online/target weight sets. + /// Must outlive the weight sets (pointers reference this memory). + /// `None` when weight sets are zero-copy views into `params_buf`. + #[allow(dead_code)] + _weight_backings: Option<( + gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking, + gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking, + )>, + /// Backing GPU allocations for ensemble extra head weight sets. + #[allow(dead_code)] + _ensemble_backings: Vec<(gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking)>, stream: Arc, /// Batch size at creation time -- must match `current_batch_size` to reuse CUDA Graph. batch_size: usize, @@ -366,20 +377,21 @@ impl FusedTrainingCtx { market_dim: 42, // Always 42 base market features — OFI features bypass bottleneck via portfolio_dim }; - // Extract weight sets from VarMaps (online + target) + // Extract weight sets from VarMaps (online + target). + // The backing storage keeps GPU allocations alive; weight sets are pointer views. let online_vars = branching_net.vars(); let target_vars = branching_target.vars(); - let online_dueling = + let (online_d_backing, online_dueling) = gpu_weights::extract_dueling_weights_branching(online_vars, &stream) .map_err(|e| anyhow::anyhow!("Extract online dueling weights: {e}"))?; - let online_branching = + let (online_b_backing, online_branching) = gpu_weights::extract_branching_weights(online_vars, &stream) .map_err(|e| anyhow::anyhow!("Extract online branching weights: {e}"))?; - let target_dueling = + let (target_d_backing, target_dueling) = gpu_weights::extract_dueling_weights_branching(target_vars, &stream) .map_err(|e| anyhow::anyhow!("Extract target dueling weights: {e}"))?; - let target_branching = + let (target_b_backing, target_branching) = gpu_weights::extract_branching_weights(target_vars, &stream) .map_err(|e| anyhow::anyhow!("Extract target branching weights: {e}"))?; @@ -546,6 +558,7 @@ impl FusedTrainingCtx { ensemble_diversity_loss_buf, ensemble_kl_grad_kernel, ensemble_d_logits_buf, + ens_backings_vec, ) = if k > 1 { use crate::cuda_pipeline::gpu_dqn_trainer::compile_ensemble_kernels; @@ -555,14 +568,18 @@ impl FusedTrainingCtx { // Allocate K-1 extra head weight sets via DtoD clone of head 0 + small noise. // DtoD: head_k ← online_dueling + online_branching weights copied on GPU. + // Each head has its own backing storage + pointer view. let mut extra_heads: Vec<(DuelingWeightSet, BranchingWeightSet)> = Vec::with_capacity(k - 1); + let mut ens_backings: Vec<(gpu_weights::DuelingWeightBacking, gpu_weights::BranchingWeightBacking)> = + Vec::with_capacity(k - 1); for head_idx in 1..k { - let head_dueling = clone_dueling_weights(&online_dueling, &stream, head_idx) + let (d_backing, head_dueling) = clone_dueling_weights(&online_dueling, &stream, head_idx) .map_err(|e| anyhow::anyhow!("Clone ensemble dueling head {head_idx}: {e}"))?; - let head_branching = clone_branching_weights(&online_branching, &stream, head_idx) + let (b_backing, head_branching) = clone_branching_weights(&online_branching, &stream, head_idx) .map_err(|e| anyhow::anyhow!("Clone ensemble branching head {head_idx}: {e}"))?; extra_heads.push((head_dueling, head_branching)); + ens_backings.push((d_backing, b_backing)); } // Pre-allocate ensemble buffers. @@ -608,10 +625,11 @@ impl FusedTrainingCtx { Some(div_loss_buf), Some(kl_grad_kernel), Some(d_logits_buf), + ens_backings, ) } else { (Vec::new(), None, None, None, None, None, None, - None, None) // kl_grad_kernel, d_logits_buf + None, None, Vec::new()) // kl_grad_kernel, d_logits_buf, ens_backings }; info!( @@ -631,6 +649,8 @@ impl FusedTrainingCtx { online_branching, target_dueling, target_branching, + _weight_backings: Some((online_d_backing, online_b_backing, target_d_backing, target_b_backing)), + _ensemble_backings: ens_backings_vec, stream, batch_size, steps_since_varmap_sync: 0, @@ -1136,7 +1156,7 @@ impl FusedTrainingCtx { ); self.trainer.target_ema_update( &self.online_dueling, &self.online_branching, - &mut self.target_dueling, &mut self.target_branching, + &self.target_dueling, &self.target_branching, tau as f32, ).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?; } @@ -1574,10 +1594,10 @@ impl FusedTrainingCtx { // Full cuBLAS value head forward for each ensemble head. // h_s2 → W_v1_k → ReLU → W_v2_k → v_logits_k // Uses shared trunk activations (save_h_s2) with per-head value weights. - let w_v1_ptr = head_dueling.w_v1.raw_ptr(); - let b_v1_ptr = head_dueling.b_v1.raw_ptr(); - let w_v2_ptr = head_dueling.w_v2.raw_ptr(); - let b_v2_ptr = head_dueling.b_v2.raw_ptr(); + let w_v1_ptr = head_dueling.w_v1; + let b_v1_ptr = head_dueling.b_v1; + let w_v2_ptr = head_dueling.w_v2; + let b_v2_ptr = head_dueling.b_v2; let dst_k_ptr = logits_buf.raw_ptr() + (k_idx * b * na * std::mem::size_of::()) as u64; @@ -2099,48 +2119,49 @@ fn compute_cosine_annealed_tau( tau_base } } +/// Deep-copy the 12 weight tensors of a DuelingWeightSet from pointer sources +/// into newly-allocated CudaSlice buffers. Returns both the backing storage and +/// a pointer-view weight set. +/// +/// Used for ensemble extra heads that need independent weight copies. fn clone_dueling_weights( src: &DuelingWeightSet, stream: &Arc, head_idx: usize, -) -> Result { - let clone_bf16 = |slice: &cudarc::driver::CudaSlice| -> Result> { - let n = slice.len(); - let dst = stream.alloc_zeros::(n) - .map_err(|e| anyhow::anyhow!("Clone alloc {n}xbf16: {e}"))?; - let (src_ptr, src_guard) = slice.device_ptr(stream); - let (dst_ptr, dst_guard) = dst.device_ptr(stream); - let _no_drop_s = std::mem::ManuallyDrop::new(src_guard); - let _no_drop_d = std::mem::ManuallyDrop::new(dst_guard); - let n_bytes = n * std::mem::size_of::(); - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, src_ptr, n_bytes, stream.cu_stream() - ).map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?; - } +) -> Result<(gpu_weights::DuelingWeightBacking, DuelingWeightSet)> { + use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy; + + let clone_from_ptr = |ptr: u64, len: usize| -> Result> { + let dst = stream.alloc_zeros::(len) + .map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?; + let n_bytes = len * std::mem::size_of::(); + dtod_copy(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_dueling") + .map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?; Ok(dst) }; - let cloned = DuelingWeightSet { - w_s1: clone_bf16(&src.w_s1)?, - b_s1: clone_bf16(&src.b_s1)?, - w_s2: clone_bf16(&src.w_s2)?, - b_s2: clone_bf16(&src.b_s2)?, - w_v1: clone_bf16(&src.w_v1)?, - b_v1: clone_bf16(&src.b_v1)?, - w_v2: clone_bf16(&src.w_v2)?, - b_v2: clone_bf16(&src.b_v2)?, - w_a1: clone_bf16(&src.w_a1)?, - b_a1: clone_bf16(&src.b_a1)?, - w_a2: clone_bf16(&src.w_a2)?, - b_a2: clone_bf16(&src.b_a2)?, + let backing = gpu_weights::DuelingWeightBacking { + slices: [ + clone_from_ptr(src.w_s1, src.w_s1_len)?, + clone_from_ptr(src.b_s1, src.b_s1_len)?, + clone_from_ptr(src.w_s2, src.w_s2_len)?, + clone_from_ptr(src.b_s2, src.b_s2_len)?, + clone_from_ptr(src.w_v1, src.w_v1_len)?, + clone_from_ptr(src.b_v1, src.b_v1_len)?, + clone_from_ptr(src.w_v2, src.w_v2_len)?, + clone_from_ptr(src.b_v2, src.b_v2_len)?, + clone_from_ptr(src.w_a1, src.w_a1_len)?, + clone_from_ptr(src.b_a1, src.b_a1_len)?, + clone_from_ptr(src.w_a2, src.w_a2_len)?, + clone_from_ptr(src.b_a2, src.b_a2_len)?, + ], }; + let ws = backing.weight_set(); - // Sync so all DtoD copies complete before we return (caller uses weights immediately) let _ = head_idx; // used for future per-head noise seeding unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); } - Ok(cloned) + Ok((backing, ws)) } /// Deep-copy a BranchingWeightSet on GPU via DtoD memcpy. @@ -2148,41 +2169,38 @@ fn clone_branching_weights( src: &BranchingWeightSet, stream: &Arc, head_idx: usize, -) -> Result { - let clone_bf16 = |slice: &cudarc::driver::CudaSlice| -> Result> { - let n = slice.len(); - let dst = stream.alloc_zeros::(n) - .map_err(|e| anyhow::anyhow!("Clone alloc {n}xbf16: {e}"))?; - let (src_ptr, src_guard) = slice.device_ptr(stream); - let (dst_ptr, dst_guard) = dst.device_ptr(stream); - let _no_drop_s = std::mem::ManuallyDrop::new(src_guard); - let _no_drop_d = std::mem::ManuallyDrop::new(dst_guard); - let n_bytes = n * std::mem::size_of::(); - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, src_ptr, n_bytes, stream.cu_stream() - ).map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?; - } +) -> Result<(gpu_weights::BranchingWeightBacking, BranchingWeightSet)> { + use crate::cuda_pipeline::gpu_dqn_trainer::dtod_copy; + + let clone_from_ptr = |ptr: u64, len: usize| -> Result> { + let dst = stream.alloc_zeros::(len) + .map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?; + let n_bytes = len * std::mem::size_of::(); + dtod_copy(dst.raw_ptr(), ptr, n_bytes, stream, 0, "clone_branching") + .map_err(|e| anyhow::anyhow!("DtoD clone: {e}"))?; Ok(dst) }; - let cloned = BranchingWeightSet { - w_bo1: clone_bf16(&src.w_bo1)?, - b_bo1: clone_bf16(&src.b_bo1)?, - w_bo2: clone_bf16(&src.w_bo2)?, - b_bo2: clone_bf16(&src.b_bo2)?, - w_bu1: clone_bf16(&src.w_bu1)?, - b_bu1: clone_bf16(&src.b_bu1)?, - w_bu2: clone_bf16(&src.w_bu2)?, - b_bu2: clone_bf16(&src.b_bu2)?, - w_bg1: clone_bf16(&src.w_bg1)?, - b_bg1: clone_bf16(&src.b_bg1)?, - w_bg2: clone_bf16(&src.w_bg2)?, - b_bg2: clone_bf16(&src.b_bg2)?, + let backing = gpu_weights::BranchingWeightBacking { + slices: [ + clone_from_ptr(src.w_bo1, src.w_bo1_len)?, + clone_from_ptr(src.b_bo1, src.b_bo1_len)?, + clone_from_ptr(src.w_bo2, src.w_bo2_len)?, + clone_from_ptr(src.b_bo2, src.b_bo2_len)?, + clone_from_ptr(src.w_bu1, src.w_bu1_len)?, + clone_from_ptr(src.b_bu1, src.b_bu1_len)?, + clone_from_ptr(src.w_bu2, src.w_bu2_len)?, + clone_from_ptr(src.b_bu2, src.b_bu2_len)?, + clone_from_ptr(src.w_bg1, src.w_bg1_len)?, + clone_from_ptr(src.b_bg1, src.b_bg1_len)?, + clone_from_ptr(src.w_bg2, src.w_bg2_len)?, + clone_from_ptr(src.b_bg2, src.b_bg2_len)?, + ], }; + let ws = backing.weight_set(); let _ = head_idx; unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); } - Ok(cloned) + Ok((backing, ws)) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs index 110658401..a62cca798 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs @@ -9,7 +9,10 @@ use std::sync::Arc; use super::helpers::cuda_device; use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer}; -use crate::cuda_pipeline::gpu_weights::{DuelingWeightSet, BranchingWeightSet}; +use crate::cuda_pipeline::gpu_weights::{ + DuelingWeightSet, BranchingWeightSet, + DuelingWeightBacking, BranchingWeightBacking, +}; /// Small config for fast kernel-level tests (no training, just buffer ops). fn test_config() -> GpuDqnTrainConfig { @@ -36,49 +39,46 @@ fn test_config() -> GpuDqnTrainConfig { /// Allocate a DuelingWeightSet with random values (not zero — spectral norm /// needs non-degenerate matrices for power iteration to converge). -fn alloc_dueling(stream: &Arc, cfg: &GpuDqnTrainConfig) -> DuelingWeightSet { +/// Returns both the backing storage (keeps GPU memory alive) and the pointer view. +fn alloc_dueling(stream: &Arc, cfg: &GpuDqnTrainConfig) -> (DuelingWeightBacking, DuelingWeightSet) { let na = cfg.num_atoms; let alloc = |n: usize| -> cudarc::driver::CudaSlice { // Fill with 0.1 (non-zero, non-degenerate for spectral norm) let bf16_data: Vec = vec![0.1; n]; stream.clone_htod(&bf16_data).expect("alloc dueling weight") }; - DuelingWeightSet { - w_s1: alloc(cfg.shared_h1 * cfg.state_dim), - b_s1: alloc(cfg.shared_h1), - w_s2: alloc(cfg.shared_h2 * cfg.shared_h1), - b_s2: alloc(cfg.shared_h2), - w_v1: alloc(cfg.value_h * cfg.shared_h2), - b_v1: alloc(cfg.value_h), - w_v2: alloc(na * cfg.value_h), - b_v2: alloc(na), - w_a1: alloc(cfg.adv_h * cfg.shared_h2), - b_a1: alloc(cfg.adv_h), - w_a2: alloc(cfg.branch_0_size * na * cfg.adv_h), - b_a2: alloc(cfg.branch_0_size * na), - } + let backing = DuelingWeightBacking { + slices: [ + alloc(cfg.shared_h1 * cfg.state_dim), alloc(cfg.shared_h1), + alloc(cfg.shared_h2 * cfg.shared_h1), alloc(cfg.shared_h2), + alloc(cfg.value_h * cfg.shared_h2), alloc(cfg.value_h), + alloc(na * cfg.value_h), alloc(na), + alloc(cfg.adv_h * cfg.shared_h2), alloc(cfg.adv_h), + alloc(cfg.branch_0_size * na * cfg.adv_h), alloc(cfg.branch_0_size * na), + ], + }; + let ws = backing.weight_set(); + (backing, ws) } -fn alloc_branching(stream: &Arc, cfg: &GpuDqnTrainConfig) -> BranchingWeightSet { +fn alloc_branching(stream: &Arc, cfg: &GpuDqnTrainConfig) -> (BranchingWeightBacking, BranchingWeightSet) { let na = cfg.num_atoms; let alloc = |n: usize| -> cudarc::driver::CudaSlice { let bf16_data: Vec = vec![0.1; n]; stream.clone_htod(&bf16_data).expect("alloc branching weight") }; - BranchingWeightSet { - w_bo1: alloc(cfg.adv_h * cfg.shared_h2), - b_bo1: alloc(cfg.adv_h), - w_bo2: alloc(cfg.branch_1_size * na * cfg.adv_h), - b_bo2: alloc(cfg.branch_1_size * na), - w_bu1: alloc(cfg.adv_h * cfg.shared_h2), - b_bu1: alloc(cfg.adv_h), - w_bu2: alloc(cfg.branch_2_size * na * cfg.adv_h), - b_bu2: alloc(cfg.branch_2_size * na), - w_bg1: alloc(cfg.adv_h * cfg.shared_h2), - b_bg1: alloc(cfg.adv_h), - w_bg2: alloc(cfg.branch_3_size * na * cfg.adv_h), - b_bg2: alloc(cfg.branch_3_size * na), - } + let backing = BranchingWeightBacking { + slices: [ + alloc(cfg.adv_h * cfg.shared_h2), alloc(cfg.adv_h), + alloc(cfg.branch_1_size * na * cfg.adv_h), alloc(cfg.branch_1_size * na), + alloc(cfg.adv_h * cfg.shared_h2), alloc(cfg.adv_h), + alloc(cfg.branch_2_size * na * cfg.adv_h), alloc(cfg.branch_2_size * na), + alloc(cfg.adv_h * cfg.shared_h2), alloc(cfg.adv_h), + alloc(cfg.branch_3_size * na * cfg.adv_h), alloc(cfg.branch_3_size * na), + ], + }; + let ws = backing.weight_set(); + (backing, ws) } /// clip_grad_buf_inplace must reduce gradient norm to at most max_norm. @@ -151,8 +151,8 @@ fn test_spectral_norm_all_heads_no_panic() -> anyhow::Result<()> { let stream = Arc::clone(dev.cuda_stream().expect("cuda stream")); let cfg = test_config(); let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?; - let dueling = alloc_dueling(&stream, &cfg); - let branching = alloc_branching(&stream, &cfg); + let (d_backing, dueling) = alloc_dueling(&stream, &cfg); + let (b_backing, branching) = alloc_branching(&stream, &cfg); // Run spectral norm — should not panic with non-zero weights. // Spectral norm now writes directly to params_bf16; unflatten syncs @@ -162,7 +162,7 @@ fn test_spectral_norm_all_heads_no_panic() -> anyhow::Result<()> { // Verify a few weights are still finite after normalization (f32 -> f32 readback) let mut w_s1_bf16 = vec![0.0_f32; cfg.shared_h1 * cfg.state_dim]; - stream.memcpy_dtoh(&dueling.w_s1, &mut w_s1_bf16) + stream.memcpy_dtoh(&d_backing.slices[0], &mut w_s1_bf16) .map_err(|e| anyhow::anyhow!("{e}"))?; let w_s1_host: Vec = w_s1_bf16.to_vec(); assert!( @@ -171,7 +171,7 @@ fn test_spectral_norm_all_heads_no_panic() -> anyhow::Result<()> { ); let mut w_v1_bf16 = vec![0.0_f32; cfg.value_h * cfg.shared_h2]; - stream.memcpy_dtoh(&dueling.w_v1, &mut w_v1_bf16) + stream.memcpy_dtoh(&d_backing.slices[4], &mut w_v1_bf16) .map_err(|e| anyhow::anyhow!("{e}"))?; let w_v1_host: Vec = w_v1_bf16.to_vec(); assert!( @@ -180,7 +180,7 @@ fn test_spectral_norm_all_heads_no_panic() -> anyhow::Result<()> { ); let mut w_bo1_bf16 = vec![0.0_f32; cfg.adv_h * cfg.shared_h2]; - stream.memcpy_dtoh(&branching.w_bo1, &mut w_bo1_bf16) + stream.memcpy_dtoh(&b_backing.slices[0], &mut w_bo1_bf16) .map_err(|e| anyhow::anyhow!("{e}"))?; let w_bo1_host: Vec = w_bo1_bf16.to_vec(); assert!( @@ -207,34 +207,28 @@ fn test_spectral_norm_constrains_operator_norm() -> anyhow::Result<()> { stream.clone_htod(&bf16_data).expect("alloc large weight") }; let na = cfg.num_atoms; - let dueling = DuelingWeightSet { - w_s1: alloc_large(cfg.shared_h1 * cfg.state_dim), - b_s1: alloc_large(cfg.shared_h1), - w_s2: alloc_large(cfg.shared_h2 * cfg.shared_h1), - b_s2: alloc_large(cfg.shared_h2), - w_v1: alloc_large(cfg.value_h * cfg.shared_h2), - b_v1: alloc_large(cfg.value_h), - w_v2: alloc_large(na * cfg.value_h), - b_v2: alloc_large(na), - w_a1: alloc_large(cfg.adv_h * cfg.shared_h2), - b_a1: alloc_large(cfg.adv_h), - w_a2: alloc_large(cfg.branch_0_size * na * cfg.adv_h), - b_a2: alloc_large(cfg.branch_0_size * na), + let d_backing = DuelingWeightBacking { + slices: [ + alloc_large(cfg.shared_h1 * cfg.state_dim), alloc_large(cfg.shared_h1), + alloc_large(cfg.shared_h2 * cfg.shared_h1), alloc_large(cfg.shared_h2), + alloc_large(cfg.value_h * cfg.shared_h2), alloc_large(cfg.value_h), + alloc_large(na * cfg.value_h), alloc_large(na), + alloc_large(cfg.adv_h * cfg.shared_h2), alloc_large(cfg.adv_h), + alloc_large(cfg.branch_0_size * na * cfg.adv_h), alloc_large(cfg.branch_0_size * na), + ], }; - let branching = BranchingWeightSet { - w_bo1: alloc_large(cfg.adv_h * cfg.shared_h2), - b_bo1: alloc_large(cfg.adv_h), - w_bo2: alloc_large(cfg.branch_1_size * na * cfg.adv_h), - b_bo2: alloc_large(cfg.branch_1_size * na), - w_bu1: alloc_large(cfg.adv_h * cfg.shared_h2), - b_bu1: alloc_large(cfg.adv_h), - w_bu2: alloc_large(cfg.branch_2_size * na * cfg.adv_h), - b_bu2: alloc_large(cfg.branch_2_size * na), - w_bg1: alloc_large(cfg.adv_h * cfg.shared_h2), - b_bg1: alloc_large(cfg.adv_h), - w_bg2: alloc_large(cfg.branch_3_size * na * cfg.adv_h), - b_bg2: alloc_large(cfg.branch_3_size * na), + let dueling = d_backing.weight_set(); + let b_backing = BranchingWeightBacking { + slices: [ + alloc_large(cfg.adv_h * cfg.shared_h2), alloc_large(cfg.adv_h), + alloc_large(cfg.branch_1_size * na * cfg.adv_h), alloc_large(cfg.branch_1_size * na), + alloc_large(cfg.adv_h * cfg.shared_h2), alloc_large(cfg.adv_h), + alloc_large(cfg.branch_2_size * na * cfg.adv_h), alloc_large(cfg.branch_2_size * na), + alloc_large(cfg.adv_h * cfg.shared_h2), alloc_large(cfg.adv_h), + alloc_large(cfg.branch_3_size * na * cfg.adv_h), alloc_large(cfg.branch_3_size * na), + ], }; + let branching = b_backing.weight_set(); // Run spectral norm multiple times (power iteration converges over iterations). // Spectral norm writes directly to params_bf16; unflatten syncs back to @@ -251,7 +245,7 @@ fn test_spectral_norm_constrains_operator_norm() -> anyhow::Result<()> { let rows = cfg.shared_h1; let cols = cfg.state_dim; let mut w_bf16 = vec![0.0_f32; rows * cols]; - stream.memcpy_dtoh(&dueling.w_s1, &mut w_bf16) + stream.memcpy_dtoh(&d_backing.slices[0], &mut w_bf16) .map_err(|e| anyhow::anyhow!("{e}"))?; let w_host: Vec = w_bf16.to_vec(); diff --git a/crates/ml/tests/gpu_kernel_parity_test.rs b/crates/ml/tests/gpu_kernel_parity_test.rs index 0aa238583..f3a79293e 100644 --- a/crates/ml/tests/gpu_kernel_parity_test.rs +++ b/crates/ml/tests/gpu_kernel_parity_test.rs @@ -493,12 +493,12 @@ mod gpu_parity { let network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap(); let target = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap(); - let weights = extract_dueling_weights(network.store(), &stream) + let (backing, weights) = extract_dueling_weights(network.store(), &stream) .expect("Weight extraction failed"); // Download and verify w_s1 is non-zero and finite (f32 weights) let mut shared_0_w = vec![0.0_f32; 256 * 54]; - stream.memcpy_dtoh(&weights.w_s1, &mut shared_0_w).unwrap(); // test-only readback + stream.memcpy_dtoh(&backing.slices[0], &mut shared_0_w).unwrap(); // test-only readback assert!(!shared_0_w.iter().all(|&x| x == 0.0_f32), "w_s1 is all zeros"); assert!(shared_0_w.iter().all(|x| x.is_finite()), "w_s1 has NaN/Inf"); @@ -524,17 +524,17 @@ mod gpu_parity { let network = DistributionalDuelingQNetwork::new(dist_config(), stream.clone()).unwrap(); - let weights = extract_dueling_weights(network.vars(), &stream) + let (backing, _weights) = extract_dueling_weights(network.vars(), &stream) .expect("Distributional weight extraction failed"); // value_out: [51, 128] = 6528 elements (f32 weights) let mut value_out_w = vec![0.0_f32; 51 * 128]; - stream.memcpy_dtoh(&weights.w_v2, &mut value_out_w).unwrap(); // test-only readback + stream.memcpy_dtoh(&backing.slices[6], &mut value_out_w).unwrap(); // test-only readback (w_v2 = index 6) assert!(value_out_w.iter().any(|&x| x != 0.0_f32), "w_v2 all zeros"); // advantage_out: [255, 128] = 32640 elements (f32 weights) let mut adv_out_w = vec![0.0_f32; 255 * 128]; - stream.memcpy_dtoh(&weights.w_a2, &mut adv_out_w).unwrap(); // test-only readback + stream.memcpy_dtoh(&backing.slices[10], &mut adv_out_w).unwrap(); // test-only readback (w_a2 = index 10) assert!(adv_out_w.iter().any(|&x| x != 0.0_f32), "w_a2 all zeros"); // RMSNorm gamma should exist and be ~1.0