diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 2b6dc7ff9..8d241817c 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -30,7 +30,7 @@ pub struct GpuBatchSlices { pub actions: CudaSlice, // [batch_size] u32 on GPU pub rewards: CudaSlice, // [batch_size] bf16 on GPU pub dones: CudaSlice, // [batch_size] bf16 on GPU (0.0/1.0) - pub weights: CudaSlice, // [batch_size] bf16 on GPU (IS weights — bf16 is fine now that PER indices use u32) + pub weights: CudaSlice, // [batch_size] f32 on GPU (IS weights — f32 to avoid bf16 overflow → Inf → NaN) pub indices: CudaSlice, // [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, weights: CudaSlice)` on GPU (weights are bf16). + /// Returns `(indices: CudaSlice, weights: CudaSlice)` on GPU (weights are f32). /// Callers process indices on GPU -- zero CPU download. - pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice, CudaSlice), MLError> { + pub fn sample_indices_gpu(&mut self, bs: usize) -> Result<(CudaSlice, CudaSlice), MLError> { let b = self.sample_proportional(bs)?; Ok((b.indices, b.weights)) } @@ -1110,7 +1095,6 @@ fn a32i(s: &Arc, n: usize, nm: &str) -> Result, 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, src: &CudaSlice, n: usize, nm: &str) -> Result, MLError> { let mut dst = a32f(s, n, nm)?; let sv = src.slice(..n); diff --git a/crates/ml-dqn/src/regime_conditional.rs b/crates/ml-dqn/src/regime_conditional.rs index 1b2beb5c3..b2e1085a8 100644 --- a/crates/ml-dqn/src/regime_conditional.rs +++ b/crates/ml-dqn/src/regime_conditional.rs @@ -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, 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::(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, 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::(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)?, diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index 64eac647d..dd498d09d 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -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 ─────────────────────────── diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index b7fd82a02..5a4b01690 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -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, // [batch_size] f32 on GPU (IS weights — f32 to avoid bf16 overflow → Inf → NaN) pub indices: cudarc::driver::CudaSlice, // [batch_size] u32 on GPU (buffer indices) /// Episode IDs per transition `[batch_size]` i32 on GPU. /// diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu index 99b31f405..515e707c7 100644 --- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -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]; diff --git a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu index 69f09fcbc..5bc2f9c7c 100644 --- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index bb36424b2..871967c5a 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1a4c4e57a..2c06ce10c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -443,7 +443,7 @@ pub struct GpuDqnTrainer { actions_buf: CudaSlice, // [B] rewards_buf: CudaSlice, // [B] dones_buf: CudaSlice, // [B] - is_weights_buf: CudaSlice, // [B] bf16 (PER indices now u32, no more overflow) + is_weights_buf: CudaSlice, // [B] f32 (bf16 overflows to Inf for PER weights) // ── Activation save buffers (forward → backward) ──────────────── save_h_s1: CudaSlice, // [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, /// Total element count of the staging buffer. upload_staging_len: usize, @@ -529,7 +529,7 @@ pub struct GpuDqnTrainer { // Pre-allocated CudaSlice 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, // [B * STATE_DIM] bf16_next_states_buf: CudaSlice, // [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::(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::(); // ── 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 (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::(); + let f32_size = std::mem::size_of::(); // 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 → i32 buf. diff --git a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu index 750b01810..04f47feb7 100644 --- a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu index ef6479350..68451217e 100644 --- a/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mse_loss_kernel.cu @@ -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); diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index 483f312e1..8eeca89b6 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -86,14 +86,6 @@ fn dtoh_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc (bf16) to host as f32 for assertion. -fn dtoh_bf16_as_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { - 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 to host for assertion. fn dtoh_u32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { 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() { diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index b01e9e79b..6dccef260 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -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 = 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}");