fix(engine): CCCL overflow_cast + checked_allocator patterns for NaN/OOM

CCCL overflow_cast.h pattern applied to qwen3_5.py:
- Prefill gate: A_log.float().clamp(-20,20).exp() prevents NaN cascade
- Decode gate: same clamp before exp (was unprotected, unlike prefill path)
- Decode g_t: clamp_(-20,20) before in-place exp_() (was raw exp_())
  Docker logs show 99.98% NaN in GatedDeltaNet layers — these unprotected
  exp() calls are the root cause.

CCCL checked_allocator.cuh pattern applied to model_runner.py:
- Wrap model forward in try/except torch.cuda.OutOfMemoryError
- On OOM: empty_cache + gc.collect + retry once
- Competitor Sub168 died permanently at layernorm x.float() OOM
  during replay (docker log evidence). This recovery keeps server alive.

Source: cccl_upstream/libcudacxx/include/cuda/__numeric/overflow_cast.h
Source: cccl_upstream/c2h/include/c2h/checked_allocator.cuh
This commit is contained in:
project6
2026-08-07 08:56:36 +00:00
parent 391866785e
commit 5a3bcbc247
2 changed files with 37 additions and 13 deletions

View File

@@ -401,7 +401,11 @@ class GatedDeltaNet(nn.Module):
v = v.reshape(1, seq_len, local_num_v, self.head_v_dim)
beta = b_all[s:e].sigmoid().unsqueeze(0) # (1, seq_len, local_num_v)
g = (-self.A_log.float().exp()
# CCCL overflow_cast pattern: clamp before exp to prevent
# overflow → NaN cascade. A_log.exp() can exceed float32 range
# when A_log > ~88; clamping to [-20,20] keeps exp in safe range.
_A_safe = self.A_log.float().clamp(-20.0, 20.0)
g = (-_A_safe.exp()
* F.softplus(a_all[s:e].float() + self.dt_bias)
).unsqueeze(0) # (1, seq_len, local_num_v)
@@ -476,7 +480,9 @@ class GatedDeltaNet(nn.Module):
v = v.reshape(num_seqs, 1, local_num_v, self.head_v_dim)
beta = b_all.sigmoid().unsqueeze(1) # (num_seqs, 1, local_num_v)
g = (-self.A_log.float().exp()
# CCCL overflow_cast pattern: clamp before exp (same as prefill path)
_A_safe = self.A_log.float().clamp(-20.0, 20.0)
g = (-_A_safe.exp()
* F.softplus(a_all.float() + self.dt_bias)
).unsqueeze(1) # (num_seqs, 1, local_num_v)
@@ -494,7 +500,7 @@ class GatedDeltaNet(nn.Module):
q_t = _l2norm(q.squeeze(1)).float() * _scale # (B, H_v, k_dim)
k_t = _l2norm(k.squeeze(1)).float() # (B, H_v, k_dim)
v_t = v.squeeze(1).float() # (B, H_v, v_dim)
g_t = g.squeeze(1).float().exp_() # (B, H_v)
g_t = g.squeeze(1).float().clamp_(-20.0, 20.0).exp_() # (B, H_v) overflow_cast
bt = beta.squeeze(1).float() # (B, H_v)
# Decay state in-place: (B, H_v, k_dim, v_dim) *= scalar per head