gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cleanup: remove dead bf16 buffers + update FactoredAction docs for 4-branch

Browse Source 
- Delete 3 unused bf16 scratch buffers (exp_h_b0/b1/b2) from
  GpuExperienceCollector — allocated VRAM but never read.
  The f32 versions (exp_h_b0_f32 etc.) remain in use.
- Fix stale doc comments: 45 actions -> 81 actions, index range 0-44 -> 0-80
- Extend round-trip tests to cover all 81 action indices
- Add TODO for future 4-branch struct refactor (direction + magnitude fields)
  — 64 callers make it too invasive for this commit.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-08 19:52:20 +02:00
parent bb0e1308d5
commit 752e19656a
3 changed files with 13 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
crates/ml-core/src/common/action.rs
View File

@@ -173,8 +173,14 @@ impl fmt::Display for Urgency {
/// Factored trading action combining exposure, order type, and urgency.
///
/// 5 exposure levels x 3 order types x 3 urgency levels = 45 actions.
/// Index mapping: `index = exposure * 9 + order * 3 + urgency` (0-44).
/// 9 exposure levels x 3 order types x 3 urgency levels = 81 actions.
/// Index mapping: `index = exposure * 9 + order * 3 + urgency` (0-80).
///
/// NOTE: The CUDA training pipeline uses a 4-branch architecture
/// (direction × magnitude × order × urgency) where the first two branches
/// correspond to the `ExposureLevel` enum (9 = 3×3 combinations).
// TODO: update for 4-branch — split `exposure` into `direction` + `magnitude`
// fields once all ~64 callers are migrated in a dedicated refactor.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct FactoredAction {
pub exposure: ExposureLevel,
@@ -192,7 +198,7 @@ impl FactoredAction {
}
}
/// Map action index (0-44) to (exposure, order, urgency)
/// Map action index (0-80) to (exposure, order, urgency)
/// Index = exposure * 9 + order * 3 + urgency
pub fn from_index(idx: usize) -> Result<Self, MLError> {
if idx >= 81 {
@@ -336,8 +342,8 @@ mod tests {
use super::*;
#[test]
fn test_factored_action_round_trip_all_45() {
for idx in 0..45 {
fn test_factored_action_round_trip_all_81() {
for idx in 0..81 {
let action = FactoredAction::from_index(idx).unwrap();
assert_eq!(
action.to_index(),

4
crates/ml-dqn/src/branching.rs
View File

@@ -818,7 +818,7 @@ impl BranchingDuelingQNetwork {
Ok(actions)
}
/// Decompose a factored action index (0-44) into per-branch indices.
/// Decompose a factored action index (0-80) into per-branch indices.
///
/// Index = exposure × (`num_order` × `num_urgency`) + order × `num_urgency` + urgency
///
@@ -1252,7 +1252,7 @@ mod tests {
#[test]
fn test_decompose_compose_roundtrip() {
for idx in 0..45_u32 {
for idx in 0..81_u32 {
let (e, o, u) = BranchingDuelingQNetwork::decompose_factored_action(idx as usize, 3, 3);
let composed =
BranchingDuelingQNetwork::compose_factored_action(e as u32, o as u32, u as u32, 3, 3);

15
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -511,9 +511,6 @@ pub struct GpuExperienceCollector {
exp_h_s1: CudaSlice<half::bf16>, // [N, shared_h1]
exp_h_s2: CudaSlice<half::bf16>, // [N, shared_h2]
exp_h_v: CudaSlice<half::bf16>, // [N, value_h]
exp_h_b0: CudaSlice<half::bf16>, // [N, adv_h]
exp_h_b1: CudaSlice<half::bf16>, // [N, adv_h]
exp_h_b2: CudaSlice<half::bf16>, // [N, adv_h]
exp_v_logits: CudaSlice<f32>, // [N, num_atoms] f32 (output layer — no bf16 truncation)
exp_b_logits: CudaSlice<f32>, // [N, total_branch_atoms] f32
@@ -943,15 +940,6 @@ impl GpuExperienceCollector {
let exp_h_v = stream
.alloc_zeros::<half::bf16>(n * value_h)
.map_err(|e| MLError::ModelError(format!("alloc exp_h_v: {e}")))?;
let exp_h_b0 = stream
.alloc_zeros::<half::bf16>(n * adv_h)
.map_err(|e| MLError::ModelError(format!("alloc exp_h_b0: {e}")))?;
let exp_h_b1 = stream
.alloc_zeros::<half::bf16>(n * adv_h)
.map_err(|e| MLError::ModelError(format!("alloc exp_h_b1: {e}")))?;
let exp_h_b2 = stream
.alloc_zeros::<half::bf16>(n * adv_h)
.map_err(|e| MLError::ModelError(format!("alloc exp_h_b2: {e}")))?;
let exp_v_logits = stream
.alloc_zeros::<f32>(n * num_atoms)
.map_err(|e| MLError::ModelError(format!("alloc exp_v_logits f32: {e}")))?;
@@ -1122,9 +1110,6 @@ impl GpuExperienceCollector {
exp_h_s1,
exp_h_s2,
exp_h_v,
exp_h_b0,
exp_h_b1,
exp_h_b2,
exp_v_logits,
exp_b_logits,
online_weights,