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

fix(critical): c51_grad writes branch-major d_adv_logits layout

Browse Source 
The old sample-major interleaved layout did not match cuBLAS backward
branch-major pointer arithmetic. At batch=16384, backward GEMM read
cross-contaminated data producing zero weight gradients.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 21:56:38 +02:00
parent 4e95f5142d
commit 1f910d0aac
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -78,9 +78,12 @@ extern "C" __global__ void c51_grad_kernel(
int A_d = branch_sizes[d];
float inv_A = 1.0f / (float)A_d;
int branch_offset = 0;
/* 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. */
int branch_base = 0;
for (int dd = 0; dd < d; dd++)
branch_offset += branch_sizes[dd] * num_atoms;
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
@@ -97,7 +100,7 @@ extern "C" __global__ void c51_grad_kernel(
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 = b * total_branch_atoms + branch_offset + a * num_atoms + j;
int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
atomicAdd(&d_adv_logits[adv_idx], grad_val);
}
}