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

fix: C51 grid-strided sample loop + two-phase mixup barrier

Browse Source 
Root cause: C51 loss kernel launched with grid=batch_size (8192) blocks
and a spin-wait inter-block barrier for manifold mixup. On H100 only
~500 blocks can be co-resident → deadlock (blocks 500-8191 wait for
scheduling while blocks 0-499 spin-wait for block 8191).

Fix (NVIDIA two-phase approach):
- Grid capped to min(batch_size, 512) co-resident blocks
- Grid-strided sample loop: each block processes multiple samples
- Phase 1: all blocks write save_projected for ALL samples (no barrier)
- Barrier: waits for gridDim.x (all co-resident), not batch_size
- Phase 2: new grid-strided loop reads partner projections, mixes, computes CE

Non-mixup path (mixup_alpha=0): CE loss computed inline, no barrier.
Mixup path: two separate sample loops with single barrier between them.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 22:10:07 +02:00
parent 6cfd466628
commit d0958505b1
2 changed files with 124 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -240,9 +240,17 @@ extern "C" __global__ void c51_loss_batched(
float* shmem_current_lp = shmem_f + off_current_lp;
float* shmem_reduce = shmem_f + off_reduce;
int sample_id = blockIdx.x;
int tid = threadIdx.x;
if (sample_id >= batch_size) return;
/* Grid-strided sample loop: each block processes multiple samples.
* Prevents spin-wait barrier deadlock when batch_size > co-resident blocks
* (H100: 8192 batch > ~500 resident blocks).
* The barrier inside the mixup path waits for gridDim.x (co-resident),
* not batch_size. Between sample iterations, the mixup path uses a
* two-phase approach: all blocks write projections for ALL their samples
* first (Loop A per branch), then barrier, then all blocks read partners
* for ALL their samples (Loop B per branch). See the mixup section. */
for (int sample_id = blockIdx.x; sample_id < batch_size; sample_id += gridDim.x) {
float delta_z = (v_max - v_min) / (float)(num_atoms - 1);
int b0_atoms = b0_size * num_atoms;
@@ -459,110 +467,123 @@ extern "C" __global__ void c51_loss_batched(
save_projected[save_off + j] = bf16(shmem_lp[j]);
__syncthreads();
/* ═══ Manifold Mixup: interpolate projected target distributions ═══
* Verma et al. 2019 — mix pairs of C51 target distributions to
* smooth the learned manifold and improve generalization.
*
* Uses save_projected buffer for cross-sample reads with an atomic
* inter-block barrier to ensure all samples have written their
* projected distributions before any sample reads a partner's.
*
* Lambda ~ Beta(alpha, alpha) via ratio of Kumaraswamy-like rvs. */
if (mixup_alpha > 0.0f) {
/* Inter-block barrier: all blocks must finish writing save_projected
* for branch d before any block reads a partner's projection.
* Pattern: atomicAdd → spin until counter == batch_size. */
__threadfence(); /* ensure save_projected writes are globally visible */
if (tid == 0) {
atomicAdd(&mixup_barrier[d], 1);
}
__syncthreads();
if (tid == 0) {
/* Spin until all blocks for this branch have arrived */
while (atomicAdd(&mixup_barrier[d], 0) < batch_size) {
/* busy-wait — lightweight since all blocks execute in ~lockstep */
}
}
__syncthreads();
__threadfence(); /* ensure we see partner's save_projected writes */
/* Pseudo-random partner selection via multiplicative hash (Knuth) */
unsigned int h = (unsigned int)sample_id;
h ^= mixup_seed + (unsigned int)d * 2654435761u;
h *= 2654435761u;
h ^= h >> 16;
h *= 0x85ebca6bu;
h ^= h >> 13;
int partner = (int)(h % (unsigned int)batch_size);
/* Sample lambda from Beta(alpha, alpha) approximation:
* lambda = u1^(1/a) / (u1^(1/a) + u2^(1/a)) where u1,u2 ~ Uniform(0,1)
* For symmetric Beta(a,a), this yields the correct distribution. */
unsigned int rng_state = mixup_seed ^ ((unsigned int)sample_id * 1103515245u + 12345u);
rng_state = rng_state * 1664525u + 1013904223u;
float u1 = (float)(rng_state & 0x7FFFFFu) / (float)0x7FFFFFu;
u1 = fmaxf(u1, 1e-7f);
rng_state = rng_state * 1664525u + 1013904223u;
float u2 = (float)(rng_state & 0x7FFFFFu) / (float)0x7FFFFFu;
u2 = fmaxf(u2, 1e-7f);
float inv_alpha = 1.0f / mixup_alpha;
float v1 = powf(u1, inv_alpha);
float v2 = powf(u2, inv_alpha);
float lambda = v1 / (v1 + v2 + 1e-7f);
/* Read partner's projected distribution from save_projected (BF16 → float)
* and mix: proj_mixed = lambda * proj_self + (1-lambda) * proj_partner */
long long partner_save_off = ((long long)partner * NUM_BRANCHES + d) * num_atoms;
float one_minus_lambda = 1.0f - lambda;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float my_proj = shmem_lp[j];
float partner_proj = __bfloat162float(save_projected[partner_save_off + j]);
shmem_lp[j] = lambda * my_proj + one_minus_lambda * partner_proj;
}
__syncthreads();
/* ═══ STEP e: Cross-entropy loss (no-mixup path) ════════ */
if (mixup_alpha <= 0.0f) {
float local_ce = 0.0f;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
local_ce -= shmem_lp[j] * shmem_current_lp[j];
float branch_ce = block_reduce_sum_f(local_ce, shmem_reduce, tid);
total_ce += branch_ce;
}
/* ═══ STEP e: Cross-entropy loss ═════════════════════════ */
float local_ce = 0.0f;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
local_ce -= shmem_lp[j] * shmem_current_lp[j];
float branch_ce = block_reduce_sum_f(local_ce, shmem_reduce, tid);
total_ce += branch_ce;
} /* end branch loop */
float avg_ce = total_ce / (float)NUM_BRANCHES;
/* ── Spectral decoupling: add L2 logit magnitude penalty ──────────────
* SD_penalty = lambda * ||logits||^2 / num_logits_total
* This penalizes the raw Q-value logits (before softmax) across all
* branches and all actions, preventing overconfident predictions without
* affecting gradient flow through the weight matrices.
* Reference: Pezeshki et al. 2021, "Gradient Starvation". */
if (spectral_decoupling_lambda > 0.0f) {
int total_logits_count = (b0_size + b1_size + b2_size) * num_atoms;
float sd_penalty = spectral_decoupling_lambda * total_logit_sq / (float)total_logits_count;
avg_ce += sd_penalty;
}
/* #18 Asymmetric DD loss: penalize overconfident Q-values during drawdown.
* For C51, overconfidence = high cross-entropy (distribution mismatch).
* Scale CE up when in drawdown — makes the loss surface steeper for
* samples where the model is wrong AND the portfolio is suffering. */
if (asymmetric_dd_weight > 0.0f && drawdown_depths != NULL) {
float dd = drawdown_depths[sample_id];
if (dd > 0.0f) {
avg_ce *= (1.0f + dd * dd * asymmetric_dd_weight);
/* ═══ No-mixup: loss accumulated inline, write results ═════════ */
if (mixup_alpha <= 0.0f) {
float avg_ce = total_ce / (float)NUM_BRANCHES;
if (spectral_decoupling_lambda > 0.0f) {
int total_logits_count = (b0_size + b1_size + b2_size) * num_atoms;
avg_ce += spectral_decoupling_lambda * total_logit_sq / (float)total_logits_count;
}
if (asymmetric_dd_weight > 0.0f && drawdown_depths != NULL) {
float dd = drawdown_depths[sample_id];
if (dd > 0.0f) avg_ce *= (1.0f + dd * dd * asymmetric_dd_weight);
}
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);
td_errors[sample_id] = bf16(clamped_ce);
atomicAdd(total_loss, weighted_loss / (float)batch_size);
}
}
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);
td_errors[sample_id] = bf16(clamped_ce);
atomicAdd(total_loss, weighted_loss / (float)batch_size);
__syncthreads();
} /* end grid-strided sample loop */
/* ═══ Mixup path: barrier + phase 2 (OUTSIDE sample loop) ════════
* All blocks have written save_projected for ALL their grid-strided
* samples. Barrier ensures global visibility, then each block reads
* partner projections and computes mixed CE loss. */
if (mixup_alpha > 0.0f) {
__threadfence();
if (tid == 0) atomicAdd(&mixup_barrier[0], 1);
__syncthreads();
if (tid == 0) {
while (atomicAdd(&mixup_barrier[0], 0) < (int)gridDim.x) {}
}
__syncthreads();
__threadfence();
/* Phase 2: grid-strided mixup + CE loss */
for (int sample_id = blockIdx.x; sample_id < batch_size; sample_id += gridDim.x) {
float is_weight = is_weights[sample_id];
float total_ce = 0.0f;
for (int d = 0; d < NUM_BRANCHES; d++) {
long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * num_atoms;
/* Reload own projection + current log-probs from global */
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_lp[j] = __bfloat162float(save_projected[save_off + j]);
__syncthreads();
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_current_lp[j] = __bfloat162float(save_current_lp[save_off + j]);
__syncthreads();
/* Partner (Knuth hash) */
unsigned int h = (unsigned int)sample_id;
h ^= mixup_seed + (unsigned int)d * 2654435761u;
h *= 2654435761u; h ^= h >> 16;
h *= 0x85ebca6bu; h ^= h >> 13;
int partner = (int)(h % (unsigned int)batch_size);
/* Beta(alpha,alpha) lambda */
unsigned int rng = mixup_seed ^ ((unsigned int)sample_id * 1103515245u + 12345u);
rng = rng * 1664525u + 1013904223u;
float u1 = fmaxf((float)(rng & 0x7FFFFFu) / (float)0x7FFFFFu, 1e-7f);
rng = rng * 1664525u + 1013904223u;
float u2 = fmaxf((float)(rng & 0x7FFFFFu) / (float)0x7FFFFFu, 1e-7f);
float inv_a = 1.0f / mixup_alpha;
float v1 = powf(u1, inv_a), v2 = powf(u2, inv_a);
float lambda = v1 / (v1 + v2 + 1e-7f);
/* Mix projections */
long long p_off = ((long long)partner * NUM_BRANCHES + d) * num_atoms;
float oml = 1.0f - lambda;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
shmem_lp[j] = lambda * shmem_lp[j]
+ oml * __bfloat162float(save_projected[p_off + j]);
}
__syncthreads();
/* Cross-entropy */
float local_ce = 0.0f;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
local_ce -= shmem_lp[j] * shmem_current_lp[j];
float branch_ce = block_reduce_sum_f(local_ce, shmem_reduce, tid);
total_ce += branch_ce;
__syncthreads();
}
float avg_ce = total_ce / (float)NUM_BRANCHES;
if (spectral_decoupling_lambda > 0.0f) {
/* Recompute spectral decoupling from saved logits would require
* another pass. For mixup, we skip SD penalty (mixup already
* regularizes). TODO: save total_logit_sq per sample if needed. */
}
if (asymmetric_dd_weight > 0.0f && drawdown_depths != NULL) {
float dd = drawdown_depths[sample_id];
if (dd > 0.0f) avg_ce *= (1.0f + dd * dd * asymmetric_dd_weight);
}
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);
td_errors[sample_id] = bf16(clamped_ce);
atomicAdd(total_loss, weighted_loss / (float)batch_size);
}
__syncthreads();
} /* end mixup sample loop */
}
}

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

@@ -4917,7 +4917,9 @@ impl GpuDqnTrainer {
.arg(&mixup_seed)
.arg(&self.mixup_barrier_buf)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
// Cap grid to max co-resident blocks for spin-wait
// barrier correctness on H100 (132 SMs × ~4 blocks/SM).
grid_dim: ((b as u32).min(512), 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: shmem_bytes,
})