Files

jgrusewski 47442de686 docs: update stale rainbow_agent_impl path in bf16 plan

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

2026-03-04 18:36:28 +01:00

18 KiB

Raw Blame History

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)

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

SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20

Step 4: Commit

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:

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

SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -20

Step 3: Commit

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:69,72
crates/ml/src/dqn/rainbow_network.rs:429,450

Pattern for each site:

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

SQLX_OFFLINE=true cargo check -p ml 2>&1 | head -30

Step 3: Commit

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

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

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

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

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:

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

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

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

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:

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

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

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

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:

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

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:

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

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)