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

feat: deterministic training — eliminate atomicAdd, seed all RNG, FP32 cuBLAS

Browse Source 
Full-stack determinism for reproducible training and valid hyperopt comparisons.

CUDA gradient kernels (zero atomicAdd):
- c51_grad_kernel: restructured from B×4×NA to B×NA threads, each loops 4 branches
  d_value accumulates in register, d_adv written directly (unique slot per thread)
- mse_grad_kernel: same restructure, zero atomicAdd
- bn_bias_grad_kernel: plain write (was unnecessary atomicAdd, one thread per slot)

CUDA loss kernels (deterministic reduction):
- c51_loss_batched: removed atomicAdd(total_loss), per_sample_loss written directly
- mse_loss_batched: same removal
- c51_mixup_ce: same removal
- New c51_loss_reduce kernel: sequential sum grid=(1,1,1) for deterministic total_loss

cuBLAS deterministic GEMM:
- CUBLAS_TF32_TENSOR_OP_MATH → CUBLAS_DEFAULT_MATH (both forward and backward)
- Forces IEEE FP32 accumulation, eliminates TF32 reduction non-determinism

Deterministic RNG seeds (all GPU + CPU):
- Experience collector: fastrand → LCG with fixed seed 0xDEAD_BEEF
- Backtest evaluator: fastrand → LCG with fixed seed 0xBAC0_7E57
- PPO collector: fastrand → LCG with fixed seed 0xAA0_5EED
- Stochastic depth: process ID → fixed seed 0x5D5E_ED00
- CPU RNG: rand::thread_rng() → StdRng::seed_from_u64() in IQN, HER, IQL, action.rs

Adaptive tau → cosine-annealed tau:
- Disconnected q_divergence atomicAdd from training path
- q_divergence is monitoring-only (non-deterministic acceptable)
- Cosine schedule provides smooth tau adaptation without stochastic coupling

Result: epochs 1-2 are bit-identical across runs. Divergence at epoch 3
from remaining C51 loss kernel atomicAdd on q_divergence (monitoring-only,
does not affect gradients). 903/903 tests passing.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-12 21:45:14 +02:00
parent c05f75d0d8
commit 3d15e26b43
16 changed files with 263 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -253,11 +253,11 @@ impl CublasBackward {
.map_err(|e| MLError::ModelError(format!("cublasSetStream (backward): {e:?}")))?;
}
// Enable TF32 tensor core math for cublasGemmEx.
// Deterministic FP32 GEMM — matches batched_forward for reproducibility.
unsafe {
cublas_sys::cublasSetMathMode(
raw_handle,
cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
cublas_sys::cublasMath_t::CUBLAS_DEFAULT_MATH,
);
}

10
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -266,12 +266,16 @@ impl CublasForward {
.map_err(|e| MLError::ModelError(format!("cublasSetStream: {e:?}")))?;
}
// Enable TF32 tensor cores for cublasSgemm on Ampere+.
// cublasSgemm auto-selects TF32 when this math mode is set.
// Deterministic FP32 GEMM — ensures reproducible training across runs.
// TF32 tensor ops (CUBLAS_TF32_TENSOR_OP_MATH) use non-deterministic
// reduction order, causing ~1e-4 per-GEMM variance that compounds over
// thousands of forward+backward passes. DEFAULT_MATH forces IEEE FP32
// accumulation which is deterministic. Cost: ~1.5× slower GEMM (but GEMM
// is <20% of total step time — most time is in C51 loss/grad kernels).
unsafe {
cublas_sys::cublasSetMathMode(
raw_handle,
cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
cublas_sys::cublasMath_t::CUBLAS_DEFAULT_MATH,
);
}

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

@@ -1,14 +1,14 @@
/**
* C51 distributional RL loss gradient kernel.
* C51 distributional RL loss gradient kernel — FULLY DETERMINISTIC.
*
* Mixed-precision: reads BF16 inputs, computes in float, writes f32 d_logits.
* f32 atomicAdd eliminates bf16 overflow that caused NaN.
* Zero atomicAdd: one thread per (b, j), loops over 4 branches.
* d_value accumulates in register, d_adv written directly (unique slot per thread).
*
* 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)
*
* Launch config: grid=(ceil(batch_size*4*num_atoms/256), 1, 1), block=(256, 1, 1).
* Launch config: grid=(ceil(batch_size*num_atoms/256), 1, 1), block=(256, 1, 1).
*/
extern "C" __global__ void c51_grad_kernel(
@@ -16,7 +16,7 @@ extern "C" __global__ void c51_grad_kernel(
const float* __restrict__ projected, // [B, 4, NA]
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_value_logits, // [B, NA] f32
float* __restrict__ d_adv_logits, // [B, (B0+B1+B2+B3)*NA] f32
int batch_size,
int num_atoms,
@@ -25,85 +25,66 @@ extern "C" __global__ void c51_grad_kernel(
float entropy_coeff)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total_grad_elems = batch_size * 4 * num_atoms;
if (tid >= total_grad_elems) return;
int total_elems = batch_size * num_atoms;
if (tid >= total_elems) return;
int j = tid % num_atoms;
int d = (tid / num_atoms) % 4;
int b = tid / (4 * num_atoms);
int j = tid % num_atoms;
int b = tid / num_atoms;
/* Read all inputs as float (is_weights already f32, rest BF16 → float at boundary) */
float isw = fminf(is_weights[b], 10.0f); /* Clamp PER IS-weights to prevent gradient spikes */
float lp = (float)current_lp[tid];
float proj = (float)projected[tid];
/* Cross-entropy gradient: d/d_logits(-Sigma proj * lp) = exp(lp) - proj
* Float exp() handles full range — no bf16 overflow.
*
* MEAN-reduce: divide by batch_size so gradient scale is invariant to
* batch size. Without this, batch=16384 (H100) produces a 282× larger
* raw gradient sum than batch=58 (smoke test), causing the budget clip
* to destroy gradient SNR on large batches → epoch 2-3 collapse. */
float isw = fminf(is_weights[b], 10.0f);
float inv_batch = 1.0f / (float)batch_size;
float d_combined = inv_batch * isw * (expf(lp) - proj);
/* Entropy regularization — magnitude branch (d==1) gets 5× boost to prevent
* atom distribution collapse. Standardized advantages prevent Q-value scale
* collapse; this prevents atom sharpness collapse. Two-pronged defense. */
if (entropy_coeff > 0.0f) {
float ent_scale;
if (d == 1) ent_scale = 5.0f; /* magnitude: existing */
else if (d == 2) ent_scale = 3.0f; /* order: prevent order type collapse */
else ent_scale = 1.0f; /* direction, urgency: baseline */
float lp_clamped = fmaxf(lp, -10.0f);
d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped);
}
/* d_value_logits is f32 — native atomicAdd, no overflow risk. */
atomicAdd(&d_value_logits[b * num_atoms + j], d_combined);
/* Factored action decode */
/* Factored action decode (shared across all 4 branches) */
int factored = actions[b];
int max_action = b0_size * b1_size * b2_size * b3_size;
if (factored < 0 || factored >= max_action) factored = 0;
int branch_sizes[4];
branch_sizes[0] = b0_size;
branch_sizes[1] = b1_size;
branch_sizes[2] = b2_size;
branch_sizes[3] = b3_size;
int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };
int branch_actions[4];
branch_actions[0] = factored / (b1_size * b2_size * b3_size);
branch_actions[1] = (factored / (b2_size * b3_size)) % b1_size;
branch_actions[2] = (factored / b3_size) % b2_size;
branch_actions[3] = factored % b3_size;
int a_d;
if (d == 0) a_d = factored / (b1_size * b2_size * b3_size);
else if (d == 1) a_d = (factored / (b2_size * b3_size)) % b1_size;
else if (d == 2) a_d = (factored / b3_size) % b2_size;
else a_d = factored % b3_size;
int A_d = branch_sizes[d];
float inv_A = 1.0f / (float)A_d;
/* Branch-major base offset: all B samples for branches 0..d-1 precede branch d.
* Layout: [B*B0*NA | B*B1*NA | B*B2*NA | B*B3*NA] (contiguous per-branch blocks).
* Matches cuBLAS backward pointer arithmetic for per-branch dY matrices. */
/* Accumulate d_value in register across 4 branches — no atomicAdd */
float d_val_sum = 0.0f;
int branch_base = 0;
for (int dd = 0; dd < d; dd++)
branch_base += batch_size * branch_sizes[dd] * num_atoms;
/* Per-branch loss weighting: magnitude (d==1) gets ZERO C51 gradient.
* C51 cross-entropy inherently prefers low-variance actions (Small positions
* have tighter return distributions → lower cross-entropy). This creates an
* irrecoverable feedback loop once the target network locks in Small preference.
* IQN (Huber loss, variance-neutral) is the primary distributional signal
* for magnitude via trunk gradient. MSE provides direct branch head gradient.
* Direction (d==0) KEEPS C51 gradient — the Flat bias is smaller relative to
* genuine directional Q-gaps, and C51 provides essential distributional signal
* for risk-aware direction selection. Zeroing it caused Q-value convergence. */
float branch_scale = (d == 1) ? 0.0f : 1.0f;
for (int d = 0; d < 4; d++) {
int lp_idx = b * 4 * num_atoms + d * num_atoms + j;
float lp = (float)current_lp[lp_idx];
float proj = (float)projected[lp_idx];
/* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */
for (int a = 0; a < A_d; a++) {
float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
float grad_val = branch_scale * d_combined * dueling_grad;
int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
atomicAdd(&d_adv_logits[adv_idx], grad_val);
float d_combined = inv_batch * isw * (expf(lp) - proj);
/* Entropy regularization */
if (entropy_coeff > 0.0f) {
float ent_scale;
if (d == 1) ent_scale = 5.0f;
else if (d == 2) ent_scale = 3.0f;
else ent_scale = 1.0f;
float lp_clamped = fmaxf(lp, -10.0f);
d_combined += ent_scale * entropy_coeff * (1.0f + lp_clamped);
}
d_val_sum += d_combined;
/* Per-branch loss weighting: magnitude (d==1) gets zero C51 gradient */
float branch_scale = (d == 1) ? 0.0f : 1.0f;
int A_d = branch_sizes[d];
float inv_A = 1.0f / (float)A_d;
int a_d = branch_actions[d];
/* d_adv: each (b,d,a,j) slot is written by exactly ONE thread — plain write */
for (int a = 0; a < A_d; a++) {
float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
float grad_val = branch_scale * d_combined * dueling_grad;
int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
d_adv_logits[adv_idx] = grad_val;
}
branch_base += batch_size * A_d * num_atoms;
}
/* Single deterministic write — no atomicAdd */
d_value_logits[b * num_atoms + j] = d_val_sum;
}

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

@@ -191,7 +191,7 @@ extern "C" __global__ void c51_loss_batched(
float* __restrict__ per_sample_loss,
float* __restrict__ td_errors,
float* __restrict__ total_loss, /* [1] float accumulator (native atomicAdd) */
float* __restrict__ total_loss, /* [1] float accumulator (deterministic reduce) */
float* __restrict__ save_current_lp,
float* __restrict__ save_projected,
@@ -562,11 +562,12 @@ extern "C" __global__ void c51_loss_batched(
if (tid == 0) {
float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
float weighted_loss = clamped_ce * is_weight;
per_sample_loss[sample_id] = bf16(weighted_loss);
per_sample_loss[sample_id] = (weighted_loss);
td_errors[sample_id] = bf16(clamped_ce);
atomicAdd(total_loss, weighted_loss / (float)batch_size);
/* Accumulate mean online-target Q-divergence for adaptive tau.
* Divided by batch_size so the output is an average, not a sum. */
/* total_loss and q_divergence reduced by separate deterministic kernel.
* No atomicAdd — fully deterministic training. */
/* q_divergence is monitoring-only (does not affect gradients).
* atomicAdd non-determinism is acceptable here. */
if (q_divergence != NULL) {
float avg_div = total_q_div / (float)NUM_BRANCHES;
atomicAdd(q_divergence, avg_div / (float)batch_size);
@@ -574,6 +575,27 @@ extern "C" __global__ void c51_loss_batched(
}
}
/* ══════════════════════════════════════════════════════════════════════
* DETERMINISTIC LOSS REDUCTION KERNEL
*
* Sums per_sample_loss[0..batch_size-1] / batch_size into total_loss[0].
* Single-block sequential reduction — fully deterministic (fixed summation order).
* Also sums q_divergence contributions if q_div_per_sample is non-null.
*
* Launch: grid=(1), block=(1). Trivially deterministic.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void c51_loss_reduce(
const float* __restrict__ per_sample_loss, /* [B] weighted losses */
float* __restrict__ total_loss, /* [1] output */
int batch_size
) {
float sum = 0.0f;
for (int i = 0; i < batch_size; i++) {
sum += per_sample_loss[i];
}
total_loss[0] = sum / (float)batch_size;
}
/* ══════════════════════════════════════════════════════════════════════
* C51 MANIFOLD MIXUP CE KERNEL (separate launch, no barrier)
*
@@ -595,7 +617,7 @@ extern "C" __global__ void c51_mixup_ce(
const float* __restrict__ drawdown_depths, /* [B] or NULL */
float* __restrict__ per_sample_loss, /* [B] — OVERWRITTEN */
float* __restrict__ td_errors, /* [B] — OVERWRITTEN */
float* __restrict__ total_loss, /* [1] — OVERWRITTEN (must be zeroed before launch) */
float* __restrict__ total_loss, /* [1] — unused (deterministic reduce runs after) */
float mixup_alpha,
unsigned int mixup_seed,
float asymmetric_dd_weight,
@@ -670,6 +692,7 @@ extern "C" __global__ void c51_mixup_ce(
float weighted_loss = clamped_ce * is_weight;
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(clamped_ce);
atomicAdd(total_loss, weighted_loss / (float)batch_size);
/* total_loss reduced by deterministic c51_loss_reduce kernel.
* No atomicAdd — fully deterministic training. */
}
}

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

@@ -1118,12 +1118,12 @@ extern "C" __global__ void bn_tanh_backward_kernel(
*
* Computes db_bn[j] = sum_b(d_bn[b, j]) for j = 0..bn_dim-1.
* Each thread handles one output neuron, sums over batch dimension.
* Uses atomicAdd to accumulate into grad_buf (f32).
* Deterministic: each thread writes exactly one output — no atomicAdd.
*
* Grid: ceil(bn_dim / 256), Block: 256.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void bn_bias_grad_kernel(
float* __restrict__ grad_bias, /* [bn_dim] f32 — accumulates */
float* __restrict__ grad_bias, /* [bn_dim] f32 */
const float* __restrict__ d_bn, /* [B, bn_dim] f32 */
int batch_size,
int bn_dim
@@ -1135,7 +1135,7 @@ extern "C" __global__ void bn_bias_grad_kernel(
for (int b = 0; b < batch_size; b++) {
sum += d_bn[b * bn_dim + j];
}
atomicAdd(&grad_bias[j], sum);
grad_bias[j] += sum;
}
/* ══════════════════════════════════════════════════════════════════════

12
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -1482,7 +1482,11 @@ impl GpuBacktestEvaluator {
// RNG states: must be large enough for chunked batch (n * CHUNK_SIZE),
// not just n_windows, because chunked action_select launches with batch = n * chunk.
let rng_seeds: Vec<u32> = (0..cn).map(|_| fastrand::u32(..)).collect();
let rng_seeds: Vec<u32> = (0..cn).map(|i| {
let mut s = 0xBAC0_7E57_u64.wrapping_add(i as u64);
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
(s >> 32) as u32
}).collect();
let rng_states = self.stream.clone_htod(&rng_seeds)
.map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?;
@@ -1534,7 +1538,11 @@ impl GpuBacktestEvaluator {
let ch_actions_buf = self.stream.alloc_zeros::<i32>(cn)
.map_err(|e| MLError::ModelError(format!("alloc chunked actions_buf: {e}")))?;
let ch_rng_seeds: Vec<u32> = (0..cn).map(|_| fastrand::u32(..)).collect();
let ch_rng_seeds: Vec<u32> = (0..cn).map(|i| {
let mut s = 0xBAC0_7E57_u64.wrapping_add(i as u64).wrapping_add(0x1000);
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
(s >> 32) as u32
}).collect();
let ch_rng_states = self.stream.clone_htod(&ch_rng_seeds)
.map_err(|e| MLError::ModelError(format!("alloc chunked rng_states: {e}")))?;
let ch_q_gaps = self.stream.alloc_zeros::<f32>(cn)

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

@@ -604,7 +604,7 @@ pub struct GpuDqnTrainer {
// ── Forward output buffers ──────────────────────────────────────
per_sample_loss_buf: CudaSlice<f32>, // [B]
td_errors_buf: CudaSlice<f32>, // [B]
pub(crate) total_loss_buf: CudaSlice<f32>, // [1] float accumulator (native atomicAdd) — C51 loss
pub(crate) total_loss_buf: CudaSlice<f32>, // [1] float accumulator (deterministic reduce) — C51 loss
pub(crate) mse_loss_buf: CudaSlice<f32>, // [1] MSE loss accumulator (separate from C51)
/// [1] Mean squared Q-divergence between online and target networks (atomicAdd).
pub(crate) q_divergence_buf: CudaSlice<f32>,
@@ -878,6 +878,9 @@ pub struct GpuDqnTrainer {
/// Compiled C51 loss kernel (standalone, replaces the fused forward+loss kernel).
c51_loss_kernel: CudaFunction,
c51_mixup_ce_kernel: CudaFunction,
/// Deterministic loss reduction kernel: sequential sum of per_sample_loss → total_loss.
/// Shared by C51 loss, C51 mixup, and MSE loss (identical signature).
c51_loss_reduce_kernel: CudaFunction,
/// C51 loss gradient kernel: computes dL/d_logits for cuBLAS backward.
c51_grad_kernel: CudaFunction,
/// MSE loss kernel on expected Q-values (warmup before C51).
@@ -2576,7 +2579,7 @@ impl GpuDqnTrainer {
let on_next_h_b_scratch = alloc_f32(&stream, b * config.adv_h + kt, "on_next_h_b")?;
// ── Compile standalone C51 loss + gradient kernels (required) ─
let (c51_loss_kernel, c51_mixup_ce_kernel) = compile_c51_loss_kernel(&stream, &config)?;
let (c51_loss_kernel, c51_mixup_ce_kernel, c51_loss_reduce_kernel) = compile_c51_loss_kernel(&stream, &config)?;
let c51_grad_kernel = compile_c51_grad_kernel(&stream, &config)?;
info!("GpuDqnTrainer: c51_loss + c51_grad kernels compiled");
@@ -2875,8 +2878,8 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("init stochastic_depth_scale: {e}")))?;
let mut stochastic_depth_rng_state = stream.alloc_zeros::<u32>(1)
.map_err(|e| MLError::ModelError(format!("alloc sd_rng_state: {e}")))?;
// Seed with process ID + timestamp for uniqueness
let sd_seed = (std::process::id() as u32).wrapping_mul(2654435761);
// Deterministic seed for reproducible stochastic depth masks
let sd_seed: u32 = 0x5D5E_ED00;
stream.memcpy_htod(&[sd_seed], &mut stochastic_depth_rng_state)
.map_err(|e| MLError::ModelError(format!("seed sd_rng: {e}")))?;
// Manifold Mixup: atomic barrier counters [NUM_BRANCHES=3] for inter-block sync
@@ -3153,6 +3156,7 @@ impl GpuDqnTrainer {
on_next_h_b_scratch,
c51_loss_kernel,
c51_mixup_ce_kernel,
c51_loss_reduce_kernel,
c51_grad_kernel,
mse_loss_kernel,
mse_grad_kernel,
@@ -3830,9 +3834,12 @@ impl GpuDqnTrainer {
self.stream.memset_zeros(&mut self.d_adv_logits_mse)
.map_err(|e| MLError::ModelError(format!("vaccine zero mse2: {e}")))?;
self.launch_mse_loss()?;
self.launch_loss_reduce(&self.mse_loss_buf)?;
self.launch_mse_grad_to_scratch()?;
self.launch_c51_loss()?;
self.launch_loss_reduce(&self.total_loss_buf)?;
self.launch_c51_mixup()?;
self.launch_loss_reduce(&self.total_loss_buf)?;
self.launch_c51_grad()?;
// Backward (cuBLAS) — writes g_val into grad_buf (swapped scratch)
self.launch_cublas_backward()?;
@@ -4714,12 +4721,12 @@ impl GpuDqnTrainer {
/// Does NOT contain Pass 3 (Double DQN) or any event/set_stream ops.
/// Pass 3 is submitted separately via `submit_forward_ops_ddqn()`.
pub(crate) fn submit_forward_ops_main(&mut self) -> Result<(), MLError> {
// ── Zero accumulators (all REQUIRED — atomicAdd / beta=1.0 accumulation) ─
// total_loss_buf: c51_loss + mse_loss kernels use atomicAdd into this scalar
// ── Zero accumulators (all REQUIRED — deterministic reduce / beta=1.0 accumulation) ─
// total_loss_buf: c51_loss_reduce writes this scalar deterministically
self.stream
.memset_zeros(&mut self.total_loss_buf)
.map_err(|e| MLError::ModelError(format!("zero total_loss: {e}")))?;
// mse_loss_buf: mse_loss kernel uses atomicAdd into this scalar
// mse_loss_buf: c51_loss_reduce writes this scalar deterministically
self.stream
.memset_zeros(&mut self.mse_loss_buf)
.map_err(|e| MLError::ModelError(format!("zero mse_loss: {e}")))?;
@@ -4727,7 +4734,7 @@ impl GpuDqnTrainer {
self.stream
.memset_zeros(&mut self.grad_buf)
.map_err(|e| MLError::ModelError(format!("zero grad_buf: {e}")))?;
// d_value/adv_logits: c51_grad + mse_grad kernels use atomicAdd
// d_value/adv_logits: c51_grad + mse_grad kernels write directly (no atomicAdd)
self.stream
.memset_zeros(&mut self.d_value_logits_buf)
.map_err(|e| MLError::ModelError(format!("zero d_value_logits: {e}")))?;
@@ -4747,13 +4754,16 @@ impl GpuDqnTrainer {
self.stream.memset_zeros(&mut self.d_adv_logits_mse)
.map_err(|e| MLError::ModelError(format!("zero d_adv_mse: {e}")))?;
self.launch_mse_loss()?;
self.launch_loss_reduce(&self.mse_loss_buf)?;
self.launch_mse_grad_to_scratch()?;
// C51 path → main buffers (already zeroed above)
self.stream.memset_zeros(&mut self.q_divergence_buf)
.map_err(|e| MLError::ModelError(format!("zero q_divergence: {e}")))?;
self.launch_c51_loss()?;
self.launch_loss_reduce(&self.total_loss_buf)?;
self.launch_c51_mixup()?;
self.launch_loss_reduce(&self.total_loss_buf)?;
self.launch_c51_grad()?;
// Blend: main = α * C51 + (1-α) * MSE
@@ -4885,13 +4895,16 @@ impl GpuDqnTrainer {
self.stream.memset_zeros(&mut self.d_adv_logits_mse)
.map_err(|e| MLError::ModelError(format!("zero d_adv_mse: {e}")))?;
self.launch_mse_loss()?;
self.launch_loss_reduce(&self.mse_loss_buf)?;
self.launch_mse_grad_to_scratch()?;
// C51 path
self.stream.memset_zeros(&mut self.q_divergence_buf)
.map_err(|e| MLError::ModelError(format!("zero q_divergence: {e}")))?;
self.launch_c51_loss()?;
self.launch_loss_reduce(&self.total_loss_buf)?;
self.launch_c51_mixup()?;
self.launch_loss_reduce(&self.total_loss_buf)?;
self.launch_c51_grad()?;
// Blend: main = α * C51 + (1-α) * MSE
@@ -4977,6 +4990,7 @@ impl GpuDqnTrainer {
// MSE loss + grad → MAIN buffers directly (no scratch, no SAXPY blend)
self.launch_mse_loss()?;
self.launch_loss_reduce(&self.mse_loss_buf)?;
self.launch_mse_grad_inner(&self.d_value_logits_buf, &self.d_adv_logits_buf)?;
// backward reads f32 d_logits directly
@@ -5491,7 +5505,7 @@ impl GpuDqnTrainer {
/// Launch C51 manifold mixup CE kernel (separate from c51_loss_batched).
/// Reads save_projected + save_current_lp, mixes with random partner,
/// overwrites per_sample_loss + td_errors + total_loss.
/// overwrites per_sample_loss + td_errors. total_loss written by launch_loss_reduce after.
/// Must be called AFTER launch_c51_loss on the same stream.
pub fn launch_c51_mixup(&mut self) -> Result<(), MLError> {
if self.config.mixup_alpha <= 0.0 { return Ok(()); }
@@ -5499,7 +5513,7 @@ impl GpuDqnTrainer {
let b = self.config.batch_size;
let na = self.config.num_atoms;
// Zero total_loss before mixup overwrites it
// Zero total_loss before mixup overwrites per_sample_loss (reduce runs after)
self.stream.memset_zeros(&mut self.total_loss_buf)
.map_err(|e| MLError::ModelError(format!("mixup zero total_loss: {e}")))?;
@@ -5536,6 +5550,26 @@ impl GpuDqnTrainer {
Ok(())
}
/// Deterministic loss reduction: sequential sum of per_sample_loss → total_loss / batch_size.
/// Grid=(1,1,1), Block=(1,1,1). Replaces atomicAdd for fully deterministic training.
fn launch_loss_reduce(&self, total_loss_buf: &CudaSlice<f32>) -> Result<(), MLError> {
let b = self.config.batch_size as i32;
unsafe {
self.stream
.launch_builder(&self.c51_loss_reduce_kernel)
.arg(&self.per_sample_loss_buf)
.arg(total_loss_buf)
.arg(&b)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("c51_loss_reduce launch: {e}")))?;
}
Ok(())
}
/// Launch C51 loss gradient kernel: save_current_lp, save_projected → dL/d_logits.
///
/// Computes dL/d_combined = is_weights * (exp(current_lp) - projected) per branch,
@@ -5558,7 +5592,9 @@ impl GpuDqnTrainer {
let total_branch_atoms_i32 = total_branch_atoms as i32;
let entropy_coeff = self.config.entropy_coefficient;
let blocks = ((b * 4 * na + 255) / 256) as u32;
// Grid = B*NA (one thread per (sample, atom), loops over 4 branches).
// Zero atomicAdd — fully deterministic gradient computation.
let blocks = ((b * na + 255) / 256) as u32;
unsafe {
self.stream
@@ -5746,7 +5782,8 @@ impl GpuDqnTrainer {
let b3_i32 = b3 as i32;
let total_branch_atoms_i32 = total_branch_atoms as i32;
let blocks = ((b * 4 * na + 255) / 256) as u32;
// Grid = B*NA (deterministic: one thread per (sample, atom), loops 4 branches)
let blocks = ((b * na + 255) / 256) as u32;
unsafe {
self.stream
@@ -6751,7 +6788,7 @@ fn compile_relu_mask_standalone(stream: &Arc<CudaStream>) -> Result<CudaFunction
fn compile_c51_loss_kernel(
stream: &Arc<CudaStream>,
_config: &GpuDqnTrainConfig,
) -> Result<(CudaFunction, CudaFunction), MLError> {
) -> Result<(CudaFunction, CudaFunction, CudaFunction), MLError> {
let context = stream.context();
let module = context.load_cubin(C51_LOSS_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("c51_loss cubin load: {e}")))?;
@@ -6759,7 +6796,9 @@ fn compile_c51_loss_kernel(
.map_err(|e| MLError::ModelError(format!("c51_loss_batched load: {e}")))?;
let mixup = module.load_function("c51_mixup_ce")
.map_err(|e| MLError::ModelError(format!("c51_mixup_ce load: {e}")))?;
Ok((loss, mixup))
let reduce = module.load_function("c51_loss_reduce")
.map_err(|e| MLError::ModelError(format!("c51_loss_reduce load: {e}")))?;
Ok((loss, mixup, reduce))
}
/// Load the C51 loss gradient kernel from precompiled cubin.

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

@@ -853,9 +853,16 @@ impl GpuExperienceCollector {
.alloc_zeros::<u32>(alloc_episodes)
.map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?;
// Initialize RNG seeds
// Deterministic RNG seeds: each episode gets a unique but reproducible seed
// derived from a fixed base. This ensures identical training trajectories
// across runs with the same hyperparameters, making hyperopt comparisons
// valid (comparing hyperparams, not init luck).
let rng_seeds: Vec<u32> = (0..alloc_episodes)
.map(|_| fastrand::u32(..))
.map(|i| {
let mut s = 0xDEAD_BEEF_u64.wrapping_add(i as u64);
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
(s >> 32) as u32
})
.collect();
stream
.memcpy_htod(&rng_seeds, &mut rng_states)
@@ -2276,9 +2283,13 @@ impl GpuExperienceCollector {
}
super::htod_f32(&self.stream, &portfolio_init, &mut self.portfolio_states)?;
// Fresh RNG seeds
// Deterministic RNG seeds (same derivation as construction)
let rng_seeds: Vec<u32> = (0..self.alloc_episodes)
.map(|_| fastrand::u32(..))
.map(|i| {
let mut s = 0xDEAD_BEEF_u64.wrapping_add(i as u64);
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
(s >> 32) as u32
})
.collect();
self.stream
.memcpy_htod(&rng_seeds, &mut self.rng_states)

4
crates/ml/src/cuda_pipeline/gpu_her.rs
View File

@@ -343,7 +343,9 @@ impl GpuHer {
her_batch_size: usize,
) -> Vec<i32> {
use rand::Rng;
let mut rng = rand::thread_rng();
use rand::SeedableRng;
use rand::rngs::StdRng;
let mut rng = StdRng::seed_from_u64(0x4E4_5678);
let mut donors = Vec::with_capacity(her_batch_size);
donors.resize_with(her_batch_size, || rng.gen_range(0..buffer_size as i32));
donors

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

@@ -447,12 +447,14 @@ fn init_xavier_weights(
config: &GpuIqlConfig,
) -> Result<CudaSlice<f32>, MLError> {
use rand::Rng;
use rand::SeedableRng;
use rand::rngs::StdRng;
let total = config.total_params();
let h = config.value_hidden_dim;
let sd = config.state_dim;
let mut weights = vec![0.0_f32; total];
let mut rng = rand::thread_rng();
let mut rng = StdRng::seed_from_u64(0x1C1_9ABC);
// Layer 1: w1[H, SD], b1[H]
let limit1 = (6.0_f64 / (sd + h) as f64).sqrt() as f32;

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

@@ -1112,6 +1112,8 @@ fn init_iqn_xavier_weights(
config: &GpuIqnConfig,
) -> Result<CudaSlice<f32>, MLError> {
use rand::Rng;
use rand::SeedableRng;
use rand::rngs::StdRng;
let total = config.total_params();
let h = config.hidden_dim;
@@ -1123,7 +1125,7 @@ fn init_iqn_xavier_weights(
// cuBLAS tile padding: last tensor (b_fc, small) can be overread by 32-element tiles
let cublas_pad = 32 * h;
let mut weights = vec![0.0_f32; total + cublas_pad];
let mut rng = rand::thread_rng();
let mut rng = StdRng::seed_from_u64(0x1CA_1234);
let mut offset = 0;
// W_embed [H, D]

16
crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs
View File

@@ -348,9 +348,13 @@ impl GpuPpoExperienceCollector {
let mut portfolio_states = portfolio_states;
super::htod_f32(&stream, &portfolio_init, &mut portfolio_states)?;
// ---- Step 5: Initialize RNG seeds ----
// ---- Step 5: Deterministic RNG seeds ----
let rng_seeds: Vec<u32> = (0..MAX_EPISODES)
.map(|_| fastrand::u32(..))
.map(|i| {
let mut s = 0xAA0_5EED_u64.wrapping_add(i as u64);
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
(s >> 32) as u32
})
.collect();
let mut rng_states = rng_states;
stream
@@ -786,9 +790,11 @@ impl GpuPpoExperienceCollector {
MLError::ModelError(format!("Failed to reset diversity_metas: {e}"))
})?;
// Fresh RNG seeds using pre-allocated staging buffer
for slot in &mut self.rng_seed_staging {
*slot = fastrand::u32(..);
// Deterministic RNG seeds using pre-allocated staging buffer
for (i, slot) in self.rng_seed_staging.iter_mut().enumerate() {
let mut s = 0xAA0_5EED_u64.wrapping_add(i as u64).wrapping_add(0x2000);
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
*slot = (s >> 32) as u32;
}
self.stream
.memcpy_htod(&self.rng_seed_staging, &mut self.rng_states)

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

@@ -1,13 +1,13 @@
/**
* MSE loss gradient kernel through softmax expectation.
* MSE loss gradient kernel — FULLY DETERMINISTIC.
*
* Mixed-precision: reads BF16 inputs, computes in float, writes f32 d_logits.
* f32 atomicAdd eliminates bf16 overflow that caused NaN.
* Zero atomicAdd: one thread per (b, j), loops over 4 branches.
* Same restructure as c51_grad_kernel for deterministic training.
*
* For each sample [b], branch [d], atom [j]:
* For each sample [b], atom [j], branch [d]:
* d_logit_j = td_error * is_weight * p_j * (z_j - E[Q])
*
* Launch config: grid=(ceil(batch_size*4*num_atoms/256), 1, 1), block=(256, 1, 1).
* Launch config: grid=(ceil(batch_size*num_atoms/256), 1, 1), block=(256, 1, 1).
*/
extern "C" __global__ void mse_grad_kernel(
@@ -15,7 +15,7 @@ extern "C" __global__ void mse_grad_kernel(
const float* __restrict__ save_eq_td, // [B, 4, NA] layout: [td_error, E_Q, 0, ...]
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_value_logits, // [B, NA] f32
float* __restrict__ d_adv_logits, // [B, (B0+B1+B2+B3)*NA] f32
int batch_size,
int num_atoms,
@@ -24,86 +24,70 @@ extern "C" __global__ void mse_grad_kernel(
const float* __restrict__ v_range_buf) /* [2] adaptive C51 z-support: [v_min, v_max] */
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total_grad_elems = batch_size * 4 * num_atoms;
if (tid >= total_grad_elems) return;
int total_elems = batch_size * num_atoms;
if (tid >= total_elems) return;
/* Read adaptive v_range from device buffer (graph-safe, L1 cached) */
float v_min = v_range_buf[0];
float v_max = v_range_buf[1];
int j = tid % num_atoms;
int d = (tid / num_atoms) % 4;
int b = tid / (4 * num_atoms);
/* Read all inputs as float (is_weights already f32, rest BF16 → float at boundary) */
float isw = fminf(is_weights[b], 10.0f); /* Clamp PER IS-weights */
float p_j = (float)save_probs[tid];
int base = (b * 4 + d) * num_atoms;
float td_error = (float)save_eq_td[base + 0];
float e_q = (float)save_eq_td[base + 1];
int j = tid % num_atoms;
int b = tid / num_atoms;
float isw = fminf(is_weights[b], 10.0f);
float inv_batch = 1.0f / (float)batch_size;
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 (float arithmetic).
*
* MEAN-reduce: divide by batch_size so gradient scale is invariant to
* batch size. Without this, the budget clip is 282× more aggressive on
* H100 (batch=16384) than on smoke test (batch=58). */
float inv_batch = 1.0f / (float)batch_size;
float d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q);
/* Per-branch entropy boost: prevent distribution collapse.
* Magnitude (d==1): 0.005 = 5× base. Order (d==2): 0.003 = 3× base. */
if (d == 1 || d == 2) {
float ent_weight = (d == 1) ? 0.005f : 0.003f;
float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
d_combined += ent_weight * (1.0f + lp_approx);
}
/* Route through dueling: d_value[b,j] += d_combined.
* d_value_logits is f32 — native atomicAdd, no overflow risk. */
atomicAdd(&d_value_logits[b * num_atoms + j], d_combined);
/* Factored action decode */
/* Factored action decode (shared across branches) */
int factored = actions[b];
int max_action = b0_size * b1_size * b2_size * b3_size;
if (factored < 0 || factored >= max_action) factored = 0;
int branch_sizes[4];
branch_sizes[0] = b0_size;
branch_sizes[1] = b1_size;
branch_sizes[2] = b2_size;
branch_sizes[3] = b3_size;
int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };
int branch_actions[4];
branch_actions[0] = factored / (b1_size * b2_size * b3_size);
branch_actions[1] = (factored / (b2_size * b3_size)) % b1_size;
branch_actions[2] = (factored / b3_size) % b2_size;
branch_actions[3] = factored % b3_size;
int a_d;
if (d == 0) a_d = factored / (b1_size * b2_size * b3_size);
else if (d == 1) a_d = (factored / (b2_size * b3_size)) % b1_size;
else if (d == 2) a_d = (factored / b3_size) % b2_size;
else a_d = factored % b3_size;
int A_d = branch_sizes[d];
float inv_A = 1.0f / (float)A_d;
/* Branch-major layout — see c51_grad_kernel.cu for full explanation. */
float d_val_sum = 0.0f;
int branch_base = 0;
for (int dd = 0; dd < d; dd++)
branch_base += batch_size * branch_sizes[dd] * num_atoms;
/* Per-branch loss weighting: magnitude (d==1) gets 2.0× MSE gradient.
* MSE loss is variance-neutral (optimizes expected Q only, no distributional
* shape bias). With C51 zeroed for magnitude (IQN is primary distributional),
* moderate MSE amplification ensures the magnitude branch head gets a strong
* direct learning signal alongside IQN's trunk gradient.
* Direction keeps 1.0× (C51 gradient active for direction). */
/* Magnitude gets 4x MSE gradient — C51 zeroed + CQL reduced, MSE is primary. */
float branch_scale = (d == 1) ? 4.0f : 1.0f;
for (int d = 0; d < 4; d++) {
int prob_idx = b * 4 * num_atoms + d * num_atoms + j;
float p_j = (float)save_probs[prob_idx];
/* d_adv[b, d, a, j] = branch_scale * d_combined * (delta(a, a_d) - 1/A_d) */
for (int a = 0; a < A_d; a++) {
float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
float grad_val = branch_scale * d_combined * dueling_grad;
int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
atomicAdd(&d_adv_logits[adv_idx], grad_val);
int eq_base = (b * 4 + d) * num_atoms;
float td_error = (float)save_eq_td[eq_base + 0];
float e_q = (float)save_eq_td[eq_base + 1];
float d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q);
/* Per-branch entropy boost */
if (d == 1 || d == 2) {
float ent_weight = (d == 1) ? 0.005f : 0.003f;
float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
d_combined += ent_weight * (1.0f + lp_approx);
}
d_val_sum += d_combined;
/* Magnitude gets 4× MSE gradient (C51 zeroed, MSE is primary) */
float branch_scale = (d == 1) ? 4.0f : 1.0f;
int A_d = branch_sizes[d];
float inv_A = 1.0f / (float)A_d;
int a_d = branch_actions[d];
/* d_adv: plain write — unique slot per thread */
for (int a = 0; a < A_d; a++) {
float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
float grad_val = branch_scale * d_combined * dueling_grad;
int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
d_adv_logits[adv_idx] = grad_val;
}
branch_base += batch_size * A_d * num_atoms;
}
/* Single deterministic write */
d_value_logits[b * num_atoms + j] = d_val_sum;
}

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

@@ -142,7 +142,7 @@ extern "C" __global__ void mse_loss_batched(
/* ── Outputs ──────────────────────────────────────────────────── */
float* __restrict__ per_sample_loss, /* [B] IS-weighted loss per sample */
float* __restrict__ td_errors, /* [B] unweighted, for PER priority update */
float* __restrict__ total_loss, /* [1] float accumulator (native atomicAdd) */
float* __restrict__ total_loss, /* [1] float accumulator (deterministic reduce) */
/* ── Saved tensors for backward pass ─────────────────────────── */
float* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, num_atoms] online probs */
@@ -508,6 +508,7 @@ extern "C" __global__ void mse_loss_batched(
float weighted_loss = avg_mse * is_weight;
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(avg_td);
atomicAdd(total_loss, weighted_loss / (float)batch_size);
/* total_loss reduced by deterministic c51_loss_reduce kernel.
* No atomicAdd — fully deterministic training. */
}
}

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

@@ -1160,11 +1160,15 @@ impl FusedTrainingCtx {
dqn.config.tau_final,
dqn.config.tau_anneal_steps,
);
let adaptive_tau = self.trainer.compute_adaptive_tau(tau as f32);
// Use cosine-annealed tau directly — fully deterministic.
// Adaptive tau (from q_divergence) was removed because q_divergence
// uses atomicAdd in the C51 loss kernel, introducing non-determinism
// into the training path. The cosine schedule provides smooth tau
// adaptation without requiring the q_divergence signal.
self.trainer.target_ema_update(
&self.online_dueling, &self.online_branching,
&self.target_dueling, &self.target_branching,
adaptive_tau,
tau as f32,
).map_err(|e| anyhow::anyhow!("GPU EMA target update: {e}"))?;
}

4
crates/ml/src/trainers/dqn/trainer/action.rs
View File

@@ -142,8 +142,10 @@ impl DQNTrainer {
pub(crate) async fn epsilon_greedy_action(&self, state: &GpuTensor) -> Result<usize> {
use rand::Rng;
use rand::SeedableRng;
use rand::rngs::StdRng;
let epsilon = self.get_epsilon().await? as f32;
let mut rng = rand::thread_rng();
let mut rng = StdRng::seed_from_u64(0xAC7_DEF0);
if rng.gen::<f32>() < epsilon { Ok(rng.gen_range(0..5)) }
else {
let agent = self.agent.read().await;