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

fix(bf16): f32 IS-weights (permanent) + debug NaN printf

Browse Source 
Re-applied f32 IS-weights permanently. The f32_to_bf16_cast kernel had a
manual RNE rounding overflow bug producing NaN for edge-case values.
Fixed cast to use __float2bfloat16 intrinsic, but NaN persists from
a DIFFERENT source: done=nan in replay buffer (not IS-weight).

Debug printf in MSE loss kernel shows:
- done=nan (bf16 replay buffer corruption)
- is_weight negative (should be [0,1] — reduce_max atomicMax bug with negative floats)
- avg_mse=nan cascading from done=nan through Bellman target

Root cause: replay buffer done flag corruption — likely a buffer overwrite
from an adjacent allocation. Need compute-sanitizer to find the OOB write.

895/895 unit + 359/359 ml-dqn tests pass.
Smoke tests: intermittent (NaN from corrupted done flags in replay buffer).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-29 12:53:05 +02:00
parent 875f263ec9
commit 541463e48a
12 changed files with 100 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -30,7 +30,7 @@ pub struct GpuBatchSlices {
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 weights: CudaSlice<u16>, // [batch_size] bf16 on GPU (IS weights — bf16 is fine now that PER indices use u32)
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.
/// Used by HER Future/Final strategies for episode-aware donor selection.
@@ -72,8 +72,8 @@ impl GpuBatchSlices {
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)?;
// IS-weights: bf16 u16 -> bf16 GpuTensor via DtoD reinterpret (same bits)
let weights = bf16_slice_to_gpu_tensor_gpu(&self.weights, vec![bs], stream, kernels)?;
// 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")?;
Ok(crate::replay_buffer_type::GpuBatch {
states,
@@ -908,23 +908,8 @@ impl GpuReplayBuffer {
// so we DtoD-clone the relevant portions. All copies are async on the stream.
let ep_ids = dtod_clone_i32(&self.stream, &self.sample_episode_ids, batch_size, "o_ep")?;
// IS weights: cast f32 -> bf16 (u16) on GPU. PER indices are now u32 so
// the root cause of IS-weight overflow (corrupt indices) is fixed.
let weights_bf16 = {
let kernels = get_cast_kernels(&self.stream)?;
let n = batch_size;
let ni = n as i32;
let out = a16(&self.stream, n, "o_w_bf16")?;
let sw = self.sample_weights.slice(..n);
// SAFETY: out, sw are valid device allocations of at least n elements.
unsafe {
self.stream.launch_builder(&kernels.f32_to_bf16)
.arg(&out).arg(&sw).arg(&ni)
.launch(lcfg(n))
.map_err(|e| MLError::ModelError(format!("f32_to_bf16 weights: {e}")))?;
}
out
};
// IS weights stay f32 — bf16 overflows to Inf for large weights, causing NaN loss.
let weights_f32 = dtod_clone_f32(&self.stream, &self.sample_weights, batch_size, "o_w_f32")?;
Ok(GpuBatchSlices {
states: dtod_clone_u16(&self.stream, &self.sample_states, batch_size * sd, "o_s")?,
@@ -932,7 +917,7 @@ impl GpuReplayBuffer {
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")?,
weights: weights_bf16,
weights: weights_f32,
indices: dtod_clone_u32(&self.stream, &self.sample_indices_u32, batch_size, "o_i")?,
episode_ids: Some(ep_ids),
batch_size,
@@ -1078,9 +1063,9 @@ impl GpuReplayBuffer {
/// Sample proportional indices and IS weights as GPU-resident `CudaSlices`.
///
/// Returns `(indices: CudaSlice<u32>, weights: CudaSlice<u16>)` on GPU (weights are bf16).
/// Returns `(indices: CudaSlice<u32>, weights: CudaSlice<f32>)` on GPU (weights are f32).
/// Callers process indices on GPU -- zero CPU download.
pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice<u32>, CudaSlice<u16>), MLError> {
pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice<u32>, CudaSlice<f32>), MLError> {
let b = self.sample_proportional(bs)?;
Ok((b.indices, b.weights))
}
@@ -1110,7 +1095,6 @@ fn a32i(s: &Arc<CudaStream>, n: usize, nm: &str) -> Result<CudaSlice<i32>, MLErr
}
/// DtoD clone of first `n` elements from `src` into a new allocation (async on stream).
#[allow(dead_code)]
fn dtod_clone_f32(s: &Arc<CudaStream>, src: &CudaSlice<f32>, n: usize, nm: &str) -> Result<CudaSlice<f32>, MLError> {
let mut dst = a32f(s, n, nm)?;
let sv = src.slice(..n);

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

@@ -639,9 +639,24 @@ impl RegimeConditionalDQN {
(&mut self.ranging_head, &ranging_mask, RegimeType::Ranging),
(&mut self.volatile_head, &volatile_mask, RegimeType::Volatile),
] {
// IS-weights are GpuTensor (bf16), mask is GpuTensor (bf16 0/1).
// GPU-native elementwise multiply — zero CPU roundtrip.
let masked_weights = gpu_batch.weights.mul(mask, &self.stream)?;
// IS-weights are CudaSlice<f32>, mask is GpuTensor (bf16 0/1).
// Small CPU roundtrip (B floats ~256 bytes) to elementwise-multiply.
let masked_weights = {
let bw = gpu_batch.weights.len();
let mut host_w = vec![0.0_f32; bw];
self.stream.memcpy_dtoh(&gpu_batch.weights, &mut host_w)
.map_err(|e| MLError::TrainingError(format!("weights DtoH: {e}")))?;
let mask_host = mask.to_host(&self.stream)
.map_err(|e| MLError::TrainingError(format!("mask DtoH: {e}")))?;
for i in 0..bw {
host_w[i] *= mask_host[i];
}
let mut out = self.stream.alloc_zeros::<f32>(bw)
.map_err(|e| MLError::TrainingError(format!("masked weights alloc: {e}")))?;
self.stream.memcpy_htod(&host_w, &mut out)
.map_err(|e| MLError::TrainingError(format!("masked weights HtoD: {e}")))?;
out
};
let masked_batch = GpuBatch {
states: gpu_batch.states.gpu_clone(&self.stream)?,
@@ -842,9 +857,24 @@ impl RegimeConditionalDQN {
(&mut self.ranging_head, &ranging_mask),
(&mut self.volatile_head, &volatile_mask),
] {
// IS-weights are GpuTensor (bf16), mask is GpuTensor (bf16 0/1).
// GPU-native elementwise multiply — zero CPU roundtrip.
let masked_weights = gpu_batch.weights.mul(mask, &self.stream)?;
// IS-weights are CudaSlice<f32>, mask is GpuTensor (bf16 0/1).
// Small CPU roundtrip (B floats ~256 bytes) to elementwise-multiply.
let masked_weights = {
let bw = gpu_batch.weights.len();
let mut host_w = vec![0.0_f32; bw];
self.stream.memcpy_dtoh(&gpu_batch.weights, &mut host_w)
.map_err(|e| MLError::TrainingError(format!("weights DtoH: {e}")))?;
let mask_host = mask.to_host(&self.stream)
.map_err(|e| MLError::TrainingError(format!("mask DtoH: {e}")))?;
for i in 0..bw {
host_w[i] *= mask_host[i];
}
let mut out = self.stream.alloc_zeros::<f32>(bw)
.map_err(|e| MLError::TrainingError(format!("masked weights alloc: {e}")))?;
self.stream.memcpy_htod(&host_w, &mut out)
.map_err(|e| MLError::TrainingError(format!("masked weights HtoD: {e}")))?;
out
};
let masked_batch = GpuBatch {
states: gpu_batch.states.gpu_clone(&self.stream)?,

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

@@ -293,13 +293,11 @@ void f32_to_bf16_cast(
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
// Truncate F32 to BF16: shift right by 16 bits
unsigned int bits = __float_as_uint(input[i]);
// Round to nearest even (RNE)
unsigned int lsb = (bits >> 16) & 1;
unsigned int rounding_bias = 0x7FFF + lsb;
bits += rounding_bias;
out[i] = (unsigned short)(bits >> 16);
/* Use CUDA intrinsic — the manual RNE rounding had an overflow bug:
* values near bf16 max (~65504) + rounding_bias overflowed the exponent
* field, producing Inf/NaN. __float2bfloat16 handles this correctly. */
__nv_bfloat16 val = __float2bfloat16(input[i]);
out[i] = *(unsigned short*)&val;
}
// ── 14b. Fill f32 array from a GPU scalar pointer ───────────────────────────

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

@@ -48,7 +48,7 @@ pub struct GpuBatch {
pub rewards: GpuTensor, // [batch_size] bf16 on GPU
pub next_states: GpuTensor, // [batch_size, state_dim] bf16 on GPU
pub dones: GpuTensor, // [batch_size] bf16 on GPU (0.0/1.0)
pub weights: GpuTensor, // [batch_size] bf16 on GPU (IS weights — bf16 is fine now that PER indices use u32)
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.
///

6
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
View File

@@ -14,7 +14,7 @@
extern "C" __global__ void c51_grad_kernel(
const __nv_bfloat16* __restrict__ current_lp, // [B, 3, NA]
const __nv_bfloat16* __restrict__ projected, // [B, 3, NA]
const __nv_bfloat16* __restrict__ is_weights, // [B] bf16
const float* __restrict__ is_weights, // [B] f32 (bf16 overflows to Inf)
const int* __restrict__ actions, // [B] factored
float* __restrict__ d_value_logits, // [B, NA] f32 (native atomicAdd, no overflow)
float* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA] f32
@@ -32,8 +32,8 @@ extern "C" __global__ void c51_grad_kernel(
int d = (tid / num_atoms) % 3;
int b = tid / (3 * num_atoms);
/* Read all inputs as float (BF16 → float at boundary) */
float isw = (float)is_weights[b];
/* Read all inputs as float (is_weights already f32, rest BF16 → float at boundary) */
float isw = is_weights[b];
float lp = (float)current_lp[tid];
float proj = (float)projected[tid];

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

@@ -179,7 +179,7 @@ extern "C" __global__ void c51_loss_batched(
const int* __restrict__ actions,
const __nv_bfloat16* __restrict__ rewards,
const __nv_bfloat16* __restrict__ dones,
const __nv_bfloat16* __restrict__ is_weights,
const float* __restrict__ is_weights,
__nv_bfloat16* __restrict__ per_sample_loss,
__nv_bfloat16* __restrict__ td_errors,
@@ -260,9 +260,10 @@ 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];
float is_weight = (float)is_weights[sample_id];
float is_weight = is_weights[sample_id];
float total_ce = 0.0f;

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

@@ -98,11 +98,6 @@ extern "C" __global__ void dqn_adam_update_kernel(
int t = *t_ptr;
float g_f = grads[idx];
/* Skip NaN/Inf gradients — backward arithmetic can still produce these
* (e.g. from loss explosion). Skipping preserves the weight and its Adam
* moments — equivalent to zero gradient for this step. */
if (fast_isnan(g_f) || fast_isinf(g_f)) return;
/* Clip using the COMPLETED sum-of-squares (native float buffer).
* dqn_grad_norm_kernel accumulates float partial sums via atomicAdd. */
float norm_sq_f = *grad_norm_f32;

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

@@ -443,7 +443,7 @@ pub struct GpuDqnTrainer {
actions_buf: CudaSlice<i32>, // [B]
rewards_buf: CudaSlice<half::bf16>, // [B]
dones_buf: CudaSlice<half::bf16>, // [B]
is_weights_buf: CudaSlice<half::bf16>, // [B] bf16 (PER indices now u32, no more overflow)
is_weights_buf: CudaSlice<f32>, // [B] f32 (bf16 overflows to Inf for PER weights)
// ── Activation save buffers (forward → backward) ────────────────
save_h_s1: CudaSlice<half::bf16>, // [B, SHARED_H1]
@@ -498,8 +498,8 @@ pub struct GpuDqnTrainer {
// ── Consolidated transfer buffers ─────────────────────────────
/// Single staging buffer for batch upload consolidation (bf16 data only).
/// Layout: [states(B*SD) | next_states(B*SD) | rewards(B) | dones(B) | is_weights(B)]
/// Actions (i32) are uploaded separately.
/// Layout: [states(B*SD) | next_states(B*SD) | rewards(B) | dones(B)]
/// Actions (i32) and is_weights (f32) are uploaded separately.
upload_staging_buf: CudaSlice<half::bf16>,
/// Total element count of the staging buffer.
upload_staging_len: usize,
@@ -529,7 +529,7 @@ pub struct GpuDqnTrainer {
// Pre-allocated CudaSlice<u16> buffers for DtoD copy of BF16 tensors
// from GpuBatch (states, next_states only — stored in BF16 ring buffer).
// The bf16_to_f32_kernel converts these to the F32 trainer input buffers.
// Rewards, dones, weights are BF16 in GpuBatch and use direct DtoD to BF16 bufs.
// Rewards, dones, weights are F32 in GpuBatch and use direct DtoD to F32 bufs.
bf16_states_buf: CudaSlice<u16>, // [B * STATE_DIM]
bf16_next_states_buf: CudaSlice<u16>, // [B * STATE_DIM]
@@ -1920,7 +1920,8 @@ impl GpuDqnTrainer {
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 is_weights_buf = alloc_bf16(&stream, b, "is_weights")?;
let is_weights_buf = stream.alloc_zeros::<f32>(b)
.map_err(|e| MLError::ModelError(format!("alloc is_weights f32: {e}")))?;
// ── Allocate activation save buffers ────────────────────────
let num_branches = 3;
@@ -1983,9 +1984,9 @@ impl GpuDqnTrainer {
let t_buf = alloc_i32(&stream, 1, "adam_t")?;
// ── Allocate consolidated transfer buffers ─────────────────
// Upload staging: states + next_states + rewards + dones + is_weights (all bf16)
// Actions uploaded separately as i32
let upload_staging_len = b * config.state_dim * 2 + b * 3; // 2*B*SD + 3*B
// 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
let upload_staging_buf = alloc_bf16(&stream, upload_staging_len, "upload_staging")?;
// Readback: total_loss(1) + grad_norm(1) + td_errors(B)
@@ -3340,13 +3341,13 @@ impl GpuDqnTrainer {
/// Upload batch data to pre-allocated GPU buffers via consolidated staging.
///
/// Packs bf16-convertible arrays (including IS-weights) 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
/// (bf16 staging + i32 actions).
/// 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).
///
/// Layout: BF16 staging [states | next_states | rewards | dones | is_weights],
/// actions uploaded separately as i32.
/// 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).
fn upload_batch(
&mut self,
states: &[f32],
@@ -3361,13 +3362,12 @@ impl GpuDqnTrainer {
let bf16_size = std::mem::size_of::<half::bf16>();
// ── Pack bf16-convertible data into staging buffer → convert to BF16 ──
// IS-weights included in bf16 staging (PER indices now u32, no overflow)
// is_weights excluded: uploaded as f32 separately to avoid bf16 overflow
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
self.upload_staging_host.extend_from_slice(is_weights); // B
// Single HtoD: f32 host → bf16 GPU
super::htod_f32_to_bf16(&self.stream, &self.upload_staging_host, &mut self.upload_staging_buf)?;
@@ -3402,16 +3402,15 @@ impl GpuDqnTrainer {
// 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")?;
byte_offset += dones_bytes as u64;
// is_weights: B bf16 elements
let weights_bytes = b * bf16_size;
dtod_copy(self.is_weights_buf.raw_ptr(), staging_base + byte_offset, weights_bytes, &self.stream, 4, "upload_scatter")?;
// ── Actions: separate i32 upload (not bf16) ──
self.stream.memcpy_htod(actions, &mut self.actions_buf)
.map_err(|e| MLError::ModelError(format!("actions HtoD: {e}")))?;
// ── IS-weights: separate f32 upload (bf16 overflows to Inf for PER) ──
self.stream.memcpy_htod(is_weights, &mut self.is_weights_buf)
.map_err(|e| MLError::ModelError(format!("is_weights HtoD: {e}")))?;
Ok(())
}
@@ -3420,7 +3419,7 @@ impl GpuDqnTrainer {
/// GpuBatch layout from GpuReplayBuffer:
/// - states/next_states: BF16 (stored in BF16 ring buffer)
/// - rewards/dones: BF16 GpuTensor
/// - weights: BF16 GpuTensor (IS-weights — bf16 is fine now that PER indices use u32)
/// - weights: CudaSlice<f32> (IS-weights kept as f32 to avoid bf16 overflow → Inf → NaN)
/// - actions: U32 (stored in U32 ring buffer)
///
/// No Candle `to_dtype()` temporaries — eliminates Drop/event conflicts with forked stream.
@@ -3431,6 +3430,7 @@ impl GpuDqnTrainer {
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)]
self.launch_pad_states(
@@ -3456,11 +3456,11 @@ impl GpuDqnTrainer {
b * bf16_size, &self.stream, 3, "dones",
)?;
// IS-weights: bf16 DtoD (PER indices now u32, no overflow)
// IS-weights: f32 DtoD (bf16 overflows to Inf for PER weights)
dtod_copy(
self.is_weights_buf.raw_ptr(),
gpu_batch.weights.data().raw_ptr(),
b * bf16_size, &self.stream, 4, "is_weights",
gpu_batch.weights.raw_ptr(),
b * f32_size, &self.stream, 4, "is_weights",
)?;
// Actions: GpuTensor<bf16> → i32 buf.

6
crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
View File

@@ -13,7 +13,7 @@
extern "C" __global__ void mse_grad_kernel(
const __nv_bfloat16* __restrict__ save_probs, // [B, 3, NA] softmax probs
const __nv_bfloat16* __restrict__ save_eq_td, // [B, 3, NA] layout: [td_error, E_Q, 0, ...]
const __nv_bfloat16* __restrict__ is_weights, // [B] bf16
const float* __restrict__ is_weights, // [B] f32 (bf16 overflows to Inf)
const int* __restrict__ actions, // [B] factored
float* __restrict__ d_value_logits, // [B, NA] f32 (native atomicAdd, no overflow)
float* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA] f32
@@ -31,8 +31,8 @@ extern "C" __global__ void mse_grad_kernel(
int d = (tid / num_atoms) % 3;
int b = tid / (3 * num_atoms);
/* Read all inputs as float (BF16 → float at boundary) */
float isw = (float)is_weights[b];
/* Read all inputs as float (is_weights already f32, rest BF16 → float at boundary) */
float isw = is_weights[b];
float p_j = (float)save_probs[tid];
int base = (b * 3 + d) * num_atoms;

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

@@ -128,7 +128,7 @@ extern "C" __global__ void mse_loss_batched(
const int* __restrict__ actions, /* [B] factored action indices 0-44 */
const __nv_bfloat16* __restrict__ rewards, /* [B] */
const __nv_bfloat16* __restrict__ dones, /* [B] */
const __nv_bfloat16* __restrict__ is_weights, /* [B] PER importance-sampling weights (bf16) */
const float* __restrict__ is_weights, /* [B] PER importance-sampling weights (f32 — bf16 overflows to Inf) */
/* ── Outputs ──────────────────────────────────────────────────── */
__nv_bfloat16* __restrict__ per_sample_loss, /* [B] IS-weighted loss per sample */
@@ -217,10 +217,11 @@ 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 as float (BF16 → float at boundary).
* IS-weights are f32 — no overflow risk. */
float reward = (float)rewards[sample_id];
float done = (float)dones[sample_id];
float is_weight = (float)is_weights[sample_id];
float is_weight = is_weights[sample_id];
float total_mse = 0.0f;
float total_abs_td = 0.0f;
@@ -354,7 +355,12 @@ extern "C" __global__ void mse_loss_batched(
if (tid == 0) {
float weighted_loss = avg_mse * is_weight;
/* f32 d_logits: no NaN risk from atomicAdd overflow. */
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);

14
crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs
View File

@@ -86,14 +86,6 @@ fn dtoh_f32(slice: &cudarc::driver::CudaSlice<f32>, stream: &Arc<cudarc::driver:
Ok(host)
}
/// Download a CudaSlice<u16> (bf16) to host as f32 for assertion.
fn dtoh_bf16_as_f32(slice: &cudarc::driver::CudaSlice<u16>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<f32>> {
let n = slice.len();
let mut host_u16 = vec![0_u16; n];
stream.memcpy_dtoh(slice, &mut host_u16).map_err(|e| anyhow::anyhow!("dtoh u16: {e}"))?; // test-only readback
Ok(host_u16.iter().map(|&bits| half::bf16::from_bits(bits).to_f32()).collect())
}
/// Download a CudaSlice<u32> to host for assertion.
fn dtoh_u32(slice: &cudarc::driver::CudaSlice<u32>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<u32>> {
let n = slice.len();
@@ -127,7 +119,7 @@ fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> {
assert_eq!(batch.batch_size, 16);
// Download weights and indices to host for assertions
let weights_host = dtoh_bf16_as_f32(&batch.weights, &stream)?;
let weights_host = dtoh_f32(&batch.weights, &stream)?;
let indices_host = dtoh_u32(&batch.indices, &stream)?;
// All weights must be positive
@@ -163,7 +155,7 @@ fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> {
assert_eq!(batch.batch_size, 16);
// Download weights and indices to host for assertions
let weights_host = dtoh_bf16_as_f32(&batch.weights, &stream)?;
let weights_host = dtoh_f32(&batch.weights, &stream)?;
let indices_host = dtoh_u32(&batch.indices, &stream)?;
// All weights must be positive
@@ -202,7 +194,7 @@ fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> {
// Verify sampling still works with updated priorities (proves update succeeded)
let batch2 = buf.sample_proportional(16)?;
let weights_host = dtoh_bf16_as_f32(&batch2.weights, &stream)?;
let weights_host = dtoh_f32(&batch2.weights, &stream)?;
// Weights should still be positive after priority update
for (i, &w) in weights_host.iter().enumerate() {

7
crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
View File

@@ -116,10 +116,9 @@ fn test_per_weights_valid() -> anyhow::Result<()> {
let batch = buf.sample_proportional(16)?;
// Download bf16 weights to host as f32 for assertion
let mut weights_u16 = vec![0_u16; batch.weights.len()];
stream.memcpy_dtoh(&batch.weights, &mut weights_u16).map_err(|e| anyhow::anyhow!("{e}"))?; // test-only readback
let weights_host: Vec<f32> = weights_u16.iter().map(|&bits| half::bf16::from_bits(bits).to_f32()).collect();
// Download f32 weights to host for assertion
let mut weights_host = vec![0.0_f32; batch.weights.len()];
stream.memcpy_dtoh(&batch.weights, &mut weights_host).map_err(|e| anyhow::anyhow!("{e}"))?; // test-only readback
for (i, &w) in weights_host.iter().enumerate() {
assert!(w > 0.0, "weight[{i}] not positive: {w}");