foxhunt/docs/plans/2026-03-03-bf16-training-design.md

# BF16 Training — Design Document

**Date:** 2026-03-03
**Scope:** All 10 models (DQN, PPO, TFT, Mamba2, TGGN, TLOB, Liquid, KAN, xLSTM, Diffusion)
**Target hardware:** L40S, H100 (Ampere+ GPUs with native BF16)
**Approach:** Pure BF16 — weights, gradients, optimizer state, activations all BF16

## Motivation

Current training runs entirely in FP32. On Ampere+ GPUs (L40S, H100), BF16 unlocks:
- Tensor core utilization for all matmuls (up to 10x theoretical FLOPS vs FP32)
- 50% VRAM reduction on weights, activations, and replay buffer states
- More parallel hyperopt trials (more VRAM headroom per GPU)
- Smaller checkpoint files (50%)

## Core Principle: Cast Once at Boundaries

Zero casts in the training hot path. BF16 flows end-to-end inside the training loop.

Cast points (exhaustive list):
1. **Data ingestion** — f32 from Rust/OHLCV → BF16 at replay buffer `store()` or dataset creation
2. **Loss scalar** — BF16 → F32 before `backward()` (single number)
3. **CUDA kernel boundary** — BF16 weights → f32 at extraction for experience collector (once per epoch)

## Design

### 1. Dynamic DType Detection

A single function determines dtype from the device at runtime:

```rust
// mixed_precision.rs
pub fn training_dtype(device: &Device) -> DType {
    match device {
        Device::Cuda(_) if is_ampere_or_newer(device) => DType::BF16,
        _ => DType::F32,
    }
}
```

Reuses existing `detect_from_gpu_name()` logic for GPU identification. No config fields, no env vars. CPU always returns F32.

### 2. Model Construction (~30 VarBuilder sites)

Every model changes from hardcoded F32 to dynamic:

```rust
// Before:
VarBuilder::from_varmap(&vars, DType::F32, &device)
// After:
VarBuilder::from_varmap(&vars, training_dtype(&device), &device)
```

Affected models and key sites:
- **DQN**: `dqn.rs` Sequential, DuelingQNetwork; `dueling.rs`; `distributional_dueling.rs`; `quantile_regression.rs`
- **PPO**: `ppo.rs` PolicyNetwork, ValueNetwork
- **TFT**: `mod.rs`, `quantized_grn.rs`, `quantile_outputs.rs`
- **Mamba2**: `mod.rs`
- **Liquid/CfC**: `adapter.rs`, `candle_cfc.rs`
- **xLSTM**: `trainable.rs`, `slstm.rs`, `mlstm.rs`
- **Diffusion**: `trainable.rs`
- **TLOB**: `tlob.rs`
- **TGGN**: model construction site
- **KAN**: model construction site

### 3. Training Data Path

**RL models (DQN, PPO):**
- Replay buffer `store()`: cast f32 states/next_states to BF16 at insertion
- Replay buffer tensors pre-allocated as BF16 (states, next_states columns)
- `sample()` returns BF16 tensors directly — zero cast in training loop
- Actions stay U32, rewards/priorities/dones stay F32

**Supervised models (TFT, Mamba2, etc.):**
- Feature extraction produces f32 from OHLCV data
- Single cast to BF16 when building epoch dataset tensor
- All training iterations read BF16 directly

### 4. Loss Computation

Loss scalar cast to F32 before `backward()`:

```rust
let loss = model.compute_loss(batch_bf16)?;
let loss_f32 = loss.to_dtype(DType::F32)?;
loss_f32.backward()?;
```

Distributional/categorical DQN loss already has F32 enforcement — keep as-is. Gradients flow back in BF16 automatically via Candle's autograd.

### 5. CUDA Pipeline

| Component | Current | After |
|-----------|---------|-------|
| GPU replay buffer: states/next_states | F32 | **BF16** |
| GPU replay buffer: actions | U32 | U32 |
| GPU replay buffer: rewards/priorities/dones | F32 | F32 |
| GPU weights extraction | f32 | **BF16→f32 cast at extraction** |
| GPU experience collector | CudaSlice<f32> | CudaSlice<f32> (unchanged) |
| GPU portfolio simulator | CudaSlice<f32> | CudaSlice<f32> (unchanged) |

Experience collector and portfolio simulator CUDA kernels stay f32 — rewriting PTX is out of scope.

### 6. Checkpoint Compatibility

- Save: weights saved as BF16 in safetensors (50% smaller files)
- Load: `training_dtype(&device)` at load time — BF16 checkpoint on CPU auto-casts to F32

```rust
// Before:
VarBuilder::from_mmaped_safetensors(&[path], DType::F32, &device)
// After:
VarBuilder::from_mmaped_safetensors(&[path], training_dtype(&device), &device)
```

### 7. Mamba2 Fix

Remove BF16/F16 rejection in `mamba/mod.rs::scalar_tensor()`:

```rust
// Before: return Err(...) for BF16/F16
// After:
Tensor::new(value, device)?.to_dtype(dtype)
```

### 8. Testing

- Existing 2506 ml tests run on CPU (F32) — no behavior change
- One integration test: DQN 1-epoch training on synthetic data with BF16, verify loss decreases
- Real validation: hyperopt run on L40S comparing F32 vs BF16 Sharpe distributions

## Expected Impact

| Metric | F32 (current) | BF16 (expected) |
|--------|--------------|-----------------|
| VRAM per DQN trial | ~8-12 GB | ~5-7 GB |
| Parallel hyperopt trials (L40S 48GB) | 3 | 5-6 |
| Parallel hyperopt trials (H100 80GB) | 4-5 | 8-10 |
| Per-trial training time | baseline | ~1.5-2x faster (tensor cores) |
| Checkpoint size | baseline | ~50% smaller |

## Risks

- **Training divergence**: BF16 optimizer state has 8-bit mantissa vs F32's 24-bit. Some models may fail to converge. Mitigation: if a model diverges, fall back to F32 for that model.
- **Numerical edge cases**: Small gradients may underflow to zero in BF16. BF16's dynamic range (same exponent bits as F32) makes this unlikely, unlike FP16.
- **CUDA kernel mismatch**: Experience collector outputs f32 into a BF16 replay buffer. The cast at insertion handles this.