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:

// 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:

// 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():

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	CudaSlice (unchanged)
GPU portfolio simulator	CudaSlice	CudaSlice (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

// 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():

// 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.