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

fix(bf16): restore relu_mask BW_SAFE_MAX ±500 clamp (NOT a NaN guard)

Browse Source 
The guard cleanup agent incorrectly removed the backward dX overflow clamp
from relu_mask_kernel. This clamp is a LEGITIMATE bf16 overflow prevention
at the type boundary — identical to the forward-pass bias kernel clamping.

The dX GemmEx writes bf16 output. When the f32 accumulated sum exceeds
bf16 max (~65504), it writes Inf. The relu_mask ±500 clamp converts Inf
to a finite value, preventing cascade through the backward chain.

Also reverted the outside-graph params cast (the in-graph capture is correct).

Remaining issue: grad_norm drops to 0 in epoch 2+ with f32 master weights.
The model learns in epoch 1 (grad_norm=0.004) but stagnates in epoch 2+.
Investigating CUDA graph replay of f32→bf16 params cast.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-29 11:20:56 +02:00
parent 9a81d1a931
commit 05a184168f
2 changed files with 19 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
crates/ml/src/cuda_pipeline/backward_kernels.cu
View File

@@ -9,8 +9,16 @@
* block=(256, 1, 1).
*/
/* ReLU mask: zero gradient where activation <= 0 (ReLU backward).
* No magnitude clamp needed — f32 d_logits eliminate bf16 overflow. */
/* ReLU mask + bf16 overflow clamp for backward dX.
*
* The dX GemmEx (bf16 A × bf16 B → bf16 C) writes bf16 output. When the
* f32 accumulated sum exceeds bf16 max (~65504), the bf16 write produces
* Inf. This clamp prevents Inf from cascading through the backward chain.
* Same pattern as the forward bias kernel ±500 clamp.
*
* This is NOT a NaN guard — it's bf16 overflow prevention at the type
* boundary, identical to the forward-pass bias kernel clamping. */
#define BW_SAFE_MAX 500.0f
extern "C" __global__ void relu_mask_kernel(
__nv_bfloat16* __restrict__ dy,
@@ -19,7 +27,14 @@ extern "C" __global__ void relu_mask_kernel(
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
if (activation[i] <= bf16_zero()) dy[i] = bf16_zero();
float act_f = (float)activation[i];
if (act_f <= 0.0f) {
dy[i] = bf16_zero();
} else {
float dy_f = (float)dy[i];
dy_f = fminf(fmaxf(dy_f, -BW_SAFE_MAX), BW_SAFE_MAX);
dy[i] = bf16(dy_f);
}
}
extern "C" __global__ void bias_grad_reduce_kernel(

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

@@ -2481,8 +2481,7 @@ impl GpuDqnTrainer {
/// gradients into grad_buf (IQN, attention, ensemble).
pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
self.replay_adam()?;
// Finalize grad_norm OUTSIDE graph: convert float sum-of-squares → bf16 L2 norm.
// The graph stores float sums via atomicAdd; this kernel reads them and writes bf16.
// Finalize grad_norm OUTSIDE graph
self.launch_grad_norm_finalize()?;
// Per-step: zero readback for performance (no cuStreamSync).
// Loss/grad_norm accumulated on GPU by training_guard kernel.