docs: full BF16 rewrite spec + plan — zero F32 on GPU
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
577
docs/superpowers/plans/2026-03-27-full-bf16-rewrite.md
Normal file
577
docs/superpowers/plans/2026-03-27-full-bf16-rewrite.md
Normal file
@@ -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<u16>` 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<f32> = stream.alloc_zeros::<f32>(n)?;
|
||||
stream.memcpy_htod(&host_f32, &mut buf)?;
|
||||
|
||||
// After:
|
||||
let buf: CudaSlice<u16> = stream.alloc_zeros::<u16>(n)?;
|
||||
let host_bf16: Vec<u16> = 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<f32>` to `CudaSlice<u16>`. 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<u16> {
|
||||
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<f32> {
|
||||
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<CudaStream>, n: usize, label: &str) -> Result<CudaSlice<u16>, MLError> {
|
||||
stream.alloc_zeros::<u16>(n).map_err(|e| {
|
||||
MLError::ModelError(format!("alloc bf16 {label}: {e}"))
|
||||
})
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Convert DuelingWeightSet fields from CudaSlice<f32> to CudaSlice<u16>**
|
||||
|
||||
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<f32>` to `CudaSlice<u16>`.
|
||||
|
||||
- [ ] **Step 3: Convert BranchingWeightSet fields similarly**
|
||||
|
||||
Change all 8 fields from `CudaSlice<f32>` to `CudaSlice<u16>`.
|
||||
|
||||
- [ ] **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<f32>
|
||||
2. Convert to Vec<u16> via `f32_to_bf16_host`
|
||||
3. Upload to CudaSlice<u16>
|
||||
|
||||
- [ ] **Step 5: Fix all compilation errors from the type change**
|
||||
|
||||
This will cascade through MANY files:
|
||||
- `gpu_dqn_trainer.rs`: `params_buf` stays `CudaSlice<f32>` for now (Adam master weights) — BUT we're going full BF16, so change it to `CudaSlice<u16>` too
|
||||
- Actually: EVERY `CudaSlice<f32>` in `GpuDqnTrainer` becomes `CudaSlice<u16>`. 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::<u16>(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<f32>, ...)` to accept `CudaSlice<u16>`. Or better: make it generic:
|
||||
```rust
|
||||
fn raw_device_ptr<T>(slice: &CudaSlice<T>, 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<f32>` → `CudaSlice<u16>`.
|
||||
|
||||
**EXCEPTION:** The `actions_buf` and `t_buf` stay as `CudaSlice<i32>` — 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<u16>, 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<f32>` to `&CudaSlice<u16>`.
|
||||
|
||||
- [ ] **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<f32>` 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<f32>` → `CudaSlice<u16>` 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<u16>`. 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<u16>`. 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.
|
||||
229
docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md
Normal file
229
docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md
Normal file
@@ -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<f32>` in Rust becomes `CudaSlice<u16>` (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 <cuda_bf16.h>` 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<u16>` (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<f32>` in `GpuDqnTrainer`, `CublasForward`, `CublasBackward`, `GpuExperienceCollector`, `GpuBacktestEvaluator`, `GpuActionSelector`, `GpuIqnHead`, `GpuAttention`, etc. becomes `CudaSlice<u16>`.
|
||||
|
||||
The `alloc_f32` helper becomes `alloc_bf16`. All `memcpy_htod` of `&[f32]` becomes conversion to `Vec<u16>` (BF16 bit pattern) first.
|
||||
|
||||
The `DuelingWeightSet` and `BranchingWeightSet` fields change from `CudaSlice<f32>` to `CudaSlice<u16>`.
|
||||
|
||||
### Conversion at Boundaries
|
||||
|
||||
**Input boundary** (CPU → GPU):
|
||||
- Feature vectors arrive as `Vec<f32>` from Rust
|
||||
- Convert to `Vec<u16>` (BF16 via `f32::to_bits() >> 16` approximation or `half::bf16::from_f32`)
|
||||
- Upload `CudaSlice<u16>`
|
||||
|
||||
**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 <cuda_bf16.h>` 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)
|
||||
Reference in New Issue
Block a user