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

feat(bf16): mixed-precision kernels, f32 IS-weights, CUTLASS padding, fast_isnan

Browse Source 
Mixed-precision loss/grad kernels:
- MSE + C51 loss: float softmax/projection/TD-error (prevents bf16 exp overflow)
- MSE + C51 grad: float arithmetic + bf16 range clamp before atomicAdd
- Shared memory: float (4 bytes/elem) for numerically stable reductions
- Bias kernels: float add+clamp ±500 (prevents bf16 Inf cascade between layers)
- Noisy bias kernel: same float clamping

fast_isnan/fast_isinf (ROOT CAUSE FIX):
- nvcc --use_fast_math implies --no-nans → isnan()/isinf() compiled to false
- ALL NaN guards across ALL kernels were dead code
- Added bit-pattern IEEE 754 checks to common_device_functions.cuh
- Replaced isnan/isinf in 7 kernel files (21 occurrences)
- ml-dqn build.rs: all kernels now get common header (no more standalone)

f32 PER IS-weights:
- GpuBatchSlices.weights: CudaSlice<u16> → CudaSlice<f32>
- GpuBatch.weights: GpuTensor → CudaSlice<f32>
- Loss/grad kernel signatures: const __nv_bfloat16* → const float*
- Upload path: separate f32 memcpy instead of bf16 staging
- Eliminates bf16 overflow in IS-weight storage

CUTLASS padding:
- pad32() helper: round up to next multiple of 32
- 6 value-logit buffers: pad32(num_atoms) (51 → 64)
- 6 branch-logit buffers: +32*3 padding per branch

895/895 unit tests, 8/9 smoke tests pass.
50-epoch convergence: NaN at step ~100-200 — backward pass produces NaN
gradients within the CUDA graph replay (same atomic execution as Adam).
Root cause: bf16 backward GemmEx inputs can overflow. Needs mixed-precision
backward pass (same pattern as loss kernels) or f32 gradient output buffers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-28 21:40:09 +01:00
parent 32c1084955
commit 75f83888ca
18 changed files with 228 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -39,22 +39,18 @@ 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()));
// Kernels that need the common BF16 header prepended
// All kernels get the common header (BF16 types + fast_isnan/fast_isinf)
let bf16_kernels = [
"rmsnorm_kernels.cu",
"noisy_kernels.cu",
"residual_kernels.cu",
"cast_kernels.cu",
];
// Kernels that are standalone (already have their own types, no BF16 header needed,
// but cast_kernels uses __nv_bfloat16 so it needs the header via #include)
// replay_buffer_kernels.cu and seg_tree_kernel.cu are pure f32/u32 kernels
let standalone_kernels = [
"replay_buffer_kernels.cu",
"seg_tree_kernel.cu",
];
let standalone_kernels: [&str; 0] = [];
let mut failed: Vec<&str> = Vec::new();
for kernel_name in &bf16_kernels {

29
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)
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.
@@ -71,7 +71,9 @@ impl GpuBatchSlices {
// 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)?;
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,
@@ -906,20 +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")?;
// Cast f32 IS weights -> bf16 (u16) for output. IS computation stays f32 for precision.
let weights_bf16 = {
let wt_slice = self.sample_weights.slice(..batch_size);
let mut out = a16(&self.stream, batch_size, "o_w_bf16")?;
let n_i = batch_size as i32;
// SAFETY: out, wt_slice are valid device allocations of at least batch_size elements.
unsafe {
self.stream.launch_builder(&self.kernels.f32_to_bf16_cast)
.arg(&mut out).arg(&wt_slice).arg(&n_i)
.launch(lcfg(batch_size))
.map_err(|e| MLError::ModelError(format!("cast w bf16: {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")?,
@@ -927,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,
@@ -1073,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))
}
@@ -1105,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);

11
crates/ml-dqn/src/noisy_kernels.cu
View File

@@ -1,7 +1,12 @@
// Noisy linear bias-add kernel -- BF16-native.
// Noisy linear bias-add kernel -- mixed precision (float arithmetic, bf16 I/O).
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: __nv_bfloat16, bf16(), etc.
//
// BF16 addition can overflow to Inf when GemmEx output is near bf16 max.
// Float arithmetic with clamping prevents Inf → NaN propagation.
#define BF16_SAFE_MAX 500.0f
extern "C" __global__
void noisy_add_bias_kernel(__nv_bfloat16* __restrict__ y,
@@ -10,6 +15,8 @@ void noisy_add_bias_kernel(__nv_bfloat16* __restrict__ y,
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < rows * cols) {
int col = idx % cols;
y[idx] = y[idx] + bias[col];
float val = (float)y[idx] + (float)bias[col];
val = fminf(fmaxf(val, -BF16_SAFE_MAX), BF16_SAFE_MAX);
y[idx] = bf16(val);
}
}

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),
] {
let masked_weights = gpu_batch.weights.mul(mask, &self.stream).map_err(|e| {
MLError::TrainingError(format!("Regime weight mask mul: {e}"))
})?;
// 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),
] {
let masked_weights = gpu_batch.weights.mul(mask, &self.stream).map_err(|e| {
MLError::TrainingError(format!("Regime weight mask mul: {e}"))
})?;
// 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)?,

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

@@ -353,7 +353,7 @@ void max_of_two_f32(
}
// ── 18. NaN/Inf check kernel ────────────────────────────────────────────────
// out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f
// out[i] = (fast_isnan(v) || fast_isinf(v)) ? 1.0f : 0.0f
extern "C" __global__
void nan_inf_check_f32(
@@ -364,7 +364,7 @@ void nan_inf_check_f32(
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
float v = input[i];
out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f;
out[i] = (fast_isnan(v) || fast_isinf(v)) ? 1.0f : 0.0f;
}
// ── 19. Dead neuron (near-zero) check kernel ────────────────────────────────

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

48
crates/ml/src/cuda_pipeline/bias_kernels.cu
View File

@@ -1,13 +1,22 @@
/**
* Bias-add kernels for the batched cuBLAS forward pass.
*
* Two kernels:
* 1. add_bias_relu_kernel — fused bias-add + ReLU for hidden layers
* 2. add_bias_kernel — bias-add only for output layers (no activation)
* Mixed-precision: reads BF16, computes bias+activation in float, writes BF16.
* BF16 addition can overflow to Inf when GemmEx output is near bf16 max (~65504).
* Float arithmetic with clamping prevents Inf → NaN propagation to the next layer.
*
* Launch config: grid=(ceil(total_elements/256), 1, 1), block=(256, 1, 1).
*/
/* BF16 max safe value for clamping (below Inf threshold) */
/* Activation clamp: prevents bf16 overflow cascade between layers.
* cuBLAS GemmEx writes bf16 output — if f32 accumulator sum > 65504,
* bf16 writes Inf. With hidden_dim=128 and prev_activation=65000,
* even small weights produce f32 sums > 65504 → Inf → NaN cascade.
* Clamp to 500: well within bf16 precision, 10× Q-clip range (±50),
* prevents multi-layer overflow while preserving training dynamics. */
#define BF16_SAFE_MAX 500.0f
extern "C" __global__ void add_bias_relu_kernel(
__nv_bfloat16* __restrict__ output,
const __nv_bfloat16* __restrict__ bias,
@@ -16,8 +25,9 @@ extern "C" __global__ void add_bias_relu_kernel(
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= total_elements) return;
__nv_bfloat16 val = output[i] + bias[i % out_dim];
output[i] = (val > bf16_zero()) ? val : bf16_zero();
float val = (float)output[i] + (float)bias[i % out_dim];
val = fminf(fmaxf(val, -BF16_SAFE_MAX), BF16_SAFE_MAX);
output[i] = (val > 0.0f) ? bf16(val) : bf16_zero();
}
extern "C" __global__ void add_bias_kernel(
@@ -28,20 +38,15 @@ extern "C" __global__ void add_bias_kernel(
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= total_elements) return;
output[i] = output[i] + bias[i % out_dim];
float val = (float)output[i] + (float)bias[i % out_dim];
val = fminf(fmaxf(val, -BF16_SAFE_MAX), BF16_SAFE_MAX);
output[i] = bf16(val);
}
/* ------------------------------------------------------------------ */
/* BF16 bias+activation kernels for cublasGemmEx tensor core path */
/* ------------------------------------------------------------------ */
/**
* Fused bias-add + ReLU for BF16 hidden layers.
* Both output[] and bias[] are __nv_bfloat16.
* Native BF16 arithmetic — no F32 intermediates.
*
* Launch config: grid=(ceil(total_elements/256), 1, 1), block=(256, 1, 1).
*/
extern "C" __global__ void add_bias_relu_bf16_kernel(
__nv_bfloat16* __restrict__ output,
const __nv_bfloat16* __restrict__ bias,
@@ -50,17 +55,11 @@ extern "C" __global__ void add_bias_relu_bf16_kernel(
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= total_elements) return;
__nv_bfloat16 val = output[i] + bias[i % out_dim];
output[i] = (val > bf16_zero()) ? val : bf16_zero();
float val = (float)output[i] + (float)bias[i % out_dim];
val = fminf(fmaxf(val, -BF16_SAFE_MAX), BF16_SAFE_MAX);
output[i] = (val > 0.0f) ? bf16(val) : bf16_zero();
}
/**
* Bias-add only (no activation) for BF16 output layers.
* Both output[] and bias[] are __nv_bfloat16.
* Native BF16 arithmetic — no F32 intermediates.
*
* Launch config: grid=(ceil(total_elements/256), 1, 1), block=(256, 1, 1).
*/
extern "C" __global__ void add_bias_bf16_kernel(
__nv_bfloat16* __restrict__ output,
const __nv_bfloat16* __restrict__ bias,
@@ -69,6 +68,7 @@ extern "C" __global__ void add_bias_bf16_kernel(
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= total_elements) return;
__nv_bfloat16 val = output[i] + bias[i % out_dim];
output[i] = val;
float val = (float)output[i] + (float)bias[i % out_dim];
val = fminf(fmaxf(val, -BF16_SAFE_MAX), BF16_SAFE_MAX);
output[i] = bf16(val);
}

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

@@ -1,22 +1,20 @@
/**
* C51 distributional RL loss gradient kernel.
*
* Computes dL/d_logits from save_current_lp and save_projected, then routes
* through the dueling architecture to produce d_value_logits and d_adv_logits.
* Mixed-precision: reads BF16, computes in float, writes BF16.
* Prevents NaN from bf16 exp() overflow and intermediate product overflow.
*
* dL/d_combined[b,d,j] = is_weights[b] * (exp(current_lp[b,d,j]) - projected[b,d,j])
* d_value[b,j] = sum_d dL/d_combined[b,d,j]
* d_adv[b,d,a,j] = dL/d_combined[b,d,j] * (delta(a,a_d) - 1/A_d)
*
* entropy_coeff is passed as a runtime parameter instead of a compile-time #define.
*
* Launch config: grid=(ceil(batch_size*3*num_atoms/256), 1, 1), block=(256, 1, 1).
*/
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]
const float* __restrict__ is_weights, // [B] f32 (bf16 overflows to Inf)
const int* __restrict__ actions, // [B] factored
__nv_bfloat16* __restrict__ d_value_logits, // [B, NA]
__nv_bfloat16* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA]
@@ -34,26 +32,29 @@ extern "C" __global__ void c51_grad_kernel(
int d = (tid / num_atoms) % 3;
int b = tid / (3 * num_atoms);
__nv_bfloat16 isw = is_weights[b];
__nv_bfloat16 lp = current_lp[tid];
__nv_bfloat16 proj = projected[tid];
/* 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];
/* Cross-entropy gradient: d/d_logits(-Sigma proj * lp) = exp(lp) - proj */
__nv_bfloat16 d_combined_bf = isw * (bf16_exp(lp) - proj);
/* Cross-entropy gradient: d/d_logits(-Sigma proj * lp) = exp(lp) - proj
* Float exp() handles full range — no bf16 overflow */
float d_combined = isw * (expf(lp) - proj);
/* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp)
* where H(p) = -Sigma p * log(p) and p = exp(lp) (log-probs).
* CRITICAL: clamp lp to [-10, 0] to prevent gradient explosion. */
__nv_bfloat16 entropy_bf = bf16(entropy_coeff);
/* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp) */
if (entropy_coeff > 0.0f) {
__nv_bfloat16 lp_clamped = bf16_fmax(lp, bf16(-10.0f)); /* floor: p >= exp(-10) ~ 4.5e-5 */
d_combined_bf = d_combined_bf + entropy_bf * (bf16_one() + lp_clamped);
float lp_clamped = fmaxf(lp, -10.0f);
d_combined += entropy_coeff * (1.0f + lp_clamped);
}
// Route through dueling: d_value[b,j] += d_combined
atomicAddBF16(&d_value_logits[b * num_atoms + j], (float)d_combined_bf);
/* Clamp to bf16 representable range before atomicAdd */
d_combined = fminf(fmaxf(d_combined, -65000.0f), 65000.0f);
if (fast_isnan(d_combined) || fast_isinf(d_combined)) d_combined = 0.0f;
// Factored action decode
/* Route through dueling: d_value[b,j] += d_combined */
atomicAddBF16(&d_value_logits[b * num_atoms + j], d_combined);
/* Factored action decode */
int factored = actions[b];
int max_action = b0_size * b1_size * b2_size;
if (factored < 0 || factored >= max_action) factored = 0;
@@ -68,17 +69,17 @@ extern "C" __global__ void c51_grad_kernel(
else a_d = factored % b2_size;
int A_d = branch_sizes[d];
__nv_bfloat16 inv_A = bf16_one() / bf16((float)A_d);
float inv_A = 1.0f / (float)A_d;
// Compute per-branch adv offset
int branch_offset = 0;
for (int dd = 0; dd < d; dd++)
branch_offset += branch_sizes[dd] * num_atoms;
// d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d)
/* d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d) */
for (int a = 0; a < A_d; a++) {
__nv_bfloat16 dueling_grad = (a == a_d) ? (bf16_one() - inv_A) : (bf16_zero() - inv_A);
float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
float grad_val = d_combined * dueling_grad;
int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
atomicAddBF16(&d_adv_logits[adv_idx], (float)(d_combined_bf * dueling_grad));
atomicAddBF16(&d_adv_logits[adv_idx], grad_val);
}
}

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 __nv_bfloat16* tg_val_row = tg_value_logits + (long long)sample_id * num_atoms;
const __nv_bfloat16* 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;

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

@@ -13,6 +13,22 @@
* 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. */

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

@@ -36,7 +36,7 @@ extern "C" __global__ void dqn_grad_norm_kernel(
* training correctness (Adam's clip_scale uses fmaxf which handles NaN→0), but
* NaN + valid = NaN poisons the norm sum. Skip them for an accurate norm. */
float g_f = (idx < total_params) ? (float)grads[idx] : 0.0f;
if (isnan(g_f) || isinf(g_f)) g_f = 0.0f;
if (fast_isnan(g_f) || fast_isinf(g_f)) g_f = 0.0f;
float g2 = g_f * g_f;
/* Warp-level reduction via shuffle (float) */
@@ -107,7 +107,7 @@ extern "C" __global__ void dqn_adam_update_kernel(
* Applying NaN to a weight corrupts it permanently. Skipping preserves
* the weight and its Adam moments — equivalent to zero gradient for this step. */
float g_check = (float)g;
if (isnan(g_check) || isinf(g_check)) return;
if (fast_isnan(g_check) || fast_isinf(g_check)) return;
/* Clip using the COMPLETED sum-of-squares (native float buffer).
* dqn_grad_norm_kernel accumulates float partial sums via atomicAdd. */

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

@@ -685,7 +685,7 @@ extern "C" __global__ void experience_env_step(
}
/* Final NaN/Inf guard on target_position. */
if (isnan(target_position) || isinf(target_position)) {
if (fast_isnan(target_position) || fast_isinf(target_position)) {
target_position = 0.0f;
}
@@ -993,7 +993,7 @@ 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 (isnan(reward) || isinf(reward)) reward = 0.0f;
if (fast_isnan(reward) || fast_isinf(reward)) reward = 0.0f;
out_rewards[out_off] = __float2bfloat16(reward);
out_dones[out_off] = __float2bfloat16((float)done);
@@ -1005,7 +1005,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 (isnan(portfolio_return) || isinf(portfolio_return)) portfolio_return = 0.0f;
if (fast_isnan(portfolio_return) || fast_isinf(portfolio_return)) portfolio_return = 0.0f;
raw_returns_out[out_off] = __float2bfloat16(portfolio_return);
}

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

@@ -433,7 +433,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]
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]
@@ -481,9 +481,9 @@ pub struct GpuDqnTrainer {
graph_adam: Option<SendSyncGraph>,
// ── Consolidated transfer buffers ─────────────────────────────
/// Single staging buffer for batch upload consolidation.
/// Layout: [states(B*SD) | next_states(B*SD) | actions_as_f32(B) | rewards(B) | dones(B) | is_weights(B)]
/// One HtoD transfer replaces 6 separate PCIe round-trips.
/// Single staging buffer for batch upload consolidation (bf16 data only).
/// 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,
@@ -1748,7 +1748,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;
@@ -1789,8 +1790,9 @@ impl GpuDqnTrainer {
let t_buf = alloc_i32(&stream, 1, "adam_t")?;
// ── Allocate consolidated transfer buffers ─────────────────
// Upload staging: states + next_states + actions(as f32) + rewards + dones + is_weights
let upload_staging_len = b * config.state_dim * 2 + b * 3; // 2*B*SD + 3*B (no actions — uploaded separately as i32)
// 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)
@@ -1821,20 +1823,20 @@ impl GpuDqnTrainer {
let tg_h_s2_buf = alloc_bf16(&stream, b * config.shared_h2, "tg_h_s2")?;
let tg_h_v_scratch = alloc_bf16(&stream, b * config.value_h, "tg_h_v")?;
let tg_h_b_scratch = alloc_bf16(&stream, b * config.adv_h, "tg_h_b")?;
let tg_v_logits_buf = alloc_bf16(&stream, b * config.num_atoms, "tg_v_logits")?;
let tg_v_logits_buf = alloc_bf16(&stream, b * pad32(config.num_atoms), "tg_v_logits")?;
let total_branch_atoms =
config.branch_0_size * config.num_atoms
+ config.branch_1_size * config.num_atoms
+ config.branch_2_size * config.num_atoms;
let tg_b_logits_buf = alloc_bf16(&stream, b * total_branch_atoms, "tg_b_logits")?;
let tg_b_logits_buf = alloc_bf16(&stream, b * total_branch_atoms + 32 * 3, "tg_b_logits")?;
// ── cuBLAS online logit buffers ────────────────────────────
let on_v_logits_buf = alloc_bf16(&stream, b * config.num_atoms, "on_v_logits")?;
let on_b_logits_buf = alloc_bf16(&stream, b * total_branch_atoms, "on_b_logits")?;
let on_v_logits_buf = alloc_bf16(&stream, b * pad32(config.num_atoms), "on_v_logits")?;
let on_b_logits_buf = alloc_bf16(&stream, b * total_branch_atoms + 32 * 3, "on_b_logits")?;
// ── cuBLAS online-on-next-states buffers (Double DQN action selector) ──
let on_next_v_logits_buf = alloc_bf16(&stream, b * config.num_atoms, "on_next_v_logits")?;
let on_next_b_logits_buf = alloc_bf16(&stream, b * total_branch_atoms, "on_next_b_logits")?;
let on_next_v_logits_buf = alloc_bf16(&stream, b * pad32(config.num_atoms), "on_next_v_logits")?;
let on_next_b_logits_buf = alloc_bf16(&stream, b * total_branch_atoms + 32 * 3, "on_next_b_logits")?;
let on_next_h_s1_scratch = alloc_bf16(&stream, b * config.shared_h1, "on_next_h_s1")?;
let on_next_h_s2_scratch = alloc_bf16(&stream, b * config.shared_h2, "on_next_h_s2")?;
let on_next_h_v_scratch = alloc_bf16(&stream, b * config.value_h, "on_next_h_v")?;
@@ -1871,15 +1873,15 @@ impl GpuDqnTrainer {
} else {
None
};
let cql_d_value_logits = alloc_bf16(&stream, b * config.num_atoms, "cql_d_value_logits")?;
let cql_d_adv_logits = alloc_bf16(&stream, b * total_branch_atoms, "cql_d_adv_logits")?;
let cql_d_value_logits = alloc_bf16(&stream, b * pad32(config.num_atoms), "cql_d_value_logits")?;
let cql_d_adv_logits = alloc_bf16(&stream, b * total_branch_atoms + 32 * 3, "cql_d_adv_logits")?;
// ── Gradient output buffers for cuBLAS backward ──────────────
let d_value_logits_buf = alloc_bf16(&stream, b * config.num_atoms, "d_value_logits")?;
let d_adv_logits_buf = alloc_bf16(&stream, b * total_branch_atoms, "d_adv_logits")?;
let d_value_logits_buf = alloc_bf16(&stream, b * pad32(config.num_atoms), "d_value_logits")?;
let d_adv_logits_buf = alloc_bf16(&stream, b * total_branch_atoms + 32 * 3, "d_adv_logits")?;
// Scratch buffers for blended MSE+C51 loss (MSE grad stored here, then blended)
let d_value_logits_mse = alloc_bf16(&stream, b * config.num_atoms, "d_value_logits_mse")?;
let d_adv_logits_mse = alloc_bf16(&stream, b * total_branch_atoms, "d_adv_logits_mse")?;
let d_value_logits_mse = alloc_bf16(&stream, b * pad32(config.num_atoms), "d_value_logits_mse")?;
let d_adv_logits_mse = alloc_bf16(&stream, b * total_branch_atoms + 32 * 3, "d_adv_logits_mse")?;
// ── Spectral normalization singular vectors ─────────────────
// Initialize with random unit vectors for proper power iteration convergence.
@@ -3110,12 +3112,13 @@ impl GpuDqnTrainer {
/// Upload batch data to pre-allocated GPU buffers via consolidated staging.
///
/// Packs all 6 arrays into a single contiguous host buffer, performs one
/// 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 1,
/// saving ~12-30us per batch on H100 PCIe Gen5.
/// 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],
@@ -3129,13 +3132,13 @@ impl GpuDqnTrainer {
let sd = self.config.state_dim;
let bf16_size = std::mem::size_of::<half::bf16>();
// ── Pack f32 data (no actions) into staging buffer → convert to BF16 ──
// ── Pack bf16-convertible data into staging buffer → convert to BF16 ──
// 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)?;
@@ -3162,16 +3165,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 is_weights_bytes = b * bf16_size;
dtod_copy(self.is_weights_buf.raw_ptr(), staging_base + byte_offset, is_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(())
}
@@ -3179,13 +3181,10 @@ impl GpuDqnTrainer {
///
/// GpuBatch layout from GpuReplayBuffer:
/// - states/next_states: BF16 (stored in BF16 ring buffer)
/// - rewards/dones/weights: F32 (stored in F32 ring buffer / computed as F32)
/// - rewards/dones: BF16 GpuTensor
/// - weights: CudaSlice<f32> (IS-weights kept as f32 to avoid bf16 overflow → Inf → NaN)
/// - actions: U32 (stored in U32 ring buffer)
///
/// For BF16 tensors (states, next_states): DtoD → BF16 staging → bf16_to_f32 kernel → F32 buf.
/// For F32 tensors (rewards, dones, weights): DtoD → F32 buf directly.
/// For U32 actions: DtoD → I32 buf (bit-compatible for values in [0, 44]).
///
/// No Candle `to_dtype()` temporaries — eliminates Drop/event conflicts with forked stream.
fn upload_batch_gpu(
&mut self,
@@ -3194,10 +3193,8 @@ impl GpuDqnTrainer {
let b = self.config.batch_size;
let sd = self.config.state_dim;
// GpuBatch stores GpuTensor (CudaSlice<half::bf16>) for states/next_states/rewards/dones/weights.
// Actions are CudaSlice<u32> — separate from GpuTensor.
// All data is already BF16 on GPU — direct DtoD copy, no conversion.
let bf16_size = std::mem::size_of::<half::bf16>();
let f32_size = std::mem::size_of::<f32>();
// States + next_states: bf16 DtoD
dtod_copy(
@@ -3211,7 +3208,7 @@ impl GpuDqnTrainer {
b * sd * bf16_size, &self.stream, 1, "next_states",
)?;
// Rewards, dones, IS-weights: bf16 DtoD
// Rewards, dones: bf16 DtoD
dtod_copy(
self.rewards_buf.raw_ptr(),
gpu_batch.rewards.data().raw_ptr(),
@@ -3222,10 +3219,12 @@ impl GpuDqnTrainer {
gpu_batch.dones.data().raw_ptr(),
b * bf16_size, &self.stream, 3, "dones",
)?;
// 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.
@@ -4285,6 +4284,18 @@ fn dtod_from_u32_to_i32(
// ── Allocation helpers ──────────────────────────────────────────────────────
/// Round up to the next multiple of 32.
///
/// cuBLAS CUTLASS kernels (e.g. `cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_128x2_tn_align8`)
/// read in 32-element M-tiles. When the M dimension (num_atoms = 51) is not a
/// multiple of 32 the kernel reads past the logical buffer end. Padding the
/// **allocation** to a 32-element boundary prevents OOB reads without changing
/// the GemmEx logical dimensions or any kernel launch parameters.
#[inline(always)]
fn pad32(n: usize) -> usize {
(n + 31) & !31
}
fn alloc_bf16(
stream: &Arc<CudaStream>,
n: usize,

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

@@ -1,14 +1,11 @@
/**
* MSE loss gradient kernel through softmax expectation.
*
* Mixed-precision: reads BF16, computes in float, writes BF16.
* Prevents NaN from bf16 intermediate product overflow.
*
* For each sample [b], branch [d], atom [j]:
* d_logit_j = td_error * is_weight * p_j * (z_j - E[Q])
* where p_j = softmax prob (from save_current_lp), z_j = support value,
* E[Q] = expected Q (from save_projected), td_error (from save_projected).
*
* v_min and v_max are passed as runtime kernel parameters.
*
* Routes through dueling: same pattern as c51_grad_kernel.
*
* Launch config: grid=(ceil(batch_size*3*num_atoms/256), 1, 1), block=(256, 1, 1).
*/
@@ -16,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]
const float* __restrict__ is_weights, // [B] f32 (bf16 overflows to Inf)
const int* __restrict__ actions, // [B] factored
__nv_bfloat16* __restrict__ d_value_logits, // [B, NA]
__nv_bfloat16* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA]
@@ -34,31 +31,28 @@ extern "C" __global__ void mse_grad_kernel(
int d = (tid / num_atoms) % 3;
int b = tid / (3 * num_atoms);
__nv_bfloat16 isw = is_weights[b];
__nv_bfloat16 p_j = save_probs[tid]; /* softmax prob for atom j */
/* 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];
/* Read td_error and E[Q] from save_eq_td.
* Layout: for each [b, d, :], element [0] = td_error, element [1] = E[Q] */
int base = (b * 3 + d) * num_atoms;
__nv_bfloat16 td_error = save_eq_td[base + 0];
__nv_bfloat16 e_q = save_eq_td[base + 1];
float td_error = (float)save_eq_td[base + 0];
float e_q = (float)save_eq_td[base + 1];
/* C51 support: z_j = v_min + j * delta_z */
__nv_bfloat16 v_min_bf = bf16(v_min);
__nv_bfloat16 v_max_bf = bf16(v_max);
__nv_bfloat16 delta_z = (num_atoms > 1)
? (v_max_bf - v_min_bf) / bf16((float)(num_atoms - 1))
: bf16_zero();
__nv_bfloat16 z_j = v_min_bf + bf16((float)j) * delta_z;
float delta_z = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;
float z_j = v_min + (float)j * delta_z;
/* Gradient of MSE loss through softmax expectation:
* dL/d_logit_j = td * p_j * (z_j - E[Q]) */
__nv_bfloat16 d_combined_bf = isw * td_error * p_j * (z_j - e_q);
/* Gradient of MSE loss through softmax expectation (float arithmetic) */
float d_combined = isw * td_error * p_j * (z_j - e_q);
// Route through dueling: d_value[b,j] += d_combined
atomicAddBF16(&d_value_logits[b * num_atoms + j], (float)d_combined_bf);
/* Clamp to bf16 representable range before atomicAdd */
d_combined = fminf(fmaxf(d_combined, -65000.0f), 65000.0f);
if (fast_isnan(d_combined) || fast_isinf(d_combined)) d_combined = 0.0f;
// Factored action decode
/* Route through dueling: d_value[b,j] += d_combined */
atomicAddBF16(&d_value_logits[b * num_atoms + j], d_combined);
/* Factored action decode */
int factored = actions[b];
int max_action = b0_size * b1_size * b2_size;
if (factored < 0 || factored >= max_action) factored = 0;
@@ -73,17 +67,17 @@ extern "C" __global__ void mse_grad_kernel(
else a_d = factored % b2_size;
int A_d = branch_sizes[d];
__nv_bfloat16 inv_A = bf16_one() / bf16((float)A_d);
float inv_A = 1.0f / (float)A_d;
// Compute per-branch adv offset
int branch_offset = 0;
for (int dd = 0; dd < d; dd++)
branch_offset += branch_sizes[dd] * num_atoms;
// d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d)
/* d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d) */
for (int a = 0; a < A_d; a++) {
__nv_bfloat16 dueling_grad = (a == a_d) ? (bf16_one() - inv_A) : (bf16_zero() - inv_A);
float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
float grad_val = d_combined * dueling_grad;
int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j;
atomicAddBF16(&d_adv_logits[adv_idx], (float)(d_combined_bf * dueling_grad));
atomicAddBF16(&d_adv_logits[adv_idx], grad_val);
}
}

7
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 */
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 __nv_bfloat16* tg_val_row = tg_value_logits + (long long)sample_id * num_atoms;
const __nv_bfloat16* 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;

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

@@ -52,8 +52,8 @@ extern "C" __global__ void training_guard_check_and_accumulate(
/* -- Guard check -- */
int loss_nan = isnan(loss) || isinf(loss);
int grad_nan = isnan(grad_norm) || isinf(grad_norm);
int loss_nan = fast_isnan(loss) || fast_isinf(loss);
int grad_nan = fast_isnan(grad_norm) || fast_isinf(grad_norm);
int halt_nan = (loss_nan || grad_nan) ? 1 : 0;
int halt_loss_clip = (!loss_nan && loss > clip_threshold) ? 1 : 0;
@@ -100,8 +100,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 = isnan(loss) || isinf(loss);
int grad_nan = isnan(grad_norm) || isinf(grad_norm);
int loss_nan = fast_isnan(loss) || fast_isinf(loss);
int grad_nan = fast_isnan(grad_norm) || fast_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;
@@ -126,10 +126,10 @@ extern "C" __global__ void training_guard_accumulate(
) {
float loss = (float)loss_scalar[0];
float grad_norm = (float)grad_norm_scalar[0];
if (!isnan(loss) && !isinf(loss)) {
if (!fast_isnan(loss) && !fast_isinf(loss)) {
acc_buf[0] = acc_buf[0] + bf16(loss);
}
if (!isnan(grad_norm) && !isinf(grad_norm)) {
if (!fast_isnan(grad_norm) && !fast_isinf(grad_norm)) {
acc_buf[1] = acc_buf[1] + bf16(grad_norm);
}
acc_buf[2] = acc_buf[2] + bf16_one();
@@ -178,7 +178,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 (!isnan(qv_f) && !isinf(qv_f)) {
if (!fast_isnan(qv_f) && !fast_isinf(qv_f)) {
sample_max = bf16_fmax(sample_max, qv);
local_all = local_all + qv;
}
@@ -269,7 +269,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 (isnan(qv_f) || isinf(qv_f)) continue;
if (fast_isnan(qv_f) || fast_isinf(qv_f)) continue;
q_min = bf16_fmin(q_min, qv);
q_max = bf16_fmax(q_max, qv);

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 = vec![0_u16; n];
stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh u16: {e}"))?;
Ok(host.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 for assertion
let mut weights_bf16 = vec![0_u16; batch.weights.len()];
stream.memcpy_dtoh(&batch.weights, &mut weights_bf16).map_err(|e| anyhow::anyhow!("{e}"))?; // test-only readback
let weights_host: Vec<f32> = weights_bf16.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}");