Reduced Triton sigmoid focal loss returns NaN gradients for gamma=0 when logits saturate #575

Summary

triton_sigmoid_focal_loss_reduce can return NaN gradients for finite inputs when gamma=0 and logits saturate in the correct direction. This matters for SAM3 training because IABCEMdetr defaults presence_gamma=0.0, so the presence loss can hit this reduced Triton focal path.

For gamma=0, sigmoid focal loss should reduce to alpha-weighted BCE, whose gradient is finite:

dL/dx = alpha_t * (sigmoid(x) - target)

However, the Triton backward computes (1 - p_t) ** (gamma - 1) before simplification. When gamma=0 and p_t == 1, this becomes 0 ** -1 = inf, and later terms produce 0 * inf = NaN.

Minimal Reproduction

import torch
from sam3.train.loss.sigmoid_focal_loss import triton_sigmoid_focal_loss_reduce

x = torch.tensor([[18.0]], device="cuda", requires_grad=True)
y = torch.tensor([[1.0]], device="cuda")

loss = triton_sigmoid_focal_loss_reduce(x, y, 0.5, 0.0)
loss.backward()

print(loss, x.grad)

Observed:

loss = tensor(0., device='cuda:0', grad_fn=<SigmoidFocalLossReducedBackward>)
x.grad = tensor([[nan]], device='cuda:0')

The same issue occurs for x=-18.0, y=0.0.

Expected Behavior

The gradient should be finite. With gamma=0, the loss is equivalent to alpha-weighted BCE, so saturated correct predictions should produce a zero or very small finite gradient, not NaN.

Related Presence Logit Clamp Issue

There is also an ineffective clamp in sam3/model/decoder.py:

intermediate_layer_presence_logits.clamp(...)

Since the return value is not assigned and this is not clamp_, the intended presence-logit clamp is not applied. This can let presence logits reach the saturation range that triggers the gamma=0 backward issue.

Environment

torch 2.12.0+cu132
triton 3.7.0
CUDA 13.2
GPU: NVIDIA RTX 6000 Ada Generation

Proposed Fix

A small fix is to route sigmoid_focal_loss(..., gamma=0) through the existing PyTorch fallback path, where the expression is numerically stable and mathematically equivalent to BCE, and to assign the presence clamp result.