From 136696d8b8eacac90607f508ecf175fe4baba93d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 27 Mar 2026 20:42:37 +0100 Subject: [PATCH] =?UTF-8?q?docs:=20full=20BF16=20rewrite=20spec=20+=20plan?= =?UTF-8?q?=20=E2=80=94=20zero=20F32=20on=20GPU?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Spec: every CUDA kernel, every GPU buffer, every cuBLAS call → BF16. 37 HOT kernels + 11 WARM + 7 COLD = 55 total kernel files. 12 implementation tasks across 4 phases. Phase 1: Foundation (buffer types, weight sets) Phase 2: cuBLAS GemmEx BF16×BF16→BF16 (forward + backward) Phase 3: Optimizer + loss kernels (Adam, C51, MSE, spectral norm) Phase 4: Environment + auxiliary (experience, backtest, IQN, attention) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../plans/2026-03-27-full-bf16-rewrite.md | 577 ++++++++++++++++++ .../2026-03-27-full-bf16-rewrite-design.md | 229 +++++++ 2 files changed, 806 insertions(+) create mode 100644 docs/superpowers/plans/2026-03-27-full-bf16-rewrite.md create mode 100644 docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md diff --git a/docs/superpowers/plans/2026-03-27-full-bf16-rewrite.md b/docs/superpowers/plans/2026-03-27-full-bf16-rewrite.md new file mode 100644 index 000000000..0bb742327 --- /dev/null +++ b/docs/superpowers/plans/2026-03-27-full-bf16-rewrite.md @@ -0,0 +1,577 @@ +# Full BF16 Rewrite — Zero F32 On GPU + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Convert every CUDA kernel, every GPU buffer, and every cuBLAS call from F32 to BF16. Zero F32 on GPU. No mixed precision. + +**Architecture:** 4-phase conversion. Phase 1 lays the foundation (BF16 type helpers, common header already has `cuda_bf16.h`). Phase 2 converts cuBLAS GEMMs + activation/weight buffers. Phase 3 converts optimizer and loss kernels. Phase 4 converts environment and auxiliary kernels. Each phase compiles and passes tests independently. + +**Tech Stack:** Rust, CUDA (`__nv_bfloat16`, `cuda_bf16.h`), cudarc 0.19.3 (`CudaSlice` for BF16), cuBLAS (`cublasGemmEx`, `CUDA_R_16BF`), `half` crate (Rust-side BF16 conversion) + +**Spec:** `docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md` + +**INVARIANT:** After every task, `SQLX_OFFLINE=true cargo check -p ml && cargo test -p ml --lib -- dqn && cargo test -p ml-dqn --lib` must pass. + +--- + +## BF16 Conversion Pattern + +Every kernel follows the same pattern: + +**CUDA side:** +```cuda +// Before: +extern "C" __global__ void my_kernel(float* data, int n) { + float val = data[i]; + val = val * 2.0f; + data[i] = val; +} + +// After: +extern "C" __global__ void my_kernel(__nv_bfloat16* data, int n) { + __nv_bfloat16 val = data[i]; + float vf = __bfloat162float(val); // promote for arithmetic + vf = vf * 2.0f; + data[i] = __float2bfloat16(vf); // demote for storage +} +``` + +**Key rule:** BF16 for STORAGE (all pointers, all buffers). F32 for ARITHMETIC inside kernels (cast on load, cast on store). This is how H100 tensor cores work — BF16 inputs, F32 accumulate, BF16 output. The casts are free on SM90 (native BF16↔F32 conversion instructions). + +**Rust side:** +```rust +// Before: +let buf: CudaSlice = stream.alloc_zeros::(n)?; +stream.memcpy_htod(&host_f32, &mut buf)?; + +// After: +let buf: CudaSlice = stream.alloc_zeros::(n)?; +let host_bf16: Vec = host_f32.iter().map(|&x| half::bf16::from_f32(x).to_bits()).collect(); +stream.memcpy_htod(&host_bf16, &mut buf)?; +``` + +--- + +## Phase 1: Foundation + +### Task 1: BF16 Rust helpers + weight set type conversion + +Convert the core type infrastructure. `DuelingWeightSet` and `BranchingWeightSet` change from `CudaSlice` to `CudaSlice`. Add BF16 conversion helpers. + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_weights.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` +- Modify: `crates/ml/src/cuda_pipeline/mod.rs` + +- [ ] **Step 1: Add BF16 helper functions to gpu_weights.rs** + +```rust +use half::bf16; + +/// Convert F32 host vector to BF16 (u16) for GPU upload. +pub fn f32_to_bf16_host(data: &[f32]) -> Vec { + data.iter().map(|&x| bf16::from_f32(x).to_bits()).collect() +} + +/// Convert BF16 (u16) GPU download to F32 host vector. +pub fn bf16_to_f32_host(data: &[u16]) -> Vec { + data.iter().map(|&x| bf16::from_bits(x).to_f32()).collect() +} + +/// Allocate BF16 buffer (u16) on GPU, zero-initialized. +pub fn alloc_bf16(stream: &Arc, n: usize, label: &str) -> Result, MLError> { + stream.alloc_zeros::(n).map_err(|e| { + MLError::ModelError(format!("alloc bf16 {label}: {e}")) + }) +} +``` + +- [ ] **Step 2: Convert DuelingWeightSet fields from CudaSlice to CudaSlice** + +Change all 12 fields (w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2, w_a1, b_a1, w_a2, b_a2) from `CudaSlice` to `CudaSlice`. + +- [ ] **Step 3: Convert BranchingWeightSet fields similarly** + +Change all 8 fields from `CudaSlice` to `CudaSlice`. + +- [ ] **Step 4: Update extract_dueling_weights_branching and extract_branching_weights** + +These functions extract weights from Candle VarMaps (F32) into GPU buffers. They must now: +1. Download F32 from Candle → host Vec +2. Convert to Vec via `f32_to_bf16_host` +3. Upload to CudaSlice + +- [ ] **Step 5: Fix all compilation errors from the type change** + +This will cascade through MANY files: +- `gpu_dqn_trainer.rs`: `params_buf` stays `CudaSlice` for now (Adam master weights) — BUT we're going full BF16, so change it to `CudaSlice` too +- Actually: EVERY `CudaSlice` in `GpuDqnTrainer` becomes `CudaSlice`. ALL of them. This includes: `states_buf`, `next_states_buf`, `rewards_buf`, `dones_buf`, `is_weights_buf`, `save_h_s1`..`save_h_b2`, `grad_buf`, `params_buf`, `target_params_buf`, `m_buf`, `v_buf`, `grad_norm_buf`, `total_loss_buf`, `td_errors_buf`, `cql_grad_scratch`, and all spectral norm u/v buffers. +- The `alloc_f32` helper calls throughout the constructor become `alloc_bf16` (or `stream.alloc_zeros::(n)`) +- The `CachedPtrs` struct: pointers are u64, type-agnostic — no change needed +- Kernel launch `.arg()` calls: u64 pointers are type-agnostic — no change needed +- `memcpy_htod` calls that upload F32 host data: convert to BF16 first +- `memcpy_dtoh` calls that download to F32 host: download as u16, convert back + +- [ ] **Step 6: Update all `raw_device_ptr` functions** + +Change `raw_device_ptr(slice: &CudaSlice, ...)` to accept `CudaSlice`. Or better: make it generic: +```rust +fn raw_device_ptr(slice: &CudaSlice, stream: &CudaStream) -> u64 { + let (ptr, guard) = slice.device_ptr(stream); + let _no_drop = std::mem::ManuallyDrop::new(guard); + ptr +} +``` + +- [ ] **Step 7: Fix all remaining compilation errors** + +This is the big cascading fix. Work through every error methodically. Every `CudaSlice` → `CudaSlice`. + +**EXCEPTION:** The `actions_buf` and `t_buf` stay as `CudaSlice` — they hold integer action indices and Adam step counter, not floating point. + +- [ ] **Step 8: Verify compilation + tests** + +Run: `SQLX_OFFLINE=true cargo check -p ml` +This WILL have many errors at first. Fix them all before proceeding. + +- [ ] **Step 9: Commit** + +```bash +git commit -m "feat(bf16): phase 1 — all GPU buffers CudaSlice, weight sets BF16" +``` + +--- + +## Phase 2: cuBLAS BF16 + +### Task 2: Convert cuBLAS forward to cublasGemmEx BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` + +- [ ] **Step 1: Set TF32 math mode on cuBLAS handle** + +In `CublasForward::new()`, after workspace allocation: +```rust +unsafe { + cublas_sys::cublasSetMathMode( + handle.0, + cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH, + ); +} +``` + +- [ ] **Step 2: Replace sgemm_layer with gemmex_bf16** + +Replace the body of `sgemm_layer`: +```rust +fn gemmex_bf16( + &self, w_ptr: u64, a_ptr: u64, c_ptr: u64, + n: usize, b: usize, k: usize, _label: &str, +) -> Result<(), MLError> { + let alpha = 1.0_f32; + let beta = 0.0_f32; + unsafe { + let status = cublas_sys::cublasGemmEx( + self.handle.0, + cublas_sys::cublasOperation_t::CUBLAS_OP_T, + cublas_sys::cublasOperation_t::CUBLAS_OP_N, + n as i32, b as i32, k as i32, + &alpha as *const f32 as *const std::ffi::c_void, + w_ptr as *const std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_16BF, k as i32, + a_ptr as *const std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_16BF, k as i32, + &beta as *const f32 as *const std::ffi::c_void, + c_ptr as *mut std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_16BF, n as i32, + cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, + ); + if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + return Err(MLError::ModelError(format!("cublasGemmEx {_label}: {status:?}"))); + } + } + Ok(()) +} +``` + +Note: `Ctype=CUDA_R_16BF` with `computeType=CUBLAS_COMPUTE_32F` → BF16 inputs, F32 internal accumulate, BF16 output. Full tensor core throughput. + +- [ ] **Step 3: Update forward_online to use BF16 pointers** + +Change `w_ptrs: &[u64; 20]` to use BF16 weight buffer pointers. Change all activation buffer refs from `&CudaSlice` to `&CudaSlice`. + +- [ ] **Step 4: Convert bias kernels to BF16** + +In `bias_kernels.cu`, change both kernels: +```cuda +extern "C" __global__ void add_bias_relu_kernel( + __nv_bfloat16* __restrict__ output, + const __nv_bfloat16* __restrict__ bias, + int out_dim, int total_elements) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total_elements) return; + float val = __bfloat162float(output[i]) + __bfloat162float(bias[i % out_dim]); + output[i] = __float2bfloat16((val > 0.0f) ? val : 0.0f); +} +``` + +Same for `add_bias_kernel`. + +- [ ] **Step 5: Update forward_target and forward_online_next similarly** + +- [ ] **Step 6: Verify compilation + tests** + +- [ ] **Step 7: Commit** + +--- + +### Task 3: Convert cuBLAS backward to cublasGemmEx BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` + +- [ ] **Step 1: Replace cublasSgemm with cublasGemmEx in backward_fc_layer** + +Both dW and dX computations use `cublasGemmEx` with BF16 inputs: +- dW = d_output^T × input: both BF16, output to BF16 grad_buf +- dX = d_output × W^T: both BF16, output to BF16 upstream gradient + +- [ ] **Step 2: Update backward_full to pass BF16 pointers** + +All activation pointers, weight pointers, and gradient pointers are now BF16 (u64 addresses into u16 buffers). + +- [ ] **Step 3: Convert relu_mask_kernel to BF16** + +In `relu_mask_kernel.cu`: +```cuda +extern "C" __global__ void relu_mask_kernel( + __nv_bfloat16* __restrict__ dx, + const __nv_bfloat16* __restrict__ activation, + int n) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + if (__bfloat162float(activation[i]) <= 0.0f) + dx[i] = __float2bfloat16(0.0f); +} +``` + +- [ ] **Step 4: Verify + commit** + +--- + +## Phase 3: Optimizer + Loss Kernels + +### Task 4: Adam kernel BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` + +- [ ] **Step 1: Convert dqn_adam_update_kernel** + +All pointers (`params`, `grads`, `m`, `v`, `grad_norm_sq`) → `__nv_bfloat16*`. Internal arithmetic stays F32 (cast on load, cast on store): + +```cuda +extern "C" __global__ void dqn_adam_update_kernel( + __nv_bfloat16* __restrict__ params, + const __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ m, + __nv_bfloat16* __restrict__ v, + const __nv_bfloat16* __restrict__ grad_norm_sq, + float lr, float beta1, float beta2, float epsilon, + float weight_decay, float max_grad_norm, + const int* __restrict__ t_ptr, int total_params +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_params) return; + int t = *t_ptr; + float g = __bfloat162float(grads[idx]); + float norm = sqrtf(__bfloat162float(*grad_norm_sq) + 1e-12f); + float clip_scale = (norm > max_grad_norm) ? (max_grad_norm / norm) : 1.0f; + float clipped_g = g * clip_scale; + float beta1_t = 1.0f - powf(beta1, (float)t); + float beta2_t = 1.0f - powf(beta2, (float)t); + float m_i = beta1 * __bfloat162float(m[idx]) + (1.0f - beta1) * clipped_g; + float v_i = beta2 * __bfloat162float(v[idx]) + (1.0f - beta2) * clipped_g * clipped_g; + m[idx] = __float2bfloat16(m_i); + v[idx] = __float2bfloat16(v_i); + float m_hat = m_i / beta1_t; + float v_hat = v_i / beta2_t; + float p = __bfloat162float(params[idx]); + p -= lr * (m_hat / (sqrtf(v_hat) + epsilon) + weight_decay * p); + params[idx] = __float2bfloat16(p); +} +``` + +- [ ] **Step 2: Convert dqn_grad_norm_kernel to BF16 input** + +```cuda +extern "C" __global__ void dqn_grad_norm_kernel( + const __nv_bfloat16* __restrict__ grads, + __nv_bfloat16* __restrict__ out_grad_norm, // BF16 output + int total_params) +``` + +Internal sum stays F32 (warp reduction), final output converted to BF16. + +- [ ] **Step 3: Convert SAXPY, clipped SAXPY, clip grad kernels** + +All `float*` → `__nv_bfloat16*`. Same pattern: BF16 storage, F32 arithmetic. + +- [ ] **Step 4: Convert spectral norm kernel** + +Weight matrix, u/v vectors: `__nv_bfloat16*`. Internal matvec and reduction: F32. + +- [ ] **Step 5: Convert shrink-perturb, regime-scale kernels** + +Same pattern. + +- [ ] **Step 6: Convert EMA kernel** + +```cuda +extern "C" __global__ void ema_kernel( + __nv_bfloat16* __restrict__ target, + const __nv_bfloat16* __restrict__ online, + float tau, int n) +``` + +- [ ] **Step 7: Verify + commit** + +--- + +### Task 5: C51 loss + gradient kernels BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu` + +- [ ] **Step 1: Convert c51_loss_batched kernel** + +All logit inputs, saved tensors, IS weights, rewards, dones → `__nv_bfloat16*`. +Shared memory arrays: `__nv_bfloat16`. +Internal arithmetic (softmax, log, cross-entropy): cast to F32 for transcendentals. +Output (per_sample_loss, td_errors, total_loss): BF16. + +- [ ] **Step 2: Convert c51_grad kernel** + +d_value_logits, d_adv_logits outputs: `__nv_bfloat16*`. + +- [ ] **Step 3: Convert MSE loss + gradient kernels** + +Same pattern for `mse_loss_kernel.cu` and `mse_grad_kernel.cu`. + +- [ ] **Step 4: Convert expected_q_kernel and q_stats_kernel** + +BF16 logit inputs, BF16 outputs. + +- [ ] **Step 5: Verify + commit** + +--- + +### Task 6: CQL + IQN + ensemble gradient kernels BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/cql_grad_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/ensemble_kernels.cu` +- Modify: `crates/ml/src/cuda_pipeline/iql_value_kernel.cu` + +- [ ] **Step 1: Convert CQL logit gradient kernel** + +All logit/gradient buffers → BF16. + +- [ ] **Step 2: Convert IQN dual-head kernel** + +Tau sampling, cosine embedding, quantile loss — all BF16 storage. + +- [ ] **Step 3: Convert IQN CVaR kernel** + +- [ ] **Step 4: Convert ensemble diversity kernels** + +- [ ] **Step 5: Convert IQL value kernel** + +- [ ] **Step 6: Update Rust launchers (gpu_iqn_head.rs, gpu_iql_trainer.rs)** + +Change all `CudaSlice` buffer types in the launcher structs. + +- [ ] **Step 7: Verify + commit** + +--- + +## Phase 4: Environment + Auxiliary Kernels + +### Task 7: Experience collector + trade physics BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/trade_physics.cuh` +- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` + +- [ ] **Step 1: Convert trade_physics.cuh — all 11 device functions** + +All function parameters and locals → BF16 storage. Use `__bfloat162float` / `__float2bfloat16` for arithmetic. For sqrt, exp, log: promote to F32. + +- [ ] **Step 2: Convert experience_env_step kernel** + +All 30 parameters that are `float*` → `__nv_bfloat16*`. Portfolio state, rewards, features, targets. + +- [ ] **Step 3: Update ExperienceCollectorConfig and launch args** + +All `CudaSlice` → `CudaSlice` in the Rust launcher. + +- [ ] **Step 4: Verify + commit** + +--- + +### Task 8: Backtest evaluator BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/backtest_env_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` +- Modify: `crates/ml/src/cuda_pipeline/trade_stats_kernel.cu` + +- [ ] **Step 1: Convert backtest_env_step kernel (already uses trade_physics.cuh)** + +- [ ] **Step 2: Convert backtest_metrics_kernel (Sharpe, Calmar, CVaR, etc.)** + +- [ ] **Step 3: Convert backtest_gather_kernel** + +- [ ] **Step 4: Convert trade_stats_kernel** + +- [ ] **Step 5: Update Rust launcher** + +- [ ] **Step 6: Verify + commit** + +--- + +### Task 9: Action selector + attention + monitoring BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/gpu_action_selector.rs` +- Modify: `crates/ml/src/cuda_pipeline/attention_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/attention_backward_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/gpu_attention.rs` +- Modify: `crates/ml/src/cuda_pipeline/monitoring_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/gpu_monitoring.rs` +- Modify: `crates/ml/src/cuda_pipeline/per_update_kernel.cu` + +- [ ] **Step 1: Convert epsilon_greedy_kernel — BF16 Q-values** +- [ ] **Step 2: Convert attention forward + backward kernels** +- [ ] **Step 3: Convert monitoring kernel** +- [ ] **Step 4: Convert PER priority update kernel** +- [ ] **Step 5: Update all Rust launchers** +- [ ] **Step 6: Verify + commit** + +--- + +### Task 10: Remaining kernels BF16 + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/nstep_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/her_episode_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/her_relabel_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/training_guard_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/statistics_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/dt_kernels.cu` +- Modify: `crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/backtest_forward_supervised_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu` +- Modify: `crates/ml/src/cuda_pipeline/backward_kernels.cu` + +- [ ] **Step 1: Convert n-step, HER, curiosity kernels** +- [ ] **Step 2: Convert training guard, statistics, DT kernels** +- [ ] **Step 3: Convert PPO, supervised, signal adapter kernels** +- [ ] **Step 4: Convert backward_kernels.cu (legacy backward)** +- [ ] **Step 5: Update all Rust launchers** +- [ ] **Step 6: Verify + commit** + +--- + +### Task 11: Fused training context + weight sync BF16 + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Remove F32→BF16 conversion step from graph_adam** + +Since params_buf IS BF16 now, there's no conversion needed after Adam. The Adam kernel writes BF16 directly to params_buf. The `bf16_params_buf` separate buffer is now REDUNDANT — `params_buf` IS the BF16 buffer. + +- [ ] **Step 2: Eliminate bf16_params_buf / bf16_target_params_buf** + +Since all buffers are BF16, the separate BF16 mirrors are no longer needed. Forward reads directly from `params_buf` (which is now BF16). Delete the mirror buffers and all conversion code. + +- [ ] **Step 3: Update unflatten_online_weights for BF16** + +The unflatten copies from `params_buf` to individual weight set tensors. Both are now `CudaSlice`. The copy is `u16` → `u16`, same as before but type changes. + +- [ ] **Step 4: Update checkpoint save/load** + +Checkpoint serialization downloads params to host. Must convert BF16 → F32 for serialization (or serialize as BF16 with a version flag). + +- [ ] **Step 5: Verify + commit** + +--- + +### Task 12: BF16-native tests + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs` +- Modify: `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` + +- [ ] **Step 1: Update test helpers to construct BF16 buffers natively** + +`alloc_dueling` and `alloc_branching` now produce `CudaSlice`. Use `half::bf16::from_f32(0.1).to_bits()` for test values. + +- [ ] **Step 2: Update gradient budget tests for BF16 precision** + +Tolerance assertions may need widening (BF16 has ~0.78% relative error vs F32's ~0.00001%). + +- [ ] **Step 3: Add BF16 roundtrip correctness test** + +Upload F32 → BF16, download BF16 → F32, verify within BF16 precision bounds. + +- [ ] **Step 4: Verify ALL tests pass** + +```bash +SQLX_OFFLINE=true cargo test -p ml --lib +SQLX_OFFLINE=true cargo test -p ml-dqn --lib +``` + +- [ ] **Step 5: Final commit** + +```bash +git commit -m "feat(bf16): complete — zero F32 on GPU, all kernels BF16" +``` + +--- + +## Execution Dependencies + +``` +Task 1 (foundation) → Task 2 (forward) → Task 3 (backward) → Task 4 (Adam) + → Task 5 (C51/MSE) + → Task 6 (CQL/IQN/ensemble) + → Task 7 (experience) → Task 8 (backtest) → Task 9 (action/attn/monitoring) + → Task 10 (remaining) + → Task 11 (fused training) — after Tasks 2-6 + → Task 12 (tests) — after all +``` + +Task 1 is the foundation — changes types everywhere. Tasks 2-10 convert individual kernel groups. Task 11 wires up the training pipeline. Task 12 validates. + +## Verification After Each Phase + +**Phase 1 (Task 1):** `cargo check --workspace` — may have many errors, fix all. +**Phase 2 (Tasks 2-3):** `cargo check -p ml && cargo test -p ml --lib -- dqn` +**Phase 3 (Tasks 4-6):** Same + `cargo test -p ml --lib -- gradient_budget` +**Phase 4 (Tasks 7-12):** Full test suite including smoke tests if GPU available. diff --git a/docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md b/docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md new file mode 100644 index 000000000..bd8e4cf2c --- /dev/null +++ b/docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md @@ -0,0 +1,229 @@ +# Spec A (Revised): Full BF16 Rewrite — Zero F32 On GPU + +**Date**: 2026-03-27 +**Scope**: Convert EVERY CUDA kernel, EVERY GPU buffer, and EVERY cuBLAS call from F32 to BF16. Zero F32 on GPU. No mixed precision. No compromises. + +## Principle + +**`__nv_bfloat16` everywhere.** Every `float` in every CUDA kernel becomes `__nv_bfloat16`. Every `CudaSlice` in Rust becomes `CudaSlice` (BF16 stored as u16 in cudarc). Every `cublasSgemm` becomes `cublasGemmEx` with `CUDA_R_16BF` for ALL three matrix types (A, B, C). Every accumulator, every intermediate, every buffer — BF16. + +The only F32 that touches GPU is the `alpha`/`beta` scalar arguments to `cublasGemmEx` (required by the API to be host-side F32 pointers even when compute type is BF16). + +## Why Full BF16 + +1. **Performance**: H100 BF16 tensor cores = 989 TFLOPS vs 67 TFLOPS F32. 14.8× theoretical, ~3× practical. +2. **Memory**: Every buffer halved. 500K params × 4 bytes = 2MB F32. In BF16 = 1MB. Adam moments halved. Activations halved. More room for larger batches. +3. **Bandwidth**: HBM3 bandwidth is the bottleneck for element-wise kernels (Adam, loss, spectral norm). Half the data = 2× throughput. +4. **Consistency**: No precision boundaries between kernels. No F32→BF16 conversion kernels needed between stages. The output of one kernel IS the input of the next — both BF16. +5. **Simplicity**: One data type everywhere. No mixed-precision bookkeeping. + +## BF16 Numeric Properties + +- **Exponent**: 8 bits (same as F32) → range ±3.4×10³⁸, no underflow/overflow risk +- **Mantissa**: 7 bits → ~1/128 relative precision → 0.78% relative error +- **Adam convergence**: Weight updates of magnitude `lr × m_hat / sqrt(v_hat)` ≈ 3e-4 × 1.0 = 3e-4. BF16 can represent this relative to weights of magnitude 0.01-10.0. No stalling. +- **Loss precision**: C51 cross-entropy ≈ 0.1-50.0. BF16 resolution at 50.0 is 0.25. Adequate for training signal. +- **Gradient precision**: Per-element gradients ≈ 1e-4 to 1e+2. BF16 handles this entire range. + +## Architecture + +### Data Flow (All BF16) + +``` +Experience Collector + | + v +states [B, SD] (BF16) + | + v cublasGemmEx (BF16 × BF16 → BF16) +Forward Pass + |-- h_s1 [B, SH1] = BF16(ReLU(states @ W_s1^T + b_s1)) + |-- h_s2 [B, SH2] = BF16(ReLU(h_s1 @ W_s2^T + b_s2)) + |-- v_logits [B, NA] = BF16(h_v @ W_v2^T + b_v2) + v +C51/MSE Loss (BF16 logits → BF16 d_logits) + | + v cublasGemmEx (BF16 × BF16 → BF16) +Backward Pass + |-- dW = d_output^T × input → BF16 (accumulated into grad_buf) + |-- dX = d_output × W^T → BF16 (upstream gradient) + v +grad_buf [TOTAL_PARAMS] (BF16) + | + v +Adam (BF16 params, BF16 grads, BF16 m, BF16 v) + | + v +params_buf [TOTAL_PARAMS] (BF16) → next forward pass +``` + +**No conversion kernels between stages.** Output of one is input of next. All BF16. + +### What Changes + +#### CUDA Shared Header (`common_device_functions.cuh`) +- Add `#include ` at the top +- Add BF16 helper macros: `__BF16_TO_FLOAT(x)`, `__FLOAT_TO_BF16(x)` +- All `float` state/param typedefs → `__nv_bfloat16` + +#### CUDA Trade Physics (`trade_physics.cuh`) +- All 11 device functions: `float` args/locals → `__nv_bfloat16` +- Intermediate arithmetic: use `__nv_bfloat16` native ops (H100 has BF16 ALU) +- For complex math (sqrt, exp, log): cast to float, compute, cast back + +#### cuBLAS Forward (`batched_forward.rs`) +- `cublasSgemm` → `cublasGemmEx` with ALL types `CUDA_R_16BF` +- `cublasComputeType_t::CUBLAS_COMPUTE_32F` (internal F32 accumulate, BF16 output) +- Wait — for full BF16 output: use `CUBLAS_COMPUTE_16BF` to get BF16 C output + - Actually: `cublasGemmEx` with `Ctype=CUDA_R_16BF` and `computeType=CUBLAS_COMPUTE_32F` gives BF16 output with F32 internal accumulation. Best of both worlds. +- All activation buffers: `CudaSlice` (BF16) +- Weight pointers: from `bf16_params_buf` (already exists) +- Bias kernels: BF16 input/output + +#### cuBLAS Backward (`batched_backward.rs`) +- Same `cublasGemmEx` pattern for dW and dX +- `grad_buf`: BF16 (was F32) — Adam reads BF16 gradients +- All intermediate activation gradients: BF16 + +#### Adam Kernel (`dqn_utility_kernels.cu`) +- `params`, `grads`, `m`, `v`: all `__nv_bfloat16*` +- Bias correction: BF16 arithmetic (cast to float for powf if needed) +- Weight decay: BF16 +- Grad norm: BF16 input, BF16 sum → final sqrt + +#### C51 Loss (`c51_loss_kernel.cu`) +- All logit inputs: BF16 +- Shared memory: BF16 +- Softmax, log, exp: cast to float for transcendentals, result back to BF16 +- Cross-entropy accumulation: BF16 +- Loss output: BF16 + +#### MSE Loss (`mse_loss_kernel.cu`) +- Same pattern — BF16 logits, BF16 loss + +#### C51/MSE Gradient Kernels +- d_logits output: BF16 + +#### Experience Kernels (`experience_kernels.cu`) +- Portfolio state: BF16 +- Rewards: BF16 +- Trade physics: BF16 + +#### Backtest Kernel (`backtest_env_kernel.cu`) +- Same as experience — all BF16 + +#### Spectral Norm (`dqn_utility_kernels.cu`) +- Weight matrix: BF16 +- u/v vectors: BF16 +- sigma computation: BF16 (cast to float for sqrt) + +#### All Other Kernels +- EMA kernel: BF16 +- PER update: BF16 +- Epsilon greedy: BF16 Q-values +- IQN head: BF16 +- Attention: BF16 +- CQL gradient: BF16 +- Ensemble: BF16 +- Every. Single. One. + +### Rust-Side Changes + +Every `CudaSlice` in `GpuDqnTrainer`, `CublasForward`, `CublasBackward`, `GpuExperienceCollector`, `GpuBacktestEvaluator`, `GpuActionSelector`, `GpuIqnHead`, `GpuAttention`, etc. becomes `CudaSlice`. + +The `alloc_f32` helper becomes `alloc_bf16`. All `memcpy_htod` of `&[f32]` becomes conversion to `Vec` (BF16 bit pattern) first. + +The `DuelingWeightSet` and `BranchingWeightSet` fields change from `CudaSlice` to `CudaSlice`. + +### Conversion at Boundaries + +**Input boundary** (CPU → GPU): +- Feature vectors arrive as `Vec` from Rust +- Convert to `Vec` (BF16 via `f32::to_bits() >> 16` approximation or `half::bf16::from_f32`) +- Upload `CudaSlice` + +**Output boundary** (GPU → CPU): +- Loss scalar, grad_norm scalar: download as BF16 u16, convert to f32 on CPU for logging +- Q-value stats: same +- Checkpoint save: download BF16 params, convert to f32 for serialization (or serialize as BF16) + +### Implementation Phases + +This is too large for one commit. Split into 4 phases, each producing compilable + testable code: + +**Phase 1: BF16 type foundation** +- Add `half` crate dependency for `bf16` type +- Create `alloc_bf16`, `bf16_from_f32_vec`, `f32_from_bf16_vec` helpers +- Add `#include ` to common_device_functions.cuh +- All tests still pass (no behavioral change yet) + +**Phase 2: cuBLAS BF16 (forward + backward)** +- Convert `cublasSgemm` → `cublasGemmEx` BF16×BF16→BF16 +- Convert activation buffers to BF16 +- Convert weight buffers to BF16 +- Bias kernels to BF16 +- Forward pass fully BF16 +- Backward pass fully BF16 + +**Phase 3: Optimizer + loss kernels BF16** +- Adam kernel: all BF16 +- Gradient norm: BF16 +- Clipping kernels: BF16 +- C51 loss + grad: BF16 +- MSE loss + grad: BF16 +- Spectral norm: BF16 +- EMA: BF16 + +**Phase 4: Environment + auxiliary kernels BF16** +- Experience collector: BF16 +- Backtest evaluator: BF16 +- Trade physics: BF16 +- Action selector: BF16 +- IQN head: BF16 +- Attention: BF16 +- CQL: BF16 +- Ensemble: BF16 +- PER: BF16 +- All remaining kernels + +### Files Changed (Complete List) + +**CUDA Kernels (37 files + 2 headers):** +Every `.cu` and `.cuh` file in `crates/ml/src/cuda_pipeline/` + +**Rust Launchers (18 files):** +Every `gpu_*.rs` and `batched_*.rs` in `crates/ml/src/cuda_pipeline/` + +**Rust Weight Types:** +`crates/ml/src/cuda_pipeline/gpu_weights.rs` — all weight set structs + +**Training Orchestration:** +`crates/ml/src/trainers/dqn/fused_training.rs` + +**Tests:** +All smoke tests, gradient budget tests + +### Success Criteria + +1. **Zero `float` in CUDA hot path**: `grep -rn "float\b" *.cu` in hot-path kernels returns 0 hits (except for cublasGemmEx alpha/beta host scalars and transcendental casts) +2. **All 488+ tests pass** +3. **Epoch time ≤3s on H100** (was ~30s F32, target: BF16 tensor cores + halved bandwidth) +4. **Training converges**: Sharpe ratio within 10% of F32 baseline over 50 epochs +5. **Peak VRAM reduced ~40%**: all buffers halved except scalar counters + +### Risks + +| Risk | Mitigation | +|---|---| +| BF16 Adam stalling (updates too small) | Monitor weight update magnitude; BF16 range covers lr=3e-4 updates on weights 0.01-10 | +| C51 log-softmax overflow | Max logit ~50 in BF16 → exp(50) = 5e21, fits BF16 range (3.4e38) | +| Spectral norm sigma precision | sigma_max=3.0, BF16 resolution at 3.0 is 0.015 — adequate | +| cuBLAS BF16 GEMM not capturable in CUDA Graph | Test on CUDA 12.x; fallback: launch outside graph (adds ~2ms) | +| Backtest metrics precision | Sharpe/Calmar use division — BF16 division is fine for 2-decimal metrics | + +### Non-Goals + +- F16 support (BF16 has superior dynamic range) +- Multi-GPU (single H100 target) +- CPU fallback path (GPU-only) +- Backward compatibility with F32 checkpoints (convert on load)