fix: restore accidentally deleted bf16 plan files
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
146
docs/plans/2026-03-03-bf16-training-design.md
Normal file
146
docs/plans/2026-03-03-bf16-training-design.md
Normal file
@@ -0,0 +1,146 @@
|
||||
# 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.
|
||||
644
docs/plans/2026-03-03-bf16-training.md
Normal file
644
docs/plans/2026-03-03-bf16-training.md
Normal file
@@ -0,0 +1,644 @@
|
||||
# BF16 Training Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Switch all 10 ML models from FP32 to BF16 training on Ampere+ GPUs with dynamic detection and zero casts in the training hot path.
|
||||
|
||||
**Architecture:** Add a single `training_dtype(device) -> DType` function that returns BF16 on Ampere+ CUDA, F32 elsewhere. Thread it through all ~150 VarBuilder sites, training tensor creation, GPU replay buffer, and checkpoint loading. Loss stays F32 (single scalar cast). CUDA experience collector kernels untouched (stay f32).
|
||||
|
||||
**Tech Stack:** Candle (v0.9.1 git pin), half crate (2.6.0), cudarc, safetensors
|
||||
|
||||
---
|
||||
|
||||
## Phase 1: Core Infrastructure
|
||||
|
||||
### Task 1: Add `training_dtype()` function
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/dqn/mixed_precision.rs`
|
||||
|
||||
**Step 1: Add the public function after `detect_from_gpu_name()` (~line 180)**
|
||||
|
||||
```rust
|
||||
/// Returns the optimal training DType for the given device.
|
||||
/// Ampere+ CUDA GPUs → BF16 (tensor core acceleration).
|
||||
/// Everything else (CPU, older GPUs) → F32.
|
||||
pub fn training_dtype(device: &candle_core::Device) -> candle_core::DType {
|
||||
match device {
|
||||
candle_core::Device::Cuda(_) => {
|
||||
// Check if GPU supports BF16 natively
|
||||
if let Some(config) = detect_from_gpu_name_auto() {
|
||||
match config.dtype {
|
||||
DTypeSelection::BF16 => candle_core::DType::BF16,
|
||||
DTypeSelection::F16 => candle_core::DType::F32, // F16 needs loss scaling, stay F32
|
||||
}
|
||||
} else {
|
||||
candle_core::DType::F32
|
||||
}
|
||||
}
|
||||
_ => candle_core::DType::F32,
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Note: `detect_from_gpu_name_auto()` already exists at line 182 — it reads the GPU name from the CUDA device and calls `detect_from_gpu_name()`. Reuse it.
|
||||
|
||||
**Step 2: Re-export from the module's public API**
|
||||
|
||||
Ensure `training_dtype` is accessible as `crate::dqn::mixed_precision::training_dtype`. Check the module's `pub use` or `mod` visibility.
|
||||
|
||||
**Step 3: Build check**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20
|
||||
```
|
||||
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/mixed_precision.rs
|
||||
git commit -m "feat(ml): add training_dtype() for dynamic BF16 detection"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Fix Mamba2 scalar_tensor BF16 rejection
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/mamba/mod.rs:605`
|
||||
|
||||
**Step 1: Change the match arm at line 605**
|
||||
|
||||
Replace the BF16/F16 rejection:
|
||||
|
||||
```rust
|
||||
// Before (line 605):
|
||||
DType::F8E4M3 | DType::U8 | DType::U32 | DType::I64 | DType::BF16 | DType::F16 => {
|
||||
Err(MLError::ModelError(format!(
|
||||
"Unsupported dtype: {:?}",
|
||||
dtype
|
||||
)))
|
||||
},
|
||||
|
||||
// After:
|
||||
DType::BF16 | DType::F16 => {
|
||||
// Create in F32 then cast — half types can't be created directly from f64
|
||||
Tensor::new(&[value as f32], device)?
|
||||
.to_dtype(dtype)?
|
||||
.reshape(())?
|
||||
.ok_or_else(|| MLError::ModelError("scalar reshape failed".into()))
|
||||
},
|
||||
DType::F8E4M3 | DType::U8 | DType::U32 | DType::I64 => {
|
||||
Err(MLError::ModelError(format!(
|
||||
"Unsupported dtype: {:?}",
|
||||
dtype
|
||||
)))
|
||||
},
|
||||
```
|
||||
|
||||
Note: Check the exact return type — `scalar_tensor` may return `Result<Tensor, MLError>`. Adjust the reshape/return accordingly. The key is: create as f32, cast to target dtype, return scalar.
|
||||
|
||||
**Step 2: Build check**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20
|
||||
```
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/mamba/mod.rs
|
||||
git commit -m "fix(ml): allow BF16/F16 in Mamba2 scalar_tensor helper"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 2: DQN Module (largest surface area)
|
||||
|
||||
### Task 3: VarBuilder sites — DQN core networks
|
||||
|
||||
Change `DType::F32` → `training_dtype(&device)` (or `training_dtype(device)` if device is a reference) in all VarBuilder::from_varmap calls across the DQN module.
|
||||
|
||||
**Files (all need the same mechanical change):**
|
||||
- `crates/ml/src/dqn/dqn.rs:669`
|
||||
- `crates/ml/src/dqn/network.rs:256,261,301,363`
|
||||
- `crates/ml/src/dqn/agent.rs:366,371,597`
|
||||
- `crates/ml/src/dqn/dueling.rs:139`
|
||||
- `crates/ml/src/dqn/distributional_dueling.rs:155`
|
||||
- `crates/ml/src/dqn/quantile_regression.rs:90`
|
||||
- `crates/ml/src/dqn/curiosity.rs:38`
|
||||
- `crates/ml/src/dqn/factored_q_network.rs:72`
|
||||
- `crates/ml/src/dqn/rainbow_agent.rs:41,45`
|
||||
- `crates/ml/src/dqn/rainbow_agent_impl.rs:69,72`
|
||||
- `crates/ml/src/dqn/rainbow_network.rs:429,450`
|
||||
|
||||
**Pattern for each site:**
|
||||
|
||||
```rust
|
||||
// Before:
|
||||
VarBuilder::from_varmap(&vars, DType::F32, &device)
|
||||
// After:
|
||||
VarBuilder::from_varmap(&vars, training_dtype(&device), &device)
|
||||
```
|
||||
|
||||
Add `use crate::dqn::mixed_precision::training_dtype;` at the top of each file that doesn't already import it.
|
||||
|
||||
**Step 1:** Apply the change to all files listed above. Use `replace_all` where `DType::F32` appears only in VarBuilder contexts. Where `DType::F32` also appears in non-VarBuilder contexts (tensor creation, loss), change only the VarBuilder lines.
|
||||
|
||||
**Step 2: Build check**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -30
|
||||
```
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/
|
||||
git commit -m "feat(ml): BF16 VarBuilder for DQN core networks"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: VarBuilder sites — DQN layer modules
|
||||
|
||||
Same pattern for the layer-level modules that have many VarBuilder sites:
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/dqn/noisy_layers.rs:314,324,346,391,431,447,487`
|
||||
- `crates/ml/src/dqn/residual.rs:178,196,222,249,276,301,330,353`
|
||||
- `crates/ml/src/dqn/attention.rs:465,479,509,549,583,614`
|
||||
- `crates/ml/src/dqn/spectral_norm.rs:240,254,275,307,330,355,376`
|
||||
- `crates/ml/src/dqn/rmsnorm.rs:239,254,269,309,353,358,406,410,458,484`
|
||||
|
||||
Same mechanical change. These files likely have `DType::F32` ONLY in VarBuilder contexts, so `replace_all` may be safe. Verify by reading each file first.
|
||||
|
||||
**Step 1:** Apply changes.
|
||||
**Step 2:** Build check.
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/
|
||||
git commit -m "feat(ml): BF16 VarBuilder for DQN layers (noisy, residual, attention, spectral, rmsnorm)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: GPU replay buffer — BF16 states
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs:74-75`
|
||||
|
||||
**Step 1: Change states/next_states allocation to use dynamic dtype**
|
||||
|
||||
```rust
|
||||
// Line 74-75, change:
|
||||
let states = Tensor::zeros(&[cap, sdim], DType::F32, device)?;
|
||||
let next_states = Tensor::zeros(&[cap, sdim], DType::F32, device)?;
|
||||
// To:
|
||||
let dtype = training_dtype(device);
|
||||
let states = Tensor::zeros(&[cap, sdim], dtype, device)?;
|
||||
let next_states = Tensor::zeros(&[cap, sdim], dtype, device)?;
|
||||
```
|
||||
|
||||
Keep rewards, dones, priorities as `DType::F32`. Keep actions as `DType::U32`.
|
||||
|
||||
**Step 2: Verify `insert_batch()` callers cast correctly**
|
||||
|
||||
The `insert_batch()` at line 167 uses `slice_scatter` which requires matching dtypes. The caller (DQN trainer) builds state tensors from f32 experience data. Add a `.to_dtype(self.states.dtype())?` cast on the incoming states/next_states args inside `insert_batch()`:
|
||||
|
||||
```rust
|
||||
// Inside insert_batch(), before slice_scatter:
|
||||
let states = states.to_dtype(self.states.dtype())?;
|
||||
let next_states = next_states.to_dtype(self.next_states.dtype())?;
|
||||
```
|
||||
|
||||
This is the ONE cast at the data ingestion boundary. After this, all `sample()` returns match the buffer dtype (BF16 on Ampere+).
|
||||
|
||||
**Step 3:** Build check.
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs
|
||||
git commit -m "feat(ml): BF16 states in GPU replay buffer (50% VRAM savings)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: GPU weights extraction — handle BF16 weights
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/gpu_weights.rs:179,204`
|
||||
|
||||
**Step 1: Cast to F32 before extraction**
|
||||
|
||||
The CUDA experience collector kernel expects f32 weights. When model weights are BF16, cast before extracting:
|
||||
|
||||
```rust
|
||||
// In extract_one() at line 179, change:
|
||||
.to_vec1::<f32>()
|
||||
// To:
|
||||
.to_dtype(candle_core::DType::F32)?
|
||||
.to_vec1::<f32>()
|
||||
```
|
||||
|
||||
Same for `sync_one()` at line 204. This is the boundary cast from BF16 model weights → f32 CUDA kernel. Happens once per epoch during experience collection, not in the training hot path.
|
||||
|
||||
**Step 2:** Build check.
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/gpu_weights.rs
|
||||
git commit -m "feat(ml): handle BF16 weights in GPU weight extraction"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: GPU data pre-upload — BF16 feature tensors
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/cuda_pipeline/mod.rs:132,135,346,353`
|
||||
|
||||
**Step 1: Cast feature/target uploads to training dtype**
|
||||
|
||||
In `DqnGpuData::upload()` (line 132):
|
||||
|
||||
```rust
|
||||
// After creating the tensor from f32 data, cast:
|
||||
let features = Tensor::from_vec(flat_features, (num_bars, feature_dim), device)?
|
||||
.to_dtype(training_dtype(device))?;
|
||||
let targets = Tensor::from_vec(flat_targets, (num_bars, target_dim), device)?
|
||||
.to_dtype(training_dtype(device))?;
|
||||
```
|
||||
|
||||
Same pattern for `GpuBufferPool::upload_dqn` (lines 346, 353) — cast after `from_slice`.
|
||||
|
||||
For PPO `PpoGpuData::upload()` (line 418) — same cast.
|
||||
|
||||
This is the data ingestion boundary cast. All downstream `build_batch_states()` and `bar_features()` calls return BF16 directly.
|
||||
|
||||
**Step 2:** Build check.
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/cuda_pipeline/mod.rs
|
||||
git commit -m "feat(ml): BF16 GPU data pre-upload for DQN and PPO"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: DQN training tensors — CPU replay buffer path
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/dqn/dqn.rs` — `compute_loss_internal()`
|
||||
|
||||
The CPU replay buffer path creates training batch tensors from `Vec<f32>`. These need to match the model's weight dtype.
|
||||
|
||||
**Step 1: Cast batch tensors at creation**
|
||||
|
||||
At lines 1540-1569, after each `Tensor::from_vec`:
|
||||
|
||||
```rust
|
||||
// States/next_states — cast to model dtype for matmul compatibility
|
||||
let states_tensor = Tensor::from_vec(states, (batch_size, self.config.state_dim), device)?
|
||||
.to_dtype(training_dtype(device))?;
|
||||
let next_states_tensor = Tensor::from_vec(next_states, (batch_size, self.config.state_dim), device)?
|
||||
.to_dtype(training_dtype(device))?;
|
||||
```
|
||||
|
||||
Actions stay U32. Rewards/dones/importance-weights stay F32 — they're used in loss math, not matmuls.
|
||||
|
||||
For the GPU replay buffer path, states already come back as BF16 from `sample()` (Task 5), so no change needed there.
|
||||
|
||||
**Step 2: Verify loss stays F32**
|
||||
|
||||
The distributional loss at line 1650/1658 already has `to_dtype(DType::F32)` enforcement. Keep as-is.
|
||||
|
||||
**Step 3:** Build check.
|
||||
**Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/dqn/dqn.rs
|
||||
git commit -m "feat(ml): BF16 training tensors in DQN compute_loss"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 9: DQN trainer auxiliary tensors
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/dqn/trainer.rs:1015,1130,2134,2141,3123,3647`
|
||||
|
||||
Same pattern — cast state tensors used in `select_actions_batch`, curiosity, and Q-value logging to `training_dtype(&self.device)`:
|
||||
|
||||
```rust
|
||||
let tensor = Tensor::from_vec(states, shape, &self.device)?
|
||||
.to_dtype(training_dtype(&self.device))?;
|
||||
```
|
||||
|
||||
These are not in the training hot path (they're action selection and logging), so the single cast is fine.
|
||||
|
||||
**Step 1:** Apply casts at listed lines.
|
||||
**Step 2:** Build check.
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/trainers/dqn/trainer.rs
|
||||
git commit -m "feat(ml): BF16 auxiliary tensors in DQN trainer"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 3: PPO Module
|
||||
|
||||
### Task 10: VarBuilder sites — PPO networks
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/ppo/ppo.rs:301,549`
|
||||
- `crates/ml/src/ppo/lstm_networks.rs:52,267`
|
||||
- `crates/ml/src/ppo/continuous_policy.rs:83`
|
||||
- `crates/ml/src/ppo/flow_policy/mod.rs:124`
|
||||
- `crates/ml/src/ppo/flow_policy/coupling_layer.rs:260`
|
||||
|
||||
Same pattern: `DType::F32` → `training_dtype(&device)`.
|
||||
|
||||
Also change checkpoint loading at lines 1795 and 1852:
|
||||
```rust
|
||||
// Before:
|
||||
VarBuilder::from_mmaped_safetensors(&[path], DType::F32, &device)
|
||||
// After:
|
||||
VarBuilder::from_mmaped_safetensors(&[path], training_dtype(&device), &device)
|
||||
```
|
||||
|
||||
**Step 1:** Apply all changes.
|
||||
**Step 2:** Build check.
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/ppo/
|
||||
git commit -m "feat(ml): BF16 VarBuilder and checkpoints for PPO networks"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 11: PPO training tensors
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/src/trainers/ppo.rs:834,895,960,1014`
|
||||
|
||||
Cast state tensors to training dtype for forward pass compatibility:
|
||||
|
||||
```rust
|
||||
let states = Tensor::from_vec(state_floats, shape, &self.device)?
|
||||
.to_dtype(training_dtype(&self.device))?;
|
||||
```
|
||||
|
||||
Lines 1156, 1346, 1348, 1372 (rewards, returns, values) — keep F32, these are loss/metric tensors not fed to the network.
|
||||
|
||||
**Step 1:** Apply casts to state tensors only.
|
||||
**Step 2:** Build check.
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/src/trainers/ppo.rs
|
||||
git commit -m "feat(ml): BF16 training tensors in PPO trainer"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 4: Supervised Models (8 models)
|
||||
|
||||
### Task 12: TFT VarBuilder sites
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/tft/mod.rs:339`
|
||||
- `crates/ml/src/tft/quantized_grn.rs:295,315`
|
||||
- `crates/ml/src/tft/quantized_attention.rs:417`
|
||||
- `crates/ml/src/tft/quantized_lstm.rs:417,442`
|
||||
- `crates/ml/src/tft/quantized_vsn.rs:61,249`
|
||||
- `crates/ml/src/tft/varmap_quantization.rs:676,722`
|
||||
|
||||
Same `DType::F32` → `training_dtype(&device)` pattern.
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for TFT"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 13: Mamba2 VarBuilder site
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/mamba/mod.rs:631`
|
||||
- `crates/ml/src/mamba/ssd_layer.rs:556`
|
||||
|
||||
Already fixed scalar_tensor in Task 2. Now change VarBuilder dtype.
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for Mamba2"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 14: Liquid/CfC VarBuilder sites
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/liquid/candle_cfc.rs:450,460,475,489,499,524,548,558,578,601,639` (11 sites)
|
||||
- `crates/ml/src/liquid/adapter.rs:58`
|
||||
- `crates/ml/src/liquid/training.rs:505`
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for Liquid/CfC"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 15: KAN VarBuilder sites
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/kan/layer.rs:138,150`
|
||||
- `crates/ml/src/kan/network.rs:94,107`
|
||||
- `crates/ml/src/kan/trainable.rs:47`
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for KAN"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 16: xLSTM VarBuilder sites
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/xlstm/slstm.rs:136,148,164,172`
|
||||
- `crates/ml/src/xlstm/mlstm.rs:229,243,259,267,282`
|
||||
- `crates/ml/src/xlstm/block.rs:127,139,152,161`
|
||||
- `crates/ml/src/xlstm/network.rs:172,185,198,214,236,248`
|
||||
- `crates/ml/src/xlstm/trainable.rs:45`
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for xLSTM"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 17: Diffusion VarBuilder sites
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/diffusion/sampler.rs:163,230`
|
||||
- `crates/ml/src/diffusion/denoiser.rs:233,244,258,270,285`
|
||||
- `crates/ml/src/diffusion/trainable.rs:45`
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for Diffusion"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 18: TGGN + TLOB VarBuilder sites
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/tgnn/trainable_adapter.rs:87`
|
||||
- `crates/ml/src/tlob/trainable_adapter.rs:116`
|
||||
- `crates/ml/src/trainers/tlob.rs:210`
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for TGGN and TLOB"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 5: Remaining Sites
|
||||
|
||||
### Task 19: Ensemble adapters + misc
|
||||
|
||||
**Files:**
|
||||
- `crates/ml/src/ensemble/adapters/liquid.rs:46,62`
|
||||
- `crates/ml/src/ensemble/adapters/diffusion.rs:49,72`
|
||||
- `crates/ml/src/ensemble/adapters/kan.rs:44,60`
|
||||
- `crates/ml/src/ensemble/adapters/tlob.rs:100,123`
|
||||
- `crates/ml/src/ensemble/adapters/tggn.rs:77,96`
|
||||
- `crates/ml/src/ensemble/adapters/xlstm.rs:57,79`
|
||||
- `crates/ml/src/portfolio_transformer.rs:190`
|
||||
- `crates/ml/src/features/multi_timeframe.rs:269,560,576`
|
||||
- `crates/ml/src/trainers/online_learning.rs:581`
|
||||
- `crates/ml/src/explainability/integrated_gradients.rs:200,247,301`
|
||||
|
||||
**Step 1:** Apply, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "feat(ml): BF16 VarBuilder for ensemble adapters and misc modules"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 6: Validation
|
||||
|
||||
### Task 20: Full workspace build and test
|
||||
|
||||
**Step 1: Workspace build**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5
|
||||
```
|
||||
|
||||
Expected: 0 errors. Fix any dtype mismatches — common issues:
|
||||
- `expected F32 but got BF16` — a tensor created without the dtype cast feeding into a module that expects matched dtypes
|
||||
- `cannot add BF16 and F32` — missing cast at a boundary
|
||||
|
||||
**Step 2: Clippy**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo clippy --workspace -- -D warnings 2>&1 | tail -10
|
||||
```
|
||||
|
||||
**Step 3: ML crate tests**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | tail -20
|
||||
```
|
||||
|
||||
Expected: 2506+ tests pass. Tests run on CPU → F32 path, no behavior change.
|
||||
|
||||
**Step 4: Commit if any fixes were needed**
|
||||
|
||||
```bash
|
||||
git commit -m "fix(ml): resolve BF16 dtype mismatches"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 21: BF16 integration test (optional — requires GPU)
|
||||
|
||||
Create a minimal integration test verifying BF16 training works end-to-end on CUDA:
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml/tests/bf16_training_integration.rs`
|
||||
|
||||
```rust
|
||||
//! Integration test: verify BF16 training on Ampere+ GPU
|
||||
//! Run with: SQLX_OFFLINE=true cargo test -p ml --test bf16_training_integration
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
mod bf16_tests {
|
||||
use ml::dqn::mixed_precision::training_dtype;
|
||||
use candle_core::{Device, DType};
|
||||
|
||||
#[test]
|
||||
fn test_training_dtype_returns_bf16_on_cuda() {
|
||||
if let Ok(device) = Device::new_cuda(0) {
|
||||
let dtype = training_dtype(&device);
|
||||
// On Ampere+ (L40S, H100), should be BF16
|
||||
// On older GPUs, F32 is fine too
|
||||
assert!(dtype == DType::BF16 || dtype == DType::F32);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_training_dtype_returns_f32_on_cpu() {
|
||||
let device = Device::Cpu;
|
||||
assert_eq!(training_dtype(&device), DType::F32);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Real validation is the hyperopt run on L40S — compare trial Sharpe distributions.
|
||||
|
||||
**Step 1:** Create test, build, commit.
|
||||
|
||||
```bash
|
||||
git commit -m "test(ml): add BF16 training dtype integration test"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
| Phase | Tasks | Sites Changed | Commit Count |
|
||||
|-------|-------|---------------|-------------|
|
||||
| 1: Infrastructure | 1-2 | training_dtype fn + mamba fix | 2 |
|
||||
| 2: DQN | 3-9 | ~80 VarBuilder + CUDA pipeline + training tensors | 7 |
|
||||
| 3: PPO | 10-11 | ~9 VarBuilder + training tensors + checkpoints | 2 |
|
||||
| 4: Supervised | 12-18 | ~60 VarBuilder across 8 models | 7 |
|
||||
| 5: Remaining | 19 | ~20 ensemble/misc sites | 1 |
|
||||
| 6: Validation | 20-21 | Build + test + integration test | 2 |
|
||||
| **Total** | **21 tasks** | **~150 sites** | **~21 commits** |
|
||||
|
||||
## Risk Checkpoints
|
||||
|
||||
After Phase 2 (DQN complete): full workspace build must pass. DQN is the most complex module — if it compiles, the rest is mechanical.
|
||||
|
||||
After Phase 6: all 2506+ tests must pass on CPU. GPU validation via hyperopt run on L40S.
|
||||
Reference in New Issue
Block a user