refactor: gpu_weights all CudaSlice<half::bf16> → CudaSlice<f32>, delete bf16 mirror infrastructure
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<half::bf16>` GPU buffers for the CUDA experience
|
||||
//! them as flat `CudaSlice<f32>` 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<half::bf16>` on GPU. Extraction and sync use
|
||||
//! 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)]
|
||||
|
||||
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<half::bf16>, // shared_0.weight [SHARED_H1, STATE_DIM]
|
||||
pub b_s1: CudaSlice<half::bf16>, // shared_0.bias [SHARED_H1]
|
||||
pub w_s2: CudaSlice<half::bf16>, // shared_1.weight [SHARED_H2, SHARED_H1]
|
||||
pub b_s2: CudaSlice<half::bf16>, // shared_1.bias [SHARED_H2]
|
||||
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]
|
||||
// Value head (distributional: NUM_ATOMS outputs, not 1)
|
||||
pub w_v1: CudaSlice<half::bf16>, // value_fc.weight [VALUE_H, SHARED_H2]
|
||||
pub b_v1: CudaSlice<half::bf16>, // value_fc.bias [VALUE_H]
|
||||
pub w_v2: CudaSlice<half::bf16>, // value_out.weight [NUM_ATOMS, VALUE_H]
|
||||
pub b_v2: CudaSlice<half::bf16>, // value_out.bias [NUM_ATOMS]
|
||||
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]
|
||||
// Exposure head / branch 0 (distributional: BRANCH_0_SIZE*NUM_ATOMS outputs)
|
||||
pub w_a1: CudaSlice<half::bf16>, // branch_0_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_a1: CudaSlice<half::bf16>, // branch_0_fc.bias [ADV_H]
|
||||
pub w_a2: CudaSlice<half::bf16>, // branch_0_out.weight [BRANCH_0_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_a2: CudaSlice<half::bf16>, // branch_0_out.bias [BRANCH_0_SIZE*NUM_ATOMS]
|
||||
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]
|
||||
}
|
||||
|
||||
/// 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<half::bf16>, // branch_1_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_bo1: CudaSlice<half::bf16>, // branch_1_fc.bias [ADV_H]
|
||||
pub w_bo2: CudaSlice<half::bf16>, // branch_1_out.weight [BRANCH_1_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_bo2: CudaSlice<half::bf16>, // branch_1_out.bias [BRANCH_1_SIZE*NUM_ATOMS]
|
||||
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]
|
||||
// Urgency head (branch 2)
|
||||
pub w_bu1: CudaSlice<half::bf16>, // branch_2_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_bu1: CudaSlice<half::bf16>, // branch_2_fc.bias [ADV_H]
|
||||
pub w_bu2: CudaSlice<half::bf16>, // branch_2_out.weight [BRANCH_2_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_bu2: CudaSlice<half::bf16>, // branch_2_out.bias [BRANCH_2_SIZE*NUM_ATOMS]
|
||||
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]
|
||||
// Magnitude head (branch 3) — 4-branch refactor
|
||||
pub w_bg1: CudaSlice<half::bf16>, // branch_3_fc.weight [ADV_H, SHARED_H2]
|
||||
pub b_bg1: CudaSlice<half::bf16>, // branch_3_fc.bias [ADV_H]
|
||||
pub w_bg2: CudaSlice<half::bf16>, // branch_3_out.weight [BRANCH_3_SIZE*NUM_ATOMS, ADV_H]
|
||||
pub b_bg2: CudaSlice<half::bf16>, // branch_3_out.bias [BRANCH_3_SIZE*NUM_ATOMS]
|
||||
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]
|
||||
}
|
||||
|
||||
/// 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<half::bf16>, // fc1.weight [CUR_HIDDEN, CUR_INPUT] = [128, 45]
|
||||
pub b1: CudaSlice<half::bf16>, // fc1.bias [CUR_HIDDEN] = [128]
|
||||
pub w2: CudaSlice<half::bf16>, // fc2.weight [CUR_OUTPUT, CUR_HIDDEN] = [42, 128]
|
||||
pub b2: CudaSlice<half::bf16>, // fc2.bias [CUR_OUTPUT] = [42]
|
||||
pub w1: CudaSlice<f32>, // fc1.weight [CUR_HIDDEN, CUR_INPUT] = [128, 45]
|
||||
pub b1: CudaSlice<f32>, // fc1.bias [CUR_HIDDEN] = [128]
|
||||
pub w2: CudaSlice<f32>, // fc2.weight [CUR_OUTPUT, CUR_HIDDEN] = [42, 128]
|
||||
pub b2: CudaSlice<f32>, // 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<CudaStream>) -> Result<Self, MLError> {
|
||||
Ok(Self {
|
||||
w1: stream.alloc_zeros::<half::bf16>(CUR_HIDDEN_DIM * CUR_INPUT_DIM).map_err(|e| {
|
||||
w1: stream.alloc_zeros::<f32>(CUR_HIDDEN_DIM * CUR_INPUT_DIM).map_err(|e| {
|
||||
MLError::ModelError(format!("Failed to alloc curiosity w1 zeros: {e}"))
|
||||
})?,
|
||||
b1: stream.alloc_zeros::<half::bf16>(CUR_HIDDEN_DIM).map_err(|e| {
|
||||
b1: stream.alloc_zeros::<f32>(CUR_HIDDEN_DIM).map_err(|e| {
|
||||
MLError::ModelError(format!("Failed to alloc curiosity b1 zeros: {e}"))
|
||||
})?,
|
||||
w2: stream.alloc_zeros::<half::bf16>(CUR_OUTPUT_DIM * CUR_HIDDEN_DIM).map_err(|e| {
|
||||
w2: stream.alloc_zeros::<f32>(CUR_OUTPUT_DIM * CUR_HIDDEN_DIM).map_err(|e| {
|
||||
MLError::ModelError(format!("Failed to alloc curiosity w2 zeros: {e}"))
|
||||
})?,
|
||||
b2: stream.alloc_zeros::<half::bf16>(CUR_OUTPUT_DIM).map_err(|e| {
|
||||
b2: stream.alloc_zeros::<f32>(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<half::bf16>, // policy_layer_0.weight [128, 54]
|
||||
pub pb1: CudaSlice<half::bf16>, // policy_layer_0.bias [128]
|
||||
pub pw2: CudaSlice<half::bf16>, // policy_layer_1.weight [64, 128]
|
||||
pub pb2: CudaSlice<half::bf16>, // policy_layer_1.bias [64]
|
||||
pub pw3: CudaSlice<half::bf16>, // policy_output.weight [45, 64]
|
||||
pub pb3: CudaSlice<half::bf16>, // policy_output.bias [45]
|
||||
pub pw1: CudaSlice<f32>, // policy_layer_0.weight [128, 54]
|
||||
pub pb1: CudaSlice<f32>, // policy_layer_0.bias [128]
|
||||
pub pw2: CudaSlice<f32>, // policy_layer_1.weight [64, 128]
|
||||
pub pb2: CudaSlice<f32>, // policy_layer_1.bias [64]
|
||||
pub pw3: CudaSlice<f32>, // policy_output.weight [45, 64]
|
||||
pub pb3: CudaSlice<f32>, // 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<half::bf16>, // value_layer_0.weight [512, 54]
|
||||
pub vb1: CudaSlice<half::bf16>, // value_layer_0.bias [512]
|
||||
pub vw2: CudaSlice<half::bf16>, // value_layer_1.weight [384, 512]
|
||||
pub vb2: CudaSlice<half::bf16>, // value_layer_1.bias [384]
|
||||
pub vw3: CudaSlice<half::bf16>, // value_layer_2.weight [256, 384]
|
||||
pub vb3: CudaSlice<half::bf16>, // value_layer_2.bias [256]
|
||||
pub vw4: CudaSlice<half::bf16>, // value_layer_3.weight [128, 256]
|
||||
pub vb4: CudaSlice<half::bf16>, // value_layer_3.bias [128]
|
||||
pub vw5: CudaSlice<half::bf16>, // value_layer_4.weight [64, 128]
|
||||
pub vb5: CudaSlice<half::bf16>, // value_layer_4.bias [64]
|
||||
pub vw6: CudaSlice<half::bf16>, // value_output.weight [1, 64]
|
||||
pub vb6: CudaSlice<half::bf16>, // value_output.bias [1]
|
||||
pub vw1: CudaSlice<f32>, // value_layer_0.weight [512, 54]
|
||||
pub vb1: CudaSlice<f32>, // value_layer_0.bias [512]
|
||||
pub vw2: CudaSlice<f32>, // value_layer_1.weight [384, 512]
|
||||
pub vb2: CudaSlice<f32>, // value_layer_1.bias [384]
|
||||
pub vw3: CudaSlice<f32>, // value_layer_2.weight [256, 384]
|
||||
pub vb3: CudaSlice<f32>, // value_layer_2.bias [256]
|
||||
pub vw4: CudaSlice<f32>, // value_layer_3.weight [128, 256]
|
||||
pub vb4: CudaSlice<f32>, // value_layer_3.bias [128]
|
||||
pub vw5: CudaSlice<f32>, // value_layer_4.weight [64, 128]
|
||||
pub vb5: CudaSlice<f32>, // value_layer_4.bias [64]
|
||||
pub vw6: CudaSlice<f32>, // value_output.weight [1, 64]
|
||||
pub vb6: CudaSlice<f32>, // 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<half::bf16>, // shared_rmsnorm_0.weight [shared_h1]
|
||||
pub gamma_s1: CudaSlice<half::bf16>, // shared_rmsnorm_1.weight [shared_h2]
|
||||
pub gamma_v: CudaSlice<half::bf16>, // value_rmsnorm.weight [value_h]
|
||||
pub gamma_a: CudaSlice<half::bf16>, // advantage_rmsnorm.weight [adv_h]
|
||||
pub gamma_s0: CudaSlice<f32>, // shared_rmsnorm_0.weight [shared_h1]
|
||||
pub gamma_s1: CudaSlice<f32>, // shared_rmsnorm_1.weight [shared_h2]
|
||||
pub gamma_v: CudaSlice<f32>, // value_rmsnorm.weight [value_h]
|
||||
pub gamma_a: CudaSlice<f32>, // advantage_rmsnorm.weight [adv_h]
|
||||
}
|
||||
|
||||
impl RmsNormWeightSet {
|
||||
@@ -283,12 +283,12 @@ impl RmsNormWeightSet {
|
||||
network_dims: (usize, usize, usize, usize),
|
||||
) -> Result<Self, MLError> {
|
||||
let (shared_h1, shared_h2, value_h, adv_h) = network_dims;
|
||||
let alloc_ones = |size: usize| -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
let alloc_ones = |size: usize| -> Result<CudaSlice<f32>, MLError> {
|
||||
let data = vec![1.0_f32; size];
|
||||
let mut buf = stream
|
||||
.alloc_zeros::<half::bf16>(size)
|
||||
.alloc_zeros::<f32>(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<half::bf16>, stream: &CudaStream) -> u64 {
|
||||
fn raw_device_ptr(slice: &CudaSlice<f32>, 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<u16> 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<u16>,
|
||||
pub b_s1: CudaSlice<u16>,
|
||||
pub w_s2: CudaSlice<u16>,
|
||||
pub b_s2: CudaSlice<u16>,
|
||||
pub w_v1: CudaSlice<u16>,
|
||||
pub b_v1: CudaSlice<u16>,
|
||||
pub w_v2: CudaSlice<u16>,
|
||||
pub b_v2: CudaSlice<u16>,
|
||||
pub w_a1: CudaSlice<u16>,
|
||||
pub b_a1: CudaSlice<u16>,
|
||||
pub w_a2: CudaSlice<u16>,
|
||||
pub b_a2: CudaSlice<u16>,
|
||||
}
|
||||
|
||||
/// BF16 mirror of BranchingWeightSet (order + urgency + magnitude heads).
|
||||
#[allow(missing_debug_implementations)]
|
||||
pub struct BranchingWeightSetBf16 {
|
||||
pub w_bo1: CudaSlice<u16>,
|
||||
pub b_bo1: CudaSlice<u16>,
|
||||
pub w_bo2: CudaSlice<u16>,
|
||||
pub b_bo2: CudaSlice<u16>,
|
||||
pub w_bu1: CudaSlice<u16>,
|
||||
pub b_bu1: CudaSlice<u16>,
|
||||
pub w_bu2: CudaSlice<u16>,
|
||||
pub b_bu2: CudaSlice<u16>,
|
||||
pub w_bg1: CudaSlice<u16>,
|
||||
pub b_bg1: CudaSlice<u16>,
|
||||
pub w_bg2: CudaSlice<u16>,
|
||||
pub b_bg2: CudaSlice<u16>,
|
||||
}
|
||||
|
||||
/// BF16 mirror of CuriosityWeightSet.
|
||||
#[allow(missing_debug_implementations)]
|
||||
pub struct CuriosityWeightSetBf16 {
|
||||
pub w1: CudaSlice<u16>,
|
||||
pub b1: CudaSlice<u16>,
|
||||
pub w2: CudaSlice<u16>,
|
||||
pub b2: CudaSlice<u16>,
|
||||
}
|
||||
|
||||
/// Allocate a BF16 buffer matching an F32 buffer's element count.
|
||||
fn alloc_bf16_mirror(
|
||||
f32_buf: &CudaSlice<half::bf16>,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> Result<CudaSlice<u16>, MLError> {
|
||||
let n = f32_buf.len();
|
||||
stream.alloc_zeros::<u16>(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<half::bf16>,
|
||||
dst: &mut CudaSlice<u16>,
|
||||
func: &cudarc::driver::CudaFunction,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> 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<CudaStream>) -> Result<Self, MLError> {
|
||||
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<CudaStream>,
|
||||
) -> 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<CudaStream>) -> Result<Self, MLError> {
|
||||
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<CudaStream>,
|
||||
) -> 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<CudaStream>) -> Result<Self, MLError> {
|
||||
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<CudaStream>,
|
||||
) -> 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<u16>, 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<CudaStream>,
|
||||
) -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
) -> Result<CudaSlice<f32>, 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::<half::bf16>(n_elems)
|
||||
.alloc_zeros::<f32>(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<half::bf16>,
|
||||
buf: &mut CudaSlice<f32>,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> 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<half::bf16>` buffers.
|
||||
/// upload them to GPU as flat `CudaSlice<f32>` 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<CudaStream>,
|
||||
) -> Result<DuelingWeightSet, MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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::<half::bf16>(),
|
||||
total_bytes = total_params * std::mem::size_of::<f32>(),
|
||||
"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<half::bf16>` buffers.
|
||||
/// and upload them to GPU as flat `CudaSlice<f32>` 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<CudaStream>,
|
||||
) -> Result<BranchingWeightSet, MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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::<half::bf16>(),
|
||||
total_bytes = total_params * std::mem::size_of::<f32>(),
|
||||
"Branching DQN weights extracted and uploaded to GPU"
|
||||
);
|
||||
|
||||
@@ -998,7 +680,7 @@ pub fn extract_dueling_weights_branching(
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> Result<DuelingWeightSet, MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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<half::bf16>` buffers.
|
||||
/// and upload them to GPU as flat `CudaSlice<f32>` 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<CudaStream>,
|
||||
) -> Result<CuriosityWeightSet, MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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::<half::bf16>(),
|
||||
total_bytes = total_params * std::mem::size_of::<f32>(),
|
||||
"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<half::bf16>,
|
||||
src: &CudaSlice<f32>,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> 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::<half::bf16>();
|
||||
let num_bytes = n_elems * std::mem::size_of::<f32>();
|
||||
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<half::bf16>` buffers.
|
||||
/// and upload them to GPU as flat `CudaSlice<f32>` 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<CudaStream>,
|
||||
) -> Result<PpoActorWeightSet, MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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::<half::bf16>(),
|
||||
total_bytes = total_params * std::mem::size_of::<f32>(),
|
||||
"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<half::bf16>` buffers.
|
||||
/// and upload them to GPU as flat `CudaSlice<f32>` 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<CudaStream>,
|
||||
) -> Result<PpoCriticWeightSet, MLError> {
|
||||
|
||||
let extract = |name: &str| -> Result<CudaSlice<half::bf16>, MLError> {
|
||||
let extract = |name: &str| -> Result<CudaSlice<f32>, 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::<half::bf16>(),
|
||||
total_bytes = total_params * std::mem::size_of::<f32>(),
|
||||
"PPO Critic weights extracted and uploaded to GPU"
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user