refactor: eliminate GpuVarStore — DuelingWeightSet is zero-copy pointer view into params_buf
DuelingWeightSet and BranchingWeightSet fields changed from owned CudaSlice<f32> to raw u64 device pointers + usize element counts. This enables zero-copy views directly into the flat params_buf for the training hot path (no D2D copies, no shadow allocations). Key changes: - Weight sets store u64 + usize pairs instead of CudaSlice<f32> - from_flat_buffer() creates zero-copy views into params_buf - from_slices() creates pointer views from owned CudaSlice allocations - DuelingWeightBacking/BranchingWeightBacking hold owned CudaSlice arrays for callers that need independent allocations (experience collector, ensemble heads, tests) - flatten/unflatten become no-ops when weight sets point into params_buf - extract functions return (Backing, WeightSet) tuples - KernelWeightPack::build() uses u64 fields directly (no raw_device_ptr) - All sync functions updated for pointer-based signatures - Ensemble clone functions return backing + pointer view pairs - gradient_budget tests updated (7/7 pass) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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);
|
||||
|
||||
|
||||
@@ -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(())
|
||||
|
||||
@@ -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<f32>; 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::<f32>()) 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::<f32>();
|
||||
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<f32>; 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::<f32>();
|
||||
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<f32>; 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::<f32>();
|
||||
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<f32>; 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::<f32>();
|
||||
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.
|
||||
|
||||
@@ -540,7 +540,7 @@ pub struct GpuExperienceCollector {
|
||||
exp_v_logits: CudaSlice<f32>, // [N, num_atoms] f32 (output layer — no f32 truncation)
|
||||
exp_b_logits: CudaSlice<f32>, // [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<f32>,
|
||||
@@ -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::<f32>();
|
||||
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(())
|
||||
}
|
||||
|
||||
@@ -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<f32>` 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<f32>`
|
||||
//! 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<f32>` 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<f32>` 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<f32>, // shared_0.weight [SHARED_H1, STATE_DIM]
|
||||
pub b_s1: CudaSlice<f32>, // shared_0.bias [SHARED_H1]
|
||||
pub w_s2: CudaSlice<f32>, // shared_1.weight [SHARED_H2, SHARED_H1]
|
||||
pub b_s2: CudaSlice<f32>, // 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<f32>, // value_fc.weight [VALUE_H, SHARED_H2]
|
||||
pub b_v1: CudaSlice<f32>, // value_fc.bias [VALUE_H]
|
||||
pub w_v2: CudaSlice<f32>, // value_out.weight [NUM_ATOMS, VALUE_H]
|
||||
pub b_v2: CudaSlice<f32>, // 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<f32>, // branch_0_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_a1: CudaSlice<f32>, // branch_0_fc.bias [ADV_H]
|
||||
pub w_a2: CudaSlice<f32>, // branch_0_out.weight [BRANCH_0_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_a2: CudaSlice<f32>, // 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::<f32>()) 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<f32>` 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<f32>, b_s1: &CudaSlice<f32>,
|
||||
w_s2: &CudaSlice<f32>, b_s2: &CudaSlice<f32>,
|
||||
w_v1: &CudaSlice<f32>, b_v1: &CudaSlice<f32>,
|
||||
w_v2: &CudaSlice<f32>, b_v2: &CudaSlice<f32>,
|
||||
w_a1: &CudaSlice<f32>, b_a1: &CudaSlice<f32>,
|
||||
w_a2: &CudaSlice<f32>, b_a2: &CudaSlice<f32>,
|
||||
) -> 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<f32>` 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<f32>, // branch_1_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_bo1: CudaSlice<f32>, // branch_1_fc.bias [ADV_H]
|
||||
pub w_bo2: CudaSlice<f32>, // branch_1_out.weight [BRANCH_1_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_bo2: CudaSlice<f32>, // 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<f32>, // branch_2_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_bu1: CudaSlice<f32>, // branch_2_fc.bias [ADV_H]
|
||||
pub w_bu2: CudaSlice<f32>, // branch_2_out.weight [BRANCH_2_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_bu2: CudaSlice<f32>, // 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<f32>, // branch_3_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_bg1: CudaSlice<f32>, // branch_3_fc.bias [ADV_H]
|
||||
pub w_bg2: CudaSlice<f32>, // branch_3_out.weight [BRANCH_3_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_bg2: CudaSlice<f32>, // 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::<f32>()) 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::<f32>()) 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<f32>` 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<f32>, b_bo1: &CudaSlice<f32>,
|
||||
w_bo2: &CudaSlice<f32>, b_bo2: &CudaSlice<f32>,
|
||||
w_bu1: &CudaSlice<f32>, b_bu1: &CudaSlice<f32>,
|
||||
w_bu2: &CudaSlice<f32>, b_bu2: &CudaSlice<f32>,
|
||||
w_bg1: &CudaSlice<f32>, b_bg1: &CudaSlice<f32>,
|
||||
w_bg2: &CudaSlice<f32>, b_bg2: &CudaSlice<f32>,
|
||||
) -> 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<f32>; 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<f32>; 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<CudaStream>,
|
||||
) -> 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::<f32>();
|
||||
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<CudaStream>,
|
||||
) -> Result<DuelingWeightSet, MLError> {
|
||||
) -> Result<(DuelingWeightBacking, DuelingWeightSet), MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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<CudaStream>,
|
||||
) -> Result<BranchingWeightSet, MLError> {
|
||||
) -> Result<(BranchingWeightBacking, BranchingWeightSet), MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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<CudaStream>,
|
||||
) -> Result<DuelingWeightSet, MLError> {
|
||||
) -> Result<(DuelingWeightBacking, DuelingWeightSet), MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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<CudaStream>,
|
||||
) -> 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<CudaStream>,
|
||||
) -> 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<CudaStream>,
|
||||
) -> 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<f32>,
|
||||
src_ptr: u64,
|
||||
src_len: usize,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> 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::<f32>();
|
||||
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<CudaStream>,
|
||||
) -> 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<CudaStream>,
|
||||
) -> 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(())
|
||||
}
|
||||
|
||||
@@ -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<cudarc::driver::CudaStream>,
|
||||
/// 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::<f32>()) 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<cudarc::driver::CudaStream>,
|
||||
head_idx: usize,
|
||||
) -> Result<DuelingWeightSet> {
|
||||
let clone_bf16 = |slice: &cudarc::driver::CudaSlice<f32>| -> Result<cudarc::driver::CudaSlice<f32>> {
|
||||
let n = slice.len();
|
||||
let dst = stream.alloc_zeros::<f32>(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::<f32>();
|
||||
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<cudarc::driver::CudaSlice<f32>> {
|
||||
let dst = stream.alloc_zeros::<f32>(len)
|
||||
.map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?;
|
||||
let n_bytes = len * std::mem::size_of::<f32>();
|
||||
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<cudarc::driver::CudaStream>,
|
||||
head_idx: usize,
|
||||
) -> Result<BranchingWeightSet> {
|
||||
let clone_bf16 = |slice: &cudarc::driver::CudaSlice<f32>| -> Result<cudarc::driver::CudaSlice<f32>> {
|
||||
let n = slice.len();
|
||||
let dst = stream.alloc_zeros::<f32>(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::<f32>();
|
||||
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<cudarc::driver::CudaSlice<f32>> {
|
||||
let dst = stream.alloc_zeros::<f32>(len)
|
||||
.map_err(|e| anyhow::anyhow!("Clone alloc {len}xf32: {e}"))?;
|
||||
let n_bytes = len * std::mem::size_of::<f32>();
|
||||
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))
|
||||
}
|
||||
|
||||
@@ -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<cudarc::driver::CudaStream>, cfg: &GpuDqnTrainConfig) -> DuelingWeightSet {
|
||||
/// Returns both the backing storage (keeps GPU memory alive) and the pointer view.
|
||||
fn alloc_dueling(stream: &Arc<cudarc::driver::CudaStream>, cfg: &GpuDqnTrainConfig) -> (DuelingWeightBacking, DuelingWeightSet) {
|
||||
let na = cfg.num_atoms;
|
||||
let alloc = |n: usize| -> cudarc::driver::CudaSlice<f32> {
|
||||
// Fill with 0.1 (non-zero, non-degenerate for spectral norm)
|
||||
let bf16_data: Vec<f32> = 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<cudarc::driver::CudaStream>, cfg: &GpuDqnTrainConfig) -> BranchingWeightSet {
|
||||
fn alloc_branching(stream: &Arc<cudarc::driver::CudaStream>, cfg: &GpuDqnTrainConfig) -> (BranchingWeightBacking, BranchingWeightSet) {
|
||||
let na = cfg.num_atoms;
|
||||
let alloc = |n: usize| -> cudarc::driver::CudaSlice<f32> {
|
||||
let bf16_data: Vec<f32> = 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<f32> = 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<f32> = 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<f32> = 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<f32> = w_bf16.to_vec();
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user