gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(cuda): remove --use_fast_math, eliminate cross-stream NaN race, f32 rewards/dones

Browse Source 
Three root causes of sporadic NaN during training:

1. --use_fast_math (nvcc) breaks IEEE 754 NaN semantics: fmaxf(NaN,x)
   returns NaN instead of x, isnan()/isinf() compile to false.
   Replaced with --ftz=true --fmad=true --prec-div=true --prec-sqrt=true
   across all 4 build.rs (ml, ml-dqn, ml-ppo, ml-core).

2. Cross-stream race: replay buffer wrote batch data on the device's
   original stream while the trainer read it on a forked stream.
   Fixed by passing the forked stream to the DQN agent via agent_device,
   so all GPU components share a single CUDA stream (zero sync overhead).

3. Rewards/dones stored as bf16 in replay buffer caused done=0xFFFF NaN.
   Converted entire rewards/dones pipeline to f32: experience collector,
   replay buffer storage, nstep kernel, loss/grad kernels.

Also:
- Removed fast_isnan/fast_isinf/fast_isfinite wrappers — standard
  isnan/isinf/isfinite work correctly without --use_fast_math
- Updated dqn-smoketest.toml: lr=1e-4, epsilon=1e-8 (f32 Adam values)
- Removed debug printfs from gather kernels
- Added curiosity_weight to training profile system
- Cleaned up smoke_params() inline overrides

11/11 smoke tests pass, 5/5 stress runs of 50-epoch test pass,
359/359 ml-dqn + 895/895 ml unit tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-29 16:18:37 +02:00
parent e4bca3fad1
commit 68804a1a51
31 changed files with 417 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
config/training/dqn-smoketest.toml
View File

@@ -1,23 +1,21 @@
# DQN Smoke Test Profile — BF16 native, small network, fast CI on RTX 3050 + H100
# DQN Smoke Test Profile — f32 master weights + bf16 tensor core shadows
# gpu_n_episodes=32 × gpu_timesteps=100 = 3200 experiences/epoch
# batch_size=64 fits in 4GB RTX 3050
#
# BF16 tuning notes:
# - lr=1e-5 (lower than F32's 3e-5 — bf16 gradients need smaller steps)
# - gradient_clip_norm=1.0 (was 10.0 — bf16 accumulation overflows with loose clip)
# - huber_delta=1.0 (was 10.0 default — large delta amplifies bf16 gradient noise)
# - q_clip ±50 (was ±200 — bf16 Q-values diverge with wide clamps)
# - reward_scale=1.0 (was 10.0 — scaling amplifies bf16 rounding into gradients)
# - noisy_sigma_init=0.3 (was 0.5 — less noise in bf16 weight perturbations)
# - spectral_norm_sigma_max=1.5 (was 3.0 — tighter weight constraint for bf16)
# Mixed-precision (NVIDIA AMP pattern):
# - f32 master weights, Adam moments, gradients, loss accumulators
# - bf16 shadow copies for cuBLAS GemmEx tensor core GEMM
# - gradient_clip_norm=1.0 (prevents large grads from overflow in bf16 backward)
# - huber_delta=1.0 (keeps TD errors bounded for stable C51 transition)
# - q_clip ±50 (prevents logit explosion through bf16 forward)
[training]
epochs = 3
batch_size = 64
learning_rate = 0.00001
learning_rate = 0.0001
gamma = 0.95
weight_decay = 0.0001
adam_epsilon = 0.001 # BF16: 1e-8 rounds to 0 → div-by-zero. Use ≥1e-4.
adam_epsilon = 1e-8
warmup_steps = 0
hidden_dim_base = 64
max_steps_per_epoch = 200
@@ -69,6 +67,7 @@ tau = 0.005
c51_warmup_epochs = 5
her_ratio = 0.2
cql_alpha = 0.1
curiosity_weight = 0.1
iqn_lambda = 0.25
spectral_norm_sigma_max = 1.5
gradient_clip_norm = 1.0

3
crates/ml-core/build.rs
View File

@@ -112,9 +112,10 @@ fn try_compile_kernel(
"-cubin",
&format!("-arch={arch}"),
"-O3",
"--use_fast_math",
"--ftz=true",
"--fmad=true",
"--prec-div=true",
"--prec-sqrt=true",
"-o",
cubin_path.to_str().unwrap(),
tmp_src.to_str().unwrap(),

5
crates/ml-dqn/build.rs
View File

@@ -39,7 +39,7 @@ fn main() {
let common_src = std::fs::read_to_string(&common_header_path)
.unwrap_or_else(|e| panic!("Failed to read {}: {e}", common_header_path.display()));
// All kernels get the common header (BF16 types + fast_isnan/fast_isinf)
// All kernels get the common header (BF16 types + math wrappers)
let bf16_kernels = [
"rmsnorm_kernels.cu",
"noisy_kernels.cu",
@@ -115,9 +115,10 @@ fn try_compile_kernel(
"-cubin",
&format!("-arch={arch}"),
"-O3",
"--use_fast_math",
"--ftz=true",
"--fmad=true",
"--prec-div=true",
"--prec-sqrt=true",
"-o",
cubin_path.to_str().unwrap(),
tmp_src.to_str().unwrap(),

31
crates/ml-dqn/src/dqn.rs
View File

@@ -2423,14 +2423,41 @@ impl DQN {
)?;
// GPU: target = reward + gamma_n * max_q_next * (1 - done)
// Convert f32 rewards/dones to bf16 for elementwise ops
let rewards_bf16 = {
let mut dst = stream.alloc_zeros::<half::bf16>(batch_size)
.map_err(|e| MLError::ModelError(format!("alloc rewards_bf16: {e}")))?;
let host: Vec<half::bf16> = {
let mut h = vec![0.0_f32; batch_size];
stream.memcpy_dtoh(&gpu.rewards, &mut h)
.map_err(|e| MLError::ModelError(format!("dtoh rewards: {e}")))?;
h.iter().map(|&v| half::bf16::from_f32(v)).collect()
};
stream.memcpy_htod(&host, &mut dst)
.map_err(|e| MLError::ModelError(format!("htod rewards_bf16: {e}")))?;
dst
};
let dones_bf16 = {
let mut dst = stream.alloc_zeros::<half::bf16>(batch_size)
.map_err(|e| MLError::ModelError(format!("alloc dones_bf16: {e}")))?;
let host: Vec<half::bf16> = {
let mut h = vec![0.0_f32; batch_size];
stream.memcpy_dtoh(&gpu.dones, &mut h)
.map_err(|e| MLError::ModelError(format!("dtoh dones: {e}")))?;
h.iter().map(|&v| half::bf16::from_f32(v)).collect()
};
stream.memcpy_htod(&host, &mut dst)
.map_err(|e| MLError::ModelError(format!("htod dones_bf16: {e}")))?;
dst
};
// (1 - done)
let one_minus_done = ew.affine(gpu.dones.data(), -1.0, 1.0, batch_size)?;
let one_minus_done = ew.affine(&dones_bf16, -1.0, 1.0, batch_size)?;
// gamma_n * max_q_next
let scaled_q = ew.affine(&max_q_next, gamma_n, 0.0, batch_size)?;
// gamma_n * max_q_next * (1 - done)
let discounted = ew.binary(&scaled_q, &one_minus_done, batch_size, 2)?; // mul
// reward + discounted
let target_vals = ew.binary(gpu.rewards.data(), &discounted, batch_size, 0)?; // add
let target_vals = ew.binary(&rewards_bf16, &discounted, batch_size, 0)?; // add
// GPU: td_error = q_sa - target
let td_error_gpu = ew.binary(&q_sa, &target_vals, batch_size, 1)?; // sub

99
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -28,8 +28,8 @@ pub struct GpuBatchSlices {
pub states: CudaSlice<u16>, // [batch_size * state_dim] bf16 on GPU
pub next_states: CudaSlice<u16>, // [batch_size * state_dim] bf16 on GPU
pub actions: CudaSlice<u32>, // [batch_size] u32 on GPU
pub rewards: CudaSlice<u16>, // [batch_size] bf16 on GPU
pub dones: CudaSlice<u16>, // [batch_size] bf16 on GPU (0.0/1.0)
pub rewards: CudaSlice<f32>, // [batch_size] f32 on GPU (no bf16 NaN risk)
pub dones: CudaSlice<f32>, // [batch_size] f32 on GPU (0.0/1.0, no bf16 NaN risk)
pub weights: CudaSlice<f32>, // [batch_size] f32 on GPU (IS weights — f32 to avoid bf16 overflow → Inf → NaN)
pub indices: CudaSlice<u32>, // [batch_size] u32 on GPU (buffer indices)
/// Episode IDs for sampled transitions `[batch_size]` i32 on GPU.
@@ -68,9 +68,9 @@ impl GpuBatchSlices {
dst
};
// bf16 slices -> bf16 GpuTensor via DtoD reinterpret (zero cast needed, already bf16)
let rewards = bf16_slice_to_gpu_tensor_gpu(&self.rewards, vec![bs], stream, kernels)?;
let dones = bf16_slice_to_gpu_tensor_gpu(&self.dones, vec![bs], stream, kernels)?;
// rewards/dones are f32 CudaSlice — DtoD clone (no bf16 NaN risk)
let rewards = dtod_clone_f32(stream, &self.rewards, bs, "r_f32")?;
let dones = dtod_clone_f32(stream, &self.dones, bs, "d_f32")?;
// IS-weights stay as f32 — no bf16 conversion (bf16 overflows to Inf for large weights)
let weights = dtod_clone_f32(stream, &self.weights, bs, "w_f32")?;
@@ -401,8 +401,8 @@ pub struct GpuReplayBuffer {
stream: Arc<CudaStream>,
kernels: ReplayKernels,
states: CudaSlice<u16>, next_states: CudaSlice<u16>,
actions: CudaSlice<u32>, rewards: CudaSlice<u16>,
dones: CudaSlice<u16>, priorities: CudaSlice<f32>,
actions: CudaSlice<u32>, rewards: CudaSlice<f32>,
dones: CudaSlice<f32>, priorities: CudaSlice<f32>,
/// Episode IDs per buffer slot `[capacity]` i32 on GPU.
/// `episode_ids[i] = i / episode_length`. Written during `insert_batch_with_episode_ids`.
episode_ids: CudaSlice<i32>,
@@ -419,8 +419,8 @@ pub struct GpuReplayBuffer {
sample_states: CudaSlice<u16>,
sample_next_states: CudaSlice<u16>,
sample_actions: CudaSlice<u32>,
sample_rewards: CudaSlice<u16>,
sample_dones: CudaSlice<u16>,
sample_rewards: CudaSlice<f32>,
sample_dones: CudaSlice<f32>,
sample_priorities: CudaSlice<f32>,
sample_weights: CudaSlice<f32>,
sample_max_weight: CudaSlice<f32>,
@@ -457,8 +457,8 @@ impl GpuReplayBuffer {
let s = a16(stream, cap * sd, "s")?;
let ns = a16(stream, cap * sd, "ns")?;
let a = a32u(stream, cap, "a")?;
let r = a16(stream, cap, "r")?;
let d = a16(stream, cap, "d")?;
let r = a32f(stream, cap, "r")?;
let d = a32f(stream, cap, "d")?;
let p = a32f(stream, cap, "p")?;
let mut mp = a32f(stream, 1, "mp")?;
stream.memcpy_htod(&[1.0_f32], &mut mp).map_err(|e| MLError::ModelError(format!("mp: {e}")))?;
@@ -474,8 +474,8 @@ impl GpuReplayBuffer {
let ss = a16(stream, mbs * sd, "s_states")?;
let sns = a16(stream, mbs * sd, "s_nstates")?;
let sa = a32u(stream, mbs, "s_act")?;
let sr = a16(stream, mbs, "s_rew")?;
let sdn = a16(stream, mbs, "s_done")?;
let sr = a32f(stream, mbs, "s_rew")?;
let sdn = a32f(stream, mbs, "s_done")?;
let sp = a32f(stream, mbs, "s_pri")?;
let sw = a32f(stream, mbs, "s_wt")?;
let smw = a32f(stream, 1, "s_mw")?;
@@ -568,33 +568,16 @@ impl GpuReplayBuffer {
.arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?;
}
// Cast f32 rewards/dones to bf16 (u16) then scatter insert as bf16 (state_dim=1)
let one_i = 1_i32;
let eff_i = eff as i32;
let mut b_r = a16(&self.stream, eff, "ib_r")?;
// SAFETY: b_r, sr are valid device allocations of at least eff elements.
// Rewards/dones: f32 scatter insert directly (no bf16 cast -- eliminates NaN risk)
// SAFETY: rewards/dones buffers are CudaSlice<f32>, sr/sd2 are valid f32 device slices.
unsafe {
self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
.arg(&mut b_r).arg(&sr).arg(&eff_i).launch(lcfg(eff))
.map_err(|e| MLError::ModelError(format!("cast r: {e}")))?;
}
// SAFETY: rewards buffer (bf16) valid, b_r cast above.
unsafe {
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
.arg(&self.rewards).arg(&b_r).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
.arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r: {e}")))?;
}
let mut b_d = a16(&self.stream, eff, "ib_d")?;
// SAFETY: b_d, sd2 are valid device allocations of at least eff elements.
unsafe {
self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
.arg(&mut b_d).arg(&sd2).arg(&eff_i).launch(lcfg(eff))
.map_err(|e| MLError::ModelError(format!("cast d: {e}")))?;
}
// SAFETY: dones buffer (bf16) valid, b_d cast above.
unsafe {
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
.arg(&self.dones).arg(&b_d).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
.arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d: {e}")))?;
}
let mut pt = a32f(&self.stream, eff, "pt")?;
@@ -652,17 +635,18 @@ impl GpuReplayBuffer {
Ok(())
}
/// Insert a batch where states/rewards/dones are already `CudaSlice<half::bf16>`.
/// Insert a batch where states are already `CudaSlice<half::bf16>` and
/// rewards/dones are `CudaSlice<f32>` (no bf16 NaN risk).
///
/// Skips the f32→bf16 cast for states (already bf16), and uses DtoD copy
/// for rewards/dones via scatter_insert_bf16 (reinterpreted as bf16 scatter).
/// Skips the f32→bf16 cast for states (already bf16). Rewards/dones
/// scatter-insert as f32 directly.
pub fn insert_batch_bf16(
&mut self,
sf: &CudaSlice<half::bf16>,
nf: &CudaSlice<half::bf16>,
ac: &CudaSlice<u32>,
rw: &CudaSlice<half::bf16>,
dn: &CudaSlice<half::bf16>,
rw: &CudaSlice<f32>,
dn: &CudaSlice<f32>,
bs: usize,
) -> Result<(), MLError> {
if bs == 0 { return Ok(()); }
@@ -692,21 +676,20 @@ impl GpuReplayBuffer {
.arg(&self.actions).arg(&sa).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc a: {e}")))?;
}
// Rewards/dones: both input and storage are bf16 — scatter directly (state_dim=1)
// Rewards/dones: f32 scatter insert directly (no bf16 NaN risk)
{
let sr = if off > 0 { rw.slice(off..) } else { rw.slice(0..) };
let sd2 = if off > 0 { dn.slice(off..) } else { dn.slice(0..) };
let one_i = 1_i32;
// SAFETY: rewards/dones buffers are u16 (bf16), rw/dn CudaSlice<half::bf16> same repr.
// SAFETY: rewards/dones buffers are CudaSlice<f32>, rw/dn CudaSlice<f32>.
unsafe {
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
.arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r bf16: {e}")))?;
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
.arg(&self.rewards).arg(&sr).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc r f32: {e}")))?;
}
unsafe {
self.stream.launch_builder(&self.kernels.scatter_insert_bf16)
.arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&one_i).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d bf16: {e}")))?;
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
.arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d f32: {e}")))?;
}
}
// Priorities — same as insert_batch
@@ -817,17 +800,17 @@ impl GpuReplayBuffer {
.arg(&mut self.sample_actions).arg(&self.actions).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g a: {e}")))?;
}
// SAFETY: same context, rewards buffer valid.
// SAFETY: same context, rewards buffer valid (f32).
unsafe {
let cap_i32 = self.capacity() as i32;
self.stream.launch_builder(&self.kernels.gather_bf16)
self.stream.launch_builder(&self.kernels.gather_f32)
.arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?;
}
// SAFETY: same context, dones buffer valid.
// SAFETY: same context, dones buffer valid (f32).
unsafe {
let cap_i32 = self.capacity() as i32;
self.stream.launch_builder(&self.kernels.gather_bf16)
self.stream.launch_builder(&self.kernels.gather_f32)
.arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
.launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g d: {e}")))?;
}
@@ -919,8 +902,8 @@ impl GpuReplayBuffer {
states: dtod_clone_u16(&self.stream, &self.sample_states, batch_size * sd, "o_s")?,
next_states: dtod_clone_u16(&self.stream, &self.sample_next_states, batch_size * sd, "o_n")?,
actions: dtod_clone_u32(&self.stream, &self.sample_actions, batch_size, "o_act")?,
rewards: dtod_clone_u16(&self.stream, &self.sample_rewards, batch_size, "o_r")?,
dones: dtod_clone_u16(&self.stream, &self.sample_dones, batch_size, "o_d")?,
rewards: dtod_clone_f32(&self.stream, &self.sample_rewards, batch_size, "o_r")?,
dones: dtod_clone_f32(&self.stream, &self.sample_dones, batch_size, "o_d")?,
weights: weights_f32,
indices: dtod_clone_u32(&self.stream, &self.sample_indices_u32, batch_size, "o_i")?,
episode_ids: Some(ep_ids),
@@ -1061,8 +1044,8 @@ impl GpuReplayBuffer {
pub const fn states_slice(&self) -> &CudaSlice<u16> { &self.states }
pub const fn next_states_slice(&self) -> &CudaSlice<u16> { &self.next_states }
pub const fn actions_slice(&self) -> &CudaSlice<u32> { &self.actions }
pub const fn rewards_slice(&self) -> &CudaSlice<u16> { &self.rewards }
pub const fn dones_slice(&self) -> &CudaSlice<u16> { &self.dones }
pub const fn rewards_slice(&self) -> &CudaSlice<f32> { &self.rewards }
pub const fn dones_slice(&self) -> &CudaSlice<f32> { &self.dones }
pub const fn priorities_slice(&self) -> &CudaSlice<f32> { &self.priorities }
/// Sample proportional indices and IS weights as GPU-resident `CudaSlices`.

42
crates/ml-dqn/src/regime_conditional.rs
View File

@@ -658,12 +658,29 @@ impl RegimeConditionalDQN {
out
};
// DtoD clone f32 rewards/dones (no bf16 NaN risk)
let rewards_clone = {
let n = gpu_batch.rewards.len();
let mut dst = self.stream.alloc_zeros::<f32>(n)
.map_err(|e| MLError::TrainingError(format!("rewards clone alloc: {e}")))?;
self.stream.memcpy_dtod(&gpu_batch.rewards, &mut dst)
.map_err(|e| MLError::TrainingError(format!("rewards clone DtoD: {e}")))?;
dst
};
let dones_clone = {
let n = gpu_batch.dones.len();
let mut dst = self.stream.alloc_zeros::<f32>(n)
.map_err(|e| MLError::TrainingError(format!("dones clone alloc: {e}")))?;
self.stream.memcpy_dtod(&gpu_batch.dones, &mut dst)
.map_err(|e| MLError::TrainingError(format!("dones clone DtoD: {e}")))?;
dst
};
let masked_batch = GpuBatch {
states: gpu_batch.states.gpu_clone(&self.stream)?,
actions: gpu_batch.actions.gpu_clone(&self.stream)?,
rewards: gpu_batch.rewards.gpu_clone(&self.stream)?,
rewards: rewards_clone,
next_states: gpu_batch.next_states.gpu_clone(&self.stream)?,
dones: gpu_batch.dones.gpu_clone(&self.stream)?,
dones: dones_clone,
weights: masked_weights,
indices: gpu_batch.indices.clone(),
episode_ids: None,
@@ -876,12 +893,29 @@ impl RegimeConditionalDQN {
out
};
// DtoD clone f32 rewards/dones (no bf16 NaN risk)
let rewards_clone = {
let n = gpu_batch.rewards.len();
let mut dst = self.stream.alloc_zeros::<f32>(n)
.map_err(|e| MLError::TrainingError(format!("rewards clone alloc: {e}")))?;
self.stream.memcpy_dtod(&gpu_batch.rewards, &mut dst)
.map_err(|e| MLError::TrainingError(format!("rewards clone DtoD: {e}")))?;
dst
};
let dones_clone = {
let n = gpu_batch.dones.len();
let mut dst = self.stream.alloc_zeros::<f32>(n)
.map_err(|e| MLError::TrainingError(format!("dones clone alloc: {e}")))?;
self.stream.memcpy_dtod(&gpu_batch.dones, &mut dst)
.map_err(|e| MLError::TrainingError(format!("dones clone DtoD: {e}")))?;
dst
};
let masked_batch = GpuBatch {
states: gpu_batch.states.gpu_clone(&self.stream)?,
actions: gpu_batch.actions.gpu_clone(&self.stream)?,
rewards: gpu_batch.rewards.gpu_clone(&self.stream)?,
rewards: rewards_clone,
next_states: gpu_batch.next_states.gpu_clone(&self.stream)?,
dones: gpu_batch.dones.gpu_clone(&self.stream)?,
dones: dones_clone,
weights: masked_weights,
indices: gpu_batch.indices.clone(),
episode_ids: None,

21
crates/ml-dqn/src/replay_buffer_kernels.cu
View File

@@ -93,11 +93,7 @@ void gather_u32(
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= batch_size) return;
int idx = (int)indices[i];
if (idx < 0 || idx >= capacity) {
printf("GATHER OOB: i=%d idx=%d cap=%d raw=%lld\n", i, idx, capacity, indices[i]);
/* Don't clamp — let the NaN propagate so the loss kernel catches it */
return;
}
if (idx < 0 || idx >= capacity) return;
out[i] = src[idx];
}
@@ -135,11 +131,7 @@ void gather_bf16(
int idx = (int)indices[i];
/* Bounds check: corrupt PER segment tree can produce OOB indices.
* Clamp to [0, capacity-1] instead of reading garbage (NaN bits). */
if (idx < 0 || idx >= capacity) {
printf("GATHER OOB: i=%d idx=%d cap=%d raw=%lld\n", i, idx, capacity, indices[i]);
/* Don't clamp — let the NaN propagate so the loss kernel catches it */
return;
}
if (idx < 0 || idx >= capacity) return;
out[i] = src[idx];
}
@@ -174,11 +166,6 @@ void is_weights_f32(
if (i >= batch_size) return;
float ts = total_sum_buf[0];
float prob = fmaxf((sampled_prios[i] * (float)n_buffer) / fmaxf(ts, 1e-8f), 1e-12f);
/* Clamp weight to prevent Inf after bf16 conversion.
* IS-weights are stored as bf16 (max ~65504). powf(tiny_prob, -beta)
* can exceed this. After normalization (÷ max_weight), values are ≤1.0,
* but PRE-normalization values must fit in bf16 to avoid Inf.
* Clamp to 60000 (below bf16 Inf threshold, normalized to ≤1.0). */
float w = powf(prob, neg_beta);
weights[i] = fminf(w, 60000.0f);
}
@@ -371,7 +358,7 @@ void max_of_two_f32(
}
// ── 18. NaN/Inf check kernel ────────────────────────────────────────────────
// out[i] = (fast_isnan(v) || fast_isinf(v)) ? 1.0f : 0.0f
// out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f
extern "C" __global__
void nan_inf_check_f32(
@@ -382,7 +369,7 @@ void nan_inf_check_f32(
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
float v = input[i];
out[i] = (fast_isnan(v) || fast_isinf(v)) ? 1.0f : 0.0f;
out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f;
}
// ── 19. Dead neuron (near-zero) check kernel ────────────────────────────────

4
crates/ml-dqn/src/replay_buffer_type.rs
View File

@@ -45,9 +45,9 @@ impl BatchSample {
pub struct GpuBatch {
pub states: GpuTensor, // [batch_size, state_dim] bf16 on GPU
pub actions: GpuTensor, // [batch_size] u32 on GPU
pub rewards: GpuTensor, // [batch_size] bf16 on GPU
pub rewards: cudarc::driver::CudaSlice<f32>, // [batch_size] f32 on GPU (no bf16 NaN risk)
pub next_states: GpuTensor, // [batch_size, state_dim] bf16 on GPU
pub dones: GpuTensor, // [batch_size] bf16 on GPU (0.0/1.0)
pub dones: cudarc::driver::CudaSlice<f32>, // [batch_size] f32 on GPU (0.0/1.0, no bf16 NaN risk)
pub weights: cudarc::driver::CudaSlice<f32>, // [batch_size] f32 on GPU (IS weights — f32 to avoid bf16 overflow → Inf → NaN)
pub indices: cudarc::driver::CudaSlice<u32>, // [batch_size] u32 on GPU (buffer indices)
/// Episode IDs per transition `[batch_size]` i32 on GPU.

3
crates/ml-ppo/build.rs
View File

@@ -109,9 +109,10 @@ fn try_compile_kernel(
"-cubin",
&format!("-arch={arch}"),
"-O3",
"--use_fast_math",
"--ftz=true",
"--fmad=true",
"--prec-div=true",
"--prec-sqrt=true",
"-o",
cubin_path.to_str().unwrap(),
tmp_src.to_str().unwrap(),

3
crates/ml/build.rs
View File

@@ -137,9 +137,10 @@ fn try_compile_kernel(
"-cubin",
&format!("-arch={arch}"),
"-O3",
"--use_fast_math",
"--ftz=true",
"--fmad=true",
"--prec-div=true",
"--prec-sqrt=true",
&include_flag,
"-o", cubin_path.to_str().unwrap(),
tmp_src.to_str().unwrap(),

11
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -177,8 +177,8 @@ extern "C" __global__ void c51_loss_batched(
const float* __restrict__ on_next_adv_logits_b2,
const int* __restrict__ actions,
const __nv_bfloat16* __restrict__ rewards,
const __nv_bfloat16* __restrict__ dones,
const float* __restrict__ rewards,
const float* __restrict__ dones,
const float* __restrict__ is_weights,
__nv_bfloat16* __restrict__ per_sample_loss,
@@ -260,9 +260,9 @@ extern "C" __global__ void c51_loss_batched(
const float* tg_val_row = tg_value_logits + (long long)sample_id * num_atoms;
const float* on_next_val_row = on_next_value_logits + (long long)sample_id * num_atoms;
/* IS-weights are f32 — no overflow risk. */
float reward = (float)rewards[sample_id];
float done = (float)dones[sample_id];
/* Rewards/dones are f32 — no bf16 NaN risk. IS-weights are f32. */
float reward = rewards[sample_id];
float done = dones[sample_id];
float is_weight = is_weights[sample_id];
float total_ce = 0.0f;
@@ -411,7 +411,6 @@ extern "C" __global__ void c51_loss_batched(
if (tid == 0) {
float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
float weighted_loss = clamped_ce * is_weight;
/* f32 d_logits: no NaN risk from atomicAdd overflow. */
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(clamped_ce);
atomicAdd(total_loss, weighted_loss / (float)batch_size);

16
crates/ml/src/cuda_pipeline/common_device_functions.cuh
View File

@@ -13,22 +13,6 @@
* This matches NVIDIA's native tensor core accumulate mode. */
#include <cuda_bf16.h>
/* ── NaN/Inf detection that survives --use_fast_math ──────────────── */
/* nvcc --use_fast_math implies --no-nans, making isnan()/isinf() */
/* always return false (dead code). Use IEEE 754 bit patterns instead. */
__device__ __forceinline__ bool fast_isnan(float x) {
unsigned int bits = __float_as_uint(x);
return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0;
}
__device__ __forceinline__ bool fast_isinf(float x) {
unsigned int bits = __float_as_uint(x);
return (bits & 0x7fffffffu) == 0x7f800000u;
}
__device__ __forceinline__ bool fast_isfinite(float x) {
unsigned int bits = __float_as_uint(x);
return (bits & 0x7f800000u) != 0x7f800000u;
}
/* ── BF16 native math wrappers ─────────────────────────────────────── */
/* Thin wrappers around F32 transcendentals for BF16 arguments. */
/* The cast is hidden inside — kernel code reads as pure BF16. */

30
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -509,8 +509,8 @@ extern "C" __global__ void experience_env_step(
__nv_bfloat16* portfolio_states,
__nv_bfloat16* out_states,
int* out_actions,
__nv_bfloat16* out_rewards,
__nv_bfloat16* out_dones,
float* out_rewards,
float* out_dones,
const __nv_bfloat16* __restrict__ batch_states,
float max_position,
float tx_cost_multiplier,
@@ -554,8 +554,8 @@ extern "C" __global__ void experience_env_step(
/* ---- Out-of-data check ---- */
if (bar_idx >= total_bars - 1) {
out_rewards[out_off] = bf16_zero();
out_dones[out_off] = bf16_one();
out_rewards[out_off] = 0.0f;
out_dones[out_off] = 1.0f;
/* Do not advance timestep: episode is at the data boundary. */
return;
}
@@ -610,8 +610,8 @@ extern "C" __global__ void experience_env_step(
{
float portfolio_val = __bfloat162float(ps[2]);
if (check_capital_floor(portfolio_val, peak_equity)) {
out_rewards[out_off] = bf16(-10.0f);
out_dones[out_off] = bf16_one();
out_rewards[out_off] = -10.0f;
out_dones[out_off] = 1.0f;
/* Reset portfolio for next episode (fresh capital) */
float init_cap = peak_equity; /* use peak as initial for next episode */
ps[0] = bf16_zero(); /* position = flat */
@@ -1004,15 +1004,9 @@ extern "C" __global__ void experience_env_step(
/* ---- Write reward and done flag ---- */
/* Final NaN guard — if reward is NaN/Inf, write 0.0 instead of poisoning
* the replay buffer. This prevents gradient explosion from propagating. */
if (fast_isnan(reward) || fast_isinf(reward)) reward = 0.0f;
out_rewards[out_off] = __float2bfloat16(reward);
__nv_bfloat16 done_bf16 = __float2bfloat16((float)done);
/* DEBUG: check if done is valid before writing */
if (fast_isnan((float)done_bf16)) {
printf("EXP DONE NAN: thread=%d done_int=%d done_float=%f out_off=%d\n",
blockIdx.x * blockDim.x + threadIdx.x, done, (float)done, out_off);
}
out_dones[out_off] = done_bf16;
if (isnan(reward) || isinf(reward)) reward = 0.0f;
out_rewards[out_off] = reward;
out_dones[out_off] = (float)done;
/* ---- Write RAW portfolio return (unshaped) for accurate Sharpe/MaxDD ---- */
/* True fractional return: (equity_t - equity_{t-1}) / equity_{t-1}.
@@ -1022,7 +1016,7 @@ extern "C" __global__ void experience_env_step(
float portfolio_return = (prev_equity > 1.0f)
? (new_portfolio_value - prev_equity) / prev_equity
: 0.0f;
if (fast_isnan(portfolio_return) || fast_isinf(portfolio_return)) portfolio_return = 0.0f;
if (isnan(portfolio_return) || isinf(portfolio_return)) portfolio_return = 0.0f;
raw_returns_out[out_off] = __float2bfloat16(portfolio_return);
}
@@ -1074,7 +1068,7 @@ extern "C" __global__ void portfolio_sim_kernel(
const int* __restrict__ actions,
__nv_bfloat16* portfolio_state,
__nv_bfloat16* portfolio_out,
__nv_bfloat16* rewards_out,
float* rewards_out,
int* done_out,
int batch_start,
int batch_size,
@@ -1190,7 +1184,7 @@ extern "C" __global__ void portfolio_sim_kernel(
reward -= (abs_pos - 0.8f) * 5.0f * 0.1f;
}
rewards_out[b] = __float2bfloat16(reward);
rewards_out[b] = reward;
int step = global_idx + 1;
int time_done = (step % episode_length == 0) ? 1 : 0;

63
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -441,8 +441,8 @@ pub struct GpuDqnTrainer {
states_buf: CudaSlice<half::bf16>, // [B, STATE_DIM]
next_states_buf: CudaSlice<half::bf16>, // [B, STATE_DIM]
actions_buf: CudaSlice<i32>, // [B]
rewards_buf: CudaSlice<half::bf16>, // [B]
dones_buf: CudaSlice<half::bf16>, // [B]
rewards_buf: CudaSlice<f32>, // [B] f32 (no bf16 NaN risk)
dones_buf: CudaSlice<f32>, // [B] f32 (no bf16 NaN risk)
is_weights_buf: CudaSlice<f32>, // [B] f32 (bf16 overflows to Inf for PER weights)
// ── Activation save buffers (forward → backward) ────────────────
@@ -714,7 +714,7 @@ impl GpuDqnTrainer {
///
/// Shape: `[B]` f32 — rewards from the batch. Valid after `train_step()` or
/// `train_step_gpu()`. Used by IQN for Bellman target computation on GPU.
pub fn rewards_buf(&self) -> &CudaSlice<half::bf16> {
pub fn rewards_buf(&self) -> &CudaSlice<f32> {
&self.rewards_buf
}
@@ -722,7 +722,7 @@ impl GpuDqnTrainer {
///
/// Shape: `[B]` f32 — done flags (0.0/1.0). Valid after `train_step()` or
/// `train_step_gpu()`. Used by IQN for Bellman target masking on GPU.
pub fn dones_buf(&self) -> &CudaSlice<half::bf16> {
pub fn dones_buf(&self) -> &CudaSlice<f32> {
&self.dones_buf
}
@@ -1918,8 +1918,10 @@ impl GpuDqnTrainer {
let states_buf = alloc_bf16(&stream, b * state_dim_padded, "states")?;
let next_states_buf = alloc_bf16(&stream, b * state_dim_padded, "next_states")?;
let actions_buf = alloc_i32(&stream, b, "actions")?;
let rewards_buf = alloc_bf16(&stream, b, "rewards")?;
let dones_buf = alloc_bf16(&stream, b, "dones")?;
let rewards_buf = stream.alloc_zeros::<f32>(b)
.map_err(|e| MLError::ModelError(format!("alloc rewards f32: {e}")))?;
let dones_buf = stream.alloc_zeros::<f32>(b)
.map_err(|e| MLError::ModelError(format!("alloc dones f32: {e}")))?;
let is_weights_buf = stream.alloc_zeros::<f32>(b)
.map_err(|e| MLError::ModelError(format!("alloc is_weights f32: {e}")))?;
@@ -1984,9 +1986,10 @@ impl GpuDqnTrainer {
let t_buf = alloc_i32(&stream, 1, "adam_t")?;
// ── Allocate consolidated transfer buffers ─────────────────
// Upload staging: states + next_states + rewards + dones (bf16 only)
// Actions uploaded separately as i32, is_weights uploaded separately as f32
let upload_staging_len = b * config.state_dim * 2 + b * 2; // 2*B*SD + 2*B
// Upload staging: states + next_states (bf16 only)
// Rewards/dones uploaded separately as f32 (no bf16 NaN risk)
// Actions uploaded separately as i32, is_weights separately as f32
let upload_staging_len = b * config.state_dim * 2; // 2*B*SD
let upload_staging_buf = alloc_bf16(&stream, upload_staging_len, "upload_staging")?;
// Readback: total_loss(1) + grad_norm(1) + td_errors(B)
@@ -3341,13 +3344,12 @@ impl GpuDqnTrainer {
/// Upload batch data to pre-allocated GPU buffers via consolidated staging.
///
/// Packs bf16-convertible arrays into a single contiguous host buffer, performs one
/// `memcpy_htod` to the GPU staging buffer, then scatters to individual
/// buffers via `memcpy_dtod_async`. This turns 6 PCIe round-trips into 2+1
/// (bf16 staging + f32 is_weights + i32 actions).
/// Packs bf16-convertible arrays (states, next_states) into a single contiguous
/// host buffer, performs one `memcpy_htod` to the GPU staging buffer, then
/// scatters to individual buffers via `memcpy_dtod_async`.
///
/// Layout: BF16 staging [states | next_states | rewards | dones], actions uploaded separately as i32,
/// is_weights uploaded separately as f32 (bf16 overflows to Inf for PER weights).
/// Rewards/dones are uploaded as f32 (no bf16 NaN risk from done flags).
/// Actions uploaded separately as i32, IS-weights separately as f32.
fn upload_batch(
&mut self,
states: &[f32],
@@ -3362,12 +3364,11 @@ impl GpuDqnTrainer {
let bf16_size = std::mem::size_of::<half::bf16>();
// ── Pack bf16-convertible data into staging buffer → convert to BF16 ──
// is_weights excluded: uploaded as f32 separately to avoid bf16 overflow
// Only states + next_states go through bf16 staging.
// Rewards/dones/is_weights uploaded as f32 separately.
self.upload_staging_host.clear();
self.upload_staging_host.extend_from_slice(states); // B * SD
self.upload_staging_host.extend_from_slice(next_states); // B * SD
self.upload_staging_host.extend_from_slice(rewards); // B
self.upload_staging_host.extend_from_slice(dones); // B
// Single HtoD: f32 host → bf16 GPU
super::htod_f32_to_bf16(&self.stream, &self.upload_staging_host, &mut self.upload_staging_buf)?;
@@ -3386,22 +3387,17 @@ impl GpuDqnTrainer {
byte_offset += states_bytes;
// next_states: contiguous [B, SD] → padded [B, pad128(SD)]
let next_states_bytes = (b * sd * bf16_size) as u64;
self.launch_pad_states(
self.next_states_buf.raw_ptr(),
staging_base + byte_offset,
b,
)?;
byte_offset += next_states_bytes;
// rewards: B bf16 elements
let rewards_bytes = b * bf16_size;
dtod_copy(self.rewards_buf.raw_ptr(), staging_base + byte_offset, rewards_bytes, &self.stream, 2, "upload_scatter")?;
byte_offset += rewards_bytes as u64;
// dones: B bf16 elements
let dones_bytes = b * bf16_size;
dtod_copy(self.dones_buf.raw_ptr(), staging_base + byte_offset, dones_bytes, &self.stream, 3, "upload_scatter")?;
// ── Rewards/dones: direct f32 HtoD (no bf16 NaN risk) ──
self.stream.memcpy_htod(rewards, &mut self.rewards_buf)
.map_err(|e| MLError::ModelError(format!("rewards HtoD: {e}")))?;
self.stream.memcpy_htod(dones, &mut self.dones_buf)
.map_err(|e| MLError::ModelError(format!("dones HtoD: {e}")))?;
// ── Actions: separate i32 upload (not bf16) ──
self.stream.memcpy_htod(actions, &mut self.actions_buf)
@@ -3429,7 +3425,6 @@ impl GpuDqnTrainer {
) -> Result<(), MLError> {
let b = self.config.batch_size;
let bf16_size = std::mem::size_of::<half::bf16>();
let f32_size = std::mem::size_of::<f32>();
// States + next_states: contiguous [B, SD] in GpuBatch → padded [B, pad128(SD)]
@@ -3444,16 +3439,16 @@ impl GpuDqnTrainer {
b,
)?;
// Rewards, dones: bf16 DtoD
// Rewards, dones: f32 DtoD (no bf16 NaN risk)
dtod_copy(
self.rewards_buf.raw_ptr(),
gpu_batch.rewards.data().raw_ptr(),
b * bf16_size, &self.stream, 2, "rewards",
gpu_batch.rewards.raw_ptr(),
b * f32_size, &self.stream, 2, "rewards",
)?;
dtod_copy(
self.dones_buf.raw_ptr(),
gpu_batch.dones.data().raw_ptr(),
b * bf16_size, &self.stream, 3, "dones",
gpu_batch.dones.raw_ptr(),
b * f32_size, &self.stream, 3, "dones",
)?;
// IS-weights: f32 DtoD (bf16 overflows to Inf for PER weights)

59
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -327,16 +327,16 @@ impl ExperienceCollectorConfig {
/// The caller syncs the stream before reading or converting on another stream.
#[allow(missing_debug_implementations)]
pub struct GpuExperienceBatch {
/// States `[total * state_dim]` on GPU (f32, row-major)
/// States `[total * state_dim]` on GPU (bf16, row-major)
pub states: CudaSlice<half::bf16>,
/// Next-states `[total * state_dim]` on GPU (f32, episode-aware shift)
/// Next-states `[total * state_dim]` on GPU (bf16, episode-aware shift)
pub next_states: CudaSlice<half::bf16>,
/// Action indices `[total]` on GPU (i32, values 0..4)
pub actions: CudaSlice<i32>,
/// Rewards `[total]` on GPU (f32)
pub rewards: CudaSlice<half::bf16>,
/// Done flags `[total]` on GPU (i32, 0 or 1)
pub dones: CudaSlice<half::bf16>,
/// Rewards `[total]` on GPU (f32 -- no bf16 NaN risk)
pub rewards: CudaSlice<f32>,
/// Done flags `[total]` on GPU (f32, 0.0 or 1.0 -- no bf16 NaN risk)
pub dones: CudaSlice<f32>,
/// Episode index per transition `[total]` on GPU (i32).
///
/// `episode_ids[i] = i / timesteps_per_episode`.
@@ -457,8 +457,8 @@ pub struct GpuExperienceCollector {
// Output buffers [alloc_episodes * alloc_timesteps, ...]
states_out: CudaSlice<half::bf16>, // [alloc_episodes * alloc_timesteps * STATE_DIM]
actions_out: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps]
rewards_out: CudaSlice<half::bf16>, // [alloc_episodes * alloc_timesteps]
done_out: CudaSlice<half::bf16>, // [alloc_episodes * alloc_timesteps]
rewards_out: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps] f32 (no bf16 NaN)
done_out: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps] f32 (no bf16 NaN)
/// Raw per-bar portfolio returns (unshaped) for accurate Sharpe/MaxDD/Sortino.
/// True fractional return: (equity_t - equity_{t-1}) / equity_{t-1}.
raw_returns_out: CudaSlice<half::bf16>, // [alloc_episodes * alloc_timesteps]
@@ -731,10 +731,10 @@ impl GpuExperienceCollector {
.alloc_zeros::<i32>(total_output)
.map_err(|e| MLError::ModelError(format!("alloc actions_out: {e}")))?;
let rewards_out = stream
.alloc_zeros::<half::bf16>(total_output)
.alloc_zeros::<f32>(total_output)
.map_err(|e| MLError::ModelError(format!("alloc rewards_out: {e}")))?;
let done_out = stream
.alloc_zeros::<half::bf16>(total_output)
.alloc_zeros::<f32>(total_output)
.map_err(|e| MLError::ModelError(format!("alloc done_out: {e}")))?;
let raw_returns_out = stream
.alloc_zeros::<half::bf16>(total_output)
@@ -1002,8 +1002,10 @@ impl GpuExperienceCollector {
// Download done_out for episode boundary detection (same size, cold path).
// done_out[i] = 1.0 when the capital floor circuit breaker fired or
// the episode ended. Used to reset the equity curve in MaxDD computation.
// done_out is f32 -- direct DtoH, no bf16 conversion needed.
let mut host_dones = vec![0.0_f32; total_output];
super::dtoh_bf16_to_f32(&self.stream, &self.done_out, &mut host_dones)?;
self.stream.memcpy_dtoh(&self.done_out, &mut host_dones)
.map_err(|e| MLError::ModelError(format!("dtoh done_out: {e}")))?;
let done_flags: Vec<f64> = host_dones.iter().map(|&d| d as f64).collect();
@@ -1037,18 +1039,18 @@ impl GpuExperienceCollector {
let sd = self.state_dim;
// Clone kernel output buffers so they can be consumed independently
let states = dtod_clone_f32(&self.stream, &self.states_out, total * sd, "states")?;
let rewards = dtod_clone_f32(&self.stream, &self.rewards_out, total, "rewards")?;
let states = dtod_clone_bf16(&self.stream, &self.states_out, total * sd, "states")?;
let rewards = dtod_clone_f32_native(&self.stream, &self.rewards_out, total, "rewards")?;
let actions = dtod_clone_i32(&self.stream, &self.actions_out, total, "actions")?;
let dones = dtod_clone_f32(&self.stream, &self.done_out, total, "dones")?;
let dones = dtod_clone_f32_native(&self.stream, &self.done_out, total, "dones")?;
// ── N-step return accumulation ───────────────────────────────────
// Converts 1-step rewards into R_n = sum(gamma^i * r_i) and
// OR's done flags over n steps. Uses double-buffering to avoid races.
let n_steps = config.n_steps.max(1);
if n_steps > 1 {
let raw_rewards = dtod_clone_f32(&self.stream, &rewards, total, "raw_rewards_nstep")?;
let raw_dones = dtod_clone_f32(&self.stream, &dones, total, "raw_dones_nstep")?;
let raw_rewards = dtod_clone_f32_native(&self.stream, &rewards, total, "raw_rewards_nstep")?;
let raw_dones = dtod_clone_f32_native(&self.stream, &dones, total, "raw_dones_nstep")?;
let gamma_f32 = config.gamma;
let n_steps_i32 = n_steps;
@@ -1548,8 +1550,8 @@ impl GpuExperienceCollector {
&self.stream
}
/// Get a reference to the rewards output GPU buffer.
pub fn rewards_gpu(&self) -> &CudaSlice<half::bf16> {
/// Get a reference to the rewards output GPU buffer (f32).
pub fn rewards_gpu(&self) -> &CudaSlice<f32> {
&self.rewards_out
}
@@ -1743,14 +1745,29 @@ fn fill_episode_ids_gpu(
// Pure cudarc DtoD helpers
// ---------------------------------------------------------------------------
/// Allocate a fresh `CudaSlice<half::bf16>` and DtoD-copy `n_elems` floats from `src`.
fn dtod_clone_f32(
/// Allocate a fresh `CudaSlice<half::bf16>` and DtoD-copy `n_elems` elements from `src`.
fn dtod_clone_bf16(
stream: &Arc<CudaStream>,
src: &CudaSlice<half::bf16>,
n_elems: usize,
label: &str,
) -> Result<CudaSlice<half::bf16>, MLError> {
let mut dst = stream.alloc_zeros::<half::bf16>(n_elems)
.map_err(|e| MLError::ModelError(format!("alloc {label} bf16[{n_elems}]: {e}")))?;
let src_view = src.slice(..n_elems);
stream.memcpy_dtod(&src_view, &mut dst)
.map_err(|e| MLError::ModelError(format!("DtoD {label}: {e}")))?;
Ok(dst)
}
/// Allocate a fresh `CudaSlice<f32>` and DtoD-copy `n_elems` f32 elements from `src`.
fn dtod_clone_f32_native(
stream: &Arc<CudaStream>,
src: &CudaSlice<f32>,
n_elems: usize,
label: &str,
) -> Result<CudaSlice<f32>, MLError> {
let mut dst = stream.alloc_zeros::<f32>(n_elems)
.map_err(|e| MLError::ModelError(format!("alloc {label} f32[{n_elems}]: {e}")))?;
let src_view = src.slice(..n_elems);
stream.memcpy_dtod(&src_view, &mut dst)
@@ -1790,7 +1807,7 @@ fn build_next_states_dtod(
let shift = shift.max(1).min(timesteps); // clamp shift to valid range
if timesteps <= shift {
return dtod_clone_f32(stream, states, total_elems, "next_states_trivial");
return dtod_clone_bf16(stream, states, total_elems, "next_states_trivial");
}
let mut dst = stream.alloc_zeros::<half::bf16>(total_elems)

2
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -234,7 +234,7 @@ impl GpuIqlTrainer {
pub fn train_value_step(
&mut self,
states_f32: &CudaSlice<half::bf16>,
q_values_f32: &CudaSlice<half::bf16>,
q_values_f32: &CudaSlice<f32>,
) -> Result<f32, MLError> {
let b = self.config.batch_size;

24
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -171,10 +171,10 @@ pub struct GpuIqnHead {
branch_actions: CudaSlice<i32>,
/// Target h_s2 computed from next_states + target trunk weights [B, H]
target_h_s2: CudaSlice<half::bf16>,
/// Persistent rewards buffer [B] — avoids per-step temp allocs
rewards_buf: CudaSlice<half::bf16>,
/// Persistent dones buffer [B] — avoids per-step temp allocs
dones_buf: CudaSlice<half::bf16>,
/// Persistent rewards buffer [B] f32 — avoids per-step temp allocs, no bf16 NaN
rewards_buf: CudaSlice<f32>,
/// Persistent dones buffer [B] f32 — avoids per-step temp allocs, no bf16 NaN
dones_buf: CudaSlice<f32>,
// ── Activation save buffers (forward → backward) ─────────────────
save_embed: CudaSlice<half::bf16>, // [B, N, H]
@@ -255,9 +255,11 @@ impl GpuIqnHead {
let per_sample_loss = alloc_f32(&stream, b, "iqn_per_sample_loss")?;
let total_loss = alloc_f32(&stream, 1, "iqn_total_loss")?;
// Persistent rewards/dones buffers (avoids per-step temp allocs)
let rewards_buf = alloc_f32(&stream, b, "iqn_rewards")?;
let dones_buf = alloc_f32(&stream, b, "iqn_dones")?;
// Persistent rewards/dones buffers (f32 — no bf16 NaN risk)
let rewards_buf = stream.alloc_zeros::<f32>(b)
.map_err(|e| MLError::ModelError(format!("GpuIqn alloc iqn_rewards ({b} f32): {e}")))?;
let dones_buf = stream.alloc_zeros::<f32>(b)
.map_err(|e| MLError::ModelError(format!("GpuIqn alloc iqn_dones ({b} f32): {e}")))?;
let vram_bytes = (total_params * 6 + b * n * 2 + b * 3
+ b * h + b * n * h * 2 + b * n * tba * 2
@@ -331,8 +333,8 @@ impl GpuIqnHead {
next_states_buf: &CudaSlice<half::bf16>,
target_dueling: &DuelingWeightSet,
dqn_actions_buf: &CudaSlice<i32>,
dqn_rewards_buf: &CudaSlice<half::bf16>,
dqn_dones_buf: &CudaSlice<half::bf16>,
dqn_rewards_buf: &CudaSlice<f32>,
dqn_dones_buf: &CudaSlice<f32>,
) -> Result<f32, MLError> {
let b = self.config.batch_size;
let shared_h1_i32 = self.config.shared_h1 as i32;
@@ -365,9 +367,9 @@ impl GpuIqnHead {
.map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
}
// 3. DtoD copy rewards/dones from DQN trainer buffers
// 3. DtoD copy rewards/dones from DQN trainer buffers (f32)
use super::gpu_dqn_trainer::dtod_copy;
let f32_bytes = std::mem::size_of::<half::bf16>();
let f32_bytes = std::mem::size_of::<f32>();
let r_src = dqn_rewards_buf.raw_ptr();
let r_dst = self.rewards_buf.raw_ptr();
dtod_copy(r_dst, r_src, b * f32_bytes, &self.stream, 0, "iqn_rewards_dtod")?;

2
crates/ml/src/cuda_pipeline/gpu_monitoring.rs
View File

@@ -48,7 +48,7 @@ impl GpuMonitoringReducer {
/// Does NOT synchronize — caller must sync before reading result.
pub fn reduce(
&mut self,
rewards: &CudaSlice<half::bf16>,
rewards: &CudaSlice<f32>,
actions: &CudaSlice<i32>,
n: usize,
) -> Result<(), MLError> {

35
crates/ml/src/cuda_pipeline/gpu_portfolio.rs
View File

@@ -37,7 +37,7 @@ pub struct GpuPortfolioSimulator {
actions_buf: CudaSlice<i32>,
portfolio_state_buf: CudaSlice<half::bf16>,
portfolio_out_buf: CudaSlice<half::bf16>,
rewards_out_buf: CudaSlice<half::bf16>,
rewards_out_buf: CudaSlice<f32>,
done_out_buf: CudaSlice<i32>,
// Config
@@ -54,8 +54,8 @@ pub struct GpuPortfolioSimulator {
pub struct GpuPortfolioSimResult {
/// Normalized portfolio features [batch_size * 3] on GPU
pub portfolio_features: CudaSlice<half::bf16>,
/// Raw PnL rewards [batch_size] on GPU
pub rewards: CudaSlice<half::bf16>,
/// Raw PnL rewards [batch_size] on GPU (f32)
pub rewards: CudaSlice<f32>,
/// Episode boundary flags [batch_size] on GPU (1=done, 0=continue)
pub done_flags: CudaSlice<i32>,
/// Number of bars processed
@@ -102,7 +102,7 @@ impl GpuPortfolioSimulator {
.alloc_zeros::<half::bf16>(MAX_BATCH_SIZE * 3)
.map_err(|e| MLError::ModelError(format!("Failed to alloc portfolio output: {e}")))?;
let rewards_out_buf = stream
.alloc_zeros::<half::bf16>(MAX_BATCH_SIZE)
.alloc_zeros::<f32>(MAX_BATCH_SIZE)
.map_err(|e| MLError::ModelError(format!("Failed to alloc rewards output: {e}")))?;
let done_out_buf = stream
.alloc_zeros::<i32>(MAX_BATCH_SIZE)
@@ -155,7 +155,7 @@ impl GpuPortfolioSimulator {
return Ok(GpuPortfolioSimResult {
portfolio_features: self.stream.alloc_zeros::<half::bf16>(0)
.map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?,
rewards: self.stream.alloc_zeros::<half::bf16>(0)
rewards: self.stream.alloc_zeros::<f32>(0)
.map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?,
done_flags: self.stream.alloc_zeros::<i32>(0)
.map_err(|e| MLError::ModelError(format!("empty alloc: {e}")))?,
@@ -206,7 +206,7 @@ impl GpuPortfolioSimulator {
let portfolio_features = Self::dtod_clone_range_f32(
&self.portfolio_out_buf, batch_size * 3, &self.stream,
)?;
let rewards = Self::dtod_clone_range_f32(
let rewards = Self::dtod_clone_range_f32_native(
&self.rewards_out_buf, batch_size, &self.stream,
)?;
let done_flags = Self::dtod_clone_range_i32(
@@ -244,6 +244,29 @@ impl GpuPortfolioSimulator {
Ok(dst)
}
/// DtoD clone the first `count` f32 elements of `src` (native f32).
#[allow(unsafe_code)]
fn dtod_clone_range_f32_native(
src: &CudaSlice<f32>,
count: usize,
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<f32>, MLError> {
let mut dst = stream.alloc_zeros::<f32>(count)
.map_err(|e| MLError::ModelError(format!("dtod_clone_f32 alloc: {e}")))?;
{
let nbytes = count * std::mem::size_of::<f32>();
let src_view = src.slice(..count);
let (src_ptr, _sg) = src_view.device_ptr(stream);
let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
).map_err(|e| MLError::ModelError(format!("dtod_clone_f32 copy: {e}")))?;
}
}
Ok(dst)
}
/// DtoD clone the first `count` i32 elements of `src`.
#[allow(unsafe_code)]
fn dtod_clone_range_i32(

8
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -93,7 +93,7 @@ __device__ __forceinline__ __nv_bfloat16 silu_grad(__nv_bfloat16 x) {
extern "C" __global__
void iql_forward_loss_kernel(
const __nv_bfloat16* __restrict__ states,
const __nv_bfloat16* __restrict__ q_values,
const float* __restrict__ q_values,
const __nv_bfloat16* __restrict__ params,
__nv_bfloat16* __restrict__ v_out,
__nv_bfloat16* __restrict__ loss_out,
@@ -168,7 +168,7 @@ void iql_forward_loss_kernel(
v_out[sample] = v_val;
/* Expectile loss: L_tau(u) = |tau - 1(u<0)| * u^2 */
__nv_bfloat16 u = q_values[sample] - v_val;
__nv_bfloat16 u = bf16(q_values[sample]) - v_val;
__nv_bfloat16 bf_tau = bf16(IQL_EXPECTILE_TAU);
__nv_bfloat16 weight = (u >= bf16_zero()) ? bf_tau : (bf16_one() - bf_tau);
__nv_bfloat16 sample_loss = weight * u * u;
@@ -198,7 +198,7 @@ void iql_forward_loss_kernel(
extern "C" __global__
void iql_backward_kernel(
const __nv_bfloat16* __restrict__ states,
const __nv_bfloat16* __restrict__ q_values,
const float* __restrict__ q_values,
const __nv_bfloat16* __restrict__ v_out,
const __nv_bfloat16* __restrict__ params,
const __nv_bfloat16* __restrict__ save_pre1,
@@ -236,7 +236,7 @@ void iql_backward_kernel(
const __nv_bfloat16* pre2 = save_pre2 + sample * VALUE_HIDDEN_DIM;
__nv_bfloat16 v_val = v_out[sample];
__nv_bfloat16 q_val = q_values[sample];
__nv_bfloat16 q_val = bf16(q_values[sample]);
/* dL/dV = -2 * weight * (Q - V) / batch_size */
__nv_bfloat16 u = q_val - v_val;

8
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -191,8 +191,8 @@ void iqn_forward_loss_kernel(
const __nv_bfloat16* __restrict__ taus, /* [B, N] online τ samples ∈ (0,1) */
const __nv_bfloat16* __restrict__ target_taus, /* [B, N] target τ samples */
const int* __restrict__ actions, /* [B, 3] branch actions (exposure, order, urgency) */
const __nv_bfloat16* __restrict__ rewards, /* [B] */
const __nv_bfloat16* __restrict__ dones, /* [B] */
const float* __restrict__ rewards, /* [B] f32 */
const float* __restrict__ dones, /* [B] f32 (0.0/1.0) */
float gamma,
/* Weights */
const __nv_bfloat16* __restrict__ online_params, /* online IQN weights */
@@ -231,8 +231,8 @@ void iqn_forward_loss_kernel(
int a0 = actions[sample * 3 + 0]; /* exposure action */
int a1 = actions[sample * 3 + 1]; /* order action */
int a2 = actions[sample * 3 + 2]; /* urgency action */
__nv_bfloat16 reward = rewards[sample];
__nv_bfloat16 done = dones[sample];
__nv_bfloat16 reward = bf16(rewards[sample]);
__nv_bfloat16 done = bf16(dones[sample]);
/* Weight pointers (same for all samples) */
const __nv_bfloat16* w_embed = online_params + off[0];

5
crates/ml/src/cuda_pipeline/monitoring_kernel.cu
View File

@@ -25,8 +25,9 @@ __device__ void atomicMaxBF16(__nv_bfloat16* addr, __nv_bfloat16 val) {
// Reduce per-experience rewards and actions into a compact summary.
// One block, parallel reduction across N elements.
// Rewards are f32 (no bf16 NaN risk from done-flag contamination).
extern "C" __global__ void monitoring_reduce(
const __nv_bfloat16* __restrict__ rewards, // [N]
const float* __restrict__ rewards, // [N] f32
const int* __restrict__ actions, // [N]
__nv_bfloat16* summary, // [16]: mean, std, min, max, sharpe, counts[9], total, _pad
int N,
@@ -59,7 +60,7 @@ extern "C" __global__ void monitoring_reduce(
int local_counts[9] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
for (int i = tid; i < N; i += stride) {
__nv_bfloat16 r = rewards[i];
__nv_bfloat16 r = bf16(rewards[i]); /* f32 → bf16 at load */
local_sum = local_sum + r;
local_sq = local_sq + r * r;
local_min = bf16_fmin(local_min, r);

23
crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
View File

@@ -126,8 +126,8 @@ extern "C" __global__ void mse_loss_batched(
/* ── Batch data ───────────────────────────────────────────────── */
const int* __restrict__ actions, /* [B] factored action indices 0-44 */
const __nv_bfloat16* __restrict__ rewards, /* [B] */
const __nv_bfloat16* __restrict__ dones, /* [B] */
const float* __restrict__ rewards, /* [B] f32 */
const float* __restrict__ dones, /* [B] f32 (0.0/1.0) */
const float* __restrict__ is_weights, /* [B] PER importance-sampling weights (f32 — bf16 overflows to Inf) */
/* ── Outputs ──────────────────────────────────────────────────── */
@@ -217,19 +217,12 @@ extern "C" __global__ void mse_loss_batched(
const float* tg_val_row = tg_value_logits + (long long)sample_id * num_atoms;
const float* on_next_val_row = on_next_value_logits + (long long)sample_id * num_atoms;
/* Read batch scalars as float (BF16 → float at boundary).
/* Read batch scalars (rewards/dones are f32, no bf16 NaN risk).
* IS-weights are f32 — no overflow risk. */
float reward = (float)rewards[sample_id];
float done = (float)dones[sample_id];
float reward = rewards[sample_id];
float done = dones[sample_id];
float is_weight = is_weights[sample_id];
/* DEBUG: log done NaN with raw bf16 bits */
if (tid == 0 && fast_isnan(done)) {
unsigned short raw = *(const unsigned short*)(&dones[sample_id]);
printf("DONE_NAN s=%d raw=0x%04X rew=%f isw=%f\n",
sample_id, (unsigned int)raw, reward, is_weight);
}
float total_mse = 0.0f;
float total_abs_td = 0.0f;
@@ -362,12 +355,6 @@ extern "C" __global__ void mse_loss_batched(
if (tid == 0) {
float weighted_loss = avg_mse * is_weight;
if (fast_isnan(weighted_loss) || fast_isinf(weighted_loss)) {
printf("NaN sample=%d: mse=%f isw=%f rew=%f done=%f tmse=%f ttd=%f\n",
sample_id, avg_mse, is_weight, reward, done, total_mse, total_abs_td);
weighted_loss = 0.0f;
avg_td = 0.0f;
}
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(avg_td);
atomicAdd(total_loss, weighted_loss / (float)batch_size);

38
crates/ml/src/cuda_pipeline/nstep_kernel.cu
View File

@@ -10,14 +10,16 @@
* Reads from raw_* (copy of original 1-step data), writes to out_*
* (overwritten in-place). Double-buffering eliminates race conditions.
*
* Rewards and dones are f32 throughout (no bf16 NaN risk from done flags).
*
* Launch config: grid=(ceil(N*L/256), 1, 1), block=(256, 1, 1).
*/
extern "C" __global__ void nstep_accumulate_kernel(
const __nv_bfloat16* __restrict__ raw_rewards, /* [N * L] original 1-step rewards */
const __nv_bfloat16* __restrict__ raw_dones, /* [N * L] original 1-step dones */
__nv_bfloat16* __restrict__ out_rewards, /* [N * L] overwritten with R_n */
__nv_bfloat16* __restrict__ out_dones, /* [N * L] overwritten with done_n */
const float* __restrict__ raw_rewards, /* [N * L] original 1-step rewards */
const float* __restrict__ raw_dones, /* [N * L] original 1-step dones */
float* __restrict__ out_rewards, /* [N * L] overwritten with R_n */
float* __restrict__ out_dones, /* [N * L] overwritten with done_n */
float gamma,
int n_steps,
int L, /* timesteps per episode */
@@ -29,25 +31,23 @@ extern "C" __global__ void nstep_accumulate_kernel(
int ep = idx / L;
int t = idx % L;
__nv_bfloat16 gamma_bf = bf16(gamma);
__nv_bfloat16 R_n = bf16_zero();
__nv_bfloat16 gamma_pow = bf16_one();
__nv_bfloat16 any_done = bf16_zero();
__nv_bfloat16 half = bf16(0.5f);
float R_n = 0.0f;
float gamma_pow = 1.0f;
float any_done = 0.0f;
for (int i = 0; i < n_steps; i++) {
int step = t + i;
if (step >= L) break;
int off = ep * L + step;
__nv_bfloat16 r_i = raw_rewards[off];
__nv_bfloat16 d_i = raw_dones[off];
float r_i = raw_rewards[off];
float d_i = raw_dones[off];
R_n = R_n + gamma_pow * r_i;
gamma_pow = gamma_pow * gamma_bf;
R_n += gamma_pow * r_i;
gamma_pow *= gamma;
if (d_i > half) {
any_done = bf16_one();
if (d_i > 0.5f) {
any_done = 1.0f;
break; /* Episode terminated -- stop accumulating */
}
}
@@ -66,19 +66,19 @@ extern "C" __global__ void nstep_accumulate_kernel(
* This puts dense shaping (0.01x) and sparse trade-completion (+/-2.0)
* rewards on the same scale for C51 distributional learning.
*
* Rewards are f32 (no bf16 precision loss during normalization).
*
* Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
* ====================================================================== */
extern "C" __global__ void reward_normalize_kernel(
__nv_bfloat16* __restrict__ rewards,
float* __restrict__ rewards,
float mean,
float inv_std,
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
__nv_bfloat16 mean_bf = bf16(mean);
__nv_bfloat16 inv_std_bf = bf16(inv_std);
rewards[i] = (rewards[i] - mean_bf) * inv_std_bf;
rewards[i] = (rewards[i] - mean) * inv_std;
}
}

16
crates/ml/src/cuda_pipeline/training_guard_kernel.cu
View File

@@ -51,8 +51,8 @@ extern "C" __global__ void training_guard_check_and_accumulate(
/* -- Guard check -- */
int loss_nan = fast_isnan(loss) || fast_isinf(loss);
int grad_nan = fast_isnan(grad_norm) || fast_isinf(grad_norm);
int loss_nan = isnan(loss) || isinf(loss);
int grad_nan = isnan(grad_norm) || isinf(grad_norm);
int halt_nan = (loss_nan || grad_nan) ? 1 : 0;
int halt_loss_clip = (!loss_nan && loss > clip_threshold) ? 1 : 0;
@@ -99,8 +99,8 @@ extern "C" __global__ void training_guard_check(
) {
float loss = (float)loss_scalar[0];
float grad_norm = (float)grad_norm_scalar[0];
int loss_nan = fast_isnan(loss) || fast_isinf(loss);
int grad_nan = fast_isnan(grad_norm) || fast_isinf(grad_norm);
int loss_nan = isnan(loss) || isinf(loss);
int grad_nan = isnan(grad_norm) || isinf(grad_norm);
int halt_nan = (loss_nan || grad_nan) ? 1 : 0;
int halt_loss_clip = (!loss_nan && loss > clip_threshold) ? 1 : 0;
int halt_grad_collapse = 0;
@@ -125,10 +125,10 @@ extern "C" __global__ void training_guard_accumulate(
) {
float loss = (float)loss_scalar[0];
float grad_norm = (float)grad_norm_scalar[0];
if (!fast_isnan(loss) && !fast_isinf(loss)) {
if (!isnan(loss) && !isinf(loss)) {
acc_buf[0] = acc_buf[0] + bf16(loss);
}
if (!fast_isnan(grad_norm) && !fast_isinf(grad_norm)) {
if (!isnan(grad_norm) && !isinf(grad_norm)) {
acc_buf[1] = acc_buf[1] + bf16(grad_norm);
}
acc_buf[2] = acc_buf[2] + bf16_one();
@@ -177,7 +177,7 @@ extern "C" __global__ void qvalue_stats_reduce(
__nv_bfloat16 qv = row[a];
/* Skip NaN/Inf: cast to F32 for isnan/isinf check */
float qv_f = (float)qv;
if (!fast_isnan(qv_f) && !fast_isinf(qv_f)) {
if (!isnan(qv_f) && !isinf(qv_f)) {
sample_max = bf16_fmax(sample_max, qv);
local_all = local_all + qv;
}
@@ -268,7 +268,7 @@ extern "C" __global__ void qvalue_divergence_check(
for (int a = 0; a < num_actions; a++) {
__nv_bfloat16 qv = q_values[a];
float qv_f = (float)qv;
if (fast_isnan(qv_f) || fast_isinf(qv_f)) continue;
if (isnan(qv_f) || isinf(qv_f)) continue;
q_min = bf16_fmin(q_min, qv);
q_max = bf16_fmax(q_max, qv);

13
crates/ml/src/trainers/dqn/config.rs
View File

@@ -402,20 +402,19 @@ impl DQNAgentType {
/// Insert a batch of experience into the GPU replay buffer.
///
/// Actions are `CudaSlice<u32>` — integers, never floating point.
/// States/next_states/rewards/dones are BF16 GpuTensors.
/// States/next_states are BF16 GpuTensors. Rewards/dones are f32 CudaSlices
/// (no bf16 NaN risk from done flags).
pub fn insert_batch_tensors(
&self,
states: &GpuTensor,
next_states: &GpuTensor,
actions: &cudarc::driver::CudaSlice<u32>,
rewards: &GpuTensor,
dones: &GpuTensor,
rewards: &cudarc::driver::CudaSlice<f32>,
dones: &cudarc::driver::CudaSlice<f32>,
) -> Result<(), MLError> {
let batch_size = states.shape().first().copied().unwrap_or(0);
let s_slice = states.data();
let ns_slice = next_states.data();
let r_slice = rewards.data();
let d_slice = dones.data();
match self {
Self::Standard(agent) => {
@@ -423,7 +422,7 @@ impl DQNAgentType {
.ok_or_else(|| MLError::TrainingError(
"CPU replay buffer fallback disabled -- use GpuPrioritized when cuda is enabled".to_owned()
))?;
gpu_buf.gpu.insert_batch_bf16(s_slice, ns_slice, actions, r_slice, d_slice, batch_size)
gpu_buf.gpu.insert_batch_bf16(s_slice, ns_slice, actions, rewards, dones, batch_size)
}
Self::RegimeConditional(agent) => {
macro_rules! insert_head {
@@ -433,7 +432,7 @@ impl DQNAgentType {
.ok_or_else(|| MLError::TrainingError(
"CPU replay buffer fallback disabled -- use GpuPrioritized when cuda is enabled".to_owned()
))?;
gpu_buf.gpu.insert_batch_bf16(s_slice, ns_slice, actions, r_slice, d_slice, batch_size)?;
gpu_buf.gpu.insert_batch_bf16(s_slice, ns_slice, actions, rewards, dones, batch_size)?;
}
};
}

98
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -1265,8 +1265,15 @@ fn gpu_her_relabel_batch(
.map_err(|e| anyhow::anyhow!("HER normal states: {e}"))?;
let normal_next = gpu.next_states.narrow(0, 0, normal_count, &stream)
.map_err(|e| anyhow::anyhow!("HER normal next_states: {e}"))?;
let normal_rewards = gpu.rewards.narrow(0, 0, normal_count, &stream)
.map_err(|e| anyhow::anyhow!("HER normal rewards: {e}"))?;
// Normal rewards: CudaSlice<f32> slice + DtoD clone
let normal_rewards_f32 = {
let src = gpu.rewards.slice(..normal_count);
let mut dst = stream.alloc_zeros::<f32>(normal_count)
.map_err(|e| anyhow::anyhow!("HER normal rewards alloc: {e}"))?;
stream.memcpy_dtod(&src, &mut dst)
.map_err(|e| anyhow::anyhow!("HER normal rewards DtoD: {e}"))?;
dst
};
// HER portion: replace goal columns with donor's achieved goal, keep rest
let her_states = if goal_dim >= state_dim {
@@ -1293,16 +1300,34 @@ fn gpu_her_relabel_batch(
};
// HER rewards: +1.0 (goal = achieved by construction in Random strategy)
let her_rewards = GpuTensor::full(&[her_batch_size], 1.0, &stream)
.map_err(|e| anyhow::anyhow!("HER reward ones: {e}"))?;
let her_rewards_f32 = {
let ones = vec![1.0_f32; her_batch_size];
stream.clone_htod(&ones)
.map_err(|e| anyhow::anyhow!("HER reward ones htod: {e}"))?
};
// Reassemble full batch: [normal | her]
let new_states = GpuTensor::cat(&[&normal_states, &her_states], 0, stream)
.map_err(|e| anyhow::anyhow!("HER concat states: {e}"))?;
let new_next = GpuTensor::cat(&[&normal_next, &her_next], 0, stream)
.map_err(|e| anyhow::anyhow!("HER concat next: {e}"))?;
let new_rewards = GpuTensor::cat(&[&normal_rewards, &her_rewards], 0, stream)
.map_err(|e| anyhow::anyhow!("HER concat rewards: {e}"))?;
// Concatenate f32 rewards: [normal_rewards | her_rewards]
let new_rewards = {
let total = normal_count + her_batch_size;
let dst = stream.alloc_zeros::<f32>(total)
.map_err(|e| anyhow::anyhow!("HER concat rewards alloc: {e}"))?;
let nbytes_normal = normal_count * std::mem::size_of::<f32>();
let nbytes_her = her_batch_size * std::mem::size_of::<f32>();
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst.raw_ptr(), normal_rewards_f32.raw_ptr(), nbytes_normal, stream.cu_stream(),
).map_err(|e| anyhow::anyhow!("HER concat rewards normal DtoD: {e}"))?;
cudarc::driver::result::memcpy_dtod_async(
dst.raw_ptr() + nbytes_normal as u64, her_rewards_f32.raw_ptr(), nbytes_her, stream.cu_stream(),
).map_err(|e| anyhow::anyhow!("HER concat rewards her DtoD: {e}"))?;
}
dst
};
tracing::debug!(
her_batch_size,
@@ -1311,12 +1336,22 @@ fn gpu_her_relabel_batch(
"GPU HER intra-batch relabeling applied"
);
// Clone dones CudaSlice<f32> for new batch
let new_dones = {
let n = gpu.dones.len();
let mut dst = stream.alloc_zeros::<f32>(n)
.map_err(|e| anyhow::anyhow!("HER dones clone alloc: {e}"))?;
stream.memcpy_dtod(&gpu.dones, &mut dst)
.map_err(|e| anyhow::anyhow!("HER dones clone DtoD: {e}"))?;
dst
};
Ok(GpuBatch {
states: new_states,
next_states: new_next,
rewards: new_rewards,
actions: gpu.actions.clone(),
dones: gpu.dones.clone(),
dones: new_dones,
weights: gpu.weights.clone(),
indices: gpu.indices.clone(),
episode_ids: None,
@@ -1380,8 +1415,15 @@ fn gpu_her_relabel_batch_with_donors(
.map_err(|e| anyhow::anyhow!("HER Future normal states: {e}"))?;
let normal_next = gpu.next_states.narrow(0, 0, normal_count, stream)
.map_err(|e| anyhow::anyhow!("HER Future normal next_states: {e}"))?;
let normal_rewards = gpu.rewards.narrow(0, 0, normal_count, stream)
.map_err(|e| anyhow::anyhow!("HER Future normal rewards: {e}"))?;
// Normal rewards: CudaSlice<f32> slice + DtoD clone
let normal_rewards_f32 = {
let src = gpu.rewards.slice(..normal_count);
let mut dst = stream.alloc_zeros::<f32>(normal_count)
.map_err(|e| anyhow::anyhow!("HER Future normal rewards alloc: {e}"))?;
stream.memcpy_dtod(&src, &mut dst)
.map_err(|e| anyhow::anyhow!("HER Future normal rewards DtoD: {e}"))?;
dst
};
// HER portion: replace goal columns with donor's achieved goal
let her_states = if goal_dim >= state_dim {
@@ -1407,16 +1449,34 @@ fn gpu_her_relabel_batch_with_donors(
};
// HER rewards: +1.0 (goal = achieved by construction)
let her_rewards = GpuTensor::full(&[her_batch_size], 1.0, stream)
.map_err(|e| anyhow::anyhow!("HER Future reward ones: {e}"))?;
let her_rewards_f32 = {
let ones = vec![1.0_f32; her_batch_size];
stream.clone_htod(&ones)
.map_err(|e| anyhow::anyhow!("HER Future reward ones htod: {e}"))?
};
// Reassemble full batch: [normal | her]
let new_states = GpuTensor::cat(&[&normal_states, &her_states], 0, stream)
.map_err(|e| anyhow::anyhow!("HER Future concat states: {e}"))?;
let new_next = GpuTensor::cat(&[&normal_next, &her_next], 0, stream)
.map_err(|e| anyhow::anyhow!("HER Future concat next: {e}"))?;
let new_rewards = GpuTensor::cat(&[&normal_rewards, &her_rewards], 0, stream)
.map_err(|e| anyhow::anyhow!("HER Future concat rewards: {e}"))?;
// Concatenate f32 rewards: [normal_rewards | her_rewards]
let new_rewards = {
let total = normal_count + her_batch_size;
let dst = stream.alloc_zeros::<f32>(total)
.map_err(|e| anyhow::anyhow!("HER Future concat rewards alloc: {e}"))?;
let nbytes_normal = normal_count * std::mem::size_of::<f32>();
let nbytes_her = her_batch_size * std::mem::size_of::<f32>();
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst.raw_ptr(), normal_rewards_f32.raw_ptr(), nbytes_normal, stream.cu_stream(),
).map_err(|e| anyhow::anyhow!("HER Future concat rewards normal DtoD: {e}"))?;
cudarc::driver::result::memcpy_dtod_async(
dst.raw_ptr() + nbytes_normal as u64, her_rewards_f32.raw_ptr(), nbytes_her, stream.cu_stream(),
).map_err(|e| anyhow::anyhow!("HER Future concat rewards her DtoD: {e}"))?;
}
dst
};
tracing::debug!(
her_batch_size,
@@ -1426,12 +1486,22 @@ fn gpu_her_relabel_batch_with_donors(
"GPU HER strategy-aware relabeling applied"
);
// Clone dones CudaSlice<f32> for new batch
let new_dones = {
let n = gpu.dones.len();
let mut dst = stream.alloc_zeros::<f32>(n)
.map_err(|e| anyhow::anyhow!("HER Future dones clone alloc: {e}"))?;
stream.memcpy_dtod(&gpu.dones, &mut dst)
.map_err(|e| anyhow::anyhow!("HER Future dones clone DtoD: {e}"))?;
dst
};
Ok(GpuBatch {
states: new_states,
next_states: new_next,
rewards: new_rewards,
actions: gpu.actions.clone(),
dones: gpu.dones.clone(),
dones: new_dones,
weights: gpu.weights.clone(),
indices: gpu.indices.clone(),
episode_ids: None,

20
crates/ml/src/trainers/dqn/smoke_tests/helpers.rs
View File

@@ -26,28 +26,16 @@ fn init_tracing() {
pub(super) fn smoke_params() -> DQNHyperparameters {
let mut p = DQNHyperparameters::conservative();
// ── Smoke test sizing — loaded from config/training/dqn-smoketest.toml ──
// Sets: epochs, batch_size, warmup_steps, hidden_dim_base, max_steps_per_epoch,
// buffer_size, min_replay_size, early_stopping_enabled, gpu_n_episodes,
// gpu_timesteps_per_episode.
// ── Single source of truth: config/training/dqn-smoketest.toml ──
// All training hyperparameters come from the TOML. No inline overrides.
let profile = crate::training_profile::DqnTrainingProfile::load("dqn-smoketest");
profile.apply_to(&mut p);
// ── Production features (always on — not in the TOML profile) ──
p.curiosity_weight = 0.1; // Curiosity — always on
p.cql_alpha = 0.1;
p.enable_kelly_sizing = true; // Kelly criterion — always on
p.kelly_fractional = 0.5;
p.kelly_max_fraction = 0.25;
p.enable_action_masking = true; // Action masking — always on
p.max_position_absolute = 2.0;
p.enable_circuit_breaker = true; // Circuit breaker — always on
// ── Testing flags (not training config) ──
// ── Test-only flags (not supported by TOML profile yet) ──
p.enable_stress_testing = false;
p.enable_compliance = false;
p.enable_regime_qnetwork = false;
p.replay_buffer_vram_fraction = 0.0; // Explicit buffer_size, no AutoReplaySizer
p.replay_buffer_vram_fraction = 0.0; // Explicit buffer_size from TOML
p.checkpoint_frequency = 100;
p

16
crates/ml/src/trainers/dqn/trainer/constructor.rs
View File

@@ -326,17 +326,27 @@ impl DQNTrainer {
let curiosity_market_dim = config.curiosity_market_dim;
let curiosity_hidden_dim = config.curiosity_hidden_dim;
// Create DQN agent
// Create DQN agent using the forked stream so ALL GPU components (trainer,
// replay buffer, experience collector) share a single CUDA stream. This
// eliminates cross-stream race conditions that caused sporadic NaN.
let agent_device = if let Some(ref forked) = cuda_stream {
MlDevice::Cuda {
context: forked.context().clone(),
stream: Arc::clone(forked),
}
} else {
device.clone()
};
let agent = if hyperparams.enable_regime_qnetwork {
info!("Creating regime-conditional DQN with 3 heads (Trending, Ranging, Volatile)");
info!(" - Regime detection: ADX (raw index 40) + CUSUM direction (raw index 41)");
info!(" - Classification: Trending (ADX>0.25), Volatile (ADX≤0.25 & |CUSUM|>0.7), Ranging (otherwise)");
let regime_agent = RegimeConditionalDQN::new_on_device(config, device.clone())
let regime_agent = RegimeConditionalDQN::new_on_device(config, agent_device.clone())
.map_err(|e| anyhow::anyhow!("Failed to create regime-conditional DQN: {}", e))?;
DQNAgentType::RegimeConditional(regime_agent)
} else {
info!("Creating standard DQN with single Q-network head");
let standard_agent = DQN::new_on_device(config, device.clone())
let standard_agent = DQN::new_on_device(config, agent_device)
.map_err(|e| anyhow::anyhow!("Failed to create DQN agent: {}", e))?;
DQNAgentType::Standard(standard_agent)
};

2
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1119,7 +1119,7 @@ impl DQNTrainer {
&*(&gpu_batch.actions as *const CudaSlice<i32> as *const CudaSlice<u32>)
};
// dones: already CudaSlice<half::bf16> (kernel outputs float 0.0/1.0)
// rewards/dones: CudaSlice<f32> (kernel outputs float directly, no bf16 NaN)
let agent = self.agent.read().await;
let mut gpu_buf = agent.memory().as_gpu_buffer()
.ok_or_else(|| anyhow::anyhow!("GPU PER buffer required"))?;

5
crates/ml/src/training_profile.rs
View File

@@ -137,6 +137,8 @@ pub struct AdvancedSection {
pub spectral_norm_sigma_max: Option<f64>,
/// Gradient clipping max norm (overrides [training].gradient_clip_norm if set).
pub gradient_clip_norm: Option<f64>,
/// Curiosity-driven exploration weight (0.0 = disabled).
pub curiosity_weight: Option<f64>,
}
/// Risk management and position-control parameters.
@@ -743,6 +745,9 @@ impl DqnTrainingProfile {
if let Some(v) = a.gradient_clip_norm {
hp.gradient_clip_norm = Some(v);
}
if let Some(v) = a.curiosity_weight {
hp.curiosity_weight = v;
}
}
// [risk]