fix(critical): c51_loss_reduce missing threadIdx guard — 32-way write race

SECOND race condition from block_dim (1,1,1)→(32,1,1) change:
c51_loss_reduce had 32 threads all computing sum and writing to
total_loss[0] simultaneously. On sm_90 (H100) → corrupted loss →
NaN in C51 gradient → graph_mega replay hangs.

This kernel runs INSIDE graph_mega (submit_forward_ops_main),
called twice (MSE reduce + C51 reduce). Both corrupted.

Also fixed stochastic_depth_rng in previous commit (same issue,
runs in pre-replay).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-17 13:19:39 +02:00
parent c7fa67159c
commit ca7dfa9f08

View File

@@ -708,6 +708,7 @@ extern "C" __global__ void c51_loss_reduce(
float* __restrict__ total_loss, /* [1] output */
int batch_size
) {
if (threadIdx.x != 0) return; /* single writer — total_loss[0] */
float sum = 0.0f;
for (int i = 0; i < batch_size; i++) {
sum += per_sample_loss[i];