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:
@@ -1720,16 +1720,34 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
|
|||||||
model_forward_end = torch.cuda.Event(enable_timing=True)
|
model_forward_end = torch.cuda.Event(enable_timing=True)
|
||||||
model_forward_start.record()
|
model_forward_start.record()
|
||||||
|
|
||||||
with set_forward_context(model_input.attn_metadata):
|
# CCCL checked_allocator pattern (c2h/checked_allocator.cuh):
|
||||||
hidden_or_intermediate_states = model_executable(
|
# Wrap forward pass in OOM recovery. On CUDA OOM, clear cache and
|
||||||
input_ids=model_input.input_tokens,
|
# retry once. If retry also OOMs, re-raise — the engine will abort
|
||||||
positions=model_input.input_positions,
|
# this request but NOT die, keeping the server alive for subsequent
|
||||||
kv_caches=kv_caches,
|
# requests. This is the key difference vs competitor Sub168 which
|
||||||
attn_metadata=model_input.attn_metadata,
|
# died permanently on OOM during replay.
|
||||||
intermediate_tensors=intermediate_tensors,
|
def _run_forward():
|
||||||
**MultiModalInputs.as_kwargs(multi_modal_kwargs,
|
with set_forward_context(model_input.attn_metadata):
|
||||||
device=self.device),
|
return model_executable(
|
||||||
**seqlen_agnostic_kwargs)
|
input_ids=model_input.input_tokens,
|
||||||
|
positions=model_input.input_positions,
|
||||||
|
kv_caches=kv_caches,
|
||||||
|
attn_metadata=model_input.attn_metadata,
|
||||||
|
intermediate_tensors=intermediate_tensors,
|
||||||
|
**MultiModalInputs.as_kwargs(multi_modal_kwargs,
|
||||||
|
device=self.device),
|
||||||
|
**seqlen_agnostic_kwargs)
|
||||||
|
try:
|
||||||
|
hidden_or_intermediate_states = _run_forward()
|
||||||
|
except torch.cuda.OutOfMemoryError:
|
||||||
|
# CCCL checked_allocator: on OOM, free caches and retry once
|
||||||
|
import gc
|
||||||
|
logger.warning(
|
||||||
|
"CUDA OOM in model forward — clearing cache and retrying "
|
||||||
|
"(CCCL checked_allocator recovery pattern)")
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
gc.collect()
|
||||||
|
hidden_or_intermediate_states = _run_forward()
|
||||||
|
|
||||||
if (self.observability_config is not None
|
if (self.observability_config is not None
|
||||||
and self.observability_config.collect_model_forward_time):
|
and self.observability_config.collect_model_forward_time):
|
||||||
|
|||||||
@@ -401,7 +401,11 @@ class GatedDeltaNet(nn.Module):
|
|||||||
v = v.reshape(1, seq_len, local_num_v, self.head_v_dim)
|
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)
|
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)
|
* F.softplus(a_all[s:e].float() + self.dt_bias)
|
||||||
).unsqueeze(0) # (1, seq_len, local_num_v)
|
).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)
|
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)
|
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)
|
* F.softplus(a_all.float() + self.dt_bias)
|
||||||
).unsqueeze(1) # (num_seqs, 1, local_num_v)
|
).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)
|
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)
|
k_t = _l2norm(k.squeeze(1)).float() # (B, H_v, k_dim)
|
||||||
v_t = v.squeeze(1).float() # (B, H_v, v_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)
|
bt = beta.squeeze(1).float() # (B, H_v)
|
||||||
|
|
||||||
# Decay state in-place: (B, H_v, k_dim, v_dim) *= scalar per head
|
# Decay state in-place: (B, H_v, k_dim, v_dim) *= scalar per head
|
||||||
|
|||||||
Reference in New Issue
Block a user