facebookresearch / facebookresearch/sam3

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

Open
#575 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
11.7k
Forks
1.8k
PR merge metrics
No merged PRs in 30d

Description

### 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:

```text
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

```python
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:

```text
loss = tensor(0., device='cuda:0', grad_fn=)
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`:

```python
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

```text
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.

Contributor guide

Open the contributing guide

Research direction

Start with sam3/train/loss/sigmoid_focal_loss.py and reproduce the issue using the CUDA example with gamma=0 and saturated logits. Compare the Triton path with the existing PyTorch fallback, then inspect sam3/model/decoder.py for the ineffective presence-logit clamp. Done means finite gradients for the reported cases and an applied presence clamp.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.