diff --git a/crates/ml/src/cuda_pipeline/gpu_weights.rs b/crates/ml/src/cuda_pipeline/gpu_weights.rs index 07b52182a..dff63bbb7 100644 --- a/crates/ml/src/cuda_pipeline/gpu_weights.rs +++ b/crates/ml/src/cuda_pipeline/gpu_weights.rs @@ -1,7 +1,7 @@ //! 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 +//! them as flat `CudaSlice` GPU buffers for the CUDA experience //! collection kernels (`experience_kernels.cu` + cuBLAS). //! //! Two weight sets are managed: @@ -11,14 +11,14 @@ //! 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 +//! 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)] use std::sync::Arc; -use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DeviceRepr, PushKernelArg}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DeviceRepr}; use ml_core::cuda_autograd::GpuVarStore; use tracing::info; @@ -113,20 +113,20 @@ const PPO_CRITIC_WEIGHT_NAMES: [&str; 12] = [ #[allow(missing_debug_implementations)] // CudaSlice does not implement 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: 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] // 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: 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] // 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: 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] } /// GPU buffers holding the 3 extra advantage head weight sets for Branching DQN. @@ -147,20 +147,20 @@ pub struct DuelingWeightSet { #[allow(missing_debug_implementations)] // CudaSlice does not implement 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: 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] // 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: 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] // 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: 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] } /// GPU buffers holding all 4 weight tensors of the Curiosity Forward Model. @@ -178,10 +178,10 @@ const CUR_OUTPUT_DIM: usize = CUR_MARKET_DIM; // 42 #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct CuriosityWeightSet { - pub w1: CudaSlice, // fc1.weight [CUR_HIDDEN, CUR_INPUT] = [128, 45] - pub b1: CudaSlice, // fc1.bias [CUR_HIDDEN] = [128] - pub w2: CudaSlice, // fc2.weight [CUR_OUTPUT, CUR_HIDDEN] = [42, 128] - pub b2: CudaSlice, // fc2.bias [CUR_OUTPUT] = [42] + pub w1: CudaSlice, // fc1.weight [CUR_HIDDEN, CUR_INPUT] = [128, 45] + pub b1: CudaSlice, // fc1.bias [CUR_HIDDEN] = [128] + pub w2: CudaSlice, // fc2.weight [CUR_OUTPUT, CUR_HIDDEN] = [42, 128] + pub b2: CudaSlice, // fc2.bias [CUR_OUTPUT] = [42] } impl CuriosityWeightSet { @@ -193,16 +193,16 @@ impl CuriosityWeightSet { /// `curiosity_scale` is set to 0.0 in the config). pub fn zeros(stream: &Arc) -> Result { Ok(Self { - w1: stream.alloc_zeros::(CUR_HIDDEN_DIM * CUR_INPUT_DIM).map_err(|e| { + w1: stream.alloc_zeros::(CUR_HIDDEN_DIM * CUR_INPUT_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity w1 zeros: {e}")) })?, - b1: stream.alloc_zeros::(CUR_HIDDEN_DIM).map_err(|e| { + b1: stream.alloc_zeros::(CUR_HIDDEN_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity b1 zeros: {e}")) })?, - w2: stream.alloc_zeros::(CUR_OUTPUT_DIM * CUR_HIDDEN_DIM).map_err(|e| { + w2: stream.alloc_zeros::(CUR_OUTPUT_DIM * CUR_HIDDEN_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity w2 zeros: {e}")) })?, - b2: stream.alloc_zeros::(CUR_OUTPUT_DIM).map_err(|e| { + b2: stream.alloc_zeros::(CUR_OUTPUT_DIM).map_err(|e| { MLError::ModelError(format!("Failed to alloc curiosity b2 zeros: {e}")) })?, }) @@ -217,12 +217,12 @@ impl CuriosityWeightSet { /// - Output: `[45, 64]`, `[45]` #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct PpoActorWeightSet { - pub pw1: CudaSlice, // policy_layer_0.weight [128, 54] - pub pb1: CudaSlice, // policy_layer_0.bias [128] - pub pw2: CudaSlice, // policy_layer_1.weight [64, 128] - pub pb2: CudaSlice, // policy_layer_1.bias [64] - pub pw3: CudaSlice, // policy_output.weight [45, 64] - pub pb3: CudaSlice, // policy_output.bias [45] + pub pw1: CudaSlice, // policy_layer_0.weight [128, 54] + pub pb1: CudaSlice, // policy_layer_0.bias [128] + pub pw2: CudaSlice, // policy_layer_1.weight [64, 128] + pub pb2: CudaSlice, // policy_layer_1.bias [64] + pub pw3: CudaSlice, // policy_output.weight [45, 64] + pub pb3: CudaSlice, // policy_output.bias [45] } /// GPU buffers holding all 12 weight tensors of the PPO Critic (ValueNetwork). @@ -236,18 +236,18 @@ pub struct PpoActorWeightSet { /// - Output: `[1, 64]`, `[1]` #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct PpoCriticWeightSet { - pub vw1: CudaSlice, // value_layer_0.weight [512, 54] - pub vb1: CudaSlice, // value_layer_0.bias [512] - pub vw2: CudaSlice, // value_layer_1.weight [384, 512] - pub vb2: CudaSlice, // value_layer_1.bias [384] - pub vw3: CudaSlice, // value_layer_2.weight [256, 384] - pub vb3: CudaSlice, // value_layer_2.bias [256] - pub vw4: CudaSlice, // value_layer_3.weight [128, 256] - pub vb4: CudaSlice, // value_layer_3.bias [128] - pub vw5: CudaSlice, // value_layer_4.weight [64, 128] - pub vb5: CudaSlice, // value_layer_4.bias [64] - pub vw6: CudaSlice, // value_output.weight [1, 64] - pub vb6: CudaSlice, // value_output.bias [1] + pub vw1: CudaSlice, // value_layer_0.weight [512, 54] + pub vb1: CudaSlice, // value_layer_0.bias [512] + pub vw2: CudaSlice, // value_layer_1.weight [384, 512] + pub vb2: CudaSlice, // value_layer_1.bias [384] + pub vw3: CudaSlice, // value_layer_2.weight [256, 384] + pub vb3: CudaSlice, // value_layer_2.bias [256] + pub vw4: CudaSlice, // value_layer_3.weight [128, 256] + pub vb4: CudaSlice, // value_layer_3.bias [128] + pub vw5: CudaSlice, // value_layer_4.weight [64, 128] + pub vb5: CudaSlice, // value_layer_4.bias [64] + pub vw6: CudaSlice, // value_output.weight [1, 64] + pub vb6: CudaSlice, // value_output.bias [1] } /// RMSNorm gamma weight names for distributional dueling network (4 total). @@ -267,10 +267,10 @@ const RMSNORM_WEIGHT_NAMES: [&str; 4] = [ /// - advantage_rmsnorm:[128] (after advantage FC) #[allow(missing_debug_implementations)] pub struct RmsNormWeightSet { - pub gamma_s0: CudaSlice, // shared_rmsnorm_0.weight [shared_h1] - pub gamma_s1: CudaSlice, // shared_rmsnorm_1.weight [shared_h2] - pub gamma_v: CudaSlice, // value_rmsnorm.weight [value_h] - pub gamma_a: CudaSlice, // advantage_rmsnorm.weight [adv_h] + pub gamma_s0: CudaSlice, // shared_rmsnorm_0.weight [shared_h1] + pub gamma_s1: CudaSlice, // shared_rmsnorm_1.weight [shared_h2] + pub gamma_v: CudaSlice, // value_rmsnorm.weight [value_h] + pub gamma_a: CudaSlice, // advantage_rmsnorm.weight [adv_h] } impl RmsNormWeightSet { @@ -283,12 +283,12 @@ impl RmsNormWeightSet { network_dims: (usize, usize, usize, usize), ) -> Result { let (shared_h1, shared_h2, value_h, adv_h) = network_dims; - let alloc_ones = |size: usize| -> Result, MLError> { + let alloc_ones = |size: usize| -> Result, MLError> { let data = vec![1.0_f32; size]; let mut buf = stream - .alloc_zeros::(size) + .alloc_zeros::(size) .map_err(|e| MLError::ModelError(format!("Alloc RMSNorm ones: {e}")))?; - super::htod_f32_to_bf16(stream, &data, &mut buf)?; + super::htod_f32(stream, &data, &mut buf)?; Ok(buf) }; Ok(Self { @@ -384,7 +384,7 @@ unsafe impl DeviceRepr for KernelWeightPack {} /// - The CudaSlice is owned by the caller and outlives the kernel launch /// - The kernel launches on the same stream, so CUDA ordering is guaranteed /// - `SyncOnDrop::Record` only records events for cross-stream sync (not needed) -fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { +fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { let (ptr, guard) = slice.device_ptr(stream); // Keep guard alive without running its Drop (skip read event recording). // Named binding prevents `let_underscore_must_use`; ManuallyDrop prevents drop. @@ -479,324 +479,6 @@ impl KernelWeightPack { } } -// --------------------------------------------------------------------------- -// BF16 Weight Mirror — parallel BF16 copies of all F32 weight buffers -// --------------------------------------------------------------------------- - -/// BF16 mirror of a DuelingWeightSet. Stored as CudaSlice where each u16 -/// holds the bit pattern of an __nv_bfloat16 value. -/// -/// Updated automatically after each F32 weight sync via the `f32_to_bf16_kernel` -/// CUDA kernel. The training forward pass reads from these BF16 buffers; -/// backward/Adam operate on the F32 originals. -#[allow(missing_debug_implementations)] -pub struct DuelingWeightSetBf16 { - pub w_s1: CudaSlice, - pub b_s1: CudaSlice, - pub w_s2: CudaSlice, - pub b_s2: CudaSlice, - pub w_v1: CudaSlice, - pub b_v1: CudaSlice, - pub w_v2: CudaSlice, - pub b_v2: CudaSlice, - pub w_a1: CudaSlice, - pub b_a1: CudaSlice, - pub w_a2: CudaSlice, - pub b_a2: CudaSlice, -} - -/// BF16 mirror of BranchingWeightSet (order + urgency + magnitude heads). -#[allow(missing_debug_implementations)] -pub struct BranchingWeightSetBf16 { - pub w_bo1: CudaSlice, - pub b_bo1: CudaSlice, - pub w_bo2: CudaSlice, - pub b_bo2: CudaSlice, - pub w_bu1: CudaSlice, - pub b_bu1: CudaSlice, - pub w_bu2: CudaSlice, - pub b_bu2: CudaSlice, - pub w_bg1: CudaSlice, - pub b_bg1: CudaSlice, - pub w_bg2: CudaSlice, - pub b_bg2: CudaSlice, -} - -/// BF16 mirror of CuriosityWeightSet. -#[allow(missing_debug_implementations)] -pub struct CuriosityWeightSetBf16 { - pub w1: CudaSlice, - pub b1: CudaSlice, - pub w2: CudaSlice, - pub b2: CudaSlice, -} - -/// Allocate a BF16 buffer matching an F32 buffer's element count. -fn alloc_bf16_mirror( - f32_buf: &CudaSlice, - stream: &Arc, -) -> Result, MLError> { - let n = f32_buf.len(); - stream.alloc_zeros::(n).map_err(|e| { - MLError::ModelError(format!("BF16 mirror alloc ({n} elems): {e}")) - }) -} - -/// Convert F32 buffer → BF16 buffer on GPU using a pre-compiled CUDA kernel. -/// -/// Launches `f32_to_bf16_kernel` (defined in `common_device_functions.cuh`, -/// compiled via NVRTC as part of the training kernel module). -/// -/// `func`: pre-compiled `CudaFunction` for `f32_to_bf16_kernel`. -pub fn convert_f32_to_bf16( - src: &CudaSlice, - dst: &mut CudaSlice, - func: &cudarc::driver::CudaFunction, - stream: &Arc, -) -> Result<(), MLError> { - let n = src.len(); - if n == 0 { return Ok(()); } - let block = 256_u32; - let grid = ((n as u32) + block - 1) / block; - let cfg = cudarc::driver::LaunchConfig { - grid_dim: (grid, 1, 1), - block_dim: (block, 1, 1), - shared_mem_bytes: 0, - }; - unsafe { - stream - .launch_builder(func) - .arg(src) - .arg(dst) - .arg(&(n as i32)) - .launch(cfg) - .map_err(|e| MLError::ModelError(format!("f32_to_bf16 launch: {e}")))?; - } - Ok(()) -} - -impl DuelingWeightSetBf16 { - /// Allocate BF16 mirrors matching an existing F32 weight set. - pub fn alloc_from(f32_set: &DuelingWeightSet, stream: &Arc) -> Result { - Ok(Self { - w_s1: alloc_bf16_mirror(&f32_set.w_s1, stream)?, - b_s1: alloc_bf16_mirror(&f32_set.b_s1, stream)?, - w_s2: alloc_bf16_mirror(&f32_set.w_s2, stream)?, - b_s2: alloc_bf16_mirror(&f32_set.b_s2, stream)?, - w_v1: alloc_bf16_mirror(&f32_set.w_v1, stream)?, - b_v1: alloc_bf16_mirror(&f32_set.b_v1, stream)?, - w_v2: alloc_bf16_mirror(&f32_set.w_v2, stream)?, - b_v2: alloc_bf16_mirror(&f32_set.b_v2, stream)?, - w_a1: alloc_bf16_mirror(&f32_set.w_a1, stream)?, - b_a1: alloc_bf16_mirror(&f32_set.b_a1, stream)?, - w_a2: alloc_bf16_mirror(&f32_set.w_a2, stream)?, - b_a2: alloc_bf16_mirror(&f32_set.b_a2, stream)?, - }) - } - - /// Sync all 12 BF16 mirrors from F32 originals using GPU kernel. - pub fn sync_from_f32( - &mut self, - f32_set: &DuelingWeightSet, - func: &cudarc::driver::CudaFunction, - stream: &Arc, - ) -> Result<(), MLError> { - convert_f32_to_bf16(&f32_set.w_s1, &mut self.w_s1, func, stream)?; - convert_f32_to_bf16(&f32_set.b_s1, &mut self.b_s1, func, stream)?; - convert_f32_to_bf16(&f32_set.w_s2, &mut self.w_s2, func, stream)?; - convert_f32_to_bf16(&f32_set.b_s2, &mut self.b_s2, func, stream)?; - convert_f32_to_bf16(&f32_set.w_v1, &mut self.w_v1, func, stream)?; - convert_f32_to_bf16(&f32_set.b_v1, &mut self.b_v1, func, stream)?; - convert_f32_to_bf16(&f32_set.w_v2, &mut self.w_v2, func, stream)?; - convert_f32_to_bf16(&f32_set.b_v2, &mut self.b_v2, func, stream)?; - convert_f32_to_bf16(&f32_set.w_a1, &mut self.w_a1, func, stream)?; - convert_f32_to_bf16(&f32_set.b_a1, &mut self.b_a1, func, stream)?; - convert_f32_to_bf16(&f32_set.w_a2, &mut self.w_a2, func, stream)?; - convert_f32_to_bf16(&f32_set.b_a2, &mut self.b_a2, func, stream)?; - Ok(()) - } -} - -impl BranchingWeightSetBf16 { - /// Allocate BF16 mirrors matching an existing F32 weight set. - pub fn alloc_from(f32_set: &BranchingWeightSet, stream: &Arc) -> Result { - Ok(Self { - w_bo1: alloc_bf16_mirror(&f32_set.w_bo1, stream)?, - b_bo1: alloc_bf16_mirror(&f32_set.b_bo1, stream)?, - w_bo2: alloc_bf16_mirror(&f32_set.w_bo2, stream)?, - b_bo2: alloc_bf16_mirror(&f32_set.b_bo2, stream)?, - w_bu1: alloc_bf16_mirror(&f32_set.w_bu1, stream)?, - b_bu1: alloc_bf16_mirror(&f32_set.b_bu1, stream)?, - w_bu2: alloc_bf16_mirror(&f32_set.w_bu2, stream)?, - b_bu2: alloc_bf16_mirror(&f32_set.b_bu2, stream)?, - w_bg1: alloc_bf16_mirror(&f32_set.w_bg1, stream)?, - b_bg1: alloc_bf16_mirror(&f32_set.b_bg1, stream)?, - w_bg2: alloc_bf16_mirror(&f32_set.w_bg2, stream)?, - b_bg2: alloc_bf16_mirror(&f32_set.b_bg2, stream)?, - }) - } - - /// Sync all 12 BF16 mirrors from F32 originals. - pub fn sync_from_f32( - &mut self, - f32_set: &BranchingWeightSet, - func: &cudarc::driver::CudaFunction, - stream: &Arc, - ) -> Result<(), MLError> { - convert_f32_to_bf16(&f32_set.w_bo1, &mut self.w_bo1, func, stream)?; - convert_f32_to_bf16(&f32_set.b_bo1, &mut self.b_bo1, func, stream)?; - convert_f32_to_bf16(&f32_set.w_bo2, &mut self.w_bo2, func, stream)?; - convert_f32_to_bf16(&f32_set.b_bo2, &mut self.b_bo2, func, stream)?; - convert_f32_to_bf16(&f32_set.w_bu1, &mut self.w_bu1, func, stream)?; - convert_f32_to_bf16(&f32_set.b_bu1, &mut self.b_bu1, func, stream)?; - convert_f32_to_bf16(&f32_set.w_bu2, &mut self.w_bu2, func, stream)?; - convert_f32_to_bf16(&f32_set.b_bu2, &mut self.b_bu2, func, stream)?; - convert_f32_to_bf16(&f32_set.w_bg1, &mut self.w_bg1, func, stream)?; - convert_f32_to_bf16(&f32_set.b_bg1, &mut self.b_bg1, func, stream)?; - convert_f32_to_bf16(&f32_set.w_bg2, &mut self.w_bg2, func, stream)?; - convert_f32_to_bf16(&f32_set.b_bg2, &mut self.b_bg2, func, stream)?; - Ok(()) - } -} - -impl CuriosityWeightSetBf16 { - /// Allocate BF16 mirrors matching an existing F32 weight set. - pub fn alloc_from(f32_set: &CuriosityWeightSet, stream: &Arc) -> Result { - Ok(Self { - w1: alloc_bf16_mirror(&f32_set.w1, stream)?, - b1: alloc_bf16_mirror(&f32_set.b1, stream)?, - w2: alloc_bf16_mirror(&f32_set.w2, stream)?, - b2: alloc_bf16_mirror(&f32_set.b2, stream)?, - }) - } - - /// Sync all 4 BF16 mirrors from F32 originals. - pub fn sync_from_f32( - &mut self, - f32_set: &CuriosityWeightSet, - func: &cudarc::driver::CudaFunction, - stream: &Arc, - ) -> Result<(), MLError> { - convert_f32_to_bf16(&f32_set.w1, &mut self.w1, func, stream)?; - convert_f32_to_bf16(&f32_set.b1, &mut self.b1, func, stream)?; - convert_f32_to_bf16(&f32_set.w2, &mut self.w2, func, stream)?; - convert_f32_to_bf16(&f32_set.b2, &mut self.b2, func, stream)?; - Ok(()) - } -} - -/// Extract raw CUDA device pointer from a BF16 CudaSlice (u16 bit pattern). -fn raw_device_ptr_bf16(slice: &CudaSlice, stream: &CudaStream) -> u64 { - let (ptr, guard) = slice.device_ptr(stream); - let _no_drop = std::mem::ManuallyDrop::new(guard); - ptr -} - -/// BF16 device pointer pack for the forward pass. -/// -/// Same layout as `KernelWeightPack` but pointing to BF16 buffers. -/// The CUDA forward kernel reads these for BF16 matmul; -/// backward/Adam kernel still reads the F32 `KernelWeightPack`. -#[repr(C)] -#[derive(Clone, Copy)] -pub(super) struct KernelWeightPackBf16 { - pub online: GpuDuelingPtrs, - pub target: GpuDuelingPtrs, - pub curiosity: GpuCuriosityPtrs, - pub rmsnorm: GpuRMSNormPtrs, // RMSNorm stays F32 (gamma scaling) - pub online_br: GpuBranchPtrs, - pub target_br: GpuBranchPtrs, -} - -#[allow(unsafe_code)] -unsafe impl DeviceRepr for KernelWeightPackBf16 {} - -impl KernelWeightPackBf16 { - /// Build BF16 weight pointer pack for forward pass. - /// - /// Weight matrices use BF16 pointers; RMSNorm gammas stay F32. - pub(super) fn build( - online_bf16: &DuelingWeightSetBf16, - target_bf16: &DuelingWeightSetBf16, - curiosity_bf16: &CuriosityWeightSetBf16, - rmsnorm_f32: &RmsNormWeightSet, - online_br_bf16: &BranchingWeightSetBf16, - target_br_bf16: &BranchingWeightSetBf16, - stream: &CudaStream, - ) -> Self { - Self { - online: GpuDuelingPtrs { - w_s1: raw_device_ptr_bf16(&online_bf16.w_s1, stream), - b_s1: raw_device_ptr_bf16(&online_bf16.b_s1, stream), - w_s2: raw_device_ptr_bf16(&online_bf16.w_s2, stream), - b_s2: raw_device_ptr_bf16(&online_bf16.b_s2, stream), - w_v1: raw_device_ptr_bf16(&online_bf16.w_v1, stream), - b_v1: raw_device_ptr_bf16(&online_bf16.b_v1, stream), - w_v2: raw_device_ptr_bf16(&online_bf16.w_v2, stream), - b_v2: raw_device_ptr_bf16(&online_bf16.b_v2, stream), - w_a1: raw_device_ptr_bf16(&online_bf16.w_a1, stream), - b_a1: raw_device_ptr_bf16(&online_bf16.b_a1, stream), - w_a2: raw_device_ptr_bf16(&online_bf16.w_a2, stream), - b_a2: raw_device_ptr_bf16(&online_bf16.b_a2, stream), - }, - target: GpuDuelingPtrs { - w_s1: raw_device_ptr_bf16(&target_bf16.w_s1, stream), - b_s1: raw_device_ptr_bf16(&target_bf16.b_s1, stream), - w_s2: raw_device_ptr_bf16(&target_bf16.w_s2, stream), - b_s2: raw_device_ptr_bf16(&target_bf16.b_s2, stream), - w_v1: raw_device_ptr_bf16(&target_bf16.w_v1, stream), - b_v1: raw_device_ptr_bf16(&target_bf16.b_v1, stream), - w_v2: raw_device_ptr_bf16(&target_bf16.w_v2, stream), - b_v2: raw_device_ptr_bf16(&target_bf16.b_v2, stream), - w_a1: raw_device_ptr_bf16(&target_bf16.w_a1, stream), - b_a1: raw_device_ptr_bf16(&target_bf16.b_a1, stream), - w_a2: raw_device_ptr_bf16(&target_bf16.w_a2, stream), - b_a2: raw_device_ptr_bf16(&target_bf16.b_a2, stream), - }, - curiosity: GpuCuriosityPtrs { - w1: raw_device_ptr_bf16(&curiosity_bf16.w1, stream), - b1: raw_device_ptr_bf16(&curiosity_bf16.b1, stream), - w2: raw_device_ptr_bf16(&curiosity_bf16.w2, stream), - b2: raw_device_ptr_bf16(&curiosity_bf16.b2, stream), - }, - rmsnorm: GpuRMSNormPtrs { - s0: raw_device_ptr(&rmsnorm_f32.gamma_s0, stream), - s1: raw_device_ptr(&rmsnorm_f32.gamma_s1, stream), - v: raw_device_ptr(&rmsnorm_f32.gamma_v, stream), - a: raw_device_ptr(&rmsnorm_f32.gamma_a, stream), - }, - online_br: GpuBranchPtrs { - w_o1: raw_device_ptr_bf16(&online_br_bf16.w_bo1, stream), - b_o1: raw_device_ptr_bf16(&online_br_bf16.b_bo1, stream), - w_o2: raw_device_ptr_bf16(&online_br_bf16.w_bo2, stream), - b_o2: raw_device_ptr_bf16(&online_br_bf16.b_bo2, stream), - w_u1: raw_device_ptr_bf16(&online_br_bf16.w_bu1, stream), - b_u1: raw_device_ptr_bf16(&online_br_bf16.b_bu1, stream), - w_u2: raw_device_ptr_bf16(&online_br_bf16.w_bu2, stream), - b_u2: raw_device_ptr_bf16(&online_br_bf16.b_bu2, stream), - w_g1: raw_device_ptr_bf16(&online_br_bf16.w_bg1, stream), - b_g1: raw_device_ptr_bf16(&online_br_bf16.b_bg1, stream), - w_g2: raw_device_ptr_bf16(&online_br_bf16.w_bg2, stream), - b_g2: raw_device_ptr_bf16(&online_br_bf16.b_bg2, stream), - }, - target_br: GpuBranchPtrs { - w_o1: raw_device_ptr_bf16(&target_br_bf16.w_bo1, stream), - b_o1: raw_device_ptr_bf16(&target_br_bf16.b_bo1, stream), - w_o2: raw_device_ptr_bf16(&target_br_bf16.w_bo2, stream), - b_o2: raw_device_ptr_bf16(&target_br_bf16.b_bo2, stream), - w_u1: raw_device_ptr_bf16(&target_br_bf16.w_bu1, stream), - b_u1: raw_device_ptr_bf16(&target_br_bf16.b_bu1, stream), - w_u2: raw_device_ptr_bf16(&target_br_bf16.w_bu2, stream), - b_u2: raw_device_ptr_bf16(&target_br_bf16.b_bu2, stream), - w_g1: raw_device_ptr_bf16(&target_br_bf16.w_bg1, stream), - b_g1: raw_device_ptr_bf16(&target_br_bf16.b_bg1, stream), - w_g2: raw_device_ptr_bf16(&target_br_bf16.w_bg2, stream), - b_g2: raw_device_ptr_bf16(&target_br_bf16.b_bg2, stream), - }, - } - } -} - // --------------------------------------------------------------------------- // Extraction helpers // --------------------------------------------------------------------------- @@ -829,14 +511,14 @@ fn extract_one( vars: &GpuVarStore, name: &str, stream: &Arc, -) -> Result, MLError> { +) -> Result, MLError> { let param = vars .get(name) .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))?; let n_elems = param.data.len(); let mut buf = stream - .alloc_zeros::(n_elems) + .alloc_zeros::(n_elems) .map_err(|e| MLError::ModelError(format!("Alloc {name}: {e}")))?; stream.memcpy_dtod(¶m.data, &mut buf) @@ -850,7 +532,7 @@ fn extract_one( fn sync_one( vars: &GpuVarStore, name: &str, - buf: &mut CudaSlice, + buf: &mut CudaSlice, stream: &Arc, ) -> Result<(), MLError> { let param = vars @@ -868,7 +550,7 @@ fn sync_one( // --------------------------------------------------------------------------- /// Extract all 12 Dueling Q-Network weight tensors from a `GpuVarStore` and -/// upload them to GPU as flat `CudaSlice` buffers. +/// upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `DuelingQNetwork`: /// `shared_0.weight`, `shared_0.bias`, `shared_1.weight`, `shared_1.bias`, @@ -880,7 +562,7 @@ pub fn extract_dueling_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -906,7 +588,7 @@ pub fn extract_dueling_weights( info!( total_params, - total_bytes = total_params * std::mem::size_of::(), + total_bytes = total_params * std::mem::size_of::(), "Dueling Q-Network weights extracted and uploaded to GPU" ); @@ -927,7 +609,7 @@ pub fn extract_dueling_weights( } /// Extract all 12 Branching DQN extra head weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `BranchingDuelingQNetwork`: /// `branch_1_fc.weight`, `branch_1_fc.bias`, `branch_1_out.weight`, `branch_1_out.bias`, @@ -938,7 +620,7 @@ pub fn extract_branching_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -964,7 +646,7 @@ pub fn extract_branching_weights( info!( total_params, - total_bytes = total_params * std::mem::size_of::(), + total_bytes = total_params * std::mem::size_of::(), "Branching DQN weights extracted and uploaded to GPU" ); @@ -998,7 +680,7 @@ pub fn extract_dueling_weights_branching( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -1053,7 +735,7 @@ pub fn sync_dueling_weights_branching( } /// Extract all 4 Curiosity Forward Model weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `ForwardDynamicsModel`: /// `fc1.weight`, `fc1.bias`, `fc2.weight`, `fc2.bias`. @@ -1062,7 +744,7 @@ pub fn extract_curiosity_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -1080,7 +762,7 @@ pub fn extract_curiosity_weights( info!( total_params, - total_bytes = total_params * std::mem::size_of::(), + total_bytes = total_params * std::mem::size_of::(), "Curiosity forward model weights extracted and uploaded to GPU" ); @@ -1151,7 +833,7 @@ pub fn sync_branching_weights( fn reverse_sync_one( vars: &GpuVarStore, name: &str, - src: &CudaSlice, + src: &CudaSlice, stream: &Arc, ) -> Result<(), MLError> { let param = vars @@ -1168,7 +850,7 @@ fn reverse_sync_one( let (dst_ptr, _dst_sync) = param.data.device_ptr(stream); let (src_ptr, _src_sync) = src.device_ptr(stream); - let num_bytes = n_elems * std::mem::size_of::(); + let num_bytes = n_elems * std::mem::size_of::(); dtod_copy_checked(dst_ptr, src_ptr, num_bytes, stream, name, "reverse_sync")?; Ok(()) @@ -1301,7 +983,7 @@ pub fn sync_curiosity_weights( } /// Extract all 6 PPO Actor (PolicyNetwork) weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `PolicyNetwork`: /// `policy_layer_0.weight`, `policy_layer_0.bias`, `policy_layer_1.weight`, @@ -1311,7 +993,7 @@ pub fn extract_ppo_actor_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -1331,7 +1013,7 @@ pub fn extract_ppo_actor_weights( info!( total_params, - total_bytes = total_params * std::mem::size_of::(), + total_bytes = total_params * std::mem::size_of::(), "PPO Actor weights extracted and uploaded to GPU" ); @@ -1346,7 +1028,7 @@ pub fn extract_ppo_actor_weights( } /// Extract all 12 PPO Critic (ValueNetwork) weight tensors from a `GpuVarStore` -/// and upload them to GPU as flat `CudaSlice` buffers. +/// and upload them to GPU as flat `CudaSlice` buffers. /// /// The `GpuVarStore` must contain tensor keys matching `ValueNetwork`: /// `value_layer_0.weight`, `value_layer_0.bias`, `value_layer_1.weight`, @@ -1358,7 +1040,7 @@ pub fn extract_ppo_critic_weights( stream: &Arc, ) -> Result { - let extract = |name: &str| -> Result, MLError> { + let extract = |name: &str| -> Result, MLError> { extract_one(vars, name, stream) }; @@ -1384,7 +1066,7 @@ pub fn extract_ppo_critic_weights( info!( total_params, - total_bytes = total_params * std::mem::size_of::(), + total_bytes = total_params * std::mem::size_of::(), "PPO Critic weights extracted and uploaded to GPU" );