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:
@@ -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];
|
||||
|
||||
Reference in New Issue
Block a user